diffusers-源码解析-三十二-
diffusers 源码解析（三十二）
.\diffusers\pipelines\kandinsky\pipeline_kandinsky_inpaint.py
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# 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.
from copy import deepcopy # 从 copy 模块导入 deepcopy 函数，用于深拷贝对象
from typing import Callable, List, Optional, Union # 导入类型提示功能，用于函数签名
import numpy as np # 导入 numpy 库，通常用于数组和矩阵操作
import PIL.Image # 导入 PIL.Image 模块，用于图像处理
import torch # 导入 PyTorch 库，深度学习框架
import torch.nn.functional as F # 导入 PyTorch 的功能性神经网络模块
from packaging import version # 导入 version 模块，用于版本管理
from PIL import Image # 导入 PIL.Image 模块，用于图像操作
from transformers import ( # 从 transformers 库导入以下组件
XLMRobertaTokenizer, # 导入 XLM-RoBERTa 的分词器
)
from ... import __version__ # 导入当前模块的版本信息
from ...models import UNet2DConditionModel, VQModel # 导入模型类
from ...schedulers import DDIMScheduler # 导入 DDIM 调度器
from ...utils import ( # 从 utils 模块导入以下工具
logging, # 导入日志模块
replace_example_docstring, # 导入用于替换示例文档字符串的函数
)
from ...utils.torch_utils import randn_tensor # 从 torch_utils 导入 randn_tensor 函数
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput # 导入扩散管道和图像管道输出类
from .text_encoder import MultilingualCLIP # 导入多语言 CLIP 文本编码器
logger = logging.get_logger(__name__) # 创建日志记录器，用于记录模块日志，名称为当前模块名
# pylint: disable=invalid-name # 禁用 pylint 对无效名称的警告
EXAMPLE_DOC_STRING = """ # 定义示例文档字符串，提供代码示例
Examples:
```py
>>> from diffusers import KandinskyInpaintPipeline, KandinskyPriorPipeline # 从 diffusers 导入管道类
>>> from diffusers.utils import load_image # 从 utils 导入图像加载函数
>>> import torch # 导入 PyTorch
>>> import numpy as np # 导入 numpy
>>> pipe_prior = KandinskyPriorPipeline.from_pretrained( # 从预训练模型创建管道
... "kandinsky-community/kandinsky-2-1-prior", torch_dtype=torch.float16 # 指定模型和数据类型
... )
>>> pipe_prior.to("cuda") # 将管道移动到 CUDA 设备
>>> prompt = "a hat" # 定义生成图像的提示词
>>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False) # 获取图像嵌入和零图像嵌入
>>> pipe = KandinskyInpaintPipeline.from_pretrained( # 从预训练模型创建图像修复管道
... "kandinsky-community/kandinsky-2-1-inpaint", torch_dtype=torch.float16 # 指定模型和数据类型
... )
>>> pipe.to("cuda") # 将管道移动到 CUDA 设备
>>> init_image = load_image( # 加载初始图像
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" # 图像的 URL
... "/kandinsky/cat.png" # 图像文件名
... )
>>> mask = np.zeros((768, 768), dtype=np.float32) # 创建一个全零的掩膜
>>> mask[:250, 250:-250] = 1 # 在掩膜中设置特定区域为 1
>>> out = pipe( # 调用管道进行图像生成
... prompt, # 提示词
... 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
"""
def get_new_h_w(h, w, scale_factor=8): # 定义函数，计算新的高和宽
new_h = h // scale_factor**2 # 计算新的高度，使用整除
if h % scale_factor**2 != 0: # 如果原高度不能被 scale_factor^2 整除
new_h += 1 # 高度加一
new_w = w // scale_factor**2 # 计算新的宽度，使用整除
if w % scale_factor**2 != 0: # 如果原宽度不能被 scale_factor^2 整除
new_w += 1 # 宽度加一
return new_h * scale_factor, new_w * scale_factor # 返回新的高和宽，均乘以 scale_factor
# 准备掩码
def prepare_mask(masks):
# 初始化一个空列表以存储准备好的掩码
prepared_masks = []
# 遍历所有输入掩码
for mask in masks:
# 深拷贝当前掩码以避免修改原始数据
old_mask = deepcopy(mask)
# 遍历掩码的每一行
for i in range(mask.shape[1]):
#
# 检查 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，确保形状符合要求
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 的维度
if mask.ndim == 3:
# 如果 mask 是单个批量图像，且没有通道维度
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 二值化
mask[mask < 0.5] = 0
mask[mask >= 0.5] = 1
# 将 image 转换为 float32 类型
image = image.to(dtype=torch.float32)
# 如果 mask 是 torch.Tensor，但 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 数组，并添加一个维度
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):
# 将多个 ndarray 在第一个维度上连接起来
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 张量，并进行归一化处理
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]
# 将每个掩码转换为灰度数组，并添加两个维度
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):
# 将多个 ndarray 在第一个维度上连接起来
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
# 将 numpy 数组转换为 PyTorch 张量
mask = torch.from_numpy(mask)
# 将掩码取反
mask = 1 - mask
# 返回掩码和图像
return mask, image
# 定义Kandinsky图像修复管道，继承自DiffusionPipeline
class KandinskyInpaintPipeline(DiffusionPipeline):
"""
使用Kandinsky2.1进行文本引导的图像修复的管道
该模型继承自[`DiffusionPipeline`]。请查看超类文档以获取库为所有管道实现的通用方法
（例如下载或保存，运行在特定设备上等）。
参数：
text_encoder ([`MultilingualCLIP`]):
冻结的文本编码器。
tokenizer ([`XLMRobertaTokenizer`]):
类的分词器
scheduler ([`DDIMScheduler`]):
用于与`unet`结合生成图像潜在表示的调度器。
unet ([`UNet2DConditionModel`]):
用于去噪图像嵌入的条件U-Net架构。
movq ([`VQModel`]):
MoVQ图像编码器和解码器
"""
# 定义模型CPU卸载顺序
model_cpu_offload_seq = "text_encoder->unet->movq"
def __init__(
self,
text_encoder: MultilingualCLIP,
movq: VQModel,
tokenizer: XLMRobertaTokenizer,
unet: UNet2DConditionModel,
scheduler: DDIMScheduler,
):
# 初始化父类
super().__init__()
# 注册模型组件
self.register_modules(
text_encoder=text_encoder,
movq=movq,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
)
# 计算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复制的方法，用于准备潜在表示
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 _encode_prompt(
self,
prompt,
device,
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
):
# 定义一个可调用的类方法
def __call__(
self,
# 接受一个字符串或字符串列表作为提示
prompt: Union[str, List[str]],
# 接受一个图像张量或 PIL 图像作为输入图像
image: Union[torch.Tensor, PIL.Image.Image],
# 接受一个掩膜图像，可以是张量、PIL 图像或 NumPy 数组
mask_image: Union[torch.Tensor, PIL.Image.Image, np.ndarray],
# 接受一个图像嵌入的张量
image_embeds: torch.Tensor,
# 接受一个负图像嵌入的张量
negative_image_embeds: torch.Tensor,
# 可选参数，接受一个字符串或字符串列表作为负提示
negative_prompt: Optional[Union[str, List[str]]] = None,
# 设置输出图像的高度，默认为 512
height: int = 512,
# 设置输出图像的宽度，默认为 512
width: int = 512,
# 设置推理步骤的数量，默认为 100
num_inference_steps: int = 100,
# 设置引导比例，默认为 4.0
guidance_scale: float = 4.0,
# 设置每个提示生成的图像数量，默认为 1
num_images_per_prompt: int = 1,
# 可选参数，接受一个随机数生成器或生成器列表
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 可选参数，接受潜在空间的张量
latents: Optional[torch.Tensor] = None,
# 可选参数，设置输出类型，默认为 "pil"
output_type: Optional[str] = "pil",
# 可选参数，接受一个回调函数，用于处理图像生成的步骤
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
# 设置回调函数调用的步数，默认为 1
callback_steps: int = 1,
# 设置返回字典的布尔值，默认为 True
return_dict: bool = True,
.\diffusers\pipelines\kandinsky\pipeline_kandinsky_prior.py
# 版权信息，表明该代码的版权归 HuggingFace 团队所有
#
# 根据 Apache 许可证，版本 2.0（“许可证”），
# 除非符合许可证，否则您不得使用此文件。
# 您可以在以下网址获得许可证的副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面同意，软件按“现状”分发，
# 不提供任何形式的担保或条件，无论是明示或暗示。
# 查看许可证以了解特定语言适用的权限和
# 限制。
# 从 dataclasses 模块导入 dataclass 装饰器
from dataclasses import dataclass
# 从 typing 模块导入 List, Optional 和 Union 类型
from typing import List, Optional, Union
# 导入 numpy 库，并简化为 np
import numpy as np
# 导入 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 (
BaseOutput, # 基本输出类
logging, # 日志模块
replace_example_docstring, # 替换示例文档字符串的函数
)
# 从 torch_utils 模块导入随机张量生成函数
from ...utils.torch_utils import randn_tensor
# 从 pipeline_utils 模块导入 DiffusionPipeline 类
from ..pipeline_utils import DiffusionPipeline
# 获取当前模块的日志记录器
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，包含用法示例
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline
>>> import torch
>>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-1-prior")
>>> pipe_prior.to("cuda")
>>> prompt = "red cat, 4k photo"
>>> out = pipe_prior(prompt)
>>> image_emb = out.image_embeds
>>> negative_image_emb = out.negative_image_embeds
>>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1")
>>> pipe.to("cuda")
>>> image = pipe(
... prompt,
... 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 = """
# 示例用法
Examples:
```py
# 从 diffusers 库中导入 KandinskyPriorPipeline 和 KandinskyPipeline
>>> from diffusers import KandinskyPriorPipeline, KandinskyPipeline
# 从 diffusers.utils 导入 load_image 函数
>>> from diffusers.utils import load_image
# 导入 PIL 库
>>> import PIL
# 导入 torch 库
>>> import torch
# 从 torchvision 导入 transforms 模块
>>> from torchvision import transforms
# 加载预训练的 KandinskyPriorPipeline，设置数据类型为 float16
>>> pipe_prior = KandinskyPriorPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-1-prior", torch_dtype=torch.float16
... )
# 将管道移动到 GPU 设备
>>> pipe_prior.to("cuda")
# 从指定 URL 加载第一张图像
>>> img1 = load_image(
... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
... "/kandinsky/cat.png"
... )
# 从指定 URL 加载第二张图像
>>> 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)
# 加载预训练的 KandinskyPipeline，设置数据类型为 float16
>>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1", 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]
# 将生成的图像保存为文件 "starry_cat.png"
>>> image.save("starry_cat.png")
"""
文档字符串，描述该模块的用途
@dataclass
装饰器，定义数据类以便更方便地管理属性
class KandinskyPriorPipelineOutput(BaseOutput):
"""
Output class for KandinskyPriorPipeline.
# 文档字符串，描述 KandinskyPriorPipeline 的输出类
Args:
image_embeds (`torch.Tensor`)
clip image embeddings for text prompt
# 描述图像嵌入的参数，供文本提示使用
negative_image_embeds (`List[PIL.Image.Image]` or `np.ndarray`)
clip image embeddings for unconditional tokens
# 描述用于无条件标记的图像嵌入参数
"""
# 定义输出类的属性，包括图像嵌入和负图像嵌入
image_embeds: Union[torch.Tensor, np.ndarray]
# 图像嵌入属性，可以是 PyTorch 张量或 NumPy 数组
negative_image_embeds: Union[torch.Tensor, np.ndarray]
# 负图像嵌入属性，也可以是 PyTorch 张量或 NumPy 数组
class KandinskyPriorPipeline(DiffusionPipeline):
"""
Pipeline for generating image prior for Kandinsky
# 文档字符串，描述生成 Kandinsky 图像先验的管道
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
# 说明该模型继承自 DiffusionPipeline，参考超类文档以了解通用方法
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
Args:
prior ([`PriorTransformer`]):
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`):
Tokenizer of class
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
# CLIP 词元化器
scheduler ([`UnCLIPScheduler`]):
A scheduler to be used in combination with `prior` to generate image embedding.
# 调度器，用于结合先验生成图像嵌入
"""
# 定义不进行 CPU 卸载的模块
_exclude_from_cpu_offload = ["prior"]
# 定义模型 CPU 卸载序列
model_cpu_offload_seq = "text_encoder->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,
)
@torch.no_grad()
# 修饰函数，表示在此函数中不需要计算梯度
@replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING)
# 替换示例文档字符串，提供相关的例子
def interpolate(
self,
images_and_prompts: List[Union[str, PIL.Image.Image, torch.Tensor]],
# 输入图像和提示的列表，可以是字符串、PIL 图像或 PyTorch 张量
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,
# 潜在变量，可选，默认为 None
negative_prior_prompt: Optional[str] = None,
# 负先验提示，默认为 None
negative_prompt: str = "",
# 负提示，默认为空字符串
guidance_scale: float = 4.0,
# 引导缩放因子，影响生成图像的风格
device=None,
# 指定设备，默认为 None
# Copied from 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
# 获取零嵌入，用于生成图像的初始状态
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
# 编码提示的私有方法，处理输入提示和设备
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__(
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",
return_dict: bool = True,
# `.\diffusers\pipelines\kandinsky\text_encoder.py`
```py
# 导入 PyTorch 库
import torch
# 从 transformers 库导入预训练模型类和配置类
from transformers import PreTrainedModel, XLMRobertaConfig, XLMRobertaModel
# 定义 MCLIPConfig 类，继承自 XLMRobertaConfig
class MCLIPConfig(XLMRobertaConfig):
# 设置模型类型为 "M-CLIP"
model_type = "M-CLIP"
# 初始化方法，接受变换器和图像维度大小的参数
def __init__(self, transformerDimSize=1024, imageDimSize=768, **kwargs):
# 将变换器维度大小赋值给实例变量
self.transformerDimensions = transformerDimSize
# 将图像维度大小赋值给实例变量
self.numDims = imageDimSize
# 调用父类的初始化方法
super().__init__(**kwargs)
# 定义 MultilingualCLIP 类，继承自 PreTrainedModel
class MultilingualCLIP(PreTrainedModel):
# 指定配置类为 MCLIPConfig
config_class = MCLIPConfig
# 初始化方法，接受配置和其他参数
def __init__(self, config, *args, **kwargs):
# 调用父类的初始化方法
super().__init__(config, *args, **kwargs)
# 创建 XLMRobertaModel 实例，使用传入的配置
self.transformer = XLMRobertaModel(config)
# 创建线性变换层，输入特征为变换器维度，输出特征为图像维度
self.LinearTransformation = torch.nn.Linear(
in_features=config.transformerDimensions, out_features=config.numDims
)
# 定义前向传播方法，接受输入 ID 和注意力掩码
def forward(self, input_ids, attention_mask):
# 获取变换器的嵌入表示，使用输入 ID 和注意力掩码
embs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)[0]
# 根据注意力掩码计算加权嵌入表示
embs2 = (embs * attention_mask.unsqueeze(2)).sum(dim=1) / attention_mask.sum(dim=1)[:, None]
# 返回线性变换后的嵌入表示和原始嵌入表示
return self.LinearTransformation(embs2), embs
.\diffusers\pipelines\kandinsky\__init__.py
# 从 typing 模块导入 TYPE_CHECKING，用于类型检查
from typing import TYPE_CHECKING
# 从 utils 模块导入相关的工具和常量
from ...utils import (
DIFFUSERS_SLOW_IMPORT, # 导入用于标识慢速导入的常量
OptionalDependencyNotAvailable, # 导入表示可选依赖不可用的异常
_LazyModule, # 导入懒加载模块的类
get_objects_from_module, # 导入从模块获取对象的函数
is_torch_available, # 导入检查 Torch 库可用性的函数
is_transformers_available, # 导入检查 Transformers 库可用性的函数
)
# 用于存储虚拟对象的字典
_dummy_objects = {}
# 用于存储模块导入结构的字典
_import_structure = {}
try:
# 检查 Transformers 和 Torch 是否都可用
if not (is_transformers_available() and is_torch_available()):
# 如果不可用，则抛出可选依赖不可用异常
raise OptionalDependencyNotAvailable()
# 捕获可选依赖不可用异常
except OptionalDependencyNotAvailable:
# 从 utils 模块导入虚拟的 Torch 和 Transformers 对象，忽略 F403 警告
from ...utils import dummy_torch_and_transformers_objects # noqa F403
# 更新 _dummy_objects 字典，获取虚拟对象
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
else:
# 如果可用，则更新导入结构字典，添加各种管道的导入
_import_structure["pipeline_kandinsky"] = ["KandinskyPipeline"]
_import_structure["pipeline_kandinsky_combined"] = [
"KandinskyCombinedPipeline",
"KandinskyImg2ImgCombinedPipeline",
"KandinskyInpaintCombinedPipeline",
]
_import_structure["pipeline_kandinsky_img2img"] = ["KandinskyImg2ImgPipeline"]
_import_structure["pipeline_kandinsky_inpaint"] = ["KandinskyInpaintPipeline"]
_import_structure["pipeline_kandinsky_prior"] = ["KandinskyPriorPipeline", "KandinskyPriorPipelineOutput"]
_import_structure["text_encoder"] = ["MultilingualCLIP"]
# 如果在类型检查模式或慢速导入模式下执行
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
try:
# 再次检查 Transformers 和 Torch 是否都可用
if not (is_transformers_available() and is_torch_available()):
# 如果不可用，则抛出可选依赖不可用异常
raise OptionalDependencyNotAvailable()
# 捕获可选依赖不可用异常
except OptionalDependencyNotAvailable:
# 从虚拟对象模块导入所有内容
from ...utils.dummy_torch_and_transformers_objects import *
else:
# 从各个管道模块导入实际的类
from .pipeline_kandinsky import KandinskyPipeline
from .pipeline_kandinsky_combined import (
KandinskyCombinedPipeline,
KandinskyImg2ImgCombinedPipeline,
KandinskyInpaintCombinedPipeline,
)
from .pipeline_kandinsky_img2img import KandinskyImg2ImgPipeline
from .pipeline_kandinsky_inpaint import KandinskyInpaintPipeline
from .pipeline_kandinsky_prior import KandinskyPriorPipeline, KandinskyPriorPipelineOutput
from .text_encoder import MultilingualCLIP
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\kandinsky2_2\pipeline_kandinsky2_2.py
# 版权声明，2024年由 HuggingFace 团队保留所有权利
#
# 根据 Apache 许可证 2.0 版（"许可证"）进行许可；
# 除非遵守许可证，否则不得使用此文件。
# 可在以下地址获得许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有规定，软件
# 以“按原样”方式分发，不附带任何形式的保证或条件，
# 明示或暗示。
# 请参见许可证以获取与权限和
# 限制相关的具体信息。
from typing import Callable, Dict, List, Optional, Union # 从 typing 模块导入类型提示相关类
import torch # 导入 PyTorch 库
# 从模型模块导入 UNet2DConditionModel 和 VQModel
from ...models import UNet2DConditionModel, VQModel
# 从调度器模块导入 DDPMScheduler
from ...schedulers import DDPMScheduler
# 从工具模块导入实用功能
from ...utils import deprecate, logging, replace_example_docstring
# 从工具中的 torch_utils 模块导入 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 = """
Examples:
```py
>>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
>>> import torch
>>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior") # 从预训练模型创建管道
>>> pipe_prior.to("cuda") # 将管道转移到 CUDA 设备
>>> prompt = "red cat, 4k photo" # 设置生成图像的提示文本
>>> out = pipe_prior(prompt) # 使用提示生成图像嵌入
>>> image_emb = out.image_embeds # 提取生成的图像嵌入
>>> zero_image_emb = out.negative_image_embeds # 提取负图像嵌入
>>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder") # 创建解码器管道
>>> pipe.to("cuda") # 将解码器管道转移到 CUDA 设备
>>> image = pipe( # 使用图像嵌入生成图像
... image_embeds=image_emb,
... negative_image_embeds=zero_image_emb,
... height=768, # 设置生成图像的高度
... width=768, # 设置生成图像的宽度
... num_inference_steps=50, # 设置推理步骤数量
... ).images # 提取生成的图像
>>> image[0].save("cat.png") # 保存生成的图像为 PNG 文件
```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 # 返回缩放后的高度和宽度
# 定义 KandinskyV22Pipeline 类，继承自 DiffusionPipeline
class KandinskyV22Pipeline(DiffusionPipeline):
"""
用于使用 Kandinsky 进行文本到图像生成的管道
该模型继承自 [`DiffusionPipeline`]。有关库为所有管道实现的通用方法的文档，请参见父类文档（例如下载或保存，运行在特定设备等）。
Args:
scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]):
与 `unet` 结合使用以生成图像潜在特征的调度器。
unet ([`UNet2DConditionModel`]):
条件 U-Net 结构，用于去噪图像嵌入。
movq ([`VQModel`]):
MoVQ 解码器，用于从潜在特征生成图像。
"""
# 定义模型的 CPU 卸载顺序，这里是将 unet 的输出移动到 q 形式
model_cpu_offload_seq = "unet->movq"
# 定义输入张量的名称列表
_callback_tensor_inputs = ["latents", "image_embeds", "negative_image_embeds"]
# 初始化方法，用于创建类的实例
def __init__(
# UNet2DConditionModel 实例，用于生成模型
unet: UNet2DConditionModel,
# DDPMScheduler 实例，用于调度过程
scheduler: DDPMScheduler,
# VQModel 实例，用于处理量化
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 复制的准备潜在变量的方法
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
# 获取指导缩放因子的属性
@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()
# 替换示例文档字符串
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
# 图像嵌入，可以是张量或张量列表
image_embeds: Union[torch.Tensor, List[torch.Tensor]],
# 负图像嵌入，可以是张量或张量列表
negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]],
# 输出图像的高度，默认 512
height: int = 512,
# 输出图像的宽度，默认 512
width: int = 512,
# 推理步骤数，默认 100
num_inference_steps: int = 100,
# 指导缩放因子，默认 4.0
guidance_scale: float = 4.0,
# 每个提示生成的图像数量，默认 1
num_images_per_prompt: int = 1,
# 随机数生成器，可以是单个生成器或生成器列表
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 可选的潜在变量张量
latents: Optional[torch.Tensor] = None,
# 输出类型，默认是 "pil"
output_type: Optional[str] = "pil",
# 是否返回字典，默认是 True
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,
.\diffusers\pipelines\kandinsky2_2\pipeline_kandinsky2_2_combined.py
# 版权信息，注明该代码的所有权归 HuggingFace 团队所有
# 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.
# 导入类型相关的类，以支持类型注解
from typing import Callable, Dict, List, Optional, Union
# 导入 PIL 库中的 Image 模块，用于图像处理
import PIL.Image
# 导入 PyTorch 库
import torch
# 从 transformers 库中导入 CLIP 相关模型和处理器
from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection
# 从本地模块导入各种模型和调度器
from ...models import PriorTransformer, UNet2DConditionModel, VQModel
from ...schedulers import DDPMScheduler, UnCLIPScheduler
from ...utils import deprecate, logging, replace_example_docstring
from ..pipeline_utils import DiffusionPipeline
from .pipeline_kandinsky2_2 import KandinskyV22Pipeline
from .pipeline_kandinsky2_2_img2img import KandinskyV22Img2ImgPipeline
from .pipeline_kandinsky2_2_inpainting import KandinskyV22InpaintPipeline
from .pipeline_kandinsky2_2_prior import KandinskyV22PriorPipeline
# 创建日志记录器，用于记录模块中的日志信息
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 定义文本到图像转换的示例文档字符串
TEXT2IMAGE_EXAMPLE_DOC_STRING = """
Examples:
```py
# 从 diffusers 库中导入自动文本到图像管道
from diffusers import AutoPipelineForText2Image
# 导入 PyTorch 库
import torch
# 加载预训练的文本到图像管道，使用半精度浮点数
pipe = AutoPipelineForText2Image.from_pretrained(
"kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
)
# 启用模型的 CPU 卸载，以减少内存占用
pipe.enable_model_cpu_offload()
# 定义要生成图像的提示语
prompt = "A lion in galaxies, spirals, nebulae, stars, smoke, iridescent, intricate detail, octane render, 8k"
# 生成图像，指定推理步骤数量，并获取生成的第一张图像
image = pipe(prompt=prompt, num_inference_steps=25).images[0]
```py
"""
# 定义图像到图像转换的示例文档字符串
IMAGE2IMAGE_EXAMPLE_DOC_STRING = """
Examples:
```py
# 从 diffusers 库中导入自动图像到图像管道
from diffusers import AutoPipelineForImage2Image
# 导入 PyTorch 库
import torch
# 导入请求库用于下载图像
import requests
# 从 io 库中导入 BytesIO，用于处理字节流
from io import BytesIO
# 导入 PIL 库中的 Image 模块
from PIL import Image
# 导入 os 库
import os
# 加载预训练的图像到图像管道，使用半精度浮点数
pipe = AutoPipelineForImage2Image.from_pretrained(
"kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
)
# 启用模型的 CPU 卸载，以减少内存占用
pipe.enable_model_cpu_offload()
# 定义要生成图像的提示语
prompt = "A fantasy landscape, Cinematic lighting"
# 定义负面提示语，用于控制生成的图像质量
negative_prompt = "low quality, bad quality"
# 定义要下载的图像 URL
url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
# 发起 GET 请求以获取图像
response = requests.get(url)
# 打开获取的图像内容，并将其转换为 RGB 模式
image = Image.open(BytesIO(response.content)).convert("RGB")
# 调整图像大小，以适应后续处理
image.thumbnail((768, 768))
# 使用管道生成新图像，传入原始图像和提示语，获取生成的第一张图像
image = pipe(prompt=prompt, image=original_image, num_inference_steps=25).images[0]
```py
"""
# 定义图像修复的示例文档字符串
INPAINT_EXAMPLE_DOC_STRING = """
# 示例代码展示如何使用 Diffusers 库进行图像修复
Examples:
```py
# 从 diffusers 库导入 AutoPipelineForInpainting 类
from diffusers import AutoPipelineForInpainting
# 从 diffusers.utils 导入 load_image 函数
from diffusers.utils import load_image
# 导入 PyTorch 库
import torch
# 导入 NumPy 库
import numpy as np
# 从预训练模型加载图像修复管道，指定数据类型为 float16
pipe = AutoPipelineForInpainting.from_pretrained(
"kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16
)
# 启用模型的 CPU 内存卸载功能，以节省显存
pipe.enable_model_cpu_offload()
# 定义图像修复的提示语
prompt = "A fantasy landscape, Cinematic lighting"
# 定义负面提示，以避免生成低质量图像
negative_prompt = "low quality, bad quality"
# 加载原始图像，使用给定的 URL
original_image = load_image(
"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png"
)
# 创建一个全零的掩码，大小为 (768, 768)，数据类型为 float32
mask = np.zeros((768, 768), dtype=np.float32)
# 在猫的头部上方的区域设置掩码为 1，以进行遮挡
mask[:250, 250:-250] = 1
# 使用管道进行图像修复，传入提示、原始图像和掩码，指定推理步骤数为 25，获取生成的图像
image = pipe(prompt=prompt, image=original_image, mask_image=mask, num_inference_steps=25).images[0]
```py
# 定义一个结合管道类，用于基于 Kandinsky 的文本到图像生成
class KandinskyV22CombinedPipeline(DiffusionPipeline):
# 文档字符串，描述该类的功能和参数
"""
Combined Pipeline for text-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 (Union[`DDIMScheduler`,`DDPMScheduler`]):
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.
prior_prior ([`PriorTransformer`]):
The canonical unCLIP prior to approximate the image embedding from the text embedding.
prior_image_encoder ([`CLIPVisionModelWithProjection`]):
Frozen image-encoder.
prior_text_encoder ([`CLIPTextModelWithProjection`]):
Frozen text-encoder.
prior_tokenizer (`CLIPTokenizer`):
Tokenizer of class
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
prior_scheduler ([`UnCLIPScheduler`]):
A scheduler to be used in combination with `prior` to generate image embedding.
prior_image_processor ([`CLIPImageProcessor`]):
A image_processor to be used to preprocess image from clip.
"""
# 定义模型的 CPU 卸载顺序，指定各部分的执行顺序
model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->unet->movq"
# 表示需要加载连接的管道
_load_connected_pipes = True
# 指定从 CPU 卸载时要排除的组件
_exclude_from_cpu_offload = ["prior_prior"]
# 初始化方法，接受多个模型作为参数
def __init__(
self,
unet: UNet2DConditionModel, # 条件 U-Net 模型，用于图像去噪
scheduler: DDPMScheduler, # 调度器，用于生成图像潜在特征
movq: VQModel, # MoVQ 解码器，从潜在特征生成图像
prior_prior: PriorTransformer, # 用于近似图像嵌入的先验转换器
prior_image_encoder: CLIPVisionModelWithProjection, # 冻结的图像编码器
prior_text_encoder: CLIPTextModelWithProjection, # 冻结的文本编码器
prior_tokenizer: CLIPTokenizer, # CLIP 的分词器
prior_scheduler: UnCLIPScheduler, # 生成图像嵌入的调度器
prior_image_processor: CLIPImageProcessor, # 用于预处理图像的图像处理器
# 初始化父类
):
super().__init__()
# 注册各个模块，以便于后续使用
self.register_modules(
unet=unet,
scheduler=scheduler,
movq=movq,
prior_prior=prior_prior,
prior_image_encoder=prior_image_encoder,
prior_text_encoder=prior_text_encoder,
prior_tokenizer=prior_tokenizer,
prior_scheduler=prior_scheduler,
prior_image_processor=prior_image_processor,
)
# 创建先前处理管道，用于图像和文本处理
self.prior_pipe = KandinskyV22PriorPipeline(
prior=prior_prior,
image_encoder=prior_image_encoder,
text_encoder=prior_text_encoder,
tokenizer=prior_tokenizer,
scheduler=prior_scheduler,
image_processor=prior_image_processor,
)
# 创建解码管道，负责生成最终输出
self.decoder_pipe = KandinskyV22Pipeline(
unet=unet,
scheduler=scheduler,
movq=movq,
)
# 启用高效内存注意力机制
def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
# 调用解码管道的方法，启用内存高效注意力
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
# 启用顺序 CPU 卸载
def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"):
r"""
将所有模型卸载到 CPU，从而显著减少内存使用。当调用时，unet、text_encoder、vae 和安全检查器的状态字典保存到 CPU，然后移至
`torch.device('meta')，并在特定子模块调用其 `forward` 方法时加载到 GPU。
注意，卸载是针对子模块进行的。内存节省比 `enable_model_cpu_offload` 高，但性能较低。
"""
# 启用先前处理管道的顺序 CPU 卸载
self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device)
# 启用解码管道的顺序 CPU 卸载
self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device)
# 进度条功能
def progress_bar(self, iterable=None, total=None):
# 调用先前处理管道的进度条功能
self.prior_pipe.progress_bar(iterable=iterable, total=total)
# 调用解码管道的进度条功能
self.decoder_pipe.progress_bar(iterable=iterable, total=total)
# 启用解码管道的模型 CPU 卸载
self.decoder_pipe.enable_model_cpu_offload()
# 设置进度条的配置
def set_progress_bar_config(self, **kwargs):
# 为先前处理管道设置进度条配置
self.prior_pipe.set_progress_bar_config(**kwargs)
# 为解码管道设置进度条配置
self.decoder_pipe.set_progress_bar_config(**kwargs)
# 无梯度计算装饰器，避免计算梯度
@torch.no_grad()
# 替换示例文档字符串
@replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING)
# 定义一个可调用的类方法，接受多个参数进行处理
def __call__(
self,
# 输入的提示，可以是字符串或字符串列表
prompt: Union[str, List[str]],
# 可选的负面提示，也可以是字符串或字符串列表
negative_prompt: Optional[Union[str, List[str]]] = None,
# 推理步骤的数量，默认为100
num_inference_steps: int = 100,
# 指导比例，默认为4.0
guidance_scale: float = 4.0,
# 每个提示生成的图像数量，默认为1
num_images_per_prompt: int = 1,
# 输出图像的高度，默认为512
height: int = 512,
# 输出图像的宽度，默认为512
width: int = 512,
# 先前指导比例，默认为4.0
prior_guidance_scale: float = 4.0,
# 先前推理步骤的数量，默认为25
prior_num_inference_steps: int = 25,
# 随机数生成器，可选，支持单个或多个生成器
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 可选的潜变量，默认为None
latents: Optional[torch.Tensor] = None,
# 输出类型，默认为"pil"
output_type: Optional[str] = "pil",
# 可选的回调函数，用于处理推理过程中的步骤
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
# 回调函数触发的步骤间隔，默认为1
callback_steps: int = 1,
# 返回结果的字典形式，默认为True
return_dict: bool = True,
# 可选的先前回调函数，在步骤结束时调用
prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 用于先前回调的张量输入，默认为["latents"]
prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# 可选的回调函数，在步骤结束时调用
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 用于回调的张量输入，默认为["latents"]
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# 定义一个结合图像到图像生成的管道类 KandinskyV22Img2ImgCombinedPipeline，继承自 DiffusionPipeline
class KandinskyV22Img2ImgCombinedPipeline(DiffusionPipeline):
"""
Combined 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 (Union[`DDIMScheduler`,`DDPMScheduler`]):
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.
prior_prior ([`PriorTransformer`]):
The canonical unCLIP prior to approximate the image embedding from the text embedding.
prior_image_encoder ([`CLIPVisionModelWithProjection`]):
Frozen image-encoder.
prior_text_encoder ([`CLIPTextModelWithProjection`]):
Frozen text-encoder.
prior_tokenizer (`CLIPTokenizer`):
Tokenizer of class
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
prior_scheduler ([`UnCLIPScheduler`]):
A scheduler to be used in combination with `prior` to generate image embedding.
prior_image_processor ([`CLIPImageProcessor`]):
A image_processor to be used to preprocess image from clip.
"""
# 定义 CPU 卸载顺序的模型组件
model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->unet->movq"
# 指定是否加载连接的管道
_load_connected_pipes = True
# 指定从 CPU 卸载中排除的组件
_exclude_from_cpu_offload = ["prior_prior"]
# 初始化方法，接受多个模型组件作为参数
def __init__(
# 条件 U-Net 模型，用于图像嵌入去噪
self,
unet: UNet2DConditionModel,
# DDPMScheduler 调度器
scheduler: DDPMScheduler,
# MoVQ 解码器
movq: VQModel,
# 用于图像嵌入的 unCLIP 先验
prior_prior: PriorTransformer,
# 冻结的图像编码器
prior_image_encoder: CLIPVisionModelWithProjection,
# 冻结的文本编码器
prior_text_encoder: CLIPTextModelWithProjection,
# CLIP 令牌化器
prior_tokenizer: CLIPTokenizer,
# 先验的调度器
prior_scheduler: UnCLIPScheduler,
# 图像预处理器
prior_image_processor: CLIPImageProcessor,
# 初始化父类
):
super().__init__()
# 注册多个模块以供后续使用
self.register_modules(
unet=unet, # UNet模型
scheduler=scheduler, # 调度器
movq=movq, # MovQ模型
prior_prior=prior_prior, # 先验模型
prior_image_encoder=prior_image_encoder, # 图像编码器
prior_text_encoder=prior_text_encoder, # 文本编码器
prior_tokenizer=prior_tokenizer, # 标记器
prior_scheduler=prior_scheduler, # 先验调度器
prior_image_processor=prior_image_processor, # 图像处理器
)
# 创建先验管道实例
self.prior_pipe = KandinskyV22PriorPipeline(
prior=prior_prior, # 先验模型
image_encoder=prior_image_encoder, # 图像编码器
text_encoder=prior_text_encoder, # 文本编码器
tokenizer=prior_tokenizer, # 标记器
scheduler=prior_scheduler, # 先验调度器
image_processor=prior_image_processor, # 图像处理器
)
# 创建图像到图像的转换管道实例
self.decoder_pipe = KandinskyV22Img2ImgPipeline(
unet=unet, # UNet模型
scheduler=scheduler, # 调度器
movq=movq, # MovQ模型
)
# 启用高效的 xformers 内存注意力机制
def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
# 启用模型 CPU 卸载
def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"):
r"""
使用 accelerate 将所有模型卸载到 CPU，减少内存使用，性能影响低。与 `enable_sequential_cpu_offload` 相比，
此方法在调用 `forward` 方法时将整个模型一次性移动到 GPU，并在下一个模型运行前保持在 GPU。
内存节省较少，但性能因 UNet 的迭代执行而更好。
"""
# 启用先验管道的 CPU 卸载
self.prior_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device)
# 启用解码管道的 CPU 卸载
self.decoder_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device)
# 启用顺序 CPU 卸载
def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"):
r"""
使用 accelerate 将所有模型卸载到 CPU，显著减少内存使用。调用时，unet、text_encoder、vae 和安全检查器
的状态字典保存到 CPU，然后移动到 `torch.device('meta')`，仅在其特定子模块调用 `forward` 方法时加载到 GPU。
注意卸载是按子模块进行的。内存节省高于 `enable_model_cpu_offload`，但性能较低。
"""
# 启用先验管道的顺序 CPU 卸载
self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device)
# 启用解码管道的顺序 CPU 卸载
self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device)
# 显示进度条
def progress_bar(self, iterable=None, total=None):
# 在先验管道中显示进度条
self.prior_pipe.progress_bar(iterable=iterable, total=total)
# 在解码管道中显示进度条
self.decoder_pipe.progress_bar(iterable=iterable, total=total)
# 启用解码管道的模型 CPU 卸载
self.decoder_pipe.enable_model_cpu_offload()
# 设置进度条的配置，允许通过关键字参数自定义
def set_progress_bar_config(self, **kwargs):
# 在优先处理管道中设置进度条配置，传入关键字参数
self.prior_pipe.set_progress_bar_config(**kwargs)
# 在解码器管道中设置进度条配置，传入关键字参数
self.decoder_pipe.set_progress_bar_config(**kwargs)
# 在不计算梯度的情况下进行操作，节省内存
@torch.no_grad()
# 用于替换示例文档字符串的装饰器
@replace_example_docstring(IMAGE2IMAGE_EXAMPLE_DOC_STRING)
# 定义调用方法，支持多种输入类型和参数配置
def __call__(
# 提示文本，可以是单个字符串或字符串列表
prompt: Union[str, List[str]],
# 输入图像，可以是张量或 PIL 图像，也支持列表
image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
# 可选的负提示文本，用于限制生成
negative_prompt: Optional[Union[str, List[str]]] = None,
# 推理步骤数量，默认为 100
num_inference_steps: int = 100,
# 指导比例，控制生成的多样性，默认为 4.0
guidance_scale: float = 4.0,
# 强度参数，控制图像变换的强烈程度，默认为 0.3
strength: float = 0.3,
# 每个提示生成的图像数量，默认为 1
num_images_per_prompt: int = 1,
# 生成图像的高度，默认为 512
height: int = 512,
# 生成图像的宽度，默认为 512
width: int = 512,
# 优先指导比例，默认为 4.0
prior_guidance_scale: float = 4.0,
# 优先推理步骤数量，默认为 25
prior_num_inference_steps: int = 25,
# 随机数生成器，支持单个或多个生成器
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 预设的潜在变量，默认为 None
latents: Optional[torch.Tensor] = None,
# 输出类型，默认为 "pil"（PIL 格式）
output_type: Optional[str] = "pil",
# 回调函数，用于处理生成过程中的步骤
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
# 回调执行的步骤频率，默认为 1
callback_steps: int = 1,
# 是否返回字典格式的结果，默认为 True
return_dict: bool = True,
# 优先回调函数在步骤结束时执行
prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 在步骤结束时传入的优先回调的张量输入，默认为 ["latents"]
prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# 回调函数在步骤结束时执行
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 在步骤结束时传入的回调的张量输入，默认为 ["latents"]
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# 定义一个名为 KandinskyV22InpaintCombinedPipeline 的类，继承自 DiffusionPipeline
class KandinskyV22InpaintCombinedPipeline(DiffusionPipeline):
"""
使用 Kandinsky 进行图像修复生成的组合管道
该模型继承自 [`DiffusionPipeline`]。有关库为所有管道实现的通用方法（例如下载或保存、在特定设备上运行等），请查看父类文档。
Args:
scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]):
要与 `unet` 结合使用的调度器，用于生成图像潜在表示。
unet ([`UNet2DConditionModel`]):
条件 U-Net 结构，用于去噪图像嵌入。
movq ([`VQModel`]):
MoVQ 解码器，从潜在表示生成图像。
prior_prior ([`PriorTransformer`]):
经典的 unCLIP 先验，用于近似文本嵌入的图像嵌入。
prior_image_encoder ([`CLIPVisionModelWithProjection`]):
冻结的图像编码器。
prior_text_encoder ([`CLIPTextModelWithProjection`]):
冻结的文本编码器。
prior_tokenizer (`CLIPTokenizer`):
该类的分词器
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer)。
prior_scheduler ([`UnCLIPScheduler`]):
要与 `prior` 结合使用的调度器，用于生成图像嵌入。
prior_image_processor ([`CLIPImageProcessor`]):
用于预处理图像的图像处理器。
"""
# 定义模型的 CPU 卸载顺序
model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->unet->movq"
# 指定是否加载连接的管道
_load_connected_pipes = True
# 指定要排除在 CPU 卸载之外的组件
_exclude_from_cpu_offload = ["prior_prior"]
# 初始化方法，定义管道的参数
def __init__(
# 条件 U-Net 模型
self,
unet: UNet2DConditionModel,
# 扩散调度器
scheduler: DDPMScheduler,
# MoVQ 解码器
movq: VQModel,
# 先验变换器
prior_prior: PriorTransformer,
# 图像编码器
prior_image_encoder: CLIPVisionModelWithProjection,
# 文本编码器
prior_text_encoder: CLIPTextModelWithProjection,
# 分词器
prior_tokenizer: CLIPTokenizer,
# 先验调度器
prior_scheduler: UnCLIPScheduler,
# 图像处理器
prior_image_processor: CLIPImageProcessor,
):
# 调用父类的初始化方法
super().__init__()
# 注册各个模块，包括 UNet、调度器等
self.register_modules(
unet=unet,
scheduler=scheduler,
movq=movq,
prior_prior=prior_prior,
prior_image_encoder=prior_image_encoder,
prior_text_encoder=prior_text_encoder,
prior_tokenizer=prior_tokenizer,
prior_scheduler=prior_scheduler,
prior_image_processor=prior_image_processor,
)
# 创建 KandinskyV22PriorPipeline 实例，传入必要的组件
self.prior_pipe = KandinskyV22PriorPipeline(
prior=prior_prior,
image_encoder=prior_image_encoder,
text_encoder=prior_text_encoder,
tokenizer=prior_tokenizer,
scheduler=prior_scheduler,
image_processor=prior_image_processor,
)
# 创建 KandinskyV22InpaintPipeline 实例，传入 UNet 和调度器
self.decoder_pipe = KandinskyV22InpaintPipeline(
unet=unet,
scheduler=scheduler,
movq=movq,
)
# 启用 xformers 的内存高效注意力机制
def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
# 调用解码管道的方法启用内存高效注意力
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
# 启用顺序 CPU 卸载，减少内存使用
def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"):
r"""
使用 accelerate 将所有模型卸载到 CPU，显著减少内存使用。调用时，unet、
text_encoder、vae 和安全检查器的状态字典被保存到 CPU，然后移动到
`torch.device('meta')，仅在特定子模块调用其 `forward` 方法时加载到 GPU。
注意，卸载是基于子模块的。内存节省大于 `enable_model_cpu_offload`，但性能较低。
"""
# 在优先管道上启用顺序 CPU 卸载
self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device)
# 在解码管道上启用顺序 CPU 卸载
self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device)
# 显示进度条
def progress_bar(self, iterable=None, total=None):
# 在优先管道上显示进度条
self.prior_pipe.progress_bar(iterable=iterable, total=total)
# 在解码管道上显示进度条
self.decoder_pipe.progress_bar(iterable=iterable, total=total)
# 启用解码管道的模型 CPU 卸载
self.decoder_pipe.enable_model_cpu_offload()
# 设置进度条配置
def set_progress_bar_config(self, **kwargs):
# 在优先管道上设置进度条配置
self.prior_pipe.set_progress_bar_config(**kwargs)
# 在解码管道上设置进度条配置
self.decoder_pipe.set_progress_bar_config(**kwargs)
# 使用装饰器以禁用梯度计算
@torch.no_grad()
# 替换示例文档字符串
@replace_example_docstring(INPAINT_EXAMPLE_DOC_STRING)
# 定义可调用对象的 __call__ 方法，允许实例像函数一样被调用
def __call__(
# 输入提示，可以是字符串或字符串列表
self,
prompt: Union[str, List[str]],
# 输入图像，可以是张量、PIL 图像或它们的列表
image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
# 掩膜图像，可以是张量、PIL 图像或它们的列表
mask_image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
# 可选的负提示，可以是字符串或字符串列表
negative_prompt: Optional[Union[str, List[str]]] = None,
# 推理步骤的数量，默认为 100
num_inference_steps: int = 100,
# 指导比例，默认为 4.0
guidance_scale: float = 4.0,
# 每个提示生成的图像数量，默认为 1
num_images_per_prompt: int = 1,
# 输出图像的高度，默认为 512
height: int = 512,
# 输出图像的宽度，默认为 512
width: int = 512,
# 先验指导比例，默认为 4.0
prior_guidance_scale: float = 4.0,
# 先验推理步骤的数量，默认为 25
prior_num_inference_steps: int = 25,
# 随机数生成器，可选，可以是张量或生成器的列表
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 可选的潜在变量，默认为 None
latents: Optional[torch.Tensor] = None,
# 输出类型，默认为 "pil"
output_type: Optional[str] = "pil",
# 是否返回字典格式的结果，默认为 True
return_dict: bool = True,
# 先验回调函数，可选，接收步骤和状态的回调
prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 先验回调函数的输入张量名称列表，默认为 ["latents"]
prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# 回调函数，可选，接收步骤和状态的回调
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 回调函数的输入张量名称列表，默认为 ["latents"]
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# 其他可选参数
**kwargs,
