diffusers-源码解析-三十三-
diffusers 源码解析（三十三）
.\diffusers\pipelines\kandinsky2_2\pipeline_kandinsky2_2_controlnet.py
# 版权所有 2024 HuggingFace 团队。保留所有权利。
#
# 根据 Apache 许可证第 2.0 版（“许可证”）进行许可；
# 除非遵守许可证，否则不得使用此文件。
# 可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，
# 根据许可证分发的软件均按“原样”提供，
# 不附有任何明示或暗示的担保或条件。
# 有关许可证所管理权限和限制的具体内容，请参阅许可证。
from typing import Callable, List, Optional, Union # 从 typing 模块导入可调用、列表、可选和联合类型
import torch # 导入 PyTorch 库
from ...models import UNet2DConditionModel, VQModel # 从模型模块导入 UNet2DConditionModel 和 VQModel 类
from ...schedulers import DDPMScheduler # 从调度器模块导入 DDPMScheduler 类
from ...utils import ( # 从 utils 模块导入 logging 工具
logging,
)
from ...utils.torch_utils import randn_tensor # 从 torch_utils 导入 randn_tensor 函数
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput # 从 pipeline_utils 导入 DiffusionPipeline 和 ImagePipelineOutput 类
logger = logging.get_logger(__name__) # 获取当前模块的日志记录器，禁用 pylint 命名检查
# 示例代码，展示如何使用深度估计和图像生成管道
Examples:
```py
# 导入 PyTorch 和 NumPy 库
>>> import torch
>>> import numpy as np
# 导入 Kandinsky V22 模型的管道和其他工具
>>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline
>>> from transformers import pipeline
>>> from diffusers.utils import load_image
# 定义生成深度提示的函数
>>> def make_hint(image, depth_estimator):
# 使用深度估计器获取图像的深度信息
... image = depth_estimator(image)["depth"]
# 将深度信息转换为 NumPy 数组
... image = np.array(image)
# 为深度图像增加一个维度
... image = image[:, :, None]
# 将深度图像复制三次，形成 RGB 格式
... image = np.concatenate([image, image, image], axis=2)
# 将 NumPy 数组转换为 PyTorch 张量并归一化
... detected_map = torch.from_numpy(image).float() / 255.0
# 调整张量的维度顺序
... hint = detected_map.permute(2, 0, 1)
# 返回生成的提示
... return hint
# 创建深度估计器管道
>>> depth_estimator = pipeline("depth-estimation")
# 加载 Kandinsky V22 先验管道，并指定数据类型为 float16
>>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16
... )
# 将管道移动到 CUDA 设备
>>> pipe_prior = pipe_prior.to("cuda")
# 加载 Kandinsky V22 控制管道，并指定数据类型为 float16
>>> pipe = KandinskyV22ControlnetPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16
... )
# 将管道移动到 CUDA 设备
>>> pipe = pipe.to("cuda")
# 从 URL 加载图像并调整大小
>>> img = load_image(
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
... "/kandinsky/cat.png"
... ).resize((768, 768))
# 生成深度提示，并调整张量维度和数据类型
>>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda")
# 定义生成图像的提示内容
>>> prompt = "A robot, 4k photo"
# 定义负提示内容，用于排除不希望出现的特征
>>> negative_prior_prompt = "lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature"
# 创建随机数生成器并设置种子
>>> generator = torch.Generator(device="cuda").manual_seed(43)
# 使用先验管道生成图像嵌入和零图像嵌入
>>> image_emb, zero_image_emb = pipe_prior(
... prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator
... ).to_tuple()
# 使用控制管道生成最终图像
>>> images = pipe(
... image_embeds=image_emb,
... negative_image_embeds=zero_image_emb,
... hint=hint,
... num_inference_steps=50,
... generator=generator,
... height=768,
... width=768,
... ).images
# 保存生成的图像
>>> images[0].save("robot_cat.png")
```py
"""
# 文档字符串，说明该模块的功能和用途
# 从 diffusers.pipelines.kandinsky2_2.pipeline_kandinsky2_2.downscale_height_and_width 复制的函数
def downscale_height_and_width(height, width, scale_factor=8):
# 根据给定的高度和宽度以及缩放因子计算新的高度
new_height = height // scale_factor**2
# 如果高度不能被缩放因子平方整除，则增加高度
if height % scale_factor**2 != 0:
new_height += 1
# 根据给定的高度和宽度以及缩放因子计算新的宽度
new_width = width // scale_factor**2
# 如果宽度不能被缩放因子平方整除，则增加宽度
if width % scale_factor**2 != 0:
new_width += 1
# 返回新的高度和宽度，乘以缩放因子以恢复到原始比例
return new_height * scale_factor, new_width * scale_factor
class KandinskyV22ControlnetPipeline(DiffusionPipeline):
"""
文档字符串，描述使用 Kandinsky 进行文本到图像生成的管道
该模型继承自 [`DiffusionPipeline`]。查看超类文档以获取库实现的通用方法（例如下载或保存、在特定设备上运行等）
参数：
scheduler ([`DDIMScheduler`]):
用于与 `unet` 结合生成图像潜变量的调度器。
unet ([`UNet2DConditionModel`]):
用于去噪图像嵌入的条件 U-Net 架构。
movq ([`VQModel`]):
MoVQ 解码器，用于从潜变量生成图像。
"""
# 定义模型中 CPU 卸载的顺序
model_cpu_offload_seq = "unet->movq"
def __init__(
self,
unet: UNet2DConditionModel,
scheduler: DDPMScheduler,
movq: VQModel,
):
# 调用父类构造函数进行初始化
super().__init__()
# 注册模块以便在管道中使用
self.register_modules(
unet=unet,
scheduler=scheduler,
movq=movq,
)
# 计算 MoVQ 的缩放因子
self.movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1)
# 从 diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents 复制的函数
def prepare_latents(self, shape, dtype, device, generator, latents, scheduler):
# 如果未提供潜变量，则生成随机潜变量
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
# 如果提供的潜变量形状不匹配，则抛出异常
if latents.shape != shape:
raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
# 将潜变量移动到指定设备
latents = latents.to(device)
# 将潜变量乘以调度器的初始噪声标准差
latents = latents * scheduler.init_noise_sigma
# 返回处理后的潜变量
return latents
@torch.no_grad()
def __call__(
self,
image_embeds: Union[torch.Tensor, List[torch.Tensor]],
negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]],
hint: torch.Tensor,
height: int = 512,
width: int = 512,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
output_type: Optional[str] = "pil",
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
return_dict: bool = True,
.\diffusers\pipelines\kandinsky2_2\pipeline_kandinsky2_2_controlnet_img2img.py
# 版权声明，表示该代码的版权所有者和保留权利
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# 根据 Apache 2.0 许可证的规定进行许可
# Licensed under the Apache License, Version 2.0 (the "License");
# 你必须遵守许可证才能使用此文件
# you may not use this file except in compliance with the License.
# 许可证可以在以下地址获取
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，软件按“原样”分发
# Unless required by applicable law or agreed to in writing, software
# 分发不提供任何明示或暗示的担保或条件
# distributed under the License is distributed on an "AS IS" BASIS,
# 参见许可证以了解特定的权限和限制
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# 从 typing 模块导入类型注解
from typing import Callable, List, Optional, Union
# 导入 numpy 库，通常用于数值计算
import numpy as np
# 导入 PIL.Image 模块，用于图像处理
import PIL.Image
# 导入 torch 库，深度学习框架
import torch
# 从 PIL 导入 Image 模块，用于图像处理
from PIL import Image
# 从相对路径导入 UNet2DConditionModel 和 VQModel 模型
from ...models import UNet2DConditionModel, VQModel
# 从相对路径导入 DDPMScheduler 调度器
from ...schedulers import DDPMScheduler
# 从相对路径导入 logging 实用程序
from ...utils import (
logging,
)
# 从相对路径导入 randn_tensor 函数
from ...utils.torch_utils import randn_tensor
# 从相对路径导入 DiffusionPipeline 和 ImagePipelineOutput
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
# 创建一个日志记录器，用于记录当前模块的日志信息
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，通常用于说明代码示例
EXAMPLE_DOC_STRING = """
# 示例代码，演示如何使用 Kandinsky 和 Transformers 进行图像处理
Examples:
```py
# 导入 PyTorch 库
>>> import torch
# 导入 NumPy 库
>>> import numpy as np
# 从 diffusers 导入所需的管道类
>>> from diffusers import KandinskyV22PriorEmb2EmbPipeline, KandinskyV22ControlnetImg2ImgPipeline
# 从 transformers 导入管道函数
>>> from transformers import pipeline
# 从 diffusers.utils 导入图像加载函数
>>> from diffusers.utils import load_image
# 定义生成提示的函数，接受图像和深度估计器作为参数
>>> def make_hint(image, depth_estimator):
... # 使用深度估计器处理图像，获取深度图
... image = depth_estimator(image)["depth"]
... # 将深度图转换为 NumPy 数组
... image = np.array(image)
... # 在数组的最后一维添加一个新的维度
... image = image[:, :, None]
... # 将深度图复制三次，形成 RGB 格式
... image = np.concatenate([image, image, image], axis=2)
... # 将 NumPy 数组转换为 PyTorch 张量并标准化到 [0, 1]
... detected_map = torch.from_numpy(image).float() / 255.0
... # 重新排列张量的维度，准备为输入
... hint = detected_map.permute(2, 0, 1)
... # 返回处理后的提示张量
... return hint
# 创建深度估计器管道
>>> depth_estimator = pipeline("depth-estimation")
# 从预训练模型加载 Kandinsky V2 Prior 管道
>>> pipe_prior = KandinskyV22PriorEmb2EmbPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16
... )
# 将管道移动到 GPU
>>> pipe_prior = pipe_prior.to("cuda")
# 从预训练模型加载 Kandinsky V2 Controlnet 图像到图像管道
>>> pipe = KandinskyV22ControlnetImg2ImgPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16
... )
# 将管道移动到 GPU
>>> pipe = pipe.to("cuda")
# 从 URL 加载图像并调整大小
>>> img = load_image(
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
... "/kandinsky/cat.png"
... ).resize((768, 768))
# 生成图像的提示，扩展维度并转换为半精度格式，移动到 GPU
>>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda")
# 设置生成的提示文本
>>> prompt = "A robot, 4k photo"
# 设置负面提示文本，避免某些特征
>>> negative_prior_prompt = "lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature"
# 创建一个随机数生成器，设置种子为 43，确保可复现性
>>> generator = torch.Generator(device="cuda").manual_seed(43)
# 生成与正面提示对应的图像嵌入
>>> img_emb = pipe_prior(prompt=prompt, image=img, strength=0.85, generator=generator)
# 生成与负面提示对应的图像嵌入
>>> negative_emb = pipe_prior(prompt=negative_prior_prompt, image=img, strength=1, generator=generator)
# 使用控制管道生成图像，包含多种参数设置
>>> images = pipe(
... image=img,
... strength=0.5,
... image_embeds=img_emb.image_embeds,
... negative_image_embeds=negative_emb.image_embeds,
... hint=hint,
... num_inference_steps=50,
... generator=generator,
... height=768,
... width=768,
... ).images
# 将生成的第一张图像保存到文件
>>> images[0].save("robot_cat.png")
# Copied from diffusers.pipelines.kandinsky2_2.pipeline_kandinsky2_2.downscale_height_and_width
def downscale_height_and_width(height, width, scale_factor=8):
# 计算新的高度，按照比例因子缩小
new_height = height // scale_factor**2
# 如果高度不是比例因子的整数倍，向上取整
if height % scale_factor**2 != 0:
new_height += 1
# 计算新的宽度，按照比例因子缩小
new_width = width // scale_factor**2
# 如果宽度不是比例因子的整数倍，向上取整
if width % scale_factor**2 != 0:
new_width += 1
# 返回调整后的高度和宽度
return new_height * scale_factor, new_width * scale_factor
# Copied from diffusers.pipelines.kandinsky.pipeline_kandinsky_img2img.prepare_image
def prepare_image(pil_image, w=512, h=512):
# 调整 PIL 图像大小为指定宽度和高度，使用双三次插值法
pil_image = pil_image.resize((w, h), resample=Image.BICUBIC, reducing_gap=1)
# 将 PIL 图像转换为 RGB 数组，并归一化到 [-1, 1] 范围内的浮点数
arr = np.array(pil_image.convert("RGB"))
arr = arr.astype(np.float32) / 127.5 - 1
# 调整数组维度顺序为 [通道, 高度, 宽度]，并转换为 PyTorch 张量
arr = np.transpose(arr, [2, 0, 1])
image = torch.from_numpy(arr).unsqueeze(0)
# 返回处理后的图像张量
return image
class KandinskyV22ControlnetImg2ImgPipeline(DiffusionPipeline):
"""
Pipeline for image-to-image generation using Kandinsky
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
Args:
scheduler ([`DDIMScheduler`]):
A scheduler to be used in combination with `unet` to generate image latents.
unet ([`UNet2DConditionModel`]):
Conditional U-Net architecture to denoise the image embedding.
movq ([`VQModel`]):
MoVQ Decoder to generate the image from the latents.
"""
model_cpu_offload_seq = "unet->movq"
def __init__(
self,
unet: UNet2DConditionModel,
scheduler: DDPMScheduler,
movq: VQModel,
):
super().__init__()
# 注册模块到管道中
self.register_modules(
unet=unet,
scheduler=scheduler,
movq=movq,
)
# 计算 MoVQ 缩放因子
self.movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1)
# Copied from diffusers.pipelines.kandinsky.pipeline_kandinsky_img2img.KandinskyImg2ImgPipeline.get_timesteps
def get_timesteps(self, num_inference_steps, strength, device):
# 根据推断步数和强度计算初始时间步长
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
# 计算起始时间步，确保不超出范围
t_start = max(num_inference_steps - init_timestep, 0)
# 获取调度器的时间步长
timesteps = self.scheduler.timesteps[t_start:]
# 返回时间步长列表和有效时间步数
return timesteps, num_inference_steps - t_start
# Copied from diffusers.pipelines.kandinsky2_2.pipeline_kandinsky2_2_img2img.KandinskyV22Img2ImgPipeline.prepare_latents
# 准备潜在变量，用于图像生成模型
def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None):
# 检查输入的图像类型是否为张量、PIL图像或列表
if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
# 如果不是，则抛出值错误，说明类型不匹配
raise ValueError(
f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
)
# 将图像转换到指定设备和数据类型
image = image.to(device=device, dtype=dtype)
# 计算有效的批量大小
batch_size = batch_size * num_images_per_prompt
# 如果图像的通道数为4，初始化潜在变量为图像本身
if image.shape[1] == 4:
init_latents = image
else:
# 如果生成器是列表并且其长度与批量大小不匹配，则抛出值错误
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
# 如果生成器是列表，使用每个生成器编码图像并采样潜在变量
elif isinstance(generator, list):
init_latents = [
self.movq.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size)
]
# 将潜在变量沿指定维度连接
init_latents = torch.cat(init_latents, dim=0)
else:
# 如果生成器不是列表，直接编码图像并采样潜在变量
init_latents = self.movq.encode(image).latent_dist.sample(generator)
# 根据配置缩放因子调整潜在变量
init_latents = self.movq.config.scaling_factor * init_latents
# 再次连接潜在变量，以确保其形状正确
init_latents = torch.cat([init_latents], dim=0)
# 获取潜在变量的形状
shape = init_latents.shape
# 生成与潜在变量形状相同的噪声张量
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# 将噪声添加到潜在变量中，以获得最终的潜在变量
init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
# 将潜在变量赋值给另一个变量
latents = init_latents
# 返回最终的潜在变量
return latents
# 装饰器，表示此函数在执行时不需要梯度计算
@torch.no_grad()
def __call__(
# 定义输入参数，包括图像嵌入、图像及其其他参数
image_embeds: Union[torch.Tensor, List[torch.Tensor]],
image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]],
hint: torch.Tensor,
height: int = 512,
width: int = 512,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
strength: float = 0.3,
num_images_per_prompt: int = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
output_type: Optional[str] = "pil",
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
return_dict: bool = True,
.\diffusers\pipelines\kandinsky2_2\pipeline_kandinsky2_2_img2img.py
# 版权声明，表明此文件的所有权及使用许可
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# 在 Apache 许可证 2.0 版本下许可使用本文件
# Licensed under the Apache License, Version 2.0 (the "License");
# 除非符合该许可证，否则不得使用此文件
# you may not use this file except in compliance with the License.
# 可以在以下网址获取许可证副本
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，软件在许可证下以“原样”提供
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 不提供任何形式的明示或暗示的保证或条件
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 请查看许可证以了解有关权限和限制的具体条款
# See the License for the specific language governing permissions and
# limitations under the License.
# 从 typing 模块导入类型注解
from typing import Callable, Dict, List, Optional, Union
# 导入 numpy 库用于数组操作
import numpy as np
# 导入 PIL 库中的 Image 模块用于图像处理
import PIL.Image
# 导入 PyTorch 库
import torch
# 从 PIL 导入 Image 模块
from PIL import Image
# 从本地模型导入 UNet2DConditionModel 和 VQModel
from ...models import UNet2DConditionModel, VQModel
# 从调度器导入 DDPMScheduler
from ...schedulers import DDPMScheduler
# 从工具模块导入 deprecate 和 logging
from ...utils import deprecate, logging
# 从 torch_utils 导入 randn_tensor 函数
from ...utils.torch_utils import randn_tensor
# 从 pipeline_utils 导入 DiffusionPipeline 和 ImagePipelineOutput
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
# 创建一个记录器实例，用于记录日志信息
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，包含使用示例
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> from diffusers import KandinskyV22Img2ImgPipeline, KandinskyV22PriorPipeline
>>> from diffusers.utils import load_image
>>> import torch
>>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16
... )
>>> pipe_prior.to("cuda")
>>> prompt = "A red cartoon frog, 4k"
>>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False)
>>> pipe = KandinskyV22Img2ImgPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
... )
>>> pipe.to("cuda")
>>> init_image = load_image(
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
... "/kandinsky/frog.png"
... )
>>> image = pipe(
... image=init_image,
... image_embeds=image_emb,
... negative_image_embeds=zero_image_emb,
... height=768,
... width=768,
... num_inference_steps=100,
... strength=0.2,
... ).images
>>> image[0].save("red_frog.png")
```py
"""
# 从 diffusers.pipelines.kandinsky2_2.pipeline_kandinsky2_2 导入下采样高度和宽度的函数
def downscale_height_and_width(height, width, scale_factor=8):
# 根据比例因子计算新的高度，进行整除
new_height = height // scale_factor**2
# 如果高度不能被比例因子平方整除，增加高度
if height % scale_factor**2 != 0:
new_height += 1
# 根据比例因子计算新的宽度，进行整除
new_width = width // scale_factor**2
# 如果宽度不能被比例因子平方整除，增加宽度
if width % scale_factor**2 != 0:
new_width += 1
# 返回按比例因子调整后的高度和宽度
return new_height * scale_factor, new_width * scale_factor
# 从 diffusers.pipelines.kandinsky.pipeline_kandinsky_img2img 导入准备图像的函数
def prepare_image(pil_image, w=512, h=512):
# 将输入的 PIL 图像调整到指定的宽度和高度，使用 BICUBIC 插值
pil_image = pil_image.resize((w, h), resample=Image.BICUBIC, reducing_gap=1)
# 将 PIL 图像转换为 RGB 格式并转为 NumPy 数组
arr = np.array(pil_image.convert("RGB"))
# 将数组数据类型转换为 float32，并归一化到 [-1, 1] 范围
arr = arr.astype(np.float32) / 127.5 - 1
# 转置数组，以改变维度顺序从 (高度, 宽度, 通道) 到 (通道, 高度, 宽度)
arr = np.transpose(arr, [2, 0, 1])
# 将 NumPy 数组转换为 PyTorch 张量，并在第一个维度上增加一个维度
image = torch.from_numpy(arr).unsqueeze(0)
# 返回处理后的图像张量
return image
# 定义 Kandinsky 图像到图像生成的管道，继承自 DiffusionPipeline
class KandinskyV22Img2ImgPipeline(DiffusionPipeline):
"""
使用 Kandinsky 进行图像到图像生成的管道
此模型继承自 [`DiffusionPipeline`]。查看父类文档以了解库为所有管道实现的通用方法
（例如下载、保存、在特定设备上运行等）。
参数：
scheduler ([`DDIMScheduler`]):
与 `unet` 结合使用以生成图像潜变量的调度器。
unet ([`UNet2DConditionModel`]):
用于去噪图像嵌入的条件 U-Net 结构。
movq ([`VQModel`]):
MoVQ 解码器，用于从潜变量生成图像。
"""
# 定义模型中 CPU 卸载的顺序
model_cpu_offload_seq = "unet->movq"
# 定义需要回调的张量输入
_callback_tensor_inputs = ["latents", "image_embeds", "negative_image_embeds"]
# 初始化方法，设置 unet、scheduler 和 movq
def __init__(
self,
unet: UNet2DConditionModel, # 条件 U-Net 模型
scheduler: DDPMScheduler, # DDPM 调度器
movq: VQModel, # VQ 解码器模型
):
# 调用父类的初始化方法
super().__init__()
# 注册模块以便在管道中使用
self.register_modules(
unet=unet,
scheduler=scheduler,
movq=movq,
)
# 计算 movq 的缩放因子
self.movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1)
# 从 KandinskyImg2ImgPipeline 获取时间步的复制方法
def get_timesteps(self, num_inference_steps, strength, device):
# 使用 init_timestep 获取原始时间步
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
# 计算起始时间步，确保不小于 0
t_start = max(num_inference_steps - init_timestep, 0)
# 从调度器中获取相应的时间步
timesteps = self.scheduler.timesteps[t_start:]
# 返回计算出的时间步和剩余的推理步骤数
return timesteps, num_inference_steps - t_start
# 准备潜在变量，输入图像及其他参数进行处理
def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None):
# 检查输入图像是否为有效类型
if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
# 如果不符合，抛出类型错误
raise ValueError(
f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
)
# 将图像转换为指定设备和数据类型
image = image.to(device=device, dtype=dtype)
# 计算有效批量大小
batch_size = batch_size * num_images_per_prompt
# 如果图像有四个通道，初始化潜在变量为图像本身
if image.shape[1] == 4:
init_latents = image
else:
# 检查生成器列表长度是否与批量大小匹配
if isinstance(generator, list) and len(generator) != batch_size:
# 如果不匹配，抛出错误
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
elif isinstance(generator, list):
# 如果是生成器列表，逐个图像编码并采样
init_latents = [
self.movq.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size)
]
# 将所有潜在变量合并
init_latents = torch.cat(init_latents, dim=0)
else:
# 否则直接编码并采样
init_latents = self.movq.encode(image).latent_dist.sample(generator)
# 对潜在变量进行缩放
init_latents = self.movq.config.scaling_factor * init_latents
# 将潜在变量维度扩展为批量维度
init_latents = torch.cat([init_latents], dim=0)
# 获取潜在变量的形状
shape = init_latents.shape
# 生成噪声张量
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# 添加噪声到潜在变量中
init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
# 设置最终的潜在变量
latents = init_latents
# 返回处理后的潜在变量
return latents
# 获取引导比例的属性
@property
def guidance_scale(self):
return self._guidance_scale
# 检查是否使用无分类器引导
@property
def do_classifier_free_guidance(self):
return self._guidance_scale > 1
# 获取时间步数的属性
@property
def num_timesteps(self):
return self._num_timesteps
# 在无梯度上下文中调用
@torch.no_grad()
def __call__(
self,
image_embeds: Union[torch.Tensor, List[torch.Tensor]],
image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]],
height: int = 512,
width: int = 512,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
strength: float = 0.3,
num_images_per_prompt: int = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
**kwargs,
.\diffusers\pipelines\kandinsky2_2\pipeline_kandinsky2_2_inpainting.py
# 版权信息，表明该文件属于 HuggingFace 团队，所有权利保留
#
# 根据 Apache 许可证 2.0 版（"许可证"）授权；
# 您只能在符合许可证的情况下使用此文件。
# 您可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面协议另有约定，软件
# 按"原样"分发，不提供任何形式的保证或条件，无论是明示或暗示的。
# 有关许可证下权限和限制的具体信息，请参见许可证。
# 从 copy 模块导入 deepcopy 函数，用于深拷贝对象
from copy import deepcopy
# 导入类型提示相关的类型
from typing import Callable, Dict, List, Optional, Union
# 导入 numpy 库并赋予别名 np，常用于数值计算
import numpy as np
# 导入 PIL 库中的 Image 模块，用于图像处理
import PIL.Image
# 导入 PyTorch 库
import torch
# 导入 PyTorch 的功能模块，用于实现功能性操作
import torch.nn.functional as F
# 导入 packaging 库中的 version 模块，用于版本管理
from packaging import version
# 从 PIL 导入 Image 类，用于处理图像对象
from PIL import Image
# 从当前包中导入版本信息
from ... import __version__
# 从模型模块导入 UNet2DConditionModel 和 VQModel 类
from ...models import UNet2DConditionModel, VQModel
# 从调度器模块导入 DDPMScheduler 类
from ...schedulers import DDPMScheduler
# 从 utils 模块导入 deprecate 和 logging 功能
from ...utils import deprecate, logging
# 从 utils.torch_utils 导入 randn_tensor 函数，用于生成随机张量
from ...utils.torch_utils import randn_tensor
# 从 pipeline_utils 导入 DiffusionPipeline 和 ImagePipelineOutput 类
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
# 创建一个日志记录器，使用当前模块的名称
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，展示如何使用该模块的功能
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> from diffusers import KandinskyV22InpaintPipeline, KandinskyV22PriorPipeline
>>> from diffusers.utils import load_image
>>> import torch
>>> import numpy as np
>>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16
... )
>>> pipe_prior.to("cuda")
>>> prompt = "a hat"
>>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False)
>>> pipe = KandinskyV22InpaintPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16
... )
>>> pipe.to("cuda")
>>> init_image = load_image(
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
... "/kandinsky/cat.png"
... )
>>> mask = np.zeros((768, 768), dtype=np.float32)
>>> mask[:250, 250:-250] = 1
>>> out = pipe(
... image=init_image,
... mask_image=mask,
... image_embeds=image_emb,
... negative_image_embeds=zero_image_emb,
... height=768,
... width=768,
... num_inference_steps=50,
... )
>>> image = out.images[0]
>>> image.save("cat_with_hat.png")
```py
"""
# 从 diffusers.pipelines.kandinsky2_2.pipeline_kandinsky2_2 模块复制的函数
def downscale_height_and_width(height, width, scale_factor=8):
# 计算新的高度，将原高度除以缩放因子的平方
new_height = height // scale_factor**2
# 如果原高度不能被缩放因子的平方整除，新的高度加 1
if height % scale_factor**2 != 0:
new_height += 1
# 计算新的宽度，将原宽度除以缩放因子的平方
new_width = width // scale_factor**2
# 如果原宽度不能被缩放因子的平方整除，新的宽度加 1
if width % scale_factor**2 != 0:
new_width += 1
# 返回调整后的高度和宽度，均乘以缩放因子
return new_height * scale_factor, new_width * scale_factor
# 从 diffusers.pipelines.kandinsky.pipeline_kandinsky_inpaint.prepare_mask 复制
def prepare_mask(masks):
# 初始化一个空列表，用于存储处理后的掩码
prepared_masks = []
# 遍历输入的每个掩码
for mask in masks:
# 深拷贝当前掩码，确保不修改原始数据
old_mask = deepcopy(mask)
# 遍历掩码的每一行
for i in range(mask.shape[1]):
# 遍历掩码的每一列
for j in range(mask.shape[2]):
# 如果当前像素值为1，则跳过
if old_mask[0][i][j] == 1:
continue
# 如果不是第一行，设置上方像素为0
if i != 0:
mask[:, i - 1, j] = 0
# 如果不是第一列，设置左侧像素为0
if j != 0:
mask[:, i, j - 1] = 0
# 如果不是第一行和第一列，设置左上角像素为0
if i != 0 and j != 0:
mask[:, i - 1, j - 1] = 0
# 如果不是最后一行，设置下方像素为0
if i != mask.shape[1] - 1:
mask[:, i + 1, j] = 0
# 如果不是最后一列，设置右侧像素为0
if j != mask.shape[2] - 1:
mask[:, i, j + 1] = 0
# 如果不是最后一行和最后一列，设置右下角像素为0
if i != mask.shape[1] - 1 and j != mask.shape[2] - 1:
mask[:, i + 1, j + 1] = 0
# 将处理后的掩码添加到列表中
prepared_masks.append(mask)
# 将所有处理后的掩码堆叠成一个张量并返回
return torch.stack(prepared_masks, dim=0)
# 从 diffusers.pipelines.kandinsky.pipeline_kandinsky_inpaint.prepare_mask_and_masked_image 复制
def prepare_mask_and_masked_image(image, mask, height, width):
r"""
准备一对（掩码，图像），以便由 Kandinsky 修复管道使用。这意味着这些输入将
被转换为 ``torch.Tensor``，形状为 ``batch x channels x height x width``，其中 ``channels`` 为 ``3``，
对于 ``image`` 和 ``1``，对于 ``mask``。
``image`` 将被转换为 ``torch.float32`` 并归一化到 ``[-1, 1]``。``mask`` 将被
二值化（``mask > 0.5``）并同样转换为 ``torch.float32``。
参数：
image (Union[np.array, PIL.Image, torch.Tensor]): 要修复的图像。
可以是 ``PIL.Image``，或 ``height x width x 3`` 的 ``np.array``，或 ``channels x height x width`` 的
``torch.Tensor``，或者是 ``batch x channels x height x width`` 的 ``torch.Tensor``。
mask (_type_): 要应用于图像的掩码，即需要修复的区域。
可以是 ``PIL.Image``，或 ``height x width`` 的 ``np.array``，或 ``1 x height x width`` 的
``torch.Tensor``，或是 ``batch x 1 x height x width`` 的 ``torch.Tensor``。
height (`int`, *可选*, 默认值为 512):
生成图像的高度（以像素为单位）。
width (`int`, *可选*, 默认值为 512):
生成图像的宽度（以像素为单位）。
引发：
ValueError: ``torch.Tensor`` 图像应在 ``[-1, 1]`` 范围内。ValueError: ``torch.Tensor`` 掩码
应在 ``[0, 1]`` 范围内。ValueError: ``mask`` 和 ``image`` 应具有相同的空间维度。
TypeError: ``mask`` 是 ``torch.Tensor`` 但 ``image`` 不是
（反之亦然）。
返回：
tuple[torch.Tensor]: 将对（掩码，图像）作为 ``torch.Tensor``，具有 4
维度：``batch x channels x height x width``。
"""
# 如果输入图像为 None，抛出错误
if image is None:
raise ValueError("`image` input cannot be undefined.")
# 检查 mask 是否为 None，如果是则抛出错误
if mask is None:
raise ValueError("`mask_image` input cannot be undefined.")
# 检查 image 是否为 torch.Tensor 类型
if isinstance(image, torch.Tensor):
# 检查 mask 是否为 torch.Tensor 类型，如果不是则抛出错误
if not isinstance(mask, torch.Tensor):
raise TypeError(f"`image` is a torch.Tensor but `mask` (type: {type(mask)} is not")
# 如果 image 是单张图像（3 个维度）
if image.ndim == 3:
# 确保图像的第一个维度为 3，表示 RGB 通道
assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)"
# 在第一个维度增加一个批量维度
image = image.unsqueeze(0)
# 如果 mask 是单张图像（2 个维度）
if mask.ndim == 2:
# 在第一个和第二个维度增加批量和通道维度
mask = mask.unsqueeze(0).unsqueeze(0)
# 如果 mask 是 3 个维度
if mask.ndim == 3:
# 如果 mask 是单张图像（第一个维度为 1）
if mask.shape[0] == 1:
# 增加一个批量维度
mask = mask.unsqueeze(0)
# 如果 mask 是批量图像（没有通道维度）
else:
# 在第二个维度增加通道维度
mask = mask.unsqueeze(1)
# 确保 image 和 mask 都是 4 维张量
assert image.ndim == 4 and mask.ndim == 4, "Image and Mask must have 4 dimensions"
# 确保 image 和 mask 的空间维度相同
assert image.shape[-2:] == mask.shape[-2:], "Image and Mask must have the same spatial dimensions"
# 确保 image 和 mask 的批量大小相同
assert image.shape[0] == mask.shape[0], "Image and Mask must have the same batch size"
# 检查 image 的值是否在 [-1, 1] 范围内
if image.min() < -1 or image.max() > 1:
raise ValueError("Image should be in [-1, 1] range")
# 检查 mask 的值是否在 [0, 1] 范围内
if mask.min() < 0 or mask.max() > 1:
raise ValueError("Mask should be in [0, 1] range")
# 将 mask 二值化，低于 0.5 的值设为 0，其余设为 1
mask[mask < 0.5] = 0
mask[mask >= 0.5] = 1
# 将 image 转换为 float32 类型
image = image.to(dtype=torch.float32)
# 如果 mask 是张量但 image 不是，则抛出错误
elif isinstance(mask, torch.Tensor):
raise TypeError(f"`mask` is a torch.Tensor but `image` (type: {type(image)} is not")
else:
# 预处理图像
if isinstance(image, (PIL.Image.Image, np.ndarray)):
# 如果输入是单张图像或数组，则将其封装成列表
image = [image]
if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):
# 如果列表中的元素是图像，按照传入的高度和宽度调整所有图像的大小
image = [i.resize((width, height), resample=Image.BICUBIC, reducing_gap=1) for i in image]
# 将调整大小后的图像转换为 RGB 格式的 NumPy 数组，并增加一个维度
image = [np.array(i.convert("RGB"))[None, :] for i in image]
# 沿着第一个维度拼接所有图像数组
image = np.concatenate(image, axis=0)
elif isinstance(image, list) and isinstance(image[0], np.ndarray):
# 如果列表中的元素是 NumPy 数组，沿着第一个维度拼接这些数组
image = np.concatenate([i[None, :] for i in image], axis=0)
# 将图像数组的维度顺序从 (N, H, W, C) 转换为 (N, C, H, W)
image = image.transpose(0, 3, 1, 2)
# 将 NumPy 数组转换为 PyTorch 张量，并归一化到 [-1, 1] 的范围
image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0
# 预处理掩膜
if isinstance(mask, (PIL.Image.Image, np.ndarray)):
# 如果输入是单张掩膜或数组，则将其封装成列表
mask = [mask]
if isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image):
# 如果列表中的元素是掩膜图像，调整所有掩膜的大小
mask = [i.resize((width, height), resample=PIL.Image.LANCZOS) for i in mask]
# 将调整大小后的掩膜转换为灰度格式的 NumPy 数组，并增加两个维度
mask = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask], axis=0)
# 将掩膜的值归一化到 [0, 1] 的范围
mask = mask.astype(np.float32) / 255.0
elif isinstance(mask, list) and isinstance(mask[0], np.ndarray):
# 如果列表中的元素是 NumPy 数组，沿着第一个维度拼接这些数组
mask = np.concatenate([m[None, None, :] for m in mask], axis=0)
# 将掩膜中小于 0.5 的值设置为 0，大于等于 0.5 的值设置为 1
mask[mask < 0.5] = 0
mask[mask >= 0.5] = 1
# 将掩膜转换为 PyTorch 张量
mask = torch.from_numpy(mask)
# 将掩膜的值反转，得到 1 - mask
mask = 1 - mask
# 返回处理后的掩膜和图像
return mask, image
# Kandinsky2.1的文本引导图像修复管道类，继承自DiffusionPipeline
class KandinskyV22InpaintPipeline(DiffusionPipeline):
"""
使用Kandinsky2.1进行文本引导图像修复的管道
该模型继承自[`DiffusionPipeline`]。请查看超类文档以获取库为所有管道实现的通用方法（例如下载或保存、在特定设备上运行等）
参数：
scheduler ([`DDIMScheduler`]):
与`unet`结合使用的调度器，用于生成图像潜变量。
unet ([`UNet2DConditionModel`]):
条件U-Net架构，用于去噪图像嵌入。
movq ([`VQModel`]):
MoVQ解码器，用于从潜变量生成图像。
"""
# 定义模型的CPU卸载顺序
model_cpu_offload_seq = "unet->movq"
# 定义回调张量输入的列表
_callback_tensor_inputs = ["latents", "image_embeds", "negative_image_embeds", "masked_image", "mask_image"]
def __init__(
self,
unet: UNet2DConditionModel,
scheduler: DDPMScheduler,
movq: VQModel,
):
# 初始化父类DiffusionPipeline
super().__init__()
# 注册模型的各个模块
self.register_modules(
unet=unet,
scheduler=scheduler,
movq=movq,
)
# 计算MoVQ缩放因子
self.movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1)
# 初始化警告标志
self._warn_has_been_called = False
# 从diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline复制的prepare_latents方法
def prepare_latents(self, shape, dtype, device, generator, latents, scheduler):
# 如果潜变量为None，生成随机潜变量
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
# 检查潜变量的形状是否与期望的形状匹配
if latents.shape != shape:
raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
# 将潜变量转移到指定设备
latents = latents.to(device)
# 将潜变量乘以调度器的初始噪声标准差
latents = latents * scheduler.init_noise_sigma
# 返回处理后的潜变量
return latents
# 获取引导缩放因子的属性
@property
def guidance_scale(self):
return self._guidance_scale
# 获取是否使用分类器自由引导的属性
@property
def do_classifier_free_guidance(self):
return self._guidance_scale > 1
# 获取时间步数的属性
@property
def num_timesteps(self):
return self._num_timesteps
# 使用torch.no_grad()装饰器，避免梯度计算
@torch.no_grad()
def __call__(
self,
image_embeds: Union[torch.Tensor, List[torch.Tensor]],
image: Union[torch.Tensor, PIL.Image.Image],
mask_image: Union[torch.Tensor, PIL.Image.Image, np.ndarray],
negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]],
height: int = 512,
width: int = 512,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
**kwargs,
.\diffusers\pipelines\kandinsky2_2\pipeline_kandinsky2_2_prior.py
# 从 typing 模块导入必要的类型提示
from typing import Callable, Dict, List, Optional, Union
# 导入 PIL 库中的 Image 模块
import PIL.Image
# 导入 torch 库
import torch
# 从 transformers 库导入相关的模型和处理器
from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection
# 从本地模型文件中导入 PriorTransformer 类
from ...models import PriorTransformer
# 从调度器模块中导入 UnCLIPScheduler 类
from ...schedulers import UnCLIPScheduler
# 从工具模块中导入 logging 和 replace_example_docstring 函数
from ...utils import (
logging,
replace_example_docstring,
)
# 从 torch_utils 模块中导入 randn_tensor 函数
from ...utils.torch_utils import randn_tensor
# 从 kandinsky 模块中导入 KandinskyPriorPipelineOutput 类
from ..kandinsky import KandinskyPriorPipelineOutput
# 从 pipeline_utils 模块中导入 DiffusionPipeline 类
from ..pipeline_utils import DiffusionPipeline
# 创建一个 logger 实例用于记录日志，禁用 pylint 的命名检查
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，包含使用示例
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
>>> import torch
>>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")
>>> pipe_prior.to("cuda")
>>> prompt = "red cat, 4k photo"
>>> image_emb, negative_image_emb = pipe_prior(prompt).to_tuple()
>>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")
>>> pipe.to("cuda")
>>> image = pipe(
... image_embeds=image_emb,
... negative_image_embeds=negative_image_emb,
... height=768,
... width=768,
... num_inference_steps=50,
... ).images
>>> image[0].save("cat.png")
```py
"""
# 示例文档字符串，包含插值的使用示例
EXAMPLE_INTERPOLATE_DOC_STRING = """
Examples:
```py
>>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22Pipeline
>>> from diffusers.utils import load_image
>>> import PIL
>>> import torch
>>> from torchvision import transforms
>>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16
... )
>>> pipe_prior.to("cuda")
>>> img1 = load_image(
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
... "/kandinsky/cat.png"
... )
>>> img2 = load_image(
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
... "/kandinsky/starry_night.jpeg"
... )
>>> images_texts = ["a cat", img1, img2]
>>> weights = [0.3, 0.3, 0.4]
>>> out = pipe_prior.interpolate(images_texts, weights)
>>> pipe = KandinskyV22Pipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
... )
>>> pipe.to("cuda")
>>> image = pipe(
... image_embeds=out.image_embeds,
... negative_image_embeds=out.negative_image_embeds,
... height=768,
... width=768,
... num_inference_steps=50,
... ).images[0]
>>> image.save("starry_cat.png")
```py
"""
# 定义 KandinskyV22PriorPipeline 类，继承自 DiffusionPipeline
class KandinskyV22PriorPipeline(DiffusionPipeline):
"""
# 该文档字符串描述了生成 Kandinsky 图像先验的管道
# 该模型继承自 [`DiffusionPipeline`]。可以查看超类文档以了解库为所有管道实现的通用方法
# （例如下载、保存、在特定设备上运行等）。
# 参数：
# prior ([`PriorTransformer`]):
# 用于从文本嵌入近似图像嵌入的标准 unCLIP 先验。
# image_encoder ([`CLIPVisionModelWithProjection`]):
# 冻结的图像编码器。
# text_encoder ([`CLIPTextModelWithProjection`]):
# 冻结的文本编码器。
# tokenizer (`CLIPTokenizer`):
# 类的标记器
# [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer)。
# scheduler ([`UnCLIPScheduler`]):
# 用于与 `prior` 结合生成图像嵌入的调度器。
# image_processor ([`CLIPImageProcessor`]):
# 用于从 CLIP 预处理图像的图像处理器。
"""
# 定义 CPU 卸载序列，包括文本编码器、图像编码器和先验
model_cpu_offload_seq = "text_encoder->image_encoder->prior"
# 定义在 CPU 卸载时排除的模块，先验不参与卸载
_exclude_from_cpu_offload = ["prior"]
# 定义需要作为回调的张量输入
_callback_tensor_inputs = ["latents", "prompt_embeds", "text_encoder_hidden_states", "text_mask"]
def __init__(
# 初始化方法，接受多个参数用于模型的构建
self,
prior: PriorTransformer, # 标准 unCLIP 先验
image_encoder: CLIPVisionModelWithProjection, # 冻结的图像编码器
text_encoder: CLIPTextModelWithProjection, # 冻结的文本编码器
tokenizer: CLIPTokenizer, # 标记器
scheduler: UnCLIPScheduler, # 调度器
image_processor: CLIPImageProcessor, # 图像处理器
):
# 调用父类的初始化方法
super().__init__()
# 注册模型所需的模块
self.register_modules(
prior=prior, # 注册先验
text_encoder=text_encoder, # 注册文本编码器
tokenizer=tokenizer, # 注册标记器
scheduler=scheduler, # 注册调度器
image_encoder=image_encoder, # 注册图像编码器
image_processor=image_processor, # 注册图像处理器
)
@torch.no_grad() # 在此装饰器下，不计算梯度以节省内存和加快计算
@replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING) # 替换示例文档字符串为预定义字符串
def interpolate(
# 定义插值方法，接受多个参数以处理图像和提示
self,
images_and_prompts: List[Union[str, PIL.Image.Image, torch.Tensor]], # 输入的图像和提示列表
weights: List[float], # 与每个图像和提示对应的权重
num_images_per_prompt: int = 1, # 每个提示生成的图像数量，默认为1
num_inference_steps: int = 25, # 推理步骤的数量，默认为25
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, # 可选的随机数生成器
latents: Optional[torch.Tensor] = None, # 可选的潜在变量张量
negative_prior_prompt: Optional[str] = None, # 可选的负先验提示
negative_prompt: str = "", # 负提示，默认为空字符串
guidance_scale: float = 4.0, # 指导比例，默认为4.0
device=None, # 可选的设备参数
# 从 diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents 复制的部分
# 准备潜在变量，生成或处理输入的潜在张量
def prepare_latents(self, shape, dtype, device, generator, latents, scheduler):
# 如果没有提供潜在张量，则随机生成一个
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
# 如果提供的潜在张量形状不匹配，则抛出异常
if latents.shape != shape:
raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
# 将潜在张量移动到指定设备
latents = latents.to(device)
# 将潜在张量乘以调度器的初始化噪声标准差
latents = latents * scheduler.init_noise_sigma
# 返回处理后的潜在张量
return latents
# 从 KandinskyPriorPipeline 获取零嵌入的复制
def get_zero_embed(self, batch_size=1, device=None):
# 如果没有指定设备，则使用默认设备
device = device or self.device
# 创建一个零图像张量，大小为图像编码器配置的图像大小
zero_img = torch.zeros(1, 3, self.image_encoder.config.image_size, self.image_encoder.config.image_size).to(
device=device, dtype=self.image_encoder.dtype
)
# 将零图像传递给图像编码器以获取图像嵌入
zero_image_emb = self.image_encoder(zero_img)["image_embeds"]
# 将嵌入复制以匹配批量大小
zero_image_emb = zero_image_emb.repeat(batch_size, 1)
# 返回零图像嵌入
return zero_image_emb
# 从 KandinskyPriorPipeline 复制的提示编码方法
def _encode_prompt(
self,
prompt,
device,
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
# 分类器自由引导的属性，检查引导比例是否大于1
@property
def do_classifier_free_guidance(self):
return self._guidance_scale > 1
# 引导比例的属性，返回当前引导比例
@property
def guidance_scale(self):
return self._guidance_scale
# 时间步数的属性，返回当前时间步数
@property
def num_timesteps(self):
return self._num_timesteps
# 禁用梯度计算并替换示例文档字符串
@torch.no_grad()
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
prompt: Union[str, List[str]],
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: int = 1,
num_inference_steps: int = 25,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
guidance_scale: float = 4.0,
output_type: Optional[str] = "pt", # 仅支持 pt 格式
return_dict: bool = True,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 指定在步骤结束时的张量输入列表
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
.\diffusers\pipelines\kandinsky2_2\pipeline_kandinsky2_2_prior_emb2emb.py
# 导入类型提示模块，用于类型注释
from typing import List, Optional, Union
# 导入 PIL 库中的 Image 模块，用于图像处理
import PIL.Image
# 导入 PyTorch 库
import torch
# 从 transformers 库导入 CLIP 相关的模型和处理器
from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection
# 从本地模块导入 PriorTransformer 模型
from ...models import PriorTransformer
# 从本地模块导入 UnCLIPScheduler 调度器
from ...schedulers import UnCLIPScheduler
# 从本地工具模块导入日志记录和文档字符串替换功能
from ...utils import (
logging,
replace_example_docstring,
)
# 从本地工具模块导入随机张量生成函数
from ...utils.torch_utils import randn_tensor
# 从本地模块导入 KandinskyPriorPipelineOutput
from ..kandinsky import KandinskyPriorPipelineOutput
# 从本地模块导入 DiffusionPipeline
from ..pipeline_utils import DiffusionPipeline
# 创建一个日志记录器，使用模块的名称作为标识
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 定义示例文档字符串，展示用法示例
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorEmb2EmbPipeline
>>> import torch
>>> pipe_prior = KandinskyPriorPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16
... )
>>> pipe_prior.to("cuda")
>>> prompt = "red cat, 4k photo"
>>> img = load_image(
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
... "/kandinsky/cat.png"
... )
>>> image_emb, nagative_image_emb = pipe_prior(prompt, image=img, strength=0.2).to_tuple()
>>> pipe = KandinskyPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-decoder, torch_dtype=torch.float16"
... )
>>> pipe.to("cuda")
>>> image = pipe(
... image_embeds=image_emb,
... negative_image_embeds=negative_image_emb,
... height=768,
... width=768,
... num_inference_steps=100,
... ).images
>>> image[0].save("cat.png")
```py
"""
# 定义插值示例文档字符串，内容为空
EXAMPLE_INTERPOLATE_DOC_STRING = """
# 示例代码，演示如何使用 Kandinsky 模型进行图像处理
Examples:
```py
# 从 diffusers 库导入所需的类和函数
>>> from diffusers import KandinskyV22PriorEmb2EmbPipeline, KandinskyV22Pipeline
>>> from diffusers.utils import load_image
>>> import PIL
# 导入 PyTorch 库
>>> import torch
>>> from torchvision import transforms
# 加载 Kandinsky 的先验模型
>>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16
... )
# 将模型转移到 GPU
>>> pipe_prior.to("cuda")
# 加载第一张图片
>>> img1 = load_image(
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
... "/kandinsky/cat.png"
... )
# 加载第二张图片
>>> img2 = load_image(
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
... "/kandinsky/starry_night.jpeg"
... )
# 创建包含文本描述和图像的列表
>>> images_texts = ["a cat", img1, img2]
# 设置每个图像的权重
>>> weights = [0.3, 0.3, 0.4]
# 调用插值方法，生成图像嵌入和零图像嵌入
>>> image_emb, zero_image_emb = pipe_prior.interpolate(images_texts, weights)
# 加载 Kandinsky 的解码器模型
>>> pipe = KandinskyV22Pipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
... )
# 将模型转移到 GPU
>>> pipe.to("cuda")
# 使用图像嵌入生成新图像
>>> image = pipe(
... image_embeds=image_emb,
... negative_image_embeds=zero_image_emb,
... height=768,
... width=768,
... num_inference_steps=150,
... ).images[0]
# 将生成的图像保存为文件
>>> image.save("starry_cat.png")
```
"""
文档字符串，描述这个类的功能
class KandinskyV22PriorEmb2EmbPipeline(DiffusionPipeline):
"""
# 文档字符串，描述生成图像先验的管道
Pipeline for generating image prior for Kandinsky
# 说明此模型继承自 [`DiffusionPipeline`]，并提到上级文档中的通用方法
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
# 构造函数参数说明
Args:
prior ([`PriorTransformer`]):
# 用于近似文本嵌入生成图像嵌入的 canonical unCLIP prior
The canonical unCLIP prior to approximate the image embedding from the text embedding.
image_encoder ([`CLIPVisionModelWithProjection`]):
# 冻结的图像编码器
Frozen image-encoder.
text_encoder ([`CLIPTextModelWithProjection`]):
# 冻结的文本编码器
Frozen text-encoder.
tokenizer (`CLIPTokenizer`):
# CLIPTokenizer 的分词器
Tokenizer of class
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
scheduler ([`UnCLIPScheduler`]):
# 与 `prior` 结合使用的调度器，用于生成图像嵌入
A scheduler to be used in combination with `prior` to generate image embedding.
"""
# 定义 CPU 释放的顺序
model_cpu_offload_seq = "text_encoder->image_encoder->prior"
# 指定不参与 CPU 释放的模块
_exclude_from_cpu_offload = ["prior"]
# 构造函数，初始化各个组件
def __init__(
self,
prior: PriorTransformer,
image_encoder: CLIPVisionModelWithProjection,
text_encoder: CLIPTextModelWithProjection,
tokenizer: CLIPTokenizer,
scheduler: UnCLIPScheduler,
image_processor: CLIPImageProcessor,
):
# 调用父类构造函数
super().__init__()
# 注册模块
self.register_modules(
prior=prior,
text_encoder=text_encoder,
tokenizer=tokenizer,
scheduler=scheduler,
image_encoder=image_encoder,
image_processor=image_processor,
)
# 获取时间步长的方法
def get_timesteps(self, num_inference_steps, strength, device):
# 计算初始时间步，确保不超过总时间步数
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
# 计算开始时间步，确保不小于零
t_start = max(num_inference_steps - init_timestep, 0)
# 获取指定时间步范围
timesteps = self.scheduler.timesteps[t_start:]
# 返回时间步和剩余推理步骤
return timesteps, num_inference_steps - t_start
# 该方法不计算梯度，用于插值
@torch.no_grad()
# 替换示例文档字符串
@replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING)
def interpolate(
# 输入图像及提示，类型可以是字符串、PIL图像或张量
images_and_prompts: List[Union[str, PIL.Image.Image, torch.Tensor]],
# 权重列表
weights: List[float],
# 每个提示生成的图像数量，默认为1
num_images_per_prompt: int = 1,
# 推理步骤数量，默认为25
num_inference_steps: int = 25,
# 随机生成器，默认为None
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 潜在张量，默认为None
latents: Optional[torch.Tensor] = None,
# 负面先验提示，默认为None
negative_prior_prompt: Optional[str] = None,
# 负面提示，默认为空字符串
negative_prompt: str = "",
# 引导尺度，默认为4.0
guidance_scale: float = 4.0,
# 设备，默认为None
device=None,
def _encode_image(
# 编码图像的方法，支持张量或PIL图像列表
self,
image: Union[torch.Tensor, List[PIL.Image.Image]],
# 设备类型
device,
# 每个提示的图像数量
num_images_per_prompt,
):
# 检查输入的图像是否为 PyTorch 张量
if not isinstance(image, torch.Tensor):
# 如果不是，则通过图像处理器将其转换为张量，并将数据类型和设备设置为合适的值
image = self.image_processor(image, return_tensors="pt").pixel_values.to(
dtype=self.image_encoder.dtype, device=device
)
# 使用图像编码器对图像进行编码，提取图像嵌入
image_emb = self.image_encoder(image)["image_embeds"] # B, D
# 按每个提示的图像数量重复图像嵌入
image_emb = image_emb.repeat_interleave(num_images_per_prompt, dim=0)
# 将图像嵌入转移到指定的设备
image_emb.to(device=device)
# 返回处理后的图像嵌入
return image_emb
def prepare_latents(self, emb, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None):
# 将嵌入转移到指定设备并设置数据类型
emb = emb.to(device=device, dtype=dtype)
# 计算新的批量大小
batch_size = batch_size * num_images_per_prompt
# 初始化潜在变量
init_latents = emb
# 检查批量大小是否大于初始潜在变量的形状，并且能否均匀复制
if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
# 计算每个提示需要额外的图像数量
additional_image_per_prompt = batch_size // init_latents.shape[0]
# 复制初始潜在变量以匹配新的批量大小
init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0)
# 检查批量大小是否大于初始潜在变量的形状且不能均匀复制
elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:
# 抛出错误，提示无法复制图像
raise ValueError(
f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
)
else:
# 如果批量大小小于或等于初始潜在变量的形状，则直接使用初始潜在变量
init_latents = torch.cat([init_latents], dim=0)
# 获取初始潜在变量的形状
shape = init_latents.shape
# 生成噪声张量
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# 获取潜在变量
init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
# 将潜在变量分配给变量以便返回
latents = init_latents
# 返回潜在变量
return latents
# 从 KandinskyPriorPipeline 复制的函数，获取零嵌入
def get_zero_embed(self, batch_size=1, device=None):
# 如果没有指定设备，则使用默认设备
device = device or self.device
# 创建一个零图像，形状为 (1, 3, 高, 宽)
zero_img = torch.zeros(1, 3, self.image_encoder.config.image_size, self.image_encoder.config.image_size).to(
device=device, dtype=self.image_encoder.dtype
)
# 对零图像进行编码以获取零嵌入
zero_image_emb = self.image_encoder(zero_img)["image_embeds"]
# 按批量大小重复零嵌入
zero_image_emb = zero_image_emb.repeat(batch_size, 1)
# 返回零嵌入
return zero_image_emb
# 从 KandinskyPriorPipeline 复制的函数，编码提示
def _encode_prompt(
self,
prompt,
device,
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
@torch.no_grad()
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
# 定义输入的提示，可以是字符串或字符串列表
prompt: Union[str, List[str]],
# 输入的图像，可以是张量、图像列表或 PIL 图像
image: Union[torch.Tensor, List[torch.Tensor], PIL.Image.Image, List[PIL.Image.Image]],
# 设置强度，默认为 0.3
strength: float = 0.3,
# 可选的负面提示
negative_prompt: Optional[Union[str, List[str]]] = None,
# 每个提示的图像数量，默认为 1
num_images_per_prompt: int = 1,
# 推理步骤数量，默认为 25
num_inference_steps: int = 25,
# 随机数生成器，可以是生成器或生成器列表
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 引导缩放比例，默认为 4.0
guidance_scale: float = 4.0,
# 输出类型，默认为 "pt"
output_type: Optional[str] = "pt", # pt only
# 返回字典的标志，默认为 True
return_dict: bool = True,
# `.\diffusers\pipelines\kandinsky2_2\__init__.py`
```py
# 导入类型检查的常量
from typing import TYPE_CHECKING
# 从 utils 模块导入所需的工具和常量
from ...utils import (
DIFFUSERS_SLOW_IMPORT, # 用于判断是否进行慢速导入
OptionalDependencyNotAvailable, # 自定义异常类，用于处理缺失依赖
_LazyModule, # 用于延迟加载模块
get_objects_from_module, # 从模块中获取对象的工具函数
is_torch_available, # 检查 PyTorch 是否可用
is_transformers_available, # 检查 Transformers 是否可用
)
# 初始化一个空字典用于存储虚拟对象
_dummy_objects = {}
# 初始化一个空字典用于存储导入结构
_import_structure = {}
try:
# 检查 Transformers 和 PyTorch 是否都可用
if not (is_transformers_available() and is_torch_available()):
# 如果不可用，抛出缺失依赖异常
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
# 导入虚拟对象，以防依赖缺失
from ...utils import dummy_torch_and_transformers_objects # noqa F403
# 更新虚拟对象字典
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
else:
# 添加相关的管道到导入结构字典
_import_structure["pipeline_kandinsky2_2"] = ["KandinskyV22Pipeline"]
_import_structure["pipeline_kandinsky2_2_combined"] = [
"KandinskyV22CombinedPipeline", # 组合管道
"KandinskyV22Img2ImgCombinedPipeline", # 图像到图像的组合管道
"KandinskyV22InpaintCombinedPipeline", # 图像修复的组合管道
]
_import_structure["pipeline_kandinsky2_2_controlnet"] = ["KandinskyV22ControlnetPipeline"] # 控制网管道
_import_structure["pipeline_kandinsky2_2_controlnet_img2img"] = ["KandinskyV22ControlnetImg2ImgPipeline"] # 控制网图像到图像的管道
_import_structure["pipeline_kandinsky2_2_img2img"] = ["KandinskyV22Img2ImgPipeline"] # 图像到图像的管道
_import_structure["pipeline_kandinsky2_2_inpainting"] = ["KandinskyV22InpaintPipeline"] # 图像修复管道
_import_structure["pipeline_kandinsky2_2_prior"] = ["KandinskyV22PriorPipeline"] # 先验管道
_import_structure["pipeline_kandinsky2_2_prior_emb2emb"] = ["KandinskyV22PriorEmb2EmbPipeline"] # 先验嵌入到嵌入的管道
# 如果处于类型检查模式或需要慢速导入
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
try:
# 再次检查 Transformers 和 PyTorch 是否可用
if not (is_transformers_available() and is_torch_available()):
raise OptionalDependencyNotAvailable() # 抛出异常
except OptionalDependencyNotAvailable:
# 导入虚拟对象以避免依赖问题
from ...utils.dummy_torch_and_transformers_objects import *
else:
# 从各个管道模块导入相应的类
from .pipeline_kandinsky2_2 import KandinskyV22Pipeline
from .pipeline_kandinsky2_2_combined import (
KandinskyV22CombinedPipeline,
KandinskyV22Img2ImgCombinedPipeline,
KandinskyV22InpaintCombinedPipeline,
)
from .pipeline_kandinsky2_2_controlnet import KandinskyV22ControlnetPipeline
from .pipeline_kandinsky2_2_controlnet_img2img import KandinskyV22ControlnetImg2ImgPipeline
from .pipeline_kandinsky2_2_img2img import KandinskyV22Img2ImgPipeline
from .pipeline_kandinsky2_2_inpainting import KandinskyV22InpaintPipeline
from .pipeline_kandinsky2_2_prior import KandinskyV22PriorPipeline
from .pipeline_kandinsky2_2_prior_emb2emb import KandinskyV22PriorEmb2EmbPipeline
else:
# 如果不处于类型检查或慢速导入，则创建懒加载模块
import sys
sys.modules[__name__] = _LazyModule(
__name__, # 模块名称
globals()["__file__"], # 模块文件路径
_import_structure, # 导入结构
module_spec=__spec__, # 模块规格
)
# 将虚拟对象添加到当前模块
for name, value in _dummy_objects.items():
setattr(sys.modules[__name__], name, value)
.\diffusers\pipelines\kandinsky3\convert_kandinsky3_unet.py
#!/usr/bin/env python3
# 指定脚本使用 Python 3 运行环境
import argparse
# 导入用于解析命令行参数的 argparse 模块
import fnmatch
# 导入用于匹配文件名的 fnmatch 模块
from safetensors.torch import load_file
# 从 safetensors.torch 导入 load_file 函数，用于加载模型文件
from diffusers import Kandinsky3UNet
# 从 diffusers 库导入 Kandinsky3UNet 模型
# 定义一个字典，用于映射原 U-Net 模型中的键到 Kandinsky3UNet 模型的键
MAPPING = {
"to_time_embed.1": "time_embedding.linear_1",
"to_time_embed.3": "time_embedding.linear_2",
"in_layer": "conv_in",
"out_layer.0": "conv_norm_out",
"out_layer.2": "conv_out",
"down_samples": "down_blocks",
"up_samples": "up_blocks",
"projection_lin": "encoder_hid_proj.projection_linear",
"projection_ln": "encoder_hid_proj.projection_norm",
"feature_pooling": "add_time_condition",
"to_query": "to_q",
"to_key": "to_k",
"to_value": "to_v",
"output_layer": "to_out.0",
"self_attention_block": "attentions.0",
}
# 定义一个动态映射字典，用于处理带有通配符的键
DYNAMIC_MAP = {
"resnet_attn_blocks.*.0": "resnets_in.*",
"resnet_attn_blocks.*.1": ("attentions.*", 1),
"resnet_attn_blocks.*.2": "resnets_out.*",
}
# DYNAMIC_MAP = {} # 备用的动态映射字典
def convert_state_dict(unet_state_dict):
"""
Args:
Convert the state dict of a U-Net model to match the key format expected by Kandinsky3UNet model.
unet_model (torch.nn.Module): The original U-Net model. unet_kandi3_model (torch.nn.Module): The Kandinsky3UNet
model to match keys with.
Returns:
OrderedDict: The converted state dictionary.
"""
# 创建一个空字典，用于存储转换后的状态字典
converted_state_dict = {}
# 遍历原 U-Net 模型的状态字典中的每个键
for key in unet_state_dict:
new_key = key # 初始化新键为原键
# 遍历映射字典，将原键中的模式替换为新模式
for pattern, new_pattern in MAPPING.items():
new_key = new_key.replace(pattern, new_pattern)
# 处理动态映射
for dyn_pattern, dyn_new_pattern in DYNAMIC_MAP.items():
has_matched = False # 初始化匹配标志
# 如果新键匹配动态模式且尚未匹配
if fnmatch.fnmatch(new_key, f"*.{dyn_pattern}.*") and not has_matched:
# 提取模式中的数字部分
star = int(new_key.split(dyn_pattern.split(".")[0])[-1].split(".")[1])
# 处理动态模式为元组的情况
if isinstance(dyn_new_pattern, tuple):
new_star = star + dyn_new_pattern[-1] # 计算新数字
dyn_new_pattern = dyn_new_pattern[0] # 提取新的模式
else:
new_star = star # 保持数字不变
# 替换动态模式中的星号
pattern = dyn_pattern.replace("*", str(star))
new_pattern = dyn_new_pattern.replace("*", str(new_star))
# 更新新键
new_key = new_key.replace(pattern, new_pattern)
has_matched = True # 设置匹配标志为 True
# 将转换后的新键与原状态字典中的值存入新的字典
converted_state_dict[new_key] = unet_state_dict[key]
return converted_state_dict # 返回转换后的状态字典
def main(model_path, output_path):
# 加载原 U-Net 模型的状态字典
unet_state_dict = load_file(model_path)
# 初始化 Kandinsky3UNet 模型的配置
config = {}
# 调用转换函数，转换状态字典
converted_state_dict = convert_state_dict(unet_state_dict)
# 实例化 Kandinsky3UNet 模型
unet = Kandinsky3UNet(config)
# 加载转换后的状态字典到模型中
unet.load_state_dict(converted_state_dict)
# 保存转换后的模型到指定路径
unet.save_pretrained(output_path)
# 打印保存模型的路径
print(f"Converted model saved to {output_path}")
if __name__ == "__main__":
# 创建命令行解析器，提供描述信息
parser = argparse.ArgumentParser(description="Convert U-Net PyTorch model to Kandinsky3UNet format")
# 添加命令行参数 '--model_path'，指定原始 U-Net PyTorch 模型的路径，类型为字符串，必填
parser.add_argument("--model_path", type=str, required=True, help="Path to the original U-Net PyTorch model")
# 添加命令行参数 '--output_path'，指定转换后模型的保存路径，类型为字符串，必填
parser.add_argument("--output_path", type=str, required=True, help="Path to save the converted model")
# 解析命令行参数并将其存储在 args 对象中
args = parser.parse_args()
# 调用 main 函数，传入模型路径和输出路径参数
main(args.model_path, args.output_path)
.\diffusers\pipelines\kandinsky3\pipeline_kandinsky3.py
# 从 typing 模块导入类型注解工具
from typing import Callable, Dict, List, Optional, Union
# 导入 PyTorch 库
import torch
# 从 transformers 库导入 T5 编码器模型和分词器
from transformers import T5EncoderModel, T5Tokenizer
# 从本地模块导入混合加载器
from ...loaders import StableDiffusionLoraLoaderMixin
# 从本地模块导入 Kandinsky3UNet 和 VQModel
from ...models import Kandinsky3UNet, VQModel
# 从本地模块导入 DDPMScheduler
from ...schedulers import DDPMScheduler
# 从本地模块导入工具函数
from ...utils import (
deprecate,
logging,
replace_example_docstring,
)
# 从本地工具模块导入随机张量生成函数
from ...utils.torch_utils import randn_tensor
# 从本地管道工具模块导入 DiffusionPipeline 和 ImagePipelineOutput
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
# 创建一个记录日志的 logger 实例，名称为当前模块名
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，展示如何使用该模块的功能
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> from diffusers import AutoPipelineForText2Image
>>> import torch
>>> pipe = AutoPipelineForText2Image.from_pretrained(
... "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16
... )
>>> pipe.enable_model_cpu_offload()
>>> prompt = "A photograph of the inside of a subway train. There are raccoons sitting on the seats. One of them is reading a newspaper. The window shows the city in the background."
>>> generator = torch.Generator(device="cpu").manual_seed(0)
>>> image = pipe(prompt, num_inference_steps=25, generator=generator).images[0]
```py
"""
# 定义 downscale_height_and_width 函数，用于缩放高度和宽度
def downscale_height_and_width(height, width, scale_factor=8):
# 计算缩放后的新高度
new_height = height // scale_factor**2
# 如果高度不能被缩放因子平方整除，则向上取整
if height % scale_factor**2 != 0:
new_height += 1
# 计算缩放后的新宽度
new_width = width // scale_factor**2
# 如果宽度不能被缩放因子平方整除，则向上取整
if width % scale_factor**2 != 0:
new_width += 1
# 返回缩放后的高度和宽度
return new_height * scale_factor, new_width * scale_factor
# 定义 Kandinsky3Pipeline 类，继承自 DiffusionPipeline 和 StableDiffusionLoraLoaderMixin
class Kandinsky3Pipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin):
# 定义模型 CPU 卸载的顺序
model_cpu_offload_seq = "text_encoder->unet->movq"
# 定义需要回调的张量输入列表
_callback_tensor_inputs = [
"latents",
"prompt_embeds",
"negative_prompt_embeds",
"negative_attention_mask",
"attention_mask",
]
# 初始化方法，接收多个模型组件作为参数
def __init__(
self,
tokenizer: T5Tokenizer,
text_encoder: T5EncoderModel,
unet: Kandinsky3UNet,
scheduler: DDPMScheduler,
movq: VQModel,
):
# 调用父类初始化方法
super().__init__()
# 注册模型组件
self.register_modules(
tokenizer=tokenizer, text_encoder=text_encoder, unet=unet, scheduler=scheduler, movq=movq
)
# 定义处理嵌入的函数
def process_embeds(self, embeddings, attention_mask, cut_context):
# 如果需要裁剪上下文
if cut_context:
# 将 attention_mask 为 0 的嵌入置为零
embeddings[attention_mask == 0] = torch.zeros_like(embeddings[attention_mask == 0])
# 计算最大序列长度
max_seq_length = attention_mask.sum(-1).max() + 1
# 裁剪嵌入和 attention_mask
embeddings = embeddings[:, :max_seq_length]
attention_mask = attention_mask[:, :max_seq_length]
# 返回处理后的嵌入和 attention_mask
return embeddings, attention_mask
# 禁用梯度计算，以节省内存
@torch.no_grad()
# 定义一个用于编码提示的函数
def encode_prompt(
# 提示内容
self,
prompt,
# 是否进行无分类器引导
do_classifier_free_guidance=True,
# 每个提示生成的图像数量
num_images_per_prompt=1,
# 设备类型（CPU/GPU）
device=None,
# 负面提示内容
negative_prompt=None,
# 提示的嵌入张量（可选）
prompt_embeds: Optional[torch.Tensor] = None,
# 负面提示的嵌入张量（可选）
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 是否截断上下文
_cut_context=False,
# 注意力掩码（可选）
attention_mask: Optional[torch.Tensor] = None,
# 负面提示的注意力掩码（可选）
negative_attention_mask: Optional[torch.Tensor] = None,
# 准备潜在变量的函数
def prepare_latents(self, shape, dtype, device, generator, latents, scheduler):
# 如果没有潜在变量，则随机生成
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
# 检查潜在变量的形状是否与预期匹配
if latents.shape != shape:
raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
# 将潜在变量移动到指定设备
latents = latents.to(device)
# 通过调度器的初始噪声标准差缩放潜在变量
latents = latents * scheduler.init_noise_sigma
# 返回处理后的潜在变量
return latents
# 检查输入有效性的函数
def check_inputs(
# 提示内容
self,
prompt,
# 回调步骤
callback_steps,
# 负面提示内容（可选）
negative_prompt=None,
# 提示的嵌入张量（可选）
prompt_embeds=None,
# 负面提示的嵌入张量（可选）
negative_prompt_embeds=None,
# 结束时的回调张量输入（可选）
callback_on_step_end_tensor_inputs=None,
# 注意力掩码（可选）
attention_mask=None,
# 负面提示的注意力掩码（可选）
negative_attention_mask=None,
# 返回引导缩放因子的属性
@property
def guidance_scale(self):
# 返回内部引导缩放因子的值
return self._guidance_scale
# 返回是否进行无分类器引导的属性
@property
def do_classifier_free_guidance(self):
# 判断引导缩放因子是否大于 1
return self._guidance_scale > 1
# 返回时间步数的属性
@property
def num_timesteps(self):
# 返回内部时间步数的值
return self._num_timesteps
# 禁用梯度计算的调用函数
@torch.no_grad()
# 替换示例文档字符串的装饰器
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
# 提示内容（可以是字符串或字符串列表）
self,
prompt: Union[str, List[str]] = None,
# 推理步骤的数量
num_inference_steps: int = 25,
# 引导缩放因子的值
guidance_scale: float = 3.0,
# 负面提示内容（可选）
negative_prompt: Optional[Union[str, List[str]]] = None,
# 每个提示生成的图像数量（可选）
num_images_per_prompt: Optional[int] = 1,
# 图像高度（可选）
height: Optional[int] = 1024,
# 图像宽度（可选）
width: Optional[int] = 1024,
# 随机生成器（可选）
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 提示的嵌入张量（可选）
prompt_embeds: Optional[torch.Tensor] = None,
# 负面提示的嵌入张量（可选）
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 注意力掩码（可选）
attention_mask: Optional[torch.Tensor] = None,
# 负面提示的注意力掩码（可选）
negative_attention_mask: Optional[torch.Tensor] = None,
# 输出类型（默认为 PIL 格式）
output_type: Optional[str] = "pil",
# 是否返回字典格式的结果（默认为 True）
return_dict: bool = True,
# 潜在变量（可选）
latents=None,
# 步骤结束时的回调函数（可选）
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 结束时的回调张量输入（默认值为 ["latents"]）
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# 其他关键字参数
**kwargs,
.\diffusers\pipelines\kandinsky3\pipeline_kandinsky3_img2img.py
import inspect # 导入 inspect 模块，用于获取对象的详细信息
from typing import Callable, Dict, List, Optional, Union # 导入类型提示相关的模块
import numpy as np # 导入 numpy 库，用于数值计算
import PIL # 导入 PIL 库，用于图像处理
import PIL.Image # 从 PIL 导入 Image 模块，用于图像对象的操作
import torch # 导入 PyTorch 库，用于深度学习
from transformers import T5EncoderModel, T5Tokenizer # 从 transformers 导入 T5 模型和分词器
from ...loaders import StableDiffusionLoraLoaderMixin # 导入混合类，用于加载 Lora 模型
from ...models import Kandinsky3UNet, VQModel # 导入 Kandinsky3 UNet 和 VQModel 模型
from ...schedulers import DDPMScheduler # 导入 DDPMScheduler，用于调度器
from ...utils import ( # 导入工具函数
deprecate, # 导入弃用装饰器
logging, # 导入日志记录模块
replace_example_docstring, # 导入用于替换示例文档字符串的函数
)
from ...utils.torch_utils import randn_tensor # 从工具模块导入随机张量生成函数
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput # 导入扩散管道和图像输出模块
logger = logging.get_logger(__name__) # 获取当前模块的日志记录器，使用模块名称作为标识，禁用 pylint 的无效名称警告
EXAMPLE_DOC_STRING = """ # 示例文档字符串，用于展示如何使用该管道
Examples: # 示例部分的开始
```py # Python 示例代码块开始
>>> from diffusers import AutoPipelineForImage2Image # 导入图像转图像的自动管道
>>> from diffusers.utils import load_image # 导入加载图像的工具函数
>>> import torch # 导入 PyTorch 库
>>> pipe = AutoPipelineForImage2Image.from_pretrained( # 从预训练模型加载图像转图像的管道
... "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16 # 指定模型和数据类型
... )
>>> pipe.enable_model_cpu_offload() # 启用模型的 CPU 内存释放
>>> prompt = "A painting of the inside of a subway train with tiny raccoons." # 定义图像生成的提示
>>> image = load_image( # 加载图像
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/kandinsky3/t2i.png" # 图像的 URL
... )
>>> generator = torch.Generator(device="cpu").manual_seed(0) # 创建 CPU 设备上的随机数生成器并设置种子
>>> image = pipe(prompt, image=image, strength=0.75, num_inference_steps=25, generator=generator).images[0] # 生成图像并获取结果
```py # 示例代码块结束
"""
def downscale_height_and_width(height, width, scale_factor=8): # 定义函数以按比例缩放高度和宽度
new_height = height // scale_factor**2 # 计算新的高度
if height % scale_factor**2 != 0: # 检查高度是否可以整除缩放因子平方
new_height += 1 # 如果不能，增加高度
new_width = width // scale_factor**2 # 计算新的宽度
if width % scale_factor**2 != 0: # 检查宽度是否可以整除缩放因子平方
new_width += 1 # 如果不能，增加宽度
return new_height * scale_factor, new_width * scale_factor # 返回按缩放因子调整后的新高度和新宽度
def prepare_image(pil_image): # 定义函数以准备 PIL 图像
arr = np.array(pil_image.convert("RGB")) # 将 PIL 图像转换为 RGB 并转为 NumPy 数组
arr = arr.astype(np.float32) / 127.5 - 1 # 将数组转换为浮点数并归一化到 [-1, 1] 范围
arr = np.transpose(arr, [2, 0, 1]) # 转换数组维度，从 (H, W, C) 到 (C, H, W)
image = torch.from_numpy(arr).unsqueeze(0) # 将 NumPy 数组转换为 PyTorch 张量并增加一个维度
return image # 返回处理后的图像张量
class Kandinsky3Img2ImgPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): # 定义 Kandinsky 3 图像到图像的管道类，继承自 DiffusionPipeline 和 StableDiffusionLoraLoaderMixin
model_cpu_offload_seq = "text_encoder->movq->unet->movq" # 定义模型 CPU 内存释放的顺序
_callback_tensor_inputs = [ # 定义需要回调的张量输入
"latents", # 潜在表示
"prompt_embeds", # 提示嵌入
"negative_prompt_embeds", # 负提示嵌入
"negative_attention_mask", # 负注意力掩码
"attention_mask", # 注意力掩码
]
def __init__( # 初始化方法
self,
tokenizer: T5Tokenizer, # T5 分词器
text_encoder: T5EncoderModel, # T5 文本编码器
unet: Kandinsky3UNet, # Kandinsky3 UNet 模型
scheduler: DDPMScheduler, # DDPMScheduler 实例
movq: VQModel, # VQModel 实例
):
super().__init__() # 调用父类的初始化方法
self.register_modules( # 注册各个模块
tokenizer=tokenizer, text_encoder=text_encoder, unet=unet, scheduler=scheduler, movq=movq # 注册分词器、文本编码器、UNet、调度器和 VQModel
)
def get_timesteps(self, num_inference_steps, strength, device):
# 获取初始时间步，使用给定的推理步骤数量和强度
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
# 计算开始时间步，确保不小于0
t_start = max(num_inference_steps - init_timestep, 0)
# 从计算的开始时间步获取调度器中的时间步
timesteps = self.scheduler.timesteps[t_start:]
# 返回时间步和剩余的推理步骤数量
return timesteps, num_inference_steps - t_start
def _process_embeds(self, embeddings, attention_mask, cut_context):
# 返回处理后的嵌入和注意力掩码
if cut_context:
# 将注意力掩码为0的嵌入置为零
embeddings[attention_mask == 0] = torch.zeros_like(embeddings[attention_mask == 0])
# 计算最大序列长度，并在此基础上切片嵌入和注意力掩码
max_seq_length = attention_mask.sum(-1).max() + 1
embeddings = embeddings[:, :max_seq_length]
attention_mask = attention_mask[:, :max_seq_length]
# 返回处理后的嵌入和注意力掩码
return embeddings, attention_mask
@torch.no_grad()
def encode_prompt(
self,
prompt,
do_classifier_free_guidance=True,
num_images_per_prompt=1,
device=None,
negative_prompt=None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
_cut_context=False,
attention_mask: Optional[torch.Tensor] = None,
negative_attention_mask: Optional[torch.Tensor] = None,
):
# 准备潜变量，处理输入图像和时间步
def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None):
# 检查输入图像类型是否合法
if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
raise ValueError(
f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
)
# 将图像转移到指定设备，并设置数据类型
image = image.to(device=device, dtype=dtype)
# 计算实际批次大小
batch_size = batch_size * num_images_per_prompt
# 如果图像的通道数为4，初始化潜变量为输入图像
if image.shape[1] == 4:
init_latents = image
else:
# 如果生成器是列表且长度与批次大小不符，抛出错误
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
# 如果生成器是列表，逐个编码图像并采样
elif isinstance(generator, list):
init_latents = [
self.movq.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size)
]
# 将所有潜变量合并到一起
init_latents = torch.cat(init_latents, dim=0)
else:
# 否则直接编码并采样
init_latents = self.movq.encode(image).latent_dist.sample(generator)
# 根据配置的缩放因子调整潜变量
init_latents = self.movq.config.scaling_factor * init_latents
# 将初始化的潜变量在第一维上进行合并
init_latents = torch.cat([init_latents], dim=0)
# 获取初始化潜变量的形状
shape = init_latents.shape
# 生成噪声张量，用于后续处理
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# 将噪声添加到潜变量中，获取最终潜变量
init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
# 设置最终的潜变量
latents = init_latents
# 返回潜变量
return latents
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs 复制而来
def prepare_extra_step_kwargs(self, generator, eta):
# 为调度器步骤准备额外的关键字参数，因为并不是所有的调度器都有相同的参数签名
# eta（η）仅在 DDIMScheduler 中使用，其他调度器将忽略它。
# eta 对应于 DDIM 论文中的 η: https://arxiv.org/abs/2010.02502
# 其值应在 [0, 1] 之间
# 检查调度器的步骤方法是否接受 eta 参数
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 初始化额外步骤参数字典
extra_step_kwargs = {}
# 如果调度器接受 eta 参数，则将其添加到额外步骤参数字典中
if accepts_eta:
extra_step_kwargs["eta"] = eta
# 检查调度器的步骤方法是否接受 generator 参数
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 如果调度器接受 generator 参数，则将其添加到额外步骤参数字典中
if accepts_generator:
extra_step_kwargs["generator"] = generator
# 返回准备好的额外步骤参数字典
return extra_step_kwargs
def check_inputs(
self,
prompt,
callback_steps,
negative_prompt=None,
prompt_embeds=None,
negative_prompt_embeds=None,
callback_on_step_end_tensor_inputs=None,
attention_mask=None,
negative_attention_mask=None,
@property
# 定义属性方法，获取指导尺度
def guidance_scale(self):
# 返回当前指导尺度的值
return self._guidance_scale
@property
# 定义属性方法，检查是否进行无分类器自由指导
def do_classifier_free_guidance(self):
# 如果指导尺度大于 1，则返回 True
return self._guidance_scale > 1
@property
# 定义属性方法，获取时间步数
def num_timesteps(self):
# 返回当前时间步数的值
return self._num_timesteps
@torch.no_grad()
# 装饰器，表示该方法在推理时不计算梯度
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
# 定义可调用对象的方法，接受多个参数用于生成图像
prompt: Union[str, List[str]] = None,
image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]] = None,
strength: float = 0.3,
num_inference_steps: int = 25,
guidance_scale: float = 3.0,
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
negative_attention_mask: Optional[torch.Tensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 定义在步骤结束时调用的回调，默认为只包含 "latents" 的列表
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
**kwargs,
