diffusers 源码解析（三十一）
.\diffusers\pipelines\free_init_utils.py
# 版权声明，指定该文件由 HuggingFace 团队版权所有，所有权利保留
# 根据 Apache 许可证 2.0 版进行授权；用户需遵循许可条款
# 提供许可证的获取地址
#
# 除非适用的法律要求或书面同意，否则该软件以 "原样" 基础提供，不提供任何形式的担保或条件
# 详见许可证中有关权限和限制的条款
# 导入数学模块
import math
# 从 typing 模块导入 Tuple 和 Union 类型
from typing import Tuple, Union
# 导入 PyTorch 库
import torch
# 导入 PyTorch 的 FFT 模块
import torch.fft as fft
# 从 utils.torch_utils 导入 randn_tensor 函数
from ..utils.torch_utils import randn_tensor
# 定义一个混入类 FreeInitMixin
class FreeInitMixin:
r"""FreeInit 的混入类."""
# 启用 FreeInit 机制的方法
def enable_free_init(
self,
num_iters: int = 3, # 默认的 FreeInit 噪声重新初始化迭代次数
use_fast_sampling: bool = False, # 是否使用快速采样，默认为 False
method: str = "butterworth", # 选择过滤方法，默认为 butterworth
order: int = 4, # butterworth 方法的过滤器阶数，默认值为 4
spatial_stop_frequency: float = 0.25, # 空间维度的归一化截止频率，默认值为 0.25
temporal_stop_frequency: float = 0.25, # 时间维度的归一化截止频率，默认值为 0.25
):
"""启用 FreeInit 机制，参考文献为 https://arxiv.org/abs/2312.07537.
此实现已根据 [官方仓库](https://github.com/TianxingWu/FreeInit) 进行了调整.
参数:
num_iters (`int`, *可选*, 默认值为 `3`):
FreeInit 噪声重新初始化的迭代次数.
use_fast_sampling (`bool`, *可选*, 默认值为 `False`):
是否以牺牲质量来加速采样过程，如果设置为 `True`，启用文中提到的 "粗到细采样" 策略.
method (`str`, *可选*, 默认值为 `butterworth`):
用于 FreeInit 低通滤波器的过滤方法，必须为 `butterworth`、`ideal` 或 `gaussian` 之一.
order (`int`, *可选*, 默认值为 `4`):
在 `butterworth` 方法中使用的滤波器阶数，较大的值会导致 `ideal` 方法行为，而较小的值会导致 `gaussian` 方法行为.
spatial_stop_frequency (`float`, *可选*, 默认值为 `0.25`):
空间维度的归一化截止频率，值必须在 0 到 1 之间，原实现中称为 `d_s`.
temporal_stop_frequency (`float`, *可选*, 默认值为 `0.25`):
时间维度的归一化截止频率，值必须在 0 到 1 之间，原实现中称为 `d_t`.
"""
# 设置 FreeInit 迭代次数
self._free_init_num_iters = num_iters
# 设置是否使用快速采样
self._free_init_use_fast_sampling = use_fast_sampling
# 设置过滤方法
self._free_init_method = method
# 设置过滤器阶数
self._free_init_order = order
# 设置空间截止频率
self._free_init_spatial_stop_frequency = spatial_stop_frequency
# 设置时间截止频率
self._free_init_temporal_stop_frequency = temporal_stop_frequency
# 禁用 FreeInit 机制（如果已启用）
def disable_free_init(self):
"""Disables the FreeInit mechanism if enabled."""
# 将 FreeInit 迭代次数设置为 None，表示禁用
self._free_init_num_iters = None
@property
# 属性，检查 FreeInit 是否启用
def free_init_enabled(self):
# 返回是否存在 FreeInit 迭代次数且不为 None
return hasattr(self, "_free_init_num_iters") and self._free_init_num_iters is not None
# 获取 FreeInit 频率滤波器
def _get_free_init_freq_filter(
self,
shape: Tuple[int, ...], # 输入形状，包含时间、高度、宽度
device: Union[str, torch.dtype], # 设备类型
filter_type: str, # 滤波器类型
order: float, # 滤波器阶数
spatial_stop_frequency: float, # 空间停止频率
temporal_stop_frequency: float, # 时间停止频率
) -> torch.Tensor:
r"""Returns the FreeInit filter based on filter type and other input conditions."""
# 提取时间、高度和宽度维度
time, height, width = shape[-3], shape[-2], shape[-1]
# 初始化全零的掩码张量
mask = torch.zeros(shape)
# 如果空间或时间停止频率为零，返回全零掩码
if spatial_stop_frequency == 0 or temporal_stop_frequency == 0:
return mask
# 根据不同滤波器类型定义掩码函数
if filter_type == "butterworth":
# Butterworth 滤波器的掩码函数
def retrieve_mask(x):
return 1 / (1 + (x / spatial_stop_frequency**2) ** order)
elif filter_type == "gaussian":
# Gaussian 滤波器的掩码函数
def retrieve_mask(x):
return math.exp(-1 / (2 * spatial_stop_frequency**2) * x)
elif filter_type == "ideal":
# 理想滤波器的掩码函数
def retrieve_mask(x):
return 1 if x <= spatial_stop_frequency * 2 else 0
else:
# 如果滤波器类型未实现，抛出异常
raise NotImplementedError("`filter_type` must be one of gaussian, butterworth or ideal")
# 遍历时间、高度和宽度，计算掩码值
for t in range(time):
for h in range(height):
for w in range(width):
# 计算距离平方，用于掩码函数
d_square = (
((spatial_stop_frequency / temporal_stop_frequency) * (2 * t / time - 1)) ** 2
+ (2 * h / height - 1) ** 2
+ (2 * w / width - 1) ** 2
)
# 根据距离平方更新掩码值
mask[..., t, h, w] = retrieve_mask(d_square)
# 将掩码张量转移到指定设备
return mask.to(device)
# 应用频率滤波器
def _apply_freq_filter(self, x: torch.Tensor, noise: torch.Tensor, low_pass_filter: torch.Tensor) -> torch.Tensor:
r"""Noise reinitialization."""
# 对输入进行快速傅里叶变换（FFT）
x_freq = fft.fftn(x, dim=(-3, -2, -1))
# 将频谱中心移到频谱的中心
x_freq = fft.fftshift(x_freq, dim=(-3, -2, -1))
# 对噪声进行快速傅里叶变换（FFT）
noise_freq = fft.fftn(noise, dim=(-3, -2, -1))
# 将噪声频谱中心移到频谱的中心
noise_freq = fft.fftshift(noise_freq, dim=(-3, -2, -1))
# 频率混合操作
high_pass_filter = 1 - low_pass_filter # 计算高通滤波器
x_freq_low = x_freq * low_pass_filter # 低通滤波器作用于输入
noise_freq_high = noise_freq * high_pass_filter # 高通滤波器作用于噪声
# 在频域中混合
x_freq_mixed = x_freq_low + noise_freq_high
# 逆快速傅里叶变换（IFFT）
x_freq_mixed = fft.ifftshift(x_freq_mixed, dim=(-3, -2, -1)) # 将混合频谱中心移回
x_mixed = fft.ifftn(x_freq_mixed, dim=(-3, -2, -1)).real # 还原到时域并取实部
# 返回混合后的结果
return x_mixed
# 应用 FreeInit 机制
def _apply_free_init(
self,
latents: torch.Tensor, # 输入的潜变量
free_init_iteration: int, # FreeInit 迭代次数
num_inference_steps: int, # 推理步骤数量
device: torch.device, # 设备类型
dtype: torch.dtype, # 数据类型
generator: torch.Generator, # 随机数生成器
# 方法体开始
):
# 如果是第一次初始化
if free_init_iteration == 0:
# 克隆初始噪声，保存在属性中
self._free_init_initial_noise = latents.detach().clone()
else:
# 获取当前潜在变量的形状
latent_shape = latents.shape
# 定义过滤器的形状，保留第一维
free_init_filter_shape = (1, *latent_shape[1:])
# 获取自由初始化频率过滤器
free_init_freq_filter = self._get_free_init_freq_filter(
shape=free_init_filter_shape, # 过滤器的形状
device=device, # 设备
filter_type=self._free_init_method, # 过滤器类型
order=self._free_init_order, # 过滤器阶数
spatial_stop_frequency=self._free_init_spatial_stop_frequency, # 空间停止频率
temporal_stop_frequency=self._free_init_temporal_stop_frequency, # 时间停止频率
)
# 获取当前扩散时间步
current_diffuse_timestep = self.scheduler.config.num_train_timesteps - 1
# 创建与潜在变量数量相同的扩散时间步张量
diffuse_timesteps = torch.full((latent_shape[0],), current_diffuse_timestep).long()
# 向潜在变量添加噪声
z_t = self.scheduler.add_noise(
original_samples=latents, # 原始潜在样本
noise=self._free_init_initial_noise, # 添加的噪声
timesteps=diffuse_timesteps.to(device) # 转移到设备
).to(dtype=torch.float32) # 转换数据类型
# 创建随机张量
z_rand = randn_tensor(
shape=latent_shape, # 随机张量的形状
generator=generator, # 随机数生成器
device=device, # 设备
dtype=torch.float32, # 数据类型
)
# 应用频率过滤器于潜在变量
latents = self._apply_freq_filter(z_t, z_rand, low_pass_filter=free_init_freq_filter)
# 转换潜在变量数据类型
latents = latents.to(dtype)
# 进行粗到细采样以加速推理（可能导致质量降低）
if self._free_init_use_fast_sampling:
# 计算推理步骤的数量
num_inference_steps = max(
1, int(num_inference_steps / self._free_init_num_iters * (free_init_iteration + 1)) # 逐步减少推理步骤
)
# 如果推理步骤大于0
if num_inference_steps > 0:
# 设置调度器的时间步
self.scheduler.set_timesteps(num_inference_steps, device=device)
# 返回潜在变量和调度器的时间步
return latents, self.scheduler.timesteps
.\diffusers\pipelines\free_noise_utils.py
# 版权所有 2024 HuggingFace 团队，保留所有权利。
#
# 根据 Apache 许可证，版本 2.0（“许可证”）授权；
# 除非遵守许可证，否则不得使用此文件。
# 可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面协议另有约定，
# 否则根据许可证分发的软件按“原样”提供，
# 不提供任何形式的保证或条件，无论是明示或暗示的。
# 有关许可证下的特定语言及其权限和限制，请参阅许可证。
from typing import Optional, Union # 导入可选和联合类型，用于类型注解
import torch # 导入 PyTorch 库，以便进行深度学习操作
from ..models.attention import BasicTransformerBlock, FreeNoiseTransformerBlock # 从上级目录导入注意力模型的基础和自由噪声变换器块
from ..models.unets.unet_motion_model import ( # 从上级目录导入 UNet 动作模型的相关模块
CrossAttnDownBlockMotion, # 导入交叉注意力下采样块
DownBlockMotion, # 导入下采样块
UpBlockMotion, # 导入上采样块
)
from ..utils import logging # 从上级目录导入日志工具
from ..utils.torch_utils import randn_tensor # 从上级目录导入生成随机张量的工具
logger = logging.get_logger(__name__) # 获取当前模块的日志记录器，便于记录日志信息
class AnimateDiffFreeNoiseMixin: # 定义一个混合类，用于实现自由噪声相关功能
r"""混合类用于 [FreeNoise](https://arxiv.org/abs/2310.15169) 的实现。""" # 文档字符串，提供该类的用途说明
# 启用变换器块中的 FreeNoise 功能的辅助函数
def _enable_free_noise_in_block(self, block: Union[CrossAttnDownBlockMotion, DownBlockMotion, UpBlockMotion]):
# 文档字符串，说明该函数的目的
for motion_module in block.motion_modules:
# 遍历每个运动模块
num_transformer_blocks = len(motion_module.transformer_blocks)
# 获取当前运动模块中变换器块的数量
for i in range(num_transformer_blocks):
# 遍历每个变换器块
if isinstance(motion_module.transformer_blocks[i], FreeNoiseTransformerBlock):
# 检查当前块是否为 FreeNoise 变换器块
motion_module.transformer_blocks[i].set_free_noise_properties(
# 设置 FreeNoise 属性
self._free_noise_context_length,
self._free_noise_context_stride,
self._free_noise_weighting_scheme,
)
else:
# 确保当前块是基本变换器块
assert isinstance(motion_module.transformer_blocks[i], BasicTransformerBlock)
basic_transfomer_block = motion_module.transformer_blocks[i]
# 获取基本变换器块的引用
motion_module.transformer_blocks[i] = FreeNoiseTransformerBlock(
# 创建新的 FreeNoise 变换器块，复制基本块的参数
dim=basic_transfomer_block.dim,
num_attention_heads=basic_transfomer_block.num_attention_heads,
attention_head_dim=basic_transfomer_block.attention_head_dim,
dropout=basic_transfomer_block.dropout,
cross_attention_dim=basic_transfomer_block.cross_attention_dim,
activation_fn=basic_transfomer_block.activation_fn,
attention_bias=basic_transfomer_block.attention_bias,
only_cross_attention=basic_transfomer_block.only_cross_attention,
double_self_attention=basic_transfomer_block.double_self_attention,
positional_embeddings=basic_transfomer_block.positional_embeddings,
num_positional_embeddings=basic_transfomer_block.num_positional_embeddings,
context_length=self._free_noise_context_length,
context_stride=self._free_noise_context_stride,
weighting_scheme=self._free_noise_weighting_scheme,
).to(device=self.device, dtype=self.dtype)
# 将新创建的块赋值到当前变换器块的位置，并设置设备和数据类型
motion_module.transformer_blocks[i].load_state_dict(
# 加载基本变换器块的状态字典到新的块
basic_transfomer_block.state_dict(), strict=True
)
# 定义一个辅助函数，用于禁用变换器块中的 FreeNoise
def _disable_free_noise_in_block(self, block: Union[CrossAttnDownBlockMotion, DownBlockMotion, UpBlockMotion]):
r"""辅助函数，用于禁用变换器块中的 FreeNoise。"""
# 遍历给定块中的所有运动模块
for motion_module in block.motion_modules:
# 计算当前运动模块中的变换器块数量
num_transformer_blocks = len(motion_module.transformer_blocks)
# 遍历每个变换器块
for i in range(num_transformer_blocks):
# 检查当前变换器块是否为 FreeNoiseTransformerBlock 类型
if isinstance(motion_module.transformer_blocks[i], FreeNoiseTransformerBlock):
# 获取当前的 FreeNoise 变换器块
free_noise_transfomer_block = motion_module.transformer_blocks[i]
# 用 BasicTransformerBlock 替换 FreeNoise 变换器块，保持相应参数
motion_module.transformer_blocks[i] = BasicTransformerBlock(
dim=free_noise_transfomer_block.dim,
num_attention_heads=free_noise_transfomer_block.num_attention_heads,
attention_head_dim=free_noise_transfomer_block.attention_head_dim,
dropout=free_noise_transfomer_block.dropout,
cross_attention_dim=free_noise_transfomer_block.cross_attention_dim,
activation_fn=free_noise_transfomer_block.activation_fn,
attention_bias=free_noise_transfomer_block.attention_bias,
only_cross_attention=free_noise_transfomer_block.only_cross_attention,
double_self_attention=free_noise_transfomer_block.double_self_attention,
positional_embeddings=free_noise_transfomer_block.positional_embeddings,
num_positional_embeddings=free_noise_transfomer_block.num_positional_embeddings,
).to(device=self.device, dtype=self.dtype)
# 加载 FreeNoise 变换器块的状态字典到新的 BasicTransformerBlock
motion_module.transformer_blocks[i].load_state_dict(
free_noise_transfomer_block.state_dict(), strict=True
)
# 定义准备自由噪声潜在变量的函数
def _prepare_latents_free_noise(
self,
batch_size: int,
num_channels_latents: int,
num_frames: int,
height: int,
width: int,
dtype: torch.dtype,
device: torch.device,
generator: Optional[torch.Generator] = None,
latents: Optional[torch.Tensor] = None,
):
# 此处省略函数的具体实现
# 定义启用自由噪声的函数
def enable_free_noise(
self,
context_length: Optional[int] = 16,
context_stride: int = 4,
weighting_scheme: str = "pyramid",
noise_type: str = "shuffle_context",
):
# 此处省略函数的具体实现
) -> None:
r"""
# 启用使用 FreeNoise 生成长视频的功能
Args:
context_length (`int`, defaults to `16`, *optional*):
# 一次处理的视频帧数量。推荐设置为运动适配器训练时的最大帧数（通常为 16/24/32）。如果为 `None`，将使用运动适配器配置中的默认值。
context_stride (`int`, *optional*):
# 通过处理多个帧生成长视频。FreeNoise 使用滑动窗口处理这些帧，窗口大小为 `context_length`。上下文步幅允许指定每个窗口之间跳过多少帧。例如，`context_length` 为 16，`context_stride` 为 4 时将处理 24 帧为：[0, 15], [4, 19], [8, 23]（基于 0 的索引）。
weighting_scheme (`str`, defaults to `pyramid`):
# 加权方案，用于在 FreeNoise 块中累积后平均潜在变量。目前支持以下加权方案：
- "pyramid"
# 使用类似金字塔的权重模式进行加权平均：[1, 2, 3, 2, 1]。
noise_type (`str`, defaults to "shuffle_context"):
# TODO
"""
# 定义允许的加权方案列表
allowed_weighting_scheme = ["pyramid"]
# 定义允许的噪声类型列表
allowed_noise_type = ["shuffle_context", "repeat_context", "random"]
# 检查上下文长度是否超过最大序列长度
if context_length > self.motion_adapter.config.motion_max_seq_length:
logger.warning(
# 记录警告，表示上下文长度设置过大可能导致生成结果不佳
f"You have set {context_length=} which is greater than {self.motion_adapter.config.motion_max_seq_length=}. This can lead to bad generation results."
)
# 验证加权方案是否在允许的选项中
if weighting_scheme not in allowed_weighting_scheme:
raise ValueError(
# 抛出错误，表示加权方案不合法
f"The parameter `weighting_scheme` must be one of {allowed_weighting_scheme}, but got {weighting_scheme=}"
)
# 验证噪声类型是否在允许的选项中
if noise_type not in allowed_noise_type:
raise ValueError(f"The parameter `noise_type` must be one of {allowed_noise_type}, but got {noise_type=}")
# 设置 FreeNoise 的上下文长度，使用提供值或最大序列长度的默认值
self._free_noise_context_length = context_length or self.motion_adapter.config.motion_max_seq_length
# 设置 FreeNoise 的上下文步幅
self._free_noise_context_stride = context_stride
# 设置 FreeNoise 的加权方案
self._free_noise_weighting_scheme = weighting_scheme
# 设置 FreeNoise 的噪声类型
self._free_noise_noise_type = noise_type
# 获取 UNet 的所有块，准备启用 FreeNoise
blocks = [*self.unet.down_blocks, self.unet.mid_block, *self.unet.up_blocks]
# 对每个块启用 FreeNoise 功能
for block in blocks:
self._enable_free_noise_in_block(block)
# 禁用 FreeNoise 功能的方法
def disable_free_noise(self) -> None:
# 将上下文长度设置为 None，表示禁用
self._free_noise_context_length = None
# 获取 UNet 的所有块，准备禁用 FreeNoise
blocks = [*self.unet.down_blocks, self.unet.mid_block, *self.unet.up_blocks]
# 对每个块禁用 FreeNoise 功能
for block in blocks:
self._disable_free_noise_in_block(block)
# 属性装饰器
@property
# 检查是否启用了自由噪声功能
def free_noise_enabled(self):
# 检查对象是否具有属性"_free_noise_context_length"且其值不为None
return hasattr(self, "_free_noise_context_length") and self._free_noise_context_length is not None
.\diffusers\pipelines\hunyuandit\pipeline_hunyuandit.py
# 版权声明，标明文件归属和使用许可信息
# Copyright 2024 HunyuanDiT Authors and The HuggingFace Team. All rights reserved.
#
# 根据 Apache License 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.
# 导入 inspect 模块以便于代码检查和获取信息
import inspect
# 从 typing 模块导入所需的类型注解
from typing import Callable, Dict, List, Optional, Tuple, Union
# 导入 NumPy 库以进行数值计算
import numpy as np
# 导入 PyTorch 库以进行深度学习模型的构建和训练
import torch
# 从 transformers 库导入所需的模型和分词器
from transformers import BertModel, BertTokenizer, CLIPImageProcessor, MT5Tokenizer, T5EncoderModel
# 从 diffusers 库导入稳定扩散管道输出类型
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
# 从回调模块导入多管道回调和单管道回调
from ...callbacks import MultiPipelineCallbacks, PipelineCallback
# 从图像处理模块导入 VAE 图像处理器
from ...image_processor import VaeImageProcessor
# 从模型模块导入自编码器和 HunyuanDiT2D 模型
from ...models import AutoencoderKL, HunyuanDiT2DModel
# 从嵌入模块导入二维旋转位置嵌入的获取函数
from ...models.embeddings import get_2d_rotary_pos_embed
# 从安全检查器模块导入稳定扩散安全检查器
from ...pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
# 从调度器模块导入 DDPMScheduler
from ...schedulers import DDPMScheduler
# 从工具模块导入各种工具函数
from ...utils import (
is_torch_xla_available,
logging,
replace_example_docstring,
)
# 从 PyTorch 工具模块导入随机张量生成函数
from ...utils.torch_utils import randn_tensor
# 从管道工具模块导入 DiffusionPipeline
from ..pipeline_utils import DiffusionPipeline
# 检查是否可用 Torch XLA 库以支持 TPU 加速
if is_torch_xla_available():
# 导入 XLA 模块以进行 TPU 操作
import torch_xla.core.xla_model as xm
# 标记 XLA 可用性为 True
XLA_AVAILABLE = True
else:
# 标记 XLA 可用性为 False
XLA_AVAILABLE = False
# 获取日志记录器以进行调试和日志记录
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，展示如何使用 HunyuanDiTPipeline
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> import torch
>>> from diffusers import HunyuanDiTPipeline
>>> pipe = HunyuanDiTPipeline.from_pretrained(
... "Tencent-Hunyuan/HunyuanDiT-Diffusers", torch_dtype=torch.float16
... )
>>> pipe.to("cuda")
>>> # 也可以使用英语提示，因为 HunyuanDiT 支持英语和中文
>>> # prompt = "An astronaut riding a horse"
>>> prompt = "一个宇航员在骑马"
>>> image = pipe(prompt).images[0]
```py
"""
# 定义标准宽高比数组
STANDARD_RATIO = np.array(
[
1.0, # 1:1
4.0 / 3.0, # 4:3
3.0 / 4.0, # 3:4
16.0 / 9.0, # 16:9
9.0 / 16.0, # 9:16
]
)
# 定义标准形状的列表，包括不同宽高比的分辨率
STANDARD_SHAPE = [
[(1024, 1024), (1280, 1280)], # 1:1
[(1024, 768), (1152, 864), (1280, 960)], # 4:3
[(768, 1024), (864, 1152), (960, 1280)], # 3:4
[(1280, 768)], # 16:9
[(768, 1280)], # 9:16
]
# 根据标准形状计算面积
STANDARD_AREA = [np.array([w * h for w, h in shapes]) for shapes in STANDARD_SHAPE]
# 定义支持的形状，列出可能的输入分辨率
SUPPORTED_SHAPE = [
(1024, 1024),
(1280, 1280), # 1:1
(1024, 768),
(1152, 864),
(1280, 960), # 4:3
(768, 1024),
(864, 1152),
(960, 1280), # 3:4
(1280, 768), # 16:9
(768, 1280), # 9:16
]
# 定义一个函数，将目标宽高映射到标准形状
def map_to_standard_shapes(target_width, target_height):
# 计算目标宽高比
target_ratio = target_width / target_height
# 找到与目标宽高比最接近的标准宽高比的索引
closest_ratio_idx = np.argmin(np.abs(STANDARD_RATIO - target_ratio))
# 找到与目标面积最接近的标准面积的索引
closest_area_idx = np.argmin(np.abs(STANDARD_AREA[closest_ratio_idx] - target_width * target_height))
# 根据最接近的宽高比和面积索引获取标准形状的宽度和高度
width, height = STANDARD_SHAPE[closest_ratio_idx][closest_area_idx]
# 返回找到的宽度和高度
return width, height
# 定义获取网格的调整大小和裁剪区域的函数，接受源图像尺寸和目标尺寸
def get_resize_crop_region_for_grid(src, tgt_size):
# 将目标尺寸赋值给高度和宽度变量
th = tw = tgt_size
# 解构源图像尺寸，获取高度和宽度
h, w = src
# 计算宽高比
r = h / w
# 根据宽高比决定调整大小的方式
# 如果高度大于宽度
if r > 1:
# 将目标高度设为目标尺寸
resize_height = th
# 根据目标高度计算目标宽度
resize_width = int(round(th / h * w))
else:
# 否则将目标宽度设为目标尺寸
resize_width = tw
# 根据目标宽度计算目标高度
resize_height = int(round(tw / w * h))
# 计算裁剪区域的上边界
crop_top = int(round((th - resize_height) / 2.0))
# 计算裁剪区域的左边界
crop_left = int(round((tw - resize_width) / 2.0))
# 返回裁剪区域的坐标和调整后的尺寸
return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width)
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg 复制的函数
def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
"""
根据 `guidance_rescale` 重新调整 `noise_cfg` 的比例。基于论文 [Common Diffusion Noise Schedules and
Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf) 的发现。见第 3.4 节
"""
# 计算噪声预测文本的标准差
std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
# 计算噪声配置的标准差
std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
# 通过标准差调整噪声预测以修正过曝问题
noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
# 根据引导重缩放原始结果，以避免图像过于单调
noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
# 返回调整后的噪声配置
return noise_cfg
# 定义 HunyuanDiT 管道类，继承自 DiffusionPipeline
class HunyuanDiTPipeline(DiffusionPipeline):
r"""
使用 HunyuanDiT 进行英语/中文到图像生成的管道。
此模型继承自 [`DiffusionPipeline`]。请查看超类文档以获取库为所有管道实现的通用方法（例如下载或保存，运行在特定设备上等）。
HunyuanDiT 使用两个文本编码器：[mT5](https://huggingface.co/google/mt5-base) 和 [双语 CLIP](fine-tuned by
ourselves)
参数：
vae ([`AutoencoderKL`]):
变分自编码器 (VAE) 模型，用于将图像编码和解码为潜在表示。我们使用 `sdxl-vae-fp16-fix`。
text_encoder (Optional[`~transformers.BertModel`, `~transformers.CLIPTextModel`]):
冻结的文本编码器 ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14))。
HunyuanDiT 使用经过微调的 [双语 CLIP]。
tokenizer (Optional[`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer`]):
用于文本标记化的 `BertTokenizer` 或 `CLIPTokenizer`。
transformer ([`HunyuanDiT2DModel`]):
腾讯 Hunyuan 设计的 HunyuanDiT 模型。
text_encoder_2 (`T5EncoderModel`):
mT5 嵌入器。具体为 't5-v1_1-xxl'。
tokenizer_2 (`MT5Tokenizer`):
mT5 嵌入器的标记器。
scheduler ([`DDPMScheduler`]):
用于与 HunyuanDiT 结合使用的调度器，以去噪编码的图像潜在。
"""
# 定义模型的 CPU 卸载序列，指定组件的调用顺序
model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
# 可选组件的列表，包含模型中可能使用的其他模块
_optional_components = [
"safety_checker", # 安全检查器
"feature_extractor", # 特征提取器
"text_encoder_2", # 第二个文本编码器
"tokenizer_2", # 第二个分词器
"text_encoder", # 第一个文本编码器
"tokenizer", # 第一个分词器
]
# 指定不参与 CPU 卸载的组件
_exclude_from_cpu_offload = ["safety_checker"] # 安全检查器不参与 CPU 卸载
# 回调张量输入列表，包含输入模型的张量名称
_callback_tensor_inputs = [
"latents", # 潜在变量
"prompt_embeds", # 提示嵌入
"negative_prompt_embeds", # 负提示嵌入
"prompt_embeds_2", # 第二个提示嵌入
"negative_prompt_embeds_2", # 第二个负提示嵌入
]
# 初始化方法，设置模型的各个组件
def __init__(
self,
vae: AutoencoderKL, # 变分自编码器
text_encoder: BertModel, # 文本编码器，使用 BERT 模型
tokenizer: BertTokenizer, # 分词器，使用 BERT 分词器
transformer: HunyuanDiT2DModel, # 转换器模型
scheduler: DDPMScheduler, # 调度器
safety_checker: StableDiffusionSafetyChecker, # 稳定扩散安全检查器
feature_extractor: CLIPImageProcessor, # 特征提取器，使用 CLIP 图像处理器
requires_safety_checker: bool = True, # 是否需要安全检查器的标志
text_encoder_2=T5EncoderModel, # 第二个文本编码器，默认使用 T5 模型
tokenizer_2=MT5Tokenizer, # 第二个分词器，默认使用 MT5 分词器
):
# 调用父类的初始化方法
super().__init__()
# 注册模型的各个组件
self.register_modules(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
tokenizer_2=tokenizer_2,
transformer=transformer,
scheduler=scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
text_encoder_2=text_encoder_2,
)
# 检查安全检查器是否为 None，并根据需要发出警告
if safety_checker is None and requires_safety_checker:
logger.warning(
f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
" that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
" results in services or applications open to the public. Both the diffusers team and Hugging Face"
" strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
" it only for use-cases that involve analyzing network behavior or auditing its results. For more"
" information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
)
# 检查安全检查器和特征提取器的配置，确保正确设置
if safety_checker is not None and feature_extractor is None:
raise ValueError(
"Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
" checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
)
# 根据 VAE 的配置计算缩放因子
self.vae_scale_factor = (
2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
)
# 初始化图像处理器
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
# 注册配置，保存是否需要安全检查器的标志
self.register_to_config(requires_safety_checker=requires_safety_checker)
# 根据 transformer 的配置确定默认样本大小
self.default_sample_size = (
self.transformer.config.sample_size
if hasattr(self, "transformer") and self.transformer is not None
else 128 # 默认样本大小为 128
)
# 定义一个编码提示的函数，包含多个参数配置
def encode_prompt(
self,
# 输入的提示文本
prompt: str,
# 设备类型，可选
device: torch.device = None,
# 数据类型，可选
dtype: torch.dtype = None,
# 每个提示生成的图像数量，默认为1
num_images_per_prompt: int = 1,
# 是否进行分类器自由引导，默认为True
do_classifier_free_guidance: bool = True,
# 负提示文本，默认为None
negative_prompt: Optional[str] = None,
# 提示的嵌入表示，默认为None
prompt_embeds: Optional[torch.Tensor] = None,
# 负提示的嵌入表示，默认为None
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 提示的注意力掩码，默认为None
prompt_attention_mask: Optional[torch.Tensor] = None,
# 负提示的注意力掩码，默认为None
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
# 最大序列长度，默认为None
max_sequence_length: Optional[int] = None,
# 文本编码器索引，默认为0
text_encoder_index: int = 0,
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker 复制的函数
def run_safety_checker(self, image, device, dtype):
# 检查安全检查器是否存在
if self.safety_checker is None:
# 如果不存在，设置nsfw概念为None
has_nsfw_concept = None
else:
# 如果输入是张量，则进行后处理
if torch.is_tensor(image):
feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
else:
# 如果不是张量，则转换为PIL格式
feature_extractor_input = self.image_processor.numpy_to_pil(image)
# 获取安全检查器的输入并移动到指定设备
safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
# 运行安全检查器，检查图像的nsfw概念
image, has_nsfw_concept = self.safety_checker(
images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
)
# 返回处理后的图像和nsfw概念
return image, has_nsfw_concept
# 从 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 = {}
if accepts_eta:
# 如果接受eta，则将其添加到额外参数中
extra_step_kwargs["eta"] = eta
# 检查调度器是否接受generator参数
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
if accepts_generator:
# 如果接受generator，则将其添加到额外参数中
extra_step_kwargs["generator"] = generator
# 返回额外的关键字参数字典
return extra_step_kwargs
# 定义一个检查输入的函数，包含多个参数
def check_inputs(
self,
# 输入的提示文本
prompt,
# 图像高度
height,
# 图像宽度
width,
# 负提示文本，默认为None
negative_prompt=None,
# 提示的嵌入表示，默认为None
prompt_embeds=None,
# 负提示的嵌入表示，默认为None
negative_prompt_embeds=None,
# 提示的注意力掩码，默认为None
prompt_attention_mask=None,
# 负提示的注意力掩码，默认为None
negative_prompt_attention_mask=None,
# 第二组提示的嵌入表示，默认为None
prompt_embeds_2=None,
# 第二组负提示的嵌入表示，默认为None
negative_prompt_embeds_2=None,
# 第二组提示的注意力掩码，默认为None
prompt_attention_mask_2=None,
# 第二组负提示的注意力掩码，默认为None
negative_prompt_attention_mask_2=None,
# 用于步骤结束回调的张量输入，默认为None
callback_on_step_end_tensor_inputs=None,
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents 复制的函数
# 准备潜在变量，定义输入形状
def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
# 定义潜在变量的形状，考虑缩放因子
shape = (
batch_size,
num_channels_latents,
int(height) // self.vae_scale_factor,
int(width) // self.vae_scale_factor,
)
# 检查生成器列表长度是否与批量大小匹配
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."
)
# 如果潜在变量为空，则随机生成
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
# 否则将潜在变量移动到指定设备
latents = latents.to(device)
# 将初始噪声按调度器要求的标准差进行缩放
latents = latents * self.scheduler.init_noise_sigma
# 返回处理后的潜在变量
return latents
# 返回指导缩放因子
@property
def guidance_scale(self):
return self._guidance_scale
# 返回指导重缩放因子
@property
def guidance_rescale(self):
return self._guidance_rescale
# 定义无分类器引导的属性，依照Imagen论文中的公式定义
@property
def do_classifier_free_guidance(self):
# 判断指导缩放因子是否大于1，决定是否进行无分类器引导
return self._guidance_scale > 1
# 返回时间步数
@property
def num_timesteps(self):
return self._num_timesteps
# 返回中断状态
@property
def interrupt(self):
return self._interrupt
# 禁用梯度计算，优化内存使用
@torch.no_grad()
# 替换示例文档字符串
@replace_example_docstring(EXAMPLE_DOC_STRING)
# 定义一个可调用的类方法，允许使用不同的参数生成输出
def __call__(
self,
# 提示内容，可以是单个字符串或字符串列表
prompt: Union[str, List[str]] = None,
# 输出图像的高度
height: Optional[int] = None,
# 输出图像的宽度
width: Optional[int] = None,
# 推理步骤的数量，默认是50
num_inference_steps: Optional[int] = 50,
# 引导比例，默认值为5.0
guidance_scale: Optional[float] = 5.0,
# 负提示内容，可以是单个字符串或字符串列表
negative_prompt: Optional[Union[str, List[str]]] = None,
# 每个提示生成的图像数量，默认是1
num_images_per_prompt: Optional[int] = 1,
# 影响生成随机性的超参数，默认是0.0
eta: Optional[float] = 0.0,
# 随机数生成器，可以是单个或多个生成器
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 潜在空间的张量，表示生成过程中的潜在表示
latents: Optional[torch.Tensor] = None,
# 提示的嵌入表示
prompt_embeds: Optional[torch.Tensor] = None,
# 第二个提示的嵌入表示
prompt_embeds_2: Optional[torch.Tensor] = None,
# 负提示的嵌入表示
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 第二个负提示的嵌入表示
negative_prompt_embeds_2: Optional[torch.Tensor] = None,
# 提示的注意力掩码
prompt_attention_mask: Optional[torch.Tensor] = None,
# 第二个提示的注意力掩码
prompt_attention_mask_2: Optional[torch.Tensor] = None,
# 负提示的注意力掩码
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
# 第二个负提示的注意力掩码
negative_prompt_attention_mask_2: Optional[torch.Tensor] = None,
# 输出类型，默认为"pil"，表示图像格式
output_type: Optional[str] = "pil",
# 是否返回字典格式的输出，默认为True
return_dict: bool = True,
# 生成步骤结束时的回调函数，可以是特定格式的函数
callback_on_step_end: Optional[
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
] = None,
# 指定在步骤结束时的张量输入名称，默认为["latents"]
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# 引导重标定的比例，默认值为0.0
guidance_rescale: float = 0.0,
# 原始图像的尺寸，默认为(1024, 1024)
original_size: Optional[Tuple[int, int]] = (1024, 1024),
# 目标图像的尺寸，默认为None，表示使用原始尺寸
target_size: Optional[Tuple[int, int]] = None,
# 裁剪区域的左上角坐标，默认为(0, 0)
crops_coords_top_left: Tuple[int, int] = (0, 0),
# 是否使用分辨率分箱，默认为True
use_resolution_binning: bool = True,
.\diffusers\pipelines\hunyuandit\__init__.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:
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_hunyuandit"] = ["HunyuanDiTPipeline"]
# 如果正在进行类型检查或慢导入
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
try:
if not (is_transformers_available() and is_torch_available()): # 再次检查依赖
raise OptionalDependencyNotAvailable() # 抛出异常，如果依赖不可用
except OptionalDependencyNotAvailable:
from ...utils.dummy_torch_and_transformers_objects import * # 导入虚拟对象以防止错误
else:
from .pipeline_hunyuandit import HunyuanDiTPipeline # 导入真实的管道实现
else:
import sys # 导入 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\i2vgen_xl\pipeline_i2vgen_xl.py
# 版权所有 2024 Alibaba DAMO-VILAB 和 The HuggingFace 团队。保留所有权利。
#
# 根据 Apache 许可证第 2.0 版（“许可证”）授权；
# 除非遵守许可证，否则您不得使用此文件。
# 您可以在以下网址获取许可证的副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，否则根据许可证分发的软件
# 是按“原样”基础分发，不提供任何形式的担保或条件。
# 有关许可证的特定权限和限制，请参见许可证。
import inspect # 导入 inspect 模块，用于获取对象的信息
from dataclasses import dataclass # 从 dataclasses 导入 dataclass 装饰器
from typing import Any, Dict, List, Optional, Tuple, Union # 导入类型提示所需的各种类型
import numpy as np # 导入 numpy 库，通常用于数值计算
import PIL # 导入 PIL 库，用于图像处理
import torch # 导入 PyTorch 库，用于深度学习
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection # 从 transformers 导入相关的 CLIP 模型和处理器
from ...image_processor import PipelineImageInput, VaeImageProcessor # 从自定义模块导入图像处理类
from ...models import AutoencoderKL # 从自定义模块导入自动编码器模型
from ...models.unets.unet_i2vgen_xl import I2VGenXLUNet # 从自定义模块导入特定的 UNet 模型
from ...schedulers import DDIMScheduler # 从自定义模块导入调度器
from ...utils import ( # 从自定义模块导入各种工具
BaseOutput, # 基础输出类
logging, # 日志记录功能
replace_example_docstring, # 替换示例文档字符串的函数
)
from ...utils.torch_utils import randn_tensor # 从自定义模块导入用于生成随机张量的函数
from ...video_processor import VideoProcessor # 从自定义模块导入视频处理器
from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin # 从自定义模块导入扩散管道和稳定扩散混合类
logger = logging.get_logger(__name__) # 获取当前模块的日志记录器实例，方便记录日志信息
EXAMPLE_DOC_STRING = """ # 定义一个示例文档字符串
Examples: # 示例部分的标题
```py # Python 代码块的开始
>>> import torch # 导入 PyTorch 库
>>> from diffusers import I2VGenXLPipeline # 从 diffusers 模块导入 I2VGenXLPipeline 类
>>> from diffusers.utils import export_to_gif, load_image # 从 utils 模块导入辅助函数
>>> pipeline = I2VGenXLPipeline.from_pretrained( # 从预训练模型创建管道实例
... "ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16" # 指定模型名称和参数
... )
>>> pipeline.enable_model_cpu_offload() # 启用模型的 CPU 卸载功能以节省内存
>>> image_url = ( # 定义图像的 URL
... "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/i2vgen_xl_images/img_0009.png" # 图像的具体 URL
... )
>>> image = load_image(image_url).convert("RGB") # 加载图像并转换为 RGB 格式
>>> prompt = "Papers were floating in the air on a table in the library" # 定义正向提示
>>> negative_prompt = "Distorted, discontinuous, Ugly, blurry, low resolution, motionless, static, disfigured, disconnected limbs, Ugly faces, incomplete arms" # 定义负向提示
>>> generator = torch.manual_seed(8888) # 设置随机种子以确保结果可重现
>>> frames = pipeline( # 调用管道以生成图像帧
... prompt=prompt, # 传递正向提示
... image=image, # 传递输入图像
... num_inference_steps=50, # 设置推理步骤数
... negative_prompt=negative_prompt, # 传递负向提示
... guidance_scale=9.0, # 设置引导比例
... generator=generator, # 传递随机数生成器
... ).frames[0] # 获取生成的第一帧
>>> video_path = export_to_gif(frames, "i2v.gif") # 将生成的帧导出为 GIF 格式的视频
```py # Python 代码块的结束
"""
@dataclass # 使用 dataclass 装饰器定义一个数据类
class I2VGenXLPipelineOutput(BaseOutput): # 定义 I2VGenXLPipelineOutput 类，继承自 BaseOutput
r""" # 文档字符串，描述该类的作用
Output class for image-to-video pipeline. # 说明这是图像到视频管道的输出类
# 函数参数文档字符串，描述函数的参数及其类型
Args:
# frames 参数可以是 torch.Tensor、np.ndarray 或嵌套列表，每个子列表包含去噪的 PIL 图像序列
frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
# 说明 frames 是一个视频输出列表，长度为 batch_size，每个子列表包含 num_frames 长度的去噪图像序列
List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised
# 说明 frames 也可以是形状为 (batch_size, num_frames, channels, height, width) 的 NumPy 数组或 Torch 张量
PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`
"""
# 定义 frames 变量类型，支持多种数据类型的联合
frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]]]
# 定义图像到视频生成的管道类，继承自 DiffusionPipeline 和 StableDiffusionMixin
class I2VGenXLPipeline(
DiffusionPipeline,
StableDiffusionMixin,
):
r"""
用于图像到视频生成的管道，如 [I2VGenXL](https://i2vgen-xl.github.io/) 所提议。
该模型继承自 [`DiffusionPipeline`]。有关所有管道的通用方法（下载、保存、在特定设备上运行等），请查看超类文档。
参数：
vae ([`AutoencoderKL`]):
用于将图像编码和解码为潜在表示的变分自编码器（VAE）模型。
text_encoder ([`CLIPTextModel`]):
冻结的文本编码器（[clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)）。
tokenizer (`CLIPTokenizer`):
用于标记文本的 [`~transformers.CLIPTokenizer`]。
unet ([`I2VGenXLUNet`]):
用于去噪编码视频潜在表示的 [`I2VGenXLUNet`]。
scheduler ([`DDIMScheduler`]):
与 `unet` 结合使用以去噪编码图像潜在表示的调度器。
"""
# 定义模型各组件的 CPU 卸载顺序
model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
def __init__(
self,
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
image_encoder: CLIPVisionModelWithProjection,
feature_extractor: CLIPImageProcessor,
unet: I2VGenXLUNet,
scheduler: DDIMScheduler,
):
# 初始化父类
super().__init__()
# 注册模型的各个模块
self.register_modules(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
image_encoder=image_encoder,
feature_extractor=feature_extractor,
unet=unet,
scheduler=scheduler,
)
# 计算 VAE 的缩放因子
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
# 设置视频处理器，不进行默认调整大小
self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor, do_resize=False)
# 定义属性，获取指导缩放值
@property
def guidance_scale(self):
return self._guidance_scale
# 定义属性，确定是否执行无分类器引导
# 这里的 `guidance_scale` 定义类似于 Imagen 论文中公式（2）的指导权重 `w`：
# `guidance_scale = 1` 对应于不进行分类器自由引导。
@property
def do_classifier_free_guidance(self):
return self._guidance_scale > 1
# 编码提示的函数
def encode_prompt(
self,
prompt,
device,
num_videos_per_prompt,
negative_prompt=None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
clip_skip: Optional[int] = None,
# 定义编码图像的私有方法，接收图像、设备和每个提示的视频数量
def _encode_image(self, image, device, num_videos_per_prompt):
# 获取图像编码器参数的数据类型
dtype = next(self.image_encoder.parameters()).dtype
# 检查输入图像是否为 PyTorch 张量
if not isinstance(image, torch.Tensor):
# 将 PIL 图像转换为 NumPy 数组
image = self.video_processor.pil_to_numpy(image)
# 将 NumPy 数组转换为 PyTorch 张量
image = self.video_processor.numpy_to_pt(image)
# 使用 CLIP 训练统计信息对图像进行归一化处理
image = self.feature_extractor(
images=image,
do_normalize=True, # 是否归一化
do_center_crop=False, # 是否中心裁剪
do_resize=False, # 是否调整大小
do_rescale=False, # 是否重新缩放
return_tensors="pt", # 返回 PyTorch 张量
).pixel_values # 获取处理后的图像像素值
# 将图像移动到指定设备，并转换为指定数据类型
image = image.to(device=device, dtype=dtype)
# 使用图像编码器对图像进行编码，获取图像嵌入
image_embeddings = self.image_encoder(image).image_embeds
# 在第二个维度上添加一个维度
image_embeddings = image_embeddings.unsqueeze(1)
# 为每个提示生成的重复图像嵌入，使用兼容 MPS 的方法
bs_embed, seq_len, _ = image_embeddings.shape # 获取嵌入的批量大小和序列长度
# 重复图像嵌入以适应每个提示的视频数量
image_embeddings = image_embeddings.repeat(1, num_videos_per_prompt, 1)
# 重新调整图像嵌入的形状以合并批量和视频数量
image_embeddings = image_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1)
# 如果启用了无分类器自由引导
if self.do_classifier_free_guidance:
# 创建与图像嵌入相同形状的零张量
negative_image_embeddings = torch.zeros_like(image_embeddings)
# 将负图像嵌入和正图像嵌入拼接
image_embeddings = torch.cat([negative_image_embeddings, image_embeddings])
# 返回最终的图像嵌入
return image_embeddings
# 定义解码潜在空间的公共方法，接收潜在向量和可选的解码块大小
def decode_latents(self, latents, decode_chunk_size=None):
# 使用 VAE 配置的缩放因子对潜在向量进行缩放
latents = 1 / self.vae.config.scaling_factor * latents
# 获取潜在向量的批量大小、通道数、帧数、高度和宽度
batch_size, channels, num_frames, height, width = latents.shape
# 重新排列潜在向量的维度，以适应 VAE 解码器的输入格式
latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width)
# 如果指定了解码块大小
if decode_chunk_size is not None:
frames = [] # 用于存储解码的帧
# 按照解码块大小逐块解码潜在向量
for i in range(0, latents.shape[0], decode_chunk_size):
# 解码当前块的潜在向量，获取采样结果
frame = self.vae.decode(latents[i : i + decode_chunk_size]).sample
frames.append(frame) # 将解码的帧添加到列表中
# 将所有帧在第一个维度上拼接成一个张量
image = torch.cat(frames, dim=0)
else:
# 如果未指定块大小，直接解码所有潜在向量
image = self.vae.decode(latents).sample
# 计算解码后的形状，以适应最终视频的结构
decode_shape = (batch_size, num_frames, -1) + image.shape[2:]
# 重新调整图像的形状，以便形成视频结构
video = image[None, :].reshape(decode_shape).permute(0, 2, 1, 3, 4)
# 始终将视频转换为 float32 格式，以兼容 bfloat16 并减少开销
video = video.float()
# 返回解码后的视频
return video
# 从 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] 范围内
# 检查调度器的 step 方法是否接受 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
# 检查调度器的 step 方法是否接受 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(
# 提示文本
prompt,
# 输入图像
image,
# 图像高度
height,
# 图像宽度
width,
# 可选的负提示文本
negative_prompt=None,
# 可选的提示嵌入
prompt_embeds=None,
# 可选的负提示嵌入
negative_prompt_embeds=None,
# 结束函数参数列表
):
# 检查高度和宽度是否都是8的倍数，若不是则抛出错误
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
# 检查是否同时提供了 prompt 和 prompt_embeds，若是则抛出错误
if prompt is not None and prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
" only forward one of the two."
)
# 检查 prompt 和 prompt_embeds 是否都未定义，若是则抛出错误
elif prompt is None and prompt_embeds is None:
raise ValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
# 检查 prompt 的类型是否为字符串或列表，若不是则抛出错误
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
# 检查是否同时提供了 negative_prompt 和 negative_prompt_embeds，若是则抛出错误
if negative_prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
# 检查 prompt_embeds 和 negative_prompt_embeds 的形状是否一致，若不一致则抛出错误
if prompt_embeds is not None and negative_prompt_embeds is not None:
if prompt_embeds.shape != negative_prompt_embeds.shape:
raise ValueError(
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)
# 检查 image 的类型是否为 torch.Tensor、PIL.Image.Image 或其列表，若不是则抛出错误
if (
not isinstance(image, torch.Tensor)
and not isinstance(image, PIL.Image.Image)
and not isinstance(image, list)
):
raise ValueError(
"`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
f" {type(image)}"
)
# 定义 prepare_image_latents 函数，准备图像潜变量
def prepare_image_latents(
self,
image,
device,
num_frames,
num_videos_per_prompt,
):
# 将图像移动到指定的设备（如 GPU）
image = image.to(device=device)
# 编码图像并从变分自编码器（VAE）获取潜在分布的样本
image_latents = self.vae.encode(image).latent_dist.sample()
# 将潜在表示缩放到 VAE 配置的缩放因子
image_latents = image_latents * self.vae.config.scaling_factor
# 为潜在图像添加帧维度
image_latents = image_latents.unsqueeze(2)
# 为每个后续帧添加位置掩码，初始图像潜在帧之后
frame_position_mask = []
# 遍历除第一帧之外的所有帧
for frame_idx in range(num_frames - 1):
# 计算当前帧的缩放因子
scale = (frame_idx + 1) / (num_frames - 1)
# 将缩放因子应用于与潜在表示相同形状的张量
frame_position_mask.append(torch.ones_like(image_latents[:, :, :1]) * scale)
# 如果位置掩码非空，则连接它们
if frame_position_mask:
frame_position_mask = torch.cat(frame_position_mask, dim=2)
# 将位置掩码附加到潜在表示上
image_latents = torch.cat([image_latents, frame_position_mask], dim=2)
# 根据每个提示的生成数量复制潜在表示，使用适合 MPS 的方法
image_latents = image_latents.repeat(num_videos_per_prompt, 1, 1, 1, 1)
# 如果使用无分类器自由引导，则重复潜在表示
if self.do_classifier_free_guidance:
image_latents = torch.cat([image_latents] * 2)
# 返回处理后的潜在表示
return image_latents
# 从 diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_synth.TextToVideoSDPipeline.prepare_latents 复制
def prepare_latents(
# 准备潜在表示的参数，包括批大小、通道数、帧数等
self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None
):
# 定义潜在表示的形状
shape = (
batch_size,
num_channels_latents,
num_frames,
height // self.vae_scale_factor,
width // self.vae_scale_factor,
)
# 检查生成器列表长度是否与批大小匹配
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."
)
# 如果未提供潜在表示，则生成随机张量
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
# 如果提供了潜在表示，则将其移动到指定设备
latents = latents.to(device)
# 根据调度器所需的标准差缩放初始噪声
latents = latents * self.scheduler.init_noise_sigma
# 返回处理后的潜在表示
return latents
# 在不计算梯度的上下文中执行
@torch.no_grad()
# 用示例文档字符串替换
@replace_example_docstring(EXAMPLE_DOC_STRING)
# 定义可调用对象的方法，允许实例像函数一样被调用
def __call__(
# 输入提示，可以是字符串或字符串列表，默认为 None
self,
prompt: Union[str, List[str]] = None,
# 输入图像，类型为 PipelineImageInput，默认为 None
image: PipelineImageInput = None,
# 目标图像高度，默认为 704 像素
height: Optional[int] = 704,
# 目标图像宽度，默认为 1280 像素
width: Optional[int] = 1280,
# 目标帧率，默认为 16 帧每秒
target_fps: Optional[int] = 16,
# 要生成的帧数，默认为 16
num_frames: int = 16,
# 推理步骤数量，默认为 50
num_inference_steps: int = 50,
# 指导比例，默认为 9.0，控制生成图像的多样性
guidance_scale: float = 9.0,
# 负提示，可以是字符串或字符串列表，默认为 None
negative_prompt: Optional[Union[str, List[str]]] = None,
# 随机噪声的调节因子，默认为 0.0
eta: float = 0.0,
# 每个提示生成的视频数量，默认为 1
num_videos_per_prompt: Optional[int] = 1,
# 解码时的块大小，默认为 1
decode_chunk_size: Optional[int] = 1,
# 随机数生成器，可以是单个或多个 torch.Generator，默认为 None
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 潜在变量张量，默认为 None
latents: Optional[torch.Tensor] = None,
# 提示嵌入张量，默认为 None
prompt_embeds: Optional[torch.Tensor] = None,
# 负提示嵌入张量，默认为 None
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 输出类型，默认为 "pil"，表示生成 PIL 图像
output_type: Optional[str] = "pil",
# 是否返回字典格式的结果，默认为 True
return_dict: bool = True,
# 交叉注意力的关键字参数，默认为 None
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
# 跳过的剪辑数，默认为 1
clip_skip: Optional[int] = 1,
# 以下实用工具来自并适应于
# https://github.com/ali-vilab/i2vgen-xl/blob/main/utils/transforms.py.
# 将 PyTorch 张量或张量列表转换为 PIL 图像
def _convert_pt_to_pil(image: Union[torch.Tensor, List[torch.Tensor]]):
# 如果输入是一个张量列表，则将其沿第一个维度拼接成一个张量
if isinstance(image, list) and isinstance(image[0], torch.Tensor):
image = torch.cat(image, 0)
# 如果输入是一个张量
if isinstance(image, torch.Tensor):
# 如果张量是 3 维的，则在第一个维度增加一个维度，使其变为 4 维
if image.ndim == 3:
image = image.unsqueeze(0)
# 将 PyTorch 张量转换为 NumPy 数组
image_numpy = VaeImageProcessor.pt_to_numpy(image)
# 将 NumPy 数组转换为 PIL 图像
image_pil = VaeImageProcessor.numpy_to_pil(image_numpy)
# 更新 image 为 PIL 图像
image = image_pil
# 返回转换后的图像
return image
# 使用双线性插值调整图像大小
def _resize_bilinear(
image: Union[torch.Tensor, List[torch.Tensor], PIL.Image.Image, List[PIL.Image.Image]], resolution: Tuple[int, int]
):
# 首先将图像转换为 PIL 格式，以防它们是浮动张量（目前仅与测试相关）
image = _convert_pt_to_pil(image)
# 如果输入是图像列表，则对每个图像进行调整大小
if isinstance(image, list):
image = [u.resize(resolution, PIL.Image.BILINEAR) for u in image]
else:
# 如果是单个图像，直接调整大小
image = image.resize(resolution, PIL.Image.BILINEAR)
# 返回调整大小后的图像
return image
# 进行中心裁剪以适应给定的分辨率
def _center_crop_wide(
image: Union[torch.Tensor, List[torch.Tensor], PIL.Image.Image, List[PIL.Image.Image]], resolution: Tuple[int, int]
):
# 首先将图像转换为 PIL 格式，以防它们是浮动张量（目前仅与测试相关）
image = _convert_pt_to_pil(image)
# 如果输入是图像列表
if isinstance(image, list):
# 计算缩放比例，确保图像适应目标分辨率
scale = min(image[0].size[0] / resolution[0], image[0].size[1] / resolution[1])
# 调整每个图像的大小
image = [u.resize((round(u.width // scale), round(u.height // scale)), resample=PIL.Image.BOX) for u in image]
# 进行中心裁剪
x1 = (image[0].width - resolution[0]) // 2 # 计算裁剪区域的左上角 x 坐标
y1 = (image[0].height - resolution[1]) // 2 # 计算裁剪区域的左上角 y 坐标
# 裁剪每个图像并返回
image = [u.crop((x1, y1, x1 + resolution[0], y1 + resolution[1])) for u in image]
return image
else:
# 对于单个图像，计算缩放比例
scale = min(image.size[0] / resolution[0], image.size[1] / resolution[1])
# 调整图像大小
image = image.resize((round(image.width // scale), round(image.height // scale)), resample=PIL.Image.BOX)
# 计算裁剪区域的左上角坐标
x1 = (image.width - resolution[0]) // 2
y1 = (image.height - resolution[1]) // 2
# 裁剪图像并返回
image = image.crop((x1, y1, x1 + resolution[0], y1 + resolution[1]))
return image
.\diffusers\pipelines\i2vgen_xl\__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, # 导入检查 PyTorch 是否可用的函数
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 导入虚拟对象，避免直接使用缺失的依赖
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_i2vgen_xl"] = ["I2VGenXLPipeline"]
# 检查类型检查或慢加载标志
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 * # noqa F403
else:
# 如果依赖可用，导入实际的 I2VGenXLPipeline
from .pipeline_i2vgen_xl import I2VGenXLPipeline
else:
# 导入 sys 模块以操作模块
import sys
# 将当前模块替换为一个延迟加载的模块
sys.modules[__name__] = _LazyModule(
__name__,
globals()["__file__"],
_import_structure,
module_spec=__spec__,
)
# 将 _dummy_objects 中的每个名称和值设置为当前模块的属性
for name, value in _dummy_objects.items():
setattr(sys.modules[__name__], name, value)
.\diffusers\pipelines\kandinsky\pipeline_kandinsky.py
# 版权声明，标识该文件的版权信息
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# 使用 Apache License, Version 2.0 进行许可
# 只有在遵守许可证的情况下，您才能使用此文件
# 许可证的副本可以在以下网址获取
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，
# 否则根据许可证分发的软件是“按原样”提供的，
# 不附带任何明示或暗示的担保或条件
# 有关许可证的具体条款，请参阅许可证文档
# 导入类型提示相关的模块
from typing import Callable, List, Optional, Union
# 导入 PyTorch 库
import torch
# 从 transformers 库中导入 XLMRobertaTokenizer
from transformers import (
XLMRobertaTokenizer,
)
# 导入自定义模型和调度器
from ...models import UNet2DConditionModel, VQModel
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import (
logging,
replace_example_docstring,
)
from ...utils.torch_utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
from .text_encoder import MultilingualCLIP
# 获取日志记录器实例，供后续使用
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
"""
# 定义函数，根据输入的高度和宽度计算新的高度和宽度
def get_new_h_w(h, w, scale_factor=8):
# 计算新的高度，向下取整除以 scale_factor^2
new_h = h // scale_factor**2
# 如果 h 不能被 scale_factor^2 整除，则高度加 1
if h % scale_factor**2 != 0:
new_h += 1
# 计算新的宽度，向下取整除以 scale_factor^2
new_w = w // scale_factor**2
# 如果 w 不能被 scale_factor^2 整除，则宽度加 1
if w % scale_factor**2 != 0:
new_w += 1
# 返回新的高度和宽度，乘以 scale_factor
return new_h * scale_factor, new_w * scale_factor
# 定义 KandinskyPipeline 类，继承自 DiffusionPipeline
class KandinskyPipeline(DiffusionPipeline):
"""
用于使用 Kandinsky 进行文本到图像生成的管道
此模型继承自 [`DiffusionPipeline`]。有关所有管道实现的通用方法（例如下载、保存、在特定设备上运行等）的文档，请查看超类文档。
# 函数参数说明
Args:
text_encoder ([`MultilingualCLIP`]): # 冻结的文本编码器
Frozen text-encoder.
tokenizer ([`XLMRobertaTokenizer`]): # 类的分词器
Tokenizer of class
scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]): # 用于与 `unet` 结合生成图像潜在变量的调度器
A scheduler to be used in combination with `unet` to generate image latents.
unet ([`UNet2DConditionModel`]): # 用于去噪图像嵌入的条件 U-Net 架构
Conditional U-Net architecture to denoise the image embedding.
movq ([`VQModel`]): # MoVQ 解码器，用于从潜在变量生成图像
MoVQ Decoder to generate the image from the latents.
"""
# 定义模型的 CPU 离线加载顺序
model_cpu_offload_seq = "text_encoder->unet->movq"
# 初始化函数，设置模型的各个组件
def __init__(
self,
text_encoder: MultilingualCLIP, # 文本编码器
tokenizer: XLMRobertaTokenizer, # 分词器
unet: UNet2DConditionModel, # U-Net 模型
scheduler: Union[DDIMScheduler, DDPMScheduler], # 调度器
movq: VQModel, # MoVQ 解码器
):
super().__init__() # 调用父类初始化函数
# 注册模型模块
self.register_modules(
text_encoder=text_encoder, # 注册文本编码器
tokenizer=tokenizer, # 注册分词器
unet=unet, # 注册 U-Net 模型
scheduler=scheduler, # 注册调度器
movq=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
# 编码提示信息
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]], # 提示文本，可以是单个字符串或字符串列表
image_embeds: Union[torch.Tensor, List[torch.Tensor]], # 图像嵌入，可以是张量或张量列表
negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]], # 负图像嵌入
negative_prompt: Optional[Union[str, List[str]]] = None, # 可选的负提示文本
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", # 输出类型，默认为 PIL 图像
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, # 可选的回调函数
callback_steps: int = 1, # 每隔多少步骤调用一次回调
return_dict: bool = True, # 是否返回字典格式的结果
.\diffusers\pipelines\kandinsky\pipeline_kandinsky_combined.py
# 版权所有 2024 HuggingFace 团队。保留所有权利。
#
# 根据 Apache 许可证 2.0 版本（"许可证"）许可；
# 除非遵守许可证，否则您不得使用此文件。
# 您可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面同意，软件在许可证下分发时以“原样”基础提供，
# 不提供任何形式的保证或条件，无论是明示或暗示的。
# 请参阅许可证以获取管理权限和
# 限制的具体语言。
from typing import Callable, List, Optional, Union # 从 typing 模块导入类型注解功能
import PIL.Image # 导入 PIL.Image 模块以处理图像
import torch # 导入 PyTorch 库用于深度学习
from transformers import ( # 从 transformers 库导入多个模型和处理器
CLIPImageProcessor, # 导入 CLIP 图像处理器
CLIPTextModelWithProjection, # 导入具有投影的 CLIP 文本模型
CLIPTokenizer, # 导入 CLIP 令牌化工具
CLIPVisionModelWithProjection, # 导入具有投影的 CLIP 视觉模型
XLMRobertaTokenizer, # 导入 XLM-Roberta 令牌化工具
)
from ...models import PriorTransformer, UNet2DConditionModel, VQModel # 从相对路径导入模型
from ...schedulers import DDIMScheduler, DDPMScheduler, UnCLIPScheduler # 导入不同的调度器
from ...utils import ( # 从工具模块导入特定功能
replace_example_docstring, # 导入替换示例文档字符串的功能
)
from ..pipeline_utils import DiffusionPipeline # 从上级路径导入扩散管道工具
from .pipeline_kandinsky import KandinskyPipeline # 导入 Kandinsky 管道
from .pipeline_kandinsky_img2img import KandinskyImg2ImgPipeline # 导入 Kandinsky 图像到图像管道
from .pipeline_kandinsky_inpaint import KandinskyInpaintPipeline # 导入 Kandinsky 修复管道
from .pipeline_kandinsky_prior import KandinskyPriorPipeline # 导入 Kandinsky 先验管道
from .text_encoder import MultilingualCLIP # 导入多语言 CLIP 文本编码器
# 定义文本到图像的示例文档字符串
TEXT2IMAGE_EXAMPLE_DOC_STRING = """
示例：
```py
from diffusers import AutoPipelineForText2Image # 导入自动文本到图像管道
import torch # 导入 PyTorch 库
# 从预训练模型加载管道
pipe = AutoPipelineForText2Image.from_pretrained(
"kandinsky-community/kandinsky-2-1", 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 = """
示例：
```py
from diffusers import AutoPipelineForImage2Image # 导入自动图像到图像管道
import torch # 导入 PyTorch 库
import requests # 导入请求库用于获取图像
from io import BytesIO # 从字节流中读取数据
from PIL import Image # 导入 PIL 图像处理库
import os # 导入操作系统接口库
# 从预训练模型加载管道
pipe = AutoPipelineForImage2Image.from_pretrained(
"kandinsky-community/kandinsky-2-1", 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"
# 发送请求获取图像
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 = """
# 示例代码块
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
# 从预训练模型加载 AutoPipelineForInpainting 对象，指定数据类型为 float16
pipe = AutoPipelineForInpainting.from_pretrained(
"kandinsky-community/kandinsky-2-1-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"
)
# 创建一个全为零的掩码数组，大小为 768x768，数据类型为 float32
mask = np.zeros((768, 768), dtype=np.float32)
# 在猫的头部上方遮罩区域
mask[:250, 250:-250] = 1
# 使用管道生成新图像，输入提示词、原始图像和掩码，设定推理步骤数量为 25，并提取生成的第一张图像
image = pipe(prompt=prompt, image=original_image, mask_image=mask, num_inference_steps=25).images[0]
```
"""
# 类定义：KandinskyCombinedPipeline，继承自 DiffusionPipeline，用于文本到图像生成
class KandinskyCombinedPipeline(DiffusionPipeline):
"""
# 文档字符串：描述使用Kandinsky进行文本到图像生成的组合管道
# 文档字符串：说明此模型继承自 [`DiffusionPipeline`]，并提到其通用方法的文档（如下载、保存、在特定设备上运行等）
# 文档字符串：模型构造函数的参数说明
# text_encoder：被冻结的文本编码器，类型为 [`MultilingualCLIP`]。
# tokenizer：类的分词器，类型为 [`XLMRobertaTokenizer`]。
# scheduler：用于与 `unet` 结合生成图像潜变量的调度器，类型为 `Union[`DDIMScheduler`,`DDPMScheduler`]`。
# unet：条件U-Net架构，用于去噪图像嵌入，类型为 [`UNet2DConditionModel`]。
# movq：从潜变量生成图像的 MoVQ 解码器，类型为 [`VQModel`]。
# prior_prior：用于从文本嵌入近似图像嵌入的规范 unCLIP 先验，类型为 [`PriorTransformer`]。
# prior_image_encoder：被冻结的图像编码器，类型为 [`CLIPVisionModelWithProjection`]。
# prior_text_encoder：被冻结的文本编码器，类型为 [`CLIPTextModelWithProjection`]。
# prior_tokenizer：类的分词器，类型为 [`CLIPTokenizer`]。
# prior_scheduler：与 `prior` 结合生成图像嵌入的调度器，类型为 [`UnCLIPScheduler`]。
"""
# 设定加载连接管道的标志为真
_load_connected_pipes = True
# 定义 CPU 卸载的模型序列
model_cpu_offload_seq = "text_encoder->unet->movq->prior_prior->prior_image_encoder->prior_text_encoder"
# 排除 CPU 卸载的部分
_exclude_from_cpu_offload = ["prior_prior"]
# 构造函数定义，接收多个参数以初始化类
def __init__(
# 文本编码器参数，类型为 MultilingualCLIP
self,
text_encoder: MultilingualCLIP,
# 分词器参数，类型为 XLMRobertaTokenizer
tokenizer: XLMRobertaTokenizer,
# U-Net 参数，类型为 UNet2DConditionModel
unet: UNet2DConditionModel,
# 调度器参数，类型为 DDIMScheduler 或 DDPMScheduler
scheduler: Union[DDIMScheduler, DDPMScheduler],
# MoVQ 解码器参数，类型为 VQModel
movq: VQModel,
# 先验参数，类型为 PriorTransformer
prior_prior: PriorTransformer,
# 图像编码器参数，类型为 CLIPVisionModelWithProjection
prior_image_encoder: CLIPVisionModelWithProjection,
# 文本编码器参数，类型为 CLIPTextModelWithProjection
prior_text_encoder: CLIPTextModelWithProjection,
# 先验分词器参数，类型为 CLIPTokenizer
prior_tokenizer: CLIPTokenizer,
# 先验调度器参数，类型为 UnCLIPScheduler
prior_scheduler: UnCLIPScheduler,
# 图像处理器参数，类型为 CLIPImageProcessor
prior_image_processor: CLIPImageProcessor,
):
# 调用父类的初始化方法
super().__init__()
# 注册多个模块，传递各自的参数
self.register_modules(
# 文本编码器
text_encoder=text_encoder,
# 分词器
tokenizer=tokenizer,
# U-Net 模型
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 = KandinskyPriorPipeline(
# 传入先验模型
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 = KandinskyPipeline(
# 传入文本编码器
text_encoder=text_encoder,
# 传入分词器
tokenizer=tokenizer,
# 传入 U-Net 模型
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 卸载的方法，接收 GPU ID 参数
def enable_sequential_cpu_offload(self, gpu_id=0):
r"""
卸载所有模型（`unet`、`text_encoder`、`vae` 和 `safety checker` 状态字典）到 CPU，使用 🤗
Accelerate，显著减少内存使用。模型被移动到 `torch.device('meta')`，仅在其特定子模块的
`forward` 方法被调用时才在 GPU 上加载。卸载是基于子模块进行的。
内存节省大于使用 `enable_model_cpu_offload`，但性能较低。
"""
# 在先验管道中启用顺序 CPU 卸载
self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id)
# 在解码器管道中启用顺序 CPU 卸载
self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id)
# 进度条方法，接收可迭代对象和总数作为参数
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)
# 使用 PyTorch 的无梯度上下文管理器，避免计算梯度
@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,
# 可选的潜在变量，通常用于生成模型的输入
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,
# 定义一个结合管道类，用于使用 Kandinsky 进行图像到图像的生成
class KandinskyImg2ImgCombinedPipeline(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:
text_encoder ([`MultilingualCLIP`]):
Frozen text-encoder. # 冻结的文本编码器
tokenizer ([`XLMRobertaTokenizer`]):
Tokenizer of class # 分词器类
scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]):
A scheduler to be used in combination with `unet` to generate image latents. # 用于与 `unet` 结合生成图像潜变量的调度器
unet ([`UNet2DConditionModel`]):
Conditional U-Net architecture to denoise the image embedding. # 条件 U-Net 架构，用于去噪图像嵌入
movq ([`VQModel`]):
MoVQ Decoder to generate the image from the latents. # MoVQ 解码器，从潜变量生成图像
prior_prior ([`PriorTransformer`]):
The canonical unCLIP prior to approximate the image embedding from the text embedding. # 经典的 unCLIP 先验，用于从文本嵌入近似图像嵌入
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` 结合生成图像嵌入的调度器
"""
_load_connected_pipes = True # 指定是否加载连接的管道
model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->prior_prior->" "text_encoder->unet->movq" # 指定模型在 CPU 卸载时的顺序
_exclude_from_cpu_offload = ["prior_prior"] # 指定在 CPU 卸载时排除的组件
def __init__( # 初始化方法
self,
text_encoder: MultilingualCLIP, # 文本编码器实例
tokenizer: XLMRobertaTokenizer, # 分词器实例
unet: UNet2DConditionModel, # U-Net 模型实例
scheduler: Union[DDIMScheduler, DDPMScheduler], # 调度器实例
movq: VQModel, # MoVQ 解码器实例
prior_prior: PriorTransformer, # 先验变换器实例
prior_image_encoder: CLIPVisionModelWithProjection, # 图像编码器实例
prior_text_encoder: CLIPTextModelWithProjection, # 文本编码器实例
prior_tokenizer: CLIPTokenizer, # 先验分词器实例
prior_scheduler: UnCLIPScheduler, # 先验调度器实例
prior_image_processor: CLIPImageProcessor, # 图像处理器实例
# 定义构造函数的结束部分，调用父类的构造函数
):
super().__init__()
# 注册多个模块及其相应的组件
self.register_modules(
text_encoder=text_encoder, # 文本编码器
tokenizer=tokenizer, # 分词器
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 = KandinskyPriorPipeline(
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 = KandinskyImg2ImgPipeline(
text_encoder=text_encoder, # 文本编码器
tokenizer=tokenizer, # 分词器
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 卸载，减少 GPU 内存使用
def enable_sequential_cpu_offload(self, gpu_id=0):
r""" # 文档字符串，描述该方法的功能
Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
`torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
Note that offloading happens on a submodule basis. Memory savings are higher than with
`enable_model_cpu_offload`, but performance is lower.
"""
# 启用先验管道的顺序 CPU 卸载
self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id)
# 启用解码管道的顺序 CPU 卸载
self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id)
# 显示进度条，监控处理过程
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) # 在解码管道中显示进度条
self.decoder_pipe.enable_model_cpu_offload() # 启用解码管道的模型 CPU 卸载
# 设置进度条的配置
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__(
self,
# 输入提示，可以是字符串或字符串列表
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,
# 每个提示生成的图像数量，默认为 1
num_images_per_prompt: int = 1,
# 强度参数，默认为 0.3，控制输入图像的影响程度
strength: float = 0.3,
# 输出图像的高度，默认为 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,
# 可选的潜在张量，用于传递预定义的潜在空间
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,
# 定义一个名为 KandinskyInpaintCombinedPipeline 的类，继承自 DiffusionPipeline 类
class KandinskyInpaintCombinedPipeline(DiffusionPipeline):
"""
Combined Pipeline for generation using Kandinsky
该模型继承自 [`DiffusionPipeline`]。请查看超类文档以获取库为所有管道实现的通用方法
（例如下载或保存、在特定设备上运行等）。
参数：
text_encoder ([`MultilingualCLIP`]):
冻结的文本编码器。
tokenizer ([`XLMRobertaTokenizer`]):
令牌化器类。
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` 结合生成图像嵌入的调度器。
"""
# 指示加载连接的管道
_load_connected_pipes = True
# 定义模型 CPU 卸载的顺序
model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->prior_prior->text_encoder->unet->movq"
# 指定从 CPU 卸载时要排除的部分
_exclude_from_cpu_offload = ["prior_prior"]
# 初始化方法，用于设置类的属性
def __init__(
# 文本编码器，类型为 MultilingualCLIP
self,
text_encoder: MultilingualCLIP,
# 令牌化器，类型为 XLMRobertaTokenizer
tokenizer: XLMRobertaTokenizer,
# 条件 U-Net，类型为 UNet2DConditionModel
unet: UNet2DConditionModel,
# 调度器，类型为 DDIMScheduler 或 DDPMScheduler
scheduler: Union[DDIMScheduler, DDPMScheduler],
# MoVQ 解码器，类型为 VQModel
movq: VQModel,
# 先验转换器，类型为 PriorTransformer
prior_prior: PriorTransformer,
# 冻结的图像编码器，类型为 CLIPVisionModelWithProjection
prior_image_encoder: CLIPVisionModelWithProjection,
# 冻结的文本编码器，类型为 CLIPTextModelWithProjection
prior_text_encoder: CLIPTextModelWithProjection,
# 先验令牌化器，类型为 CLIPTokenizer
prior_tokenizer: CLIPTokenizer,
# 先验调度器，类型为 UnCLIPScheduler
prior_scheduler: UnCLIPScheduler,
# 图像处理器，类型为 CLIPImageProcessor
prior_image_processor: CLIPImageProcessor,
):
# 调用父类的初始化方法
super().__init__()
# 注册多个模块及其对应的参数
self.register_modules(
# 文本编码器模块
text_encoder=text_encoder,
# 分词器模块
tokenizer=tokenizer,
# UNet模块
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 = KandinskyPriorPipeline(
# 传入先验模块
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 = KandinskyInpaintPipeline(
# 传入文本编码器
text_encoder=text_encoder,
# 传入分词器
tokenizer=tokenizer,
# 传入 UNet 模块
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=0):
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)
# 启用解码管道的顺序 CPU 离线处理
self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id)
# 显示进度条
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)
# 定义一个可调用的类方法，接受多个参数以生成图像
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",
# 可选的回调函数，在生成过程中调用，接收步数和生成的张量
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
# 回调函数调用的步数间隔，默认为 1
callback_steps: int = 1,
# 是否返回字典格式的结果，默认为 True
return_dict: bool = True,
.\diffusers\pipelines\kandinsky\pipeline_kandinsky_img2img.py
# 版权信息，表明该文件的所有权及相关许可信息
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# 根据 Apache 许可证 2.0 版的条款提供该文件，用户需遵循此许可证
# Licensed under the Apache License, Version 2.0 (the "License");
# 仅在符合许可证的情况下才能使用此文件
# 可以在以下网址获得许可证副本
# 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, List, Optional, Union
# 导入 numpy 库
import numpy as np
# 导入 PIL 库中的 Image 模块
import PIL.Image
# 导入 PyTorch 库
import torch
# 从 PIL 导入 Image 类
from PIL import Image
# 从 transformers 库导入 XLMRobertaTokenizer
from transformers import (
XLMRobertaTokenizer,
)
# 从相对路径导入 UNet2DConditionModel 和 VQModel
from ...models import UNet2DConditionModel, VQModel
# 从相对路径导入 DDIMScheduler
from ...schedulers import DDIMScheduler
# 从相对路径导入 logging 和 replace_example_docstring
from ...utils import (
logging,
replace_example_docstring,
)
# 从相对路径导入 randn_tensor
from ...utils.torch_utils import randn_tensor
# 从相对路径导入 DiffusionPipeline 和 ImagePipelineOutput
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
# 从相对路径导入 MultilingualCLIP
from .text_encoder import MultilingualCLIP
# 创建一个日志记录器，使用模块的名称
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，展示如何使用该模块的功能
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> from diffusers import KandinskyImg2ImgPipeline, KandinskyPriorPipeline
>>> from diffusers.utils import load_image
>>> import torch
>>> pipe_prior = KandinskyPriorPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-1-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 = KandinskyImg2ImgPipeline.from_pretrained(
... "kandinsky-community/kandinsky-2-1", 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(
... prompt,
... 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
"""
# 定义一个函数，用于计算新的高度和宽度
def get_new_h_w(h, w, scale_factor=8):
# 根据给定的高度和缩放因子计算新的高度
new_h = h // scale_factor**2
# 如果高度不能被缩放因子平方整除，增加新的高度
if h % scale_factor**2 != 0:
new_h += 1
# 根据给定的宽度和缩放因子计算新的宽度
new_w = w // scale_factor**2
# 如果宽度不能被缩放因子平方整除，增加新的宽度
if w % scale_factor**2 != 0:
new_w += 1
# 返回新的高度和宽度，乘以缩放因子
return new_h * scale_factor, new_w * scale_factor
# 定义一个函数，用于准备图像
def prepare_image(pil_image, w=512, h=512):
# 调整 PIL 图像的大小，使用双三次插值法
pil_image = pil_image.resize((w, h), resample=Image.BICUBIC, reducing_gap=1)
# 将 PIL 图像转换为 NumPy 数组并转换为 RGB
arr = np.array(pil_image.convert("RGB"))
# 将数组数据类型转换为 float32，并归一化到 [-1, 1] 范围
arr = arr.astype(np.float32) / 127.5 - 1
# 调整数组维度，从 (H, W, C) 转为 (C, H, W)
arr = np.transpose(arr, [2, 0, 1])
# 将 NumPy 数组转换为 PyTorch 张量，并在第一个维度上添加一个维度
image = torch.from_numpy(arr).unsqueeze(0)
# 返回处理后的张量
return image
# Kandinsky 图像到图像生成管道类，继承自 DiffusionPipeline
class KandinskyImg2ImgPipeline(DiffusionPipeline):
"""
使用 Kandinsky 进行图像生成的管道
此模型继承自 [`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__()
# 注册模块，包括文本编码器、分词器、U-Net、调度器和 MoVQ
self.register_modules(
text_encoder=text_encoder,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
movq=movq,
)
# 设置 MoVQ 的缩放因子
self.movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1)
def get_timesteps(self, num_inference_steps, strength, device):
# 使用 init_timestep 获取原始时间步
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
def prepare_latents(self, latents, latent_timestep, shape, dtype, device, generator, 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
# 获取潜在值的形状
shape = latents.shape
# 生成与潜在值相同形状的随机噪声
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# 向潜在值添加噪声
latents = self.add_noise(latents, noise, latent_timestep)
# 返回处理后的潜在值
return latents
def _encode_prompt(
self,
prompt,
device,
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
# 添加噪声的方法覆盖调度器中的同名方法，因为它使用不同的 beta 调度进行添加噪声与采样
def add_noise(
self,
original_samples: torch.Tensor,
noise: torch.Tensor,
timesteps: torch.IntTensor,
# 返回一个张量类型的输出
) -> torch.Tensor:
# 生成 0.0001 到 0.02 之间的 1000 个等间隔的 beta 值
betas = torch.linspace(0.0001, 0.02, 1000, dtype=torch.float32)
# 计算 alpha 值，等于 1 减去 beta 值
alphas = 1.0 - betas
# 计算 alpha 的累积乘积
alphas_cumprod = torch.cumprod(alphas, dim=0)
# 将累积乘积转换为原样本的设备和数据类型
alphas_cumprod = alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
# 将时间步转换为原样本的设备
timesteps = timesteps.to(original_samples.device)
# 计算 sqrt(alpha) 的乘积，取出对应时间步的值
sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
# 将其展平为一维
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
# 如果维度少于原样本，增加维度
while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
# 计算 sqrt(1 - alpha) 的乘积，取出对应时间步的值
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
# 将其展平为一维
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
# 如果维度少于原样本，增加维度
while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
# 根据加权公式生成带噪声的样本
noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
# 返回生成的带噪声的样本
return noisy_samples
# 装饰器，表示不需要计算梯度
@torch.no_grad()
# 替换示例文档字符串
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
# 接收的提示，可以是字符串或字符串列表
self,
prompt: Union[str, List[str]],
# 接收的图像，可以是张量或 PIL 图像
image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
# 图像的嵌入表示
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,
# 强度参数，默认 0.3
strength: float = 0.3,
# 指导比例，默认 7.0
guidance_scale: float = 7.0,
# 每个提示生成的图像数量，默认 1
num_images_per_prompt: int = 1,
# 可选的生成器
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = 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,
