diffusers 源码解析（三十九）
.\diffusers\pipelines\pag\pipeline_pag_kolors.py
# 版权所有 2024 Stability AI, Kwai-Kolors Team 和 The HuggingFace Team。保留所有权利。
#
# 根据 Apache 许可证，版本 2.0（“许可证”）进行许可；
# 除非遵循许可证，否则您不得使用此文件。
# 您可以在以下位置获得许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，按照许可证分发的软件均按“原样”提供，
# 不提供任何形式的担保或条件，明示或暗示。
# 请参阅许可证以获取有关权限和限制的具体语言。
import inspect # 导入 inspect 模块以便进行对象检查
from typing import Any, Callable, Dict, List, Optional, Tuple, Union # 导入类型提示以便于类型注释
import torch # 导入 PyTorch 库
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection # 从 transformers 导入图像处理器和模型
from ...callbacks import MultiPipelineCallbacks, PipelineCallback # 从回调模块导入多管道回调和单个管道回调
from ...image_processor import PipelineImageInput, VaeImageProcessor # 从图像处理模块导入图像输入和变分自编码器图像处理器
from ...loaders import IPAdapterMixin, StableDiffusionXLLoraLoaderMixin # 从加载器模块导入适配器混合类
from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel # 从模型模块导入自编码器、图像投影和条件模型
from ...models.attention_processor import AttnProcessor2_0, FusedAttnProcessor2_0, XFormersAttnProcessor # 导入注意力处理器
from ...schedulers import KarrasDiffusionSchedulers # 从调度器模块导入 Karras 扩散调度器
from ...utils import is_torch_xla_available, logging, replace_example_docstring # 从工具模块导入工具函数
from ...utils.torch_utils import randn_tensor # 从 PyTorch 工具模块导入随机张量生成函数
from ..kolors.pipeline_output import KolorsPipelineOutput # 从 Kolors 模块导入管道输出类
from ..kolors.text_encoder import ChatGLMModel # 从 Kolors 模块导入文本编码器模型
from ..kolors.tokenizer import ChatGLMTokenizer # 从 Kolors 模块导入分词器
from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin # 从管道工具模块导入扩散管道和稳定扩散混合类
from .pag_utils import PAGMixin # 从 PAG 工具模块导入 PAG 混合类
# 检查 XLA 是否可用，如果可用，则导入相关模块
if is_torch_xla_available():
import torch_xla.core.xla_model as xm # 导入 XLA 模型相关功能
XLA_AVAILABLE = True # 标记 XLA 可用
else:
XLA_AVAILABLE = False # 标记 XLA 不可用
# 创建日志记录器实例，用于记录模块内的日志
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，提供用法示例
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> import torch # 导入 PyTorch 库
>>> from diffusers import AutoPipelineForText2Image # 从 diffusers 导入自动文本转图像管道
>>> pipe = AutoPipelineForText2Image.from_pretrained( # 从预训练模型创建管道实例
... "Kwai-Kolors/Kolors-diffusers", # 指定模型名称
... variant="fp16", # 指定变体为 fp16
... torch_dtype=torch.float16, # 设置 PyTorch 数据类型为 float16
... enable_pag=True, # 启用 PAG 功能
... pag_applied_layers=["down.block_2.attentions_1", "up.block_0.attentions_1"], # 指定 PAG 应用层
... )
>>> pipe = pipe.to("cuda") # 将管道移动到 CUDA 设备
>>> prompt = ( # 定义提示文本
... "A photo of a ladybug, macro, zoom, high quality, film, holding a wooden sign with the text 'KOLORS'" # 提示内容
... )
>>> image = pipe(prompt, guidance_scale=5.5, pag_scale=1.5).images[0] # 生成图像并提取第一张
```py
"""
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion 导入的函数
def retrieve_timesteps( # 定义函数以检索时间步
scheduler, # 调度器对象
num_inference_steps: Optional[int] = None, # 可选的推理步数
device: Optional[Union[str, torch.device]] = None, # 可选的设备类型
timesteps: Optional[List[int]] = None, # 可选的时间步列表
sigmas: Optional[List[float]] = None, # 可选的 sigma 值列表
**kwargs, # 其他可选参数
):
"""
# 调用调度器的 `set_timesteps` 方法，并在调用后从调度器中获取时间步
# 处理自定义时间步。任何关键字参数将被传递给 `scheduler.set_timesteps`。
# 参数说明：
# scheduler (`SchedulerMixin`):
# 要从中获取时间步的调度器。
# num_inference_steps (`int`):
# 在使用预训练模型生成样本时使用的扩散步骤数量。如果使用此参数，则 `timesteps`
# 必须为 `None`。
# device (`str` 或 `torch.device`, *可选*):
# 时间步应该移动到的设备。如果为 `None`，则时间步不会被移动。
# timesteps (`List[int]`, *可选*):
# 自定义时间步，用于覆盖调度器的时间步间隔策略。如果传递 `timesteps`，
# 则 `num_inference_steps` 和 `sigmas` 必须为 `None`。
# sigmas (`List[float]`, *可选*):
# 自定义 sigmas，用于覆盖调度器的时间步间隔策略。如果传递 `sigmas`，
# 则 `num_inference_steps` 和 `timesteps` 必须为 `None`。
# 返回:
# `Tuple[torch.Tensor, int]`: 一个元组，其中第一个元素是调度器的时间步调度，
# 第二个元素是推理步骤的数量。
"""
# 检查是否同时传递了 `timesteps` 和 `sigmas`
if timesteps is not None and sigmas is not None:
# 抛出错误，提示只能选择其中一个作为自定义值
raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
# 如果提供了自定义时间步
if timesteps is not None:
# 检查调度器的 `set_timesteps` 方法是否接受时间步
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
# 如果不接受，抛出错误提示
if not accepts_timesteps:
raise ValueError(
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
f" timestep schedules. Please check whether you are using the correct scheduler."
)
# 调用调度器的 `set_timesteps` 方法，传入自定义时间步及设备和关键字参数
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
# 获取调度器的时间步
timesteps = scheduler.timesteps
# 计算推理步骤的数量
num_inference_steps = len(timesteps)
# 如果提供了自定义 sigmas
elif sigmas is not None:
# 检查调度器的 `set_timesteps` 方法是否接受 sigmas
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
# 如果不接受，抛出错误提示
if not accept_sigmas:
raise ValueError(
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
f" sigmas schedules. Please check whether you are using the correct scheduler."
)
# 调用调度器的 `set_timesteps` 方法，传入自定义 sigmas 及设备和关键字参数
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
# 获取调度器的时间步
timesteps = scheduler.timesteps
# 计算推理步骤的数量
num_inference_steps = len(timesteps)
# 如果没有提供 `timesteps` 和 `sigmas`
else:
# 调用调度器的 `set_timesteps` 方法，传入推理步骤数量及设备和关键字参数
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
# 获取调度器的时间步
timesteps = scheduler.timesteps
# 返回时间步和推理步骤数量
return timesteps, num_inference_steps
# 定义一个名为 KolorsPAGPipeline 的类，继承自多个基类，用于文本到图像的生成
class KolorsPAGPipeline(
# 继承 DiffusionPipeline、StableDiffusionMixin、StableDiffusionXLLoraLoaderMixin、IPAdapterMixin 和 PAGMixin
DiffusionPipeline, StableDiffusionMixin, StableDiffusionXLLoraLoaderMixin, IPAdapterMixin, PAGMixin
):
r"""
使用 Kolors 进行文本到图像生成的管道。
该模型继承自 [`DiffusionPipeline`]。请查看超类文档以了解库为所有管道实现的通用方法（如下载或保存、在特定设备上运行等）。
该管道还继承了以下加载方法：
- [`~loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] 用于加载 LoRA 权重
- [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] 用于保存 LoRA 权重
- [`~loaders.IPAdapterMixin.load_ip_adapter`] 用于加载 IP 适配器
参数：
vae ([`AutoencoderKL`]):
变分自编码器 (VAE) 模型，用于将图像编码和解码为潜在表示。
text_encoder ([`ChatGLMModel`]):
冻结的文本编码器。Kolors 使用 [ChatGLM3-6B](https://huggingface.co/THUDM/chatglm3-6b)。
tokenizer (`ChatGLMTokenizer`):
类的标记器
[ChatGLMTokenizer](https://huggingface.co/THUDM/chatglm3-6b/blob/main/tokenization_chatglm.py)。
unet ([`UNet2DConditionModel`]): 条件 U-Net 架构，用于去噪编码的图像潜在表示。
scheduler ([`SchedulerMixin`]):
结合 `unet` 使用的调度器，用于去噪编码的图像潜在表示。可以是
[`DDIMScheduler`]、[`LMSDiscreteScheduler`] 或 [`PNDMScheduler`]。
force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"False"`):
是否强制负提示嵌入始终设置为 0。另请参见 `Kwai-Kolors/Kolors-diffusers` 的配置。
pag_applied_layers (`str` or `List[str]``, *optional*, defaults to `"mid"`):
设置要应用扰动注意力引导的变压器注意力层。可以是字符串或字符串列表，包含 "down"、"mid"、"up"，或整个变压器块或特定变压器块注意力层，例如：
["mid"] ["down", "mid"] ["down", "mid", "up.block_1"] ["down", "mid", "up.block_1.attentions_0",
"up.block_1.attentions_1"]
"""
# 定义模型在 CPU 上的卸载顺序
model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
# 定义可选组件的列表，这些组件在实例化时可以选择性提供
_optional_components = [
"image_encoder",
"feature_extractor",
]
# 定义需要作为张量输入的回调列表，这些输入将用于后续的处理
_callback_tensor_inputs = [
"latents",
"prompt_embeds",
"negative_prompt_embeds",
"add_text_embeds",
"add_time_ids",
"negative_pooled_prompt_embeds",
"negative_add_time_ids",
]
# 初始化方法，用于设置类的初始状态
def __init__(
self,
vae: AutoencoderKL, # VAE模型，用于图像编码
text_encoder: ChatGLMModel, # 文本编码器模型，用于处理输入文本
tokenizer: ChatGLMTokenizer, # 令牌化工具，将文本转为模型可处理的格式
unet: UNet2DConditionModel, # UNet模型，用于生成图像
scheduler: KarrasDiffusionSchedulers, # 调度器，用于管理采样过程
image_encoder: CLIPVisionModelWithProjection = None, # 可选的图像编码器
feature_extractor: CLIPImageProcessor = None, # 可选的特征提取器
force_zeros_for_empty_prompt: bool = False, # 是否强制为空提示时输出零
pag_applied_layers: Union[str, List[str]] = "mid", # 应用的层，默认设置为“mid”
):
super().__init__() # 调用父类的初始化方法
# 注册多个模块，以便于管理和访问
self.register_modules(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
image_encoder=image_encoder,
feature_extractor=feature_extractor,
)
# 将配置参数注册到类中
self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)
# 根据 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
)
# 创建图像处理器实例，使用 VAE 的缩放因子
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
# 设置默认的样本大小，从 UNet 配置中获取
self.default_sample_size = self.unet.config.sample_size
# 设置应用的层
self.set_pag_applied_layers(pag_applied_layers)
# 从其他管道复制的 encode_prompt 方法，用于编码提示
def encode_prompt(
self,
prompt, # 输入提示文本
device: Optional[torch.device] = None, # 可选的设备参数，用于模型计算
num_images_per_prompt: int = 1, # 每个提示生成的图像数量
do_classifier_free_guidance: bool = True, # 是否进行无分类器引导
negative_prompt=None, # 可选的负提示文本
prompt_embeds: Optional[torch.FloatTensor] = None, # 可选的提示嵌入
pooled_prompt_embeds: Optional[torch.Tensor] = None, # 可选的池化提示嵌入
negative_prompt_embeds: Optional[torch.FloatTensor] = None, # 可选的负提示嵌入
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, # 可选的负池化提示嵌入
max_sequence_length: int = 256, # 最大序列长度
# 从其他管道复制的 encode_image 方法
# 定义一个编码图像的函数，接收图像、设备、每个提示的图像数量和可选的隐藏状态输出
def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
# 获取图像编码器参数的数据类型
dtype = next(self.image_encoder.parameters()).dtype
# 如果输入的图像不是张量，则通过特征提取器处理并返回张量格式
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values
# 将图像移到指定设备并转换为适当的数据类型
image = image.to(device=device, dtype=dtype)
# 如果需要输出隐藏状态
if output_hidden_states:
# 获取编码后的隐藏状态的倒数第二个层
image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
# 复制隐藏状态以适应每个提示的图像数量
image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
# 对于无条件图像，使用全零张量获取隐藏状态
uncond_image_enc_hidden_states = self.image_encoder(
torch.zeros_like(image), output_hidden_states=True
).hidden_states[-2]
# 复制无条件隐藏状态以适应每个提示的图像数量
uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
num_images_per_prompt, dim=0
)
# 返回编码图像和无条件图像的隐藏状态
return image_enc_hidden_states, uncond_image_enc_hidden_states
else:
# 获取编码后的图像嵌入
image_embeds = self.image_encoder(image).image_embeds
# 复制图像嵌入以适应每个提示的图像数量
image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
# 创建与图像嵌入形状相同的全零张量作为无条件图像嵌入
uncond_image_embeds = torch.zeros_like(image_embeds)
# 返回编码图像和无条件图像的嵌入
return image_embeds, uncond_image_embeds
# 从 stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds 复制的函数
def prepare_ip_adapter_image_embeds(
# 定义输入适配器图像、适配器图像嵌入、设备、每个提示的图像数量和是否进行无分类器引导
self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance
):
# 初始化图像嵌入列表
image_embeds = []
# 如果启用分类器自由引导，初始化负图像嵌入列表
if do_classifier_free_guidance:
negative_image_embeds = []
# 如果输入适配器的图像嵌入为空
if ip_adapter_image_embeds is None:
# 确保输入的图像是列表形式
if not isinstance(ip_adapter_image, list):
ip_adapter_image = [ip_adapter_image]
# 检查输入图像的数量与适配器数量是否匹配
if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers):
raise ValueError(
# 抛出错误提示输入图像数量不匹配
f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
)
# 遍历每个单独的输入适配器图像和对应的投影层
for single_ip_adapter_image, image_proj_layer in zip(
ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers
):
# 确定输出是否为隐藏状态
output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
# 对图像进行编码，返回嵌入
single_image_embeds, single_negative_image_embeds = self.encode_image(
single_ip_adapter_image, device, 1, output_hidden_state
)
# 将单个图像嵌入添加到图像嵌入列表
image_embeds.append(single_image_embeds[None, :])
# 如果启用分类器自由引导，将负图像嵌入添加到列表
if do_classifier_free_guidance:
negative_image_embeds.append(single_negative_image_embeds[None, :])
else:
# 遍历已存在的图像嵌入
for single_image_embeds in ip_adapter_image_embeds:
# 如果启用分类器自由引导，将图像嵌入拆分为负和正
if do_classifier_free_guidance:
single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
negative_image_embeds.append(single_negative_image_embeds)
# 将图像嵌入添加到列表
image_embeds.append(single_image_embeds)
# 初始化适配器图像嵌入列表
ip_adapter_image_embeds = []
# 遍历图像嵌入并处理
for i, single_image_embeds in enumerate(image_embeds):
# 根据提示的图像数量重复图像嵌入
single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
# 如果启用分类器自由引导，处理负图像嵌入
if do_classifier_free_guidance:
single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0)
single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], dim=0)
# 将图像嵌入移动到指定设备
single_image_embeds = single_image_embeds.to(device=device)
# 将处理后的图像嵌入添加到适配器图像嵌入列表
ip_adapter_image_embeds.append(single_image_embeds)
# 返回适配器图像嵌入
return ip_adapter_image_embeds
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs 复制的代码
# 准备额外的参数用于调度器的步骤，因为并非所有调度器的签名相同
def prepare_extra_step_kwargs(self, generator, eta):
# eta（η）仅在 DDIMScheduler 中使用，对其他调度器将被忽略
# eta 对应于 DDIM 论文中的 η: https://arxiv.org/abs/2010.02502
# 其值应在 [0, 1] 之间
# 检查调度器的步骤方法是否接受 eta 参数
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 初始化额外步骤参数字典
extra_step_kwargs = {}
# 如果调度器接受 eta，则将其添加到额外参数中
if accepts_eta:
extra_step_kwargs["eta"] = eta
# 检查调度器的步骤方法是否接受 generator 参数
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 如果调度器接受 generator，则将其添加到额外参数中
if accepts_generator:
extra_step_kwargs["generator"] = generator
# 返回准备好的额外步骤参数字典
return extra_step_kwargs
# 从 diffusers.pipelines.kolors.pipeline_kolors.KolorsPipeline.check_inputs 复制的函数
def check_inputs(
self,
prompt,
num_inference_steps,
height,
width,
negative_prompt=None,
prompt_embeds=None,
pooled_prompt_embeds=None,
negative_prompt_embeds=None,
negative_pooled_prompt_embeds=None,
ip_adapter_image=None,
ip_adapter_image_embeds=None,
callback_on_step_end_tensor_inputs=None,
max_sequence_length=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,
)
# 检查传入的 generator 列表长度是否与批量大小一致
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"您传入的生成器列表长度为 {len(generator)}，请求的有效批量大小为 {batch_size}。"
f" 确保批量大小与生成器长度一致。"
)
# 如果潜变量为 None，则生成随机潜变量
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
# 如果潜变量不为 None，将其移至指定设备
latents = latents.to(device)
# 按调度器要求的标准差缩放初始噪声
latents = latents * self.scheduler.init_noise_sigma
# 返回处理后的潜变量
return latents
# 从 diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline._get_add_time_ids 复制的函数
def _get_add_time_ids(
self, original_size, crops_coords_top_left, target_size, dtype, text_encoder_projection_dim=None
):
# 创建一个包含原始尺寸、裁剪左上角坐标和目标尺寸的列表
add_time_ids = list(original_size + crops_coords_top_left + target_size)
# 计算传入的附加时间嵌入维度
passed_add_embed_dim = (
self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
)
# 获取模型期望的附加嵌入维度
expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
# 如果期望的嵌入维度与实际的不同，则抛出错误
if expected_add_embed_dim != passed_add_embed_dim:
raise ValueError(
f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
)
# 将附加时间 ID 转换为张量，指定数据类型
add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
# 返回附加时间 ID 张量
return add_time_ids
# 从 diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.upcast_vae 复制的函数
def upcast_vae(self):
# 获取 VAE 的数据类型
dtype = self.vae.dtype
# 将 VAE 转换为 float32 类型
self.vae.to(dtype=torch.float32)
# 检查是否使用了 torch 2.0 或 xformers，确定注意力处理器的类型
use_torch_2_0_or_xformers = isinstance(
self.vae.decoder.mid_block.attentions[0].processor,
(
AttnProcessor2_0,
XFormersAttnProcessor,
FusedAttnProcessor2_0,
),
)
# 如果使用了 xformers 或 torch 2.0，则注意力块不需要是 float32，这样可以节省大量内存
if use_torch_2_0_or_xformers:
# 将后量化卷积层转换为指定的数据类型
self.vae.post_quant_conv.to(dtype)
# 将输入卷积层转换为指定的数据类型
self.vae.decoder.conv_in.to(dtype)
# 将中间块转换为指定的数据类型
self.vae.decoder.mid_block.to(dtype)
# 从 diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding 复制的函数
def get_guidance_scale_embedding(
# 输入的张量 w 和嵌入维度以及数据类型的默认值
self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32
# 该函数返回生成的嵌入向量，形状为 (len(w), embedding_dim)
) -> torch.Tensor:
"""
# 引用文档，详细说明函数的实现和参数
See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
Args:
w (`torch.Tensor`):
# 用于生成嵌入向量的指导权重
Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
embedding_dim (`int`, *optional*, defaults to 512):
# 生成的嵌入维度，默认值为512
Dimension of the embeddings to generate.
dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
# 生成的嵌入的数据类型，默认为torch.float32
Data type of the generated embeddings.
Returns:
`torch.Tensor`: # 返回嵌入向量
Embedding vectors with shape `(len(w), embedding_dim)`.
"""
# 确保输入的 w 只有一维
assert len(w.shape) == 1
# 将 w 的值扩大1000倍
w = w * 1000.0
# 计算嵌入的一半维度
half_dim = embedding_dim // 2
# 计算缩放因子的对数
emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
# 计算嵌入的指数值
emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
# 将 w 转换为指定 dtype 并进行广播
emb = w.to(dtype)[:, None] * emb[None, :]
# 组合正弦和余弦函数的嵌入
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
# 如果嵌入维度是奇数，则进行零填充
if embedding_dim % 2 == 1: # zero pad
emb = torch.nn.functional.pad(emb, (0, 1))
# 确保生成的嵌入形状符合预期
assert emb.shape == (w.shape[0], embedding_dim)
# 返回生成的嵌入
return emb
# 返回指导尺度属性
@property
def guidance_scale(self):
return self._guidance_scale
# 定义指导尺度，与Imagen论文中的指导权重 w 类似
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# 对应于不执行分类器自由指导。
@property
def do_classifier_free_guidance(self):
# 返回是否进行分类器自由指导的布尔值
return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
# 返回交叉注意力的关键字参数
@property
def cross_attention_kwargs(self):
return self._cross_attention_kwargs
# 返回去噪结束的属性
@property
def denoising_end(self):
return self._denoising_end
# 返回时间步数的属性
@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: int = 50,
# 时间步的列表，可选
timesteps: List[int] = None,
# sigma 值的列表，可选
sigmas: List[float] = None,
# 去噪结束值，可选
denoising_end: Optional[float] = None,
# 引导缩放因子，默认为 5.0
guidance_scale: float = 5.0,
# 负面提示，可以是字符串或字符串列表，默认为 None
negative_prompt: Optional[Union[str, List[str]]] = None,
# 每个提示生成的图像数量，默认为 1
num_images_per_prompt: Optional[int] = 1,
# eta 值，默认为 0.0
eta: float = 0.0,
# 随机数生成器，可选
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 潜在表示，可选
latents: Optional[torch.Tensor] = None,
# 提示嵌入，可选
prompt_embeds: Optional[torch.Tensor] = None,
# 处理过的提示嵌入，可选
pooled_prompt_embeds: Optional[torch.Tensor] = None,
# 负面提示嵌入，可选
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 处理过的负面提示嵌入，可选
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
# IP 适配器图像，可选
ip_adapter_image: Optional[PipelineImageInput] = None,
# IP 适配器图像嵌入的列表，可选
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
# 输出类型，默认为 "pil"
output_type: Optional[str] = "pil",
# 是否返回字典，默认为 True
return_dict: bool = True,
# 跨注意力参数的字典，可选
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
# 原始图像大小，可选
original_size: Optional[Tuple[int, int]] = None,
# 裁剪坐标的左上角，默认为 (0, 0)
crops_coords_top_left: Tuple[int, int] = (0, 0),
# 目标大小，可选
target_size: Optional[Tuple[int, int]] = None,
# 负面图像的原始大小，可选
negative_original_size: Optional[Tuple[int, int]] = None,
# 负面裁剪坐标的左上角，默认为 (0, 0)
negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
# 负面目标大小，可选
negative_target_size: Optional[Tuple[int, int]] = None,
# 在步骤结束时的回调函数，可选
callback_on_step_end: Optional[
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
] = None,
# 在步骤结束时的张量输入回调名称，默认为 ["latents"]
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# PAG 缩放因子，默认为 3.0
pag_scale: float = 3.0,
# 自适应 PAG 缩放，默认为 0.0
pag_adaptive_scale: float = 0.0,
# 最大序列长度，默认为 256
max_sequence_length: int = 256,
.\diffusers\pipelines\pag\pipeline_pag_pixart_sigma.py
# 版权声明，说明本文件的版权归 PixArt-Sigma 作者和 HuggingFace 团队所有，保留所有权利。
#
# 根据 Apache 许可证第 2.0 版（"许可证"）进行许可；
# 除非遵循许可证，否则您不得使用此文件。
# 您可以在以下地址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面同意，否则根据许可证分发的软件在 "按原样" 基础上提供，
# 不附带任何明示或暗示的担保或条件。
# 有关许可证下特定权限和限制的信息，请参见许可证。
import html # 导入 HTML 处理库
import inspect # 导入用于检查对象的库
import re # 导入正则表达式库
import urllib.parse as ul # 导入用于解析和构建 URL 的库
from typing import Callable, List, Optional, Tuple, Union # 导入类型提示
import torch # 导入 PyTorch 库
from transformers import T5EncoderModel, T5Tokenizer # 从 transformers 库导入 T5 编码器模型和分词器
from ...image_processor import PixArtImageProcessor # 从相对路径导入 PixArt 图像处理器
from ...models import AutoencoderKL, PixArtTransformer2DModel # 从相对路径导入模型
from ...schedulers import KarrasDiffusionSchedulers # 从相对路径导入 Karras 扩散调度器
from ...utils import ( # 从相对路径导入多个实用工具
BACKENDS_MAPPING, # 后端映射
deprecate, # 废弃标记
is_bs4_available, # 检查 BeautifulSoup 是否可用的函数
is_ftfy_available, # 检查 ftfy 是否可用的函数
logging, # 日志记录工具
replace_example_docstring, # 替换示例文档字符串的工具
)
from ...utils.torch_utils import randn_tensor # 从相对路径导入生成随机张量的工具
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput # 从相对路径导入扩散管道和图像输出类
from ..pixart_alpha.pipeline_pixart_alpha import ( # 从相对路径导入 PixArt Alpha 管道中的多个常量
ASPECT_RATIO_256_BIN, # 256 维的长宽比常量
ASPECT_RATIO_512_BIN, # 512 维的长宽比常量
ASPECT_RATIO_1024_BIN, # 1024 维的长宽比常量
)
from ..pixart_alpha.pipeline_pixart_sigma import ASPECT_RATIO_2048_BIN # 从相对路径导入 2048 维的长宽比常量
from .pag_utils import PAGMixin # 从当前模块导入 PAGMixin 类
logger = logging.get_logger(__name__) # 获取当前模块的日志记录器实例，禁用 pylint 的无效名称警告
if is_bs4_available(): # 检查是否可用 BeautifulSoup 库
from bs4 import BeautifulSoup # 从 bs4 导入 BeautifulSoup 类
if is_ftfy_available(): # 检查是否可用 ftfy 库
import ftfy # 导入 ftfy 库
EXAMPLE_DOC_STRING = """ # 示例文档字符串，用于展示如何使用该管道
Examples: # 示例部分的标题
```py # Python 代码块开始
>>> import torch # 导入 PyTorch 库
>>> from diffusers import AutoPipelineForText2Image # 从 diffusers 库导入自动文本到图像的管道
>>> pipe = AutoPipelineForText2Image.from_pretrained( # 从预训练模型加载管道
... "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", # 指定模型路径
... torch_dtype=torch.float16, # 设置使用的 PyTorch 数据类型
... pag_applied_layers=["blocks.14"], # 指定应用 PAG 的层
... enable_pag=True, # 启用 PAG
... ) # 结束管道初始化
>>> pipe = pipe.to("cuda") # 将管道移动到 CUDA 设备
>>> prompt = "A small cactus with a happy face in the Sahara desert" # 定义提示语
>>> image = pipe(prompt, pag_scale=4.0, guidance_scale=1.0).images[0] # 生成图像并获取第一张图像
```py # Python 代码块结束
"""
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion 中复制的函数
def retrieve_timesteps( # 定义函数以检索时间步
scheduler, # 传入调度器对象
num_inference_steps: Optional[int] = None, # 可选参数，推理步骤数量
device: Optional[Union[str, torch.device]] = None, # 可选参数，设备信息
timesteps: Optional[List[int]] = None, # 可选参数，时间步列表
sigmas: Optional[List[float]] = None, # 可选参数，sigma 值列表
**kwargs, # 额外的关键字参数
):
""" # 函数文档字符串
调用调度器的 `set_timesteps` 方法，并在调用后从调度器中检索时间步。处理
自定义时间步。任何 kwargs 将被传递给 `scheduler.set_timesteps`。
``` # 函数文档字符串结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
```py # 代码块结束
``` # 代码块结束
# 参数说明部分，解释函数输入的参数类型和含义
Args:
scheduler (`SchedulerMixin`): # 调度器，用于获取时间步
The scheduler to get timesteps from. # 从调度器中获取时间步
num_inference_steps (`int`): # 生成样本时使用的扩散步骤数量
The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` # 如果使用此参数，则 `timesteps` 必须为 `None`
must be `None`.
device (`str` or `torch.device`, *optional*): # 目标设备，时间步将被移动到该设备上
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. # 如果为 `None`，则不移动时间步
timesteps (`List[int]`, *optional*): # 自定义时间步，用于覆盖调度器的时间步间隔策略
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, # 如果传入 `timesteps`，则 `num_inference_steps` 和 `sigmas` 必须为 `None`
`num_inference_steps` and `sigmas` must be `None`.
sigmas (`List[float]`, *optional*): # 自定义 sigma，用于覆盖调度器的时间步间隔策略
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, # 如果传入 `sigmas`，则 `num_inference_steps` 和 `timesteps` 必须为 `None`
`num_inference_steps` and `timesteps` must be `None`.
Returns: # 返回值说明部分
`Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the # 返回一个元组，第一个元素是来自调度器的时间步调度
second element is the number of inference steps. # 第二个元素是推理步骤的数量
"""
if timesteps is not None and sigmas is not None: # 检查是否同时提供了时间步和 sigma
raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") # 如果同时提供，抛出值错误
if timesteps is not None: # 如果提供了时间步
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) # 检查调度器是否支持自定义时间步
if not accepts_timesteps: # 如果不支持
raise ValueError( # 抛出值错误，提示不支持自定义时间步
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
f" timestep schedules. Please check whether you are using the correct scheduler."
)
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) # 设置调度器的时间步
timesteps = scheduler.timesteps # 获取调度器的时间步
num_inference_steps = len(timesteps) # 计算推理步骤的数量
elif sigmas is not None: # 如果提供了 sigma
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) # 检查调度器是否支持自定义 sigma
if not accept_sigmas: # 如果不支持
raise ValueError( # 抛出值错误，提示不支持自定义 sigma
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
f" sigmas schedules. Please check whether you are using the correct scheduler."
)
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) # 设置调度器的 sigma
timesteps = scheduler.timesteps # 获取调度器的时间步
num_inference_steps = len(timesteps) # 计算推理步骤的数量
else: # 如果既没有提供时间步也没有提供 sigma
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) # 使用推理步骤数量设置时间步
timesteps = scheduler.timesteps # 获取调度器的时间步
return timesteps, num_inference_steps # 返回时间步和推理步骤的数量
# 定义一个名为 PixArtSigmaPAGPipeline 的类，继承自 DiffusionPipeline 和 PAGMixin
class PixArtSigmaPAGPipeline(DiffusionPipeline, PAGMixin):
r"""
[PAG pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/pag) for text-to-image generation
using PixArt-Sigma.
""" # 文档字符串，描述该管道的用途
# 定义一个用于匹配不良标点符号的正则表达式
bad_punct_regex = re.compile(
r"["
+ "#®•©™&@·º½¾¿¡§~" # 匹配特定的不良字符
+ r"\)" # 匹配右括号
+ r"\(" # 匹配左括号
+ r"\]" # 匹配右方括号
+ r"\[" # 匹配左方括号
+ r"\}" # 匹配右花括号
+ r"\{" # 匹配左花括号
+ r"\|" # 匹配竖线
+ "\\" # 匹配反斜杠
+ r"\/" # 匹配斜杠
+ r"\*" # 匹配星号
+ r"]{1,}" # 匹配一个或多个上述字符
) # noqa
# 定义可选组件的列表，包括 tokenizer 和 text_encoder
_optional_components = ["tokenizer", "text_encoder"]
# 定义模型的 CPU 离线顺序
model_cpu_offload_seq = "text_encoder->transformer->vae"
# 初始化方法，接收多个参数以构建管道
def __init__(
self,
tokenizer: T5Tokenizer, # 传入的 tokenizer 对象
text_encoder: T5EncoderModel, # 传入的文本编码器对象
vae: AutoencoderKL, # 传入的变分自编码器对象
transformer: PixArtTransformer2DModel, # 传入的变换器模型对象
scheduler: KarrasDiffusionSchedulers, # 传入的调度器对象
pag_applied_layers: Union[str, List[str]] = "blocks.1", # 应用 PAG 的层，默认为第一个变换器块
):
# 调用父类的初始化方法
super().__init__()
# 注册各个模块以供后续使用
self.register_modules(
tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
)
# 计算 VAE 的缩放因子
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
# 创建 PixArt 图像处理器对象，使用 VAE 缩放因子
self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor)
# 设置 PAG 应用的层
self.set_pag_applied_layers(pag_applied_layers)
# 从 PixArtAlphaPipeline 复制的 encode_prompt 方法，修改了最大序列长度
def encode_prompt(
self,
prompt: Union[str, List[str]], # 输入的提示文本，可以是字符串或字符串列表
do_classifier_free_guidance: bool = True, # 是否进行无分类器引导
negative_prompt: str = "", # 可选的负面提示文本
num_images_per_prompt: int = 1, # 每个提示生成的图像数量
device: Optional[torch.device] = None, # 可选的设备参数
prompt_embeds: Optional[torch.Tensor] = None, # 可选的提示嵌入
negative_prompt_embeds: Optional[torch.Tensor] = None, # 可选的负面提示嵌入
prompt_attention_mask: Optional[torch.Tensor] = None, # 可选的提示注意力掩码
negative_prompt_attention_mask: Optional[torch.Tensor] = None, # 可选的负面提示注意力掩码
clean_caption: bool = False, # 是否清理标题
max_sequence_length: int = 300, # 最大序列长度，默认为 300
**kwargs, # 其他可选参数
# 从 StableDiffusionPipeline 复制的 prepare_extra_step_kwargs 方法
# 准备调度器步骤的额外参数，因为并非所有调度器都有相同的参数签名
def prepare_extra_step_kwargs(self, generator, eta):
# eta (η) 仅用于 DDIMScheduler，对于其他调度器将被忽略
# eta 对应于 DDIM 论文中的 η: https://arxiv.org/abs/2010.02502
# eta 应在 [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
# 从 diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.check_inputs 复制而来
def check_inputs(
self,
# 输入的提示信息
prompt,
# 输出图像的高度
height,
# 输出图像的宽度
width,
# 负提示信息
negative_prompt,
# 回调步骤
callback_steps,
# 可选的提示嵌入
prompt_embeds=None,
# 可选的负提示嵌入
negative_prompt_embeds=None,
# 可选的提示注意力掩码
prompt_attention_mask=None,
# 可选的负提示注意力掩码
negative_prompt_attention_mask=None,
# 从 diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing 复制而来
def _text_preprocessing(self, text, clean_caption=False):
# 如果需要清理标题但缺少 bs4 库，则发出警告
if clean_caption and not is_bs4_available():
logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
logger.warning("Setting `clean_caption` to False...")
# 将 clean_caption 设置为 False
clean_caption = False
# 如果需要清理标题但缺少 ftfy 库，则发出警告
if clean_caption and not is_ftfy_available():
logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
logger.warning("Setting `clean_caption` to False...")
# 将 clean_caption 设置为 False
clean_caption = False
# 如果输入的文本不是元组或列表，则将其转换为列表
if not isinstance(text, (tuple, list)):
text = [text]
# 定义处理文本的函数
def process(text: str):
# 如果需要清理标题，则执行清理操作
if clean_caption:
text = self._clean_caption(text)
text = self._clean_caption(text)
else:
# 否则将文本转换为小写并去掉首尾空白
text = text.lower().strip()
# 返回处理后的文本
return text
# 返回处理后的所有文本
return [process(t) for t in text]
# 从 diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption 复制而来
# 从 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
# 禁用梯度计算以提高性能
@torch.no_grad()
# 替换文档字符串
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
# 接受提示，可以是单个字符串或字符串列表
prompt: Union[str, List[str]] = None,
# 负提示，默认为空字符串
negative_prompt: str = "",
# 推理步骤的数量，默认为20
num_inference_steps: int = 20,
# 时间步列表，默认为None
timesteps: List[int] = None,
# sigma值列表，默认为None
sigmas: List[float] = None,
# 指导比例，默认为4.5
guidance_scale: float = 4.5,
# 每个提示生成的图像数量，默认为1
num_images_per_prompt: Optional[int] = 1,
# 图像高度，默认为None
height: Optional[int] = None,
# 图像宽度，默认为None
width: Optional[int] = None,
# eta值，默认为0.0
eta: float = 0.0,
# 随机生成器，默认为None
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 潜在变量，默认为None
latents: Optional[torch.Tensor] = None,
# 提示嵌入，默认为None
prompt_embeds: Optional[torch.Tensor] = None,
# 提示注意力掩码，默认为None
prompt_attention_mask: Optional[torch.Tensor] = None,
# 负提示嵌入，默认为None
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 负提示注意力掩码，默认为None
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
# 输出类型，默认为"pil"
output_type: Optional[str] = "pil",
# 是否返回字典，默认为True
return_dict: bool = True,
# 回调函数，默认为None
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
# 回调步骤，默认为1
callback_steps: int = 1,
# 是否清理提示，默认为True
clean_caption: bool = True,
# 是否使用分辨率分箱，默认为True
use_resolution_binning: bool = True,
# 最大序列长度，默认为300
max_sequence_length: int = 300,
# pag_scale，默认为3.0
pag_scale: float = 3.0,
# pag_adaptive_scale，默认为0.0
pag_adaptive_scale: float = 0.0,
.\diffusers\pipelines\pag\pipeline_pag_sd.py
# 版权声明，指明版权归 HuggingFace 团队所有
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# 根据 Apache 许可证第 2.0 版（“许可证”）进行授权；
# 除非遵循许可证，否则不得使用此文件。
# 可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，否则根据许可证分发的软件是按“现状”基础分发的，
# 不附有任何形式的明示或暗示的担保或条件。
# 有关许可证下特定权限和限制的更多信息，请参见许可证。
import inspect # 导入 inspect 模块，用于获取活跃对象的详细信息
from typing import Any, Callable, Dict, List, Optional, Union # 导入类型提示相关的工具
import torch # 导入 PyTorch 库
from packaging import version # 导入版本管理工具
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection # 导入 CLIP 相关的处理器和模型
from ...configuration_utils import FrozenDict # 导入 FrozenDict，可能用于不可变字典配置
from ...image_processor import PipelineImageInput, VaeImageProcessor # 导入图像处理相关的类
from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin # 导入不同的加载器混合类
from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel # 导入模型类
from ...models.lora import adjust_lora_scale_text_encoder # 导入用于调整 Lora 模型的函数
from ...schedulers import KarrasDiffusionSchedulers # 导入 Karras 扩散调度器
from ...utils import ( # 从工具模块导入多个实用函数
USE_PEFT_BACKEND, # 指示是否使用 PEFT 后端
deprecate, # 用于标记弃用的功能
logging, # 导入日志记录工具
replace_example_docstring, # 用于替换示例文档字符串的函数
scale_lora_layers, # 用于缩放 Lora 层的函数
unscale_lora_layers, # 用于取消缩放 Lora 层的函数
)
from ...utils.torch_utils import randn_tensor # 导入随机张量生成工具
from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin # 导入扩散管道和稳定扩散混合类
from ..stable_diffusion.pipeline_output import StableDiffusionPipelineOutput # 导入稳定扩散管道的输出类
from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker # 导入稳定扩散安全检查器
from .pag_utils import PAGMixin # 导入 PAG 相关的混合类
logger = logging.get_logger(__name__) # 创建一个日志记录器，记录当前模块的信息；禁用 pylint 的无效名称警告
EXAMPLE_DOC_STRING = """ # 定义一个示例文档字符串
Examples: # 示例部分
```py # 示例代码块开始
>>> import torch # 导入 PyTorch 库
>>> from diffusers import AutoPipelineForText2Image # 从 diffusers 导入自动文本到图像管道
>>> pipe = AutoPipelineForText2Image.from_pretrained( # 从预训练模型加载管道
... "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, enable_pag=True # 指定模型路径、数据类型及启用 PAG
... )
>>> pipe = pipe.to("cuda") # 将管道移动到 CUDA 设备
>>> prompt = "a photo of an astronaut riding a horse on mars" # 定义生成图像的提示
>>> image = pipe(prompt, pag_scale=0.3).images[0] # 生成图像并获取第一张图像
```py # 示例代码块结束
"""
# 从稳定扩散管道复制的函数，调整噪声配置
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
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion 中复制的
def retrieve_timesteps(
# 调度器，用于获取时间步
scheduler,
# 可选的推理步骤数量
num_inference_steps: Optional[int] = None,
# 可选的设备，指定时间步移动的目标设备
device: Optional[Union[str, torch.device]] = None,
# 可选的自定义时间步
timesteps: Optional[List[int]] = None,
# 可选的自定义 sigmas
sigmas: Optional[List[float]] = None,
# 其他关键字参数
**kwargs,
):
"""
调用调度器的 `set_timesteps` 方法并在调用后从调度器获取时间步。处理
自定义时间步。所有关键字参数将传递给 `scheduler.set_timesteps`。
参数:
scheduler (`SchedulerMixin`):
用于获取时间步的调度器。
num_inference_steps (`int`):
生成样本时使用的扩散步骤数量。如果使用，则 `timesteps`
必须为 `None`。
device (`str` 或 `torch.device`, *可选*):
时间步应移动到的设备。如果为 `None`，则时间步不移动。
timesteps (`List[int]`, *可选*):
用于覆盖调度器的时间步间隔策略的自定义时间步。如果传入 `timesteps`，
则 `num_inference_steps` 和 `sigmas` 必须为 `None`。
sigmas (`List[float]`, *可选*):
用于覆盖调度器的时间步间隔策略的自定义 sigmas。如果传入 `sigmas`，
则 `num_inference_steps` 和 `timesteps` 必须为 `None`。
返回:
`Tuple[torch.Tensor, int]`: 一个元组，第一个元素是来自调度器的时间步计划，
第二个元素是推理步骤的数量。
"""
# 如果同时传入了自定义时间步和 sigmas，抛出异常
if timesteps is not None and sigmas is not None:
raise ValueError("只能传入 `timesteps` 或 `sigmas` 中的一个。请选择一个以设置自定义值")
# 如果传入了自定义时间步
if timesteps is not None:
# 检查调度器的 `set_timesteps` 方法是否接受自定义时间步
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
# 如果不支持自定义时间步，抛出异常
if not accepts_timesteps:
raise ValueError(
f"当前调度器类 {scheduler.__class__} 的 `set_timesteps` 不支持自定义"
f" 时间步计划。请检查是否使用了正确的调度器。"
)
# 设置调度器的时间步
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
# 从调度器获取设置后的时间步
timesteps = scheduler.timesteps
# 计算推理步骤的数量
num_inference_steps = len(timesteps)
# 如果传入了自定义 sigmas
elif sigmas is not None:
# 检查调度器的 `set_timesteps` 方法是否接受自定义 sigmas
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
# 如果不支持自定义 sigmas，抛出异常
if not accept_sigmas:
raise ValueError(
f"当前调度器类 {scheduler.__class__} 的 `set_timesteps` 不支持自定义"
f" sigmas 时间步计划。请检查是否使用了正确的调度器。"
)
# 设置调度器的 sigmas
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
# 从调度器获取设置后的时间步
timesteps = scheduler.timesteps
# 计算推理步骤的数量
num_inference_steps = len(timesteps)
# 如果不是满足条件的情况，执行以下操作
else:
# 设置推理步骤的时间步，指定设备并传递额外参数
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
# 获取调度器中的时间步列表
timesteps = scheduler.timesteps
# 返回时间步和推理步骤的数量
return timesteps, num_inference_steps
# 定义一个类 StableDiffusionPAGPipeline，继承自多个基类以实现特定功能
class StableDiffusionPAGPipeline(
# 继承自 DiffusionPipeline 基类，提供基本的扩散管道功能
DiffusionPipeline,
# 继承自 StableDiffusionMixin，添加 Stable Diffusion 特有的功能
StableDiffusionMixin,
# 继承自 TextualInversionLoaderMixin，支持文本反演加载功能
TextualInversionLoaderMixin,
# 继承自 StableDiffusionLoraLoaderMixin，支持加载 LoRA 权重
StableDiffusionLoraLoaderMixin,
# 继承自 IPAdapterMixin，支持加载 IP 适配器
IPAdapterMixin,
# 继承自 FromSingleFileMixin，支持从单一文件加载模型
FromSingleFileMixin,
# 继承自 PAGMixin，添加特定于 PAG 的功能
PAGMixin,
):
# 类文档字符串，描述该管道的功能
r"""
Pipeline for text-to-image generation using Stable Diffusion.
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
implemented for all pipelines (downloading, saving, running on a particular device, etc.).
The pipeline also inherits the following loading methods:
- [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
- [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for loading LoRA weights
- [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for saving LoRA weights
- [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
- [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
Args:
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
text_encoder ([`~transformers.CLIPTextModel`]):
Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
tokenizer ([`~transformers.CLIPTokenizer`]):
A `CLIPTokenizer` to tokenize text.
unet ([`UNet2DConditionModel`]):
A `UNet2DConditionModel` to denoise the encoded image latents.
scheduler ([`SchedulerMixin`]):
A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
[`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
safety_checker ([`StableDiffusionSafetyChecker`]):
Classification module that estimates whether generated images could be considered offensive or harmful.
Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
about a model's potential harms.
feature_extractor ([`~transformers.CLIPImageProcessor`]):
A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
"""
# 定义模型在 CPU 上的卸载顺序
model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
# 定义可选组件的列表，包含可能的附加功能
_optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
# 定义不进行 CPU 卸载的组件列表，确保安全检查器始终在 CPU 上
_exclude_from_cpu_offload = ["safety_checker"]
# 定义回调张量输入的列表，指定需要监控的输入数据
_callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
# 初始化类的构造函数，接收多个参数用于模型的配置
def __init__(
self,
# VAE模型，负责生成和重构图像
vae: AutoencoderKL,
# 文本编码器，使用CLIP进行文本嵌入
text_encoder: CLIPTextModel,
# 分词器，将文本转换为模型可处理的格式
tokenizer: CLIPTokenizer,
# UNet模型，用于图像生成
unet: UNet2DConditionModel,
# 调度器，控制生成过程中的采样
scheduler: KarrasDiffusionSchedulers,
# 安全检查器，用于确保生成内容的安全性
safety_checker: StableDiffusionSafetyChecker,
# 特征提取器，处理输入图像
feature_extractor: CLIPImageProcessor,
# 可选的图像编码器，用于图像的额外处理
image_encoder: CLIPVisionModelWithProjection = None,
# 是否需要安全检查的标志
requires_safety_checker: bool = True,
# 应用层的配置，控制图像处理的层次
pag_applied_layers: Union[str, List[str]] = "mid",
# 从稳定扩散管道中复制的函数，编码提示文本
def encode_prompt(
self,
# 输入的文本提示
prompt,
# 设备类型（CPU或GPU）
device,
# 每个提示生成的图像数量
num_images_per_prompt,
# 是否执行分类器自由引导
do_classifier_free_guidance,
# 可选的负面提示
negative_prompt=None,
# 可选的提示嵌入，若已预先计算
prompt_embeds: Optional[torch.Tensor] = None,
# 可选的负面提示嵌入，若已预先计算
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 可选的LoRA缩放参数
lora_scale: Optional[float] = None,
# 可选的跳过CLIP层的参数
clip_skip: Optional[int] = None,
# 从稳定扩散管道中复制的函数，编码输入图像
def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
# 获取图像编码器参数的数据类型
dtype = next(self.image_encoder.parameters()).dtype
# 如果输入不是张量，则通过特征提取器处理
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values
# 将图像转移到指定设备，并设置数据类型
image = image.to(device=device, dtype=dtype)
# 如果需要输出隐藏状态，则进行隐藏状态的编码
if output_hidden_states:
# 编码图像并获取倒数第二层的隐藏状态
image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
# 重复嵌入以匹配图像数量
image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
# 对无条件图像编码
uncond_image_enc_hidden_states = self.image_encoder(
torch.zeros_like(image), output_hidden_states=True
).hidden_states[-2]
# 重复无条件图像嵌入以匹配图像数量
uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
num_images_per_prompt, dim=0
)
# 返回图像和无条件图像的隐藏状态
return image_enc_hidden_states, uncond_image_enc_hidden_states
else:
# 对图像进行编码，获取图像嵌入
image_embeds = self.image_encoder(image).image_embeds
# 重复图像嵌入以匹配图像数量
image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
# 创建与图像嵌入形状相同的零张量作为无条件图像嵌入
uncond_image_embeds = torch.zeros_like(image_embeds)
# 返回图像嵌入和无条件图像嵌入
return image_embeds, uncond_image_embeds
# 从稳定扩散管道中复制的函数，准备图像嵌入
def prepare_ip_adapter_image_embeds(
# 输入的IP适配器图像
self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance
): # 方法结束括号，开始代码块
image_embeds = [] # 初始化图像嵌入列表
if do_classifier_free_guidance: # 检查是否进行分类器自由引导
negative_image_embeds = [] # 初始化负图像嵌入列表
if ip_adapter_image_embeds is None: # 检查 IP 适配器图像嵌入是否为空
if not isinstance(ip_adapter_image, list): # 确保 ip_adapter_image 是列表
ip_adapter_image = [ip_adapter_image] # 将其转换为单元素列表
if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers): # 检查列表长度
raise ValueError( # 引发值错误异常
f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." # 异常信息，显示不匹配的长度
)
for single_ip_adapter_image, image_proj_layer in zip( # 遍历 IP 适配器图像和图像投影层
ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers # 将两者打包在一起
):
output_hidden_state = not isinstance(image_proj_layer, ImageProjection) # 判断是否输出隐藏状态
single_image_embeds, single_negative_image_embeds = self.encode_image( # 编码单个图像并获取嵌入
single_ip_adapter_image, device, 1, output_hidden_state # 调用图像编码方法
)
image_embeds.append(single_image_embeds[None, :]) # 将单个图像嵌入添加到列表
if do_classifier_free_guidance: # 如果进行分类器自由引导
negative_image_embeds.append(single_negative_image_embeds[None, :]) # 添加负图像嵌入
else: # 如果 ip_adapter_image_embeds 不是 None
for single_image_embeds in ip_adapter_image_embeds: # 遍历已有的图像嵌入
if do_classifier_free_guidance: # 如果进行分类器自由引导
single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) # 分割嵌入
negative_image_embeds.append(single_negative_image_embeds) # 添加负图像嵌入
image_embeds.append(single_image_embeds) # 添加图像嵌入
ip_adapter_image_embeds = [] # 初始化最终的 IP 适配器图像嵌入列表
for i, single_image_embeds in enumerate(image_embeds): # 遍历图像嵌入
single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) # 重复嵌入以匹配每个提示的数量
if do_classifier_free_guidance: # 如果进行分类器自由引导
single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) # 重复负图像嵌入
single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], dim=0) # 连接负图像和正图像嵌入
single_image_embeds = single_image_embeds.to(device=device) # 将嵌入移动到指定设备
ip_adapter_image_embeds.append(single_image_embeds) # 添加到最终嵌入列表
return ip_adapter_image_embeds # 返回最终的 IP 适配器图像嵌入
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker # 复制自其他模块的安全检查器方法
# 运行安全检查器，验证图像的内容是否安全
def run_safety_checker(self, image, device, dtype):
# 如果安全检查器未定义，则将不安全内容标志设为 None
if self.safety_checker is None:
has_nsfw_concept = None
else:
# 如果输入的图像是张量，进行后处理以转换为 PIL 格式
if torch.is_tensor(image):
feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
# 如果输入的图像不是张量，将其从 numpy 数组转换为 PIL 格式
else:
feature_extractor_input = self.image_processor.numpy_to_pil(image)
# 使用特征提取器处理图像，并将其转换为设备上的张量
safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
# 运行安全检查器，检查图像并返回处理后的图像和不安全内容标志
image, has_nsfw_concept = self.safety_checker(
images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
)
# 返回处理后的图像和不安全内容标志
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 = {}
# 如果调度器接受 eta，则将其添加到字典中
if accepts_eta:
extra_step_kwargs["eta"] = eta
# 检查调度器的步骤方法是否接受 generator 参数
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 如果调度器接受 generator，则将其添加到字典中
if accepts_generator:
extra_step_kwargs["generator"] = generator
# 返回准备好的额外步骤参数字典
return extra_step_kwargs
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs 复制
def check_inputs(
self,
prompt,
height,
width,
callback_steps,
negative_prompt=None,
prompt_embeds=None,
negative_prompt_embeds=None,
ip_adapter_image=None,
ip_adapter_image_embeds=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
# 从 diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img 导入的函数
def get_guidance_scale_embedding(
self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32
) -> torch.Tensor:
"""
参见 https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
参数:
w (`torch.Tensor`):
使用指定的引导比例生成嵌入向量，以随后丰富时间步嵌入。
embedding_dim (`int`, *可选*, 默认为 512):
生成的嵌入的维度。
dtype (`torch.dtype`, *可选*, 默认为 `torch.float32`):
生成嵌入的数据类型。
返回:
`torch.Tensor`: 形状为 `(len(w), embedding_dim)` 的嵌入向量。
"""
# 确保输入张量是一维的
assert len(w.shape) == 1
# 将 w 乘以 1000.0 进行缩放
w = w * 1000.0
# 计算嵌入的半维度
half_dim = embedding_dim // 2
# 计算用于生成嵌入的基数
emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
# 生成衰减因子
emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
# 根据输入张量和衰减因子生成嵌入
emb = w.to(dtype)[:, None] * emb[None, :]
# 将正弦和余弦值连接到嵌入向量中
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
# 如果嵌入维度为奇数，则进行零填充
if embedding_dim % 2 == 1: # zero pad
emb = torch.nn.functional.pad(emb, (0, 1))
# 确保生成的嵌入形状正确
assert emb.shape == (w.shape[0], embedding_dim)
# 返回生成的嵌入向量
return emb
# 返回引导比例的属性
@property
def guidance_scale(self):
return self._guidance_scale
# 返回引导重标定的属性
@property
def guidance_rescale(self):
return self._guidance_rescale
# 返回跳过剪辑的属性
@property
def clip_skip(self):
return self._clip_skip
# 这里的 `guidance_scale` 按照 Imagen 论文中的引导权重 `w` 定义
# `guidance_scale = 1` 对应于不进行无分类器引导
@property
# 定义一个方法用于无分类器自由引导
def do_classifier_free_guidance(self):
# 检查引导比例是否大于1且时间条件投影维度为None
return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
# 定义一个属性，用于获取交叉注意力的参数
@property
def cross_attention_kwargs(self):
# 返回交叉注意力的参数
return self._cross_attention_kwargs
# 定义一个属性，用于获取时间步数
@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__(
# 输入的提示文本，可以是字符串或字符串列表
prompt: Union[str, List[str]] = None,
# 高度参数，默认为None
height: Optional[int] = None,
# 宽度参数，默认为None
width: Optional[int] = None,
# 推理步骤数量，默认为50
num_inference_steps: int = 50,
# 时间步列表，默认为None
timesteps: List[int] = None,
# Sigma值列表，默认为None
sigmas: List[float] = None,
# 引导比例，默认为7.5
guidance_scale: float = 7.5,
# 负提示文本，可以是字符串或字符串列表，默认为None
negative_prompt: Optional[Union[str, List[str]]] = None,
# 每个提示生成的图像数量，默认为1
num_images_per_prompt: Optional[int] = 1,
# Eta值，默认为0.0
eta: float = 0.0,
# 随机数生成器，默认为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,
# 图像适配器输入，默认为None
ip_adapter_image: Optional[PipelineImageInput] = None,
# 图像适配器嵌入，默认为None
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
# 输出类型，默认为"pil"
output_type: Optional[str] = "pil",
# 是否返回字典，默认为True
return_dict: bool = True,
# 交叉注意力的参数，默认为None
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
# 引导重标定值，默认为0.0
guidance_rescale: float = 0.0,
# 跳过的剪辑次数，默认为None
clip_skip: Optional[int] = None,
# 每步结束时的回调函数，默认为None
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 每步结束时的张量输入回调列表，默认为["latents"]
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# PAG比例，默认为3.0
pag_scale: float = 3.0,
# 自适应PAG比例，默认为0.0
pag_adaptive_scale: float = 0.0,
.\diffusers\pipelines\pag\pipeline_pag_sd_3.py
# 版权声明，指明版权归属及相关许可信息
# Copyright 2024 Stability AI and 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.
# 导入用于检查对象信息的模块
import inspect
# 导入类型提示相关的类型
from typing import Any, Callable, Dict, List, Optional, Union
# 导入 PyTorch 库
import torch
# 导入 transformers 库中的相关模型和分词器
from transformers import (
CLIPTextModelWithProjection, # 导入 CLIP 文本模型
CLIPTokenizer, # 导入 CLIP 分词器
T5EncoderModel, # 导入 T5 编码器模型
T5TokenizerFast, # 导入快速 T5 分词器
)
# 导入自定义图像处理器
from ...image_processor import VaeImageProcessor
# 导入自定义加载器
from ...loaders import FromSingleFileMixin, SD3LoraLoaderMixin
# 导入自定义注意力处理器
from ...models.attention_processor import PAGCFGJointAttnProcessor2_0, PAGJointAttnProcessor2_0
# 导入自定义自动编码器
from ...models.autoencoders import AutoencoderKL
# 导入自定义变换器模型
from ...models.transformers import SD3Transformer2DModel
# 导入自定义调度器
from ...schedulers import FlowMatchEulerDiscreteScheduler
# 导入各种实用工具函数
from ...utils import (
USE_PEFT_BACKEND, # 导入是否使用 PEFT 后端的标识
is_torch_xla_available, # 导入检查 XLA 是否可用的函数
logging, # 导入日志记录模块
replace_example_docstring, # 导入替换示例文档字符串的函数
scale_lora_layers, # 导入缩放 LoRA 层的函数
unscale_lora_layers, # 导入取消缩放 LoRA 层的函数
)
# 导入 PyTorch 相关的实用工具函数
from ...utils.torch_utils import randn_tensor
# 导入扩散管道的实用工具
from ..pipeline_utils import DiffusionPipeline
# 导入稳定扩散的管道输出
from ..stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput
# 导入 PAG 相关的实用工具
from .pag_utils import PAGMixin
# 检查 XLA 是否可用，如果可用则导入 XLA 相关模块
if is_torch_xla_available():
import torch_xla.core.xla_model as xm # 导入 XLA 核心模型
XLA_AVAILABLE = True # 设置 XLA 可用标识为真
else:
XLA_AVAILABLE = False # 设置 XLA 可用标识为假
# 获取当前模块的日志记录器
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，包含使用示例
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> import torch
>>> from diffusers import AutoPipelineForText2Image
>>> pipe = AutoPipelineForText2Image.from_pretrained(
... "stabilityai/stable-diffusion-3-medium-diffusers",
... torch_dtype=torch.float16,
... enable_pag=True,
... pag_applied_layers=["blocks.13"],
... )
>>> pipe.to("cuda") # 将管道移动到 GPU
>>> prompt = "A cat holding a sign that says hello world" # 定义生成图像的提示
>>> image = pipe(prompt, guidance_scale=5.0, pag_scale=0.7).images[0] # 生成图像
>>> image.save("sd3_pag.png") # 保存生成的图像
```py
"""
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion 中复制的函数
def retrieve_timesteps(
scheduler, # 调度器实例
num_inference_steps: Optional[int] = None, # 可选的推理步骤数
device: Optional[Union[str, torch.device]] = None, # 可选的设备参数
timesteps: Optional[List[int]] = None, # 可选的时间步列表
sigmas: Optional[List[float]] = None, # 可选的 sigma 列表
**kwargs, # 其他可选参数
):
"""
调用调度器的 `set_timesteps` 方法，并在调用后从调度器检索时间步。处理自定义时间步。
Any kwargs will be supplied to `scheduler.set_timesteps`.
# 参数说明
Args:
# 调度器，用于获取时间步的类
scheduler (`SchedulerMixin`):
# 从调度器获取时间步
The scheduler to get timesteps from.
# 生成样本时的扩散步骤数量
num_inference_steps (`int`):
# 使用预训练模型生成样本时的扩散步骤数。如果使用，`timesteps` 必须为 `None`
The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
must be `None`.
# 可选参数，指定时间步移动的设备
device (`str` or `torch.device`, *optional*):
# 如果为 `None`，则时间步不会被移动
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
# 可选参数，自定义时间步以覆盖调度器的时间步间隔策略
timesteps (`List[int]`, *optional*):
# 如果传递 `timesteps`，则 `num_inference_steps` 和 `sigmas` 必须为 `None`
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
# 可选参数，自定义 sigmas 以覆盖调度器的时间步间隔策略
sigmas (`List[float]`, *optional*):
# 如果传递 `sigmas`，则 `num_inference_steps` 和 `timesteps` 必须为 `None`
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
# 返回一个元组，包含时间步调度和推理步骤数量
Returns:
`Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
# 如果同时传递了时间步和 sigmas，抛出错误
if timesteps is not None and sigmas is not None:
raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
# 如果传递了时间步
if timesteps is not None:
# 检查当前调度器是否支持时间步
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
# 如果不支持，抛出错误
if not accepts_timesteps:
raise ValueError(
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
f" timestep schedules. Please check whether you are using the correct scheduler."
)
# 设置调度器的时间步
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
# 从调度器获取当前时间步
timesteps = scheduler.timesteps
# 计算推理步骤数量
num_inference_steps = len(timesteps)
# 如果传递了 sigmas
elif sigmas is not None:
# 检查当前调度器是否支持 sigmas
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
# 如果不支持，抛出错误
if not accept_sigmas:
raise ValueError(
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
f" sigmas schedules. Please check whether you are using the correct scheduler."
)
# 设置调度器的 sigmas
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
# 从调度器获取当前时间步
timesteps = scheduler.timesteps
# 计算推理步骤数量
num_inference_steps = len(timesteps)
# 如果没有传递时间步或 sigmas
else:
# 设置调度器的时间步数量
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
# 从调度器获取当前时间步
timesteps = scheduler.timesteps
# 返回时间步和推理步骤数量
return timesteps, num_inference_steps
# 定义一个名为 StableDiffusion3PAGPipeline 的类，继承自多个基类
class StableDiffusion3PAGPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, PAGMixin):
r"""
# 文档字符串，描述该类的功能和参数
[PAG pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/pag) for text-to-image generation
using Stable Diffusion 3.
Args:
transformer ([`SD3Transformer2DModel`]):
Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
scheduler ([`FlowMatchEulerDiscreteScheduler`]):
A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
text_encoder ([`CLIPTextModelWithProjection`]):
[CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant,
with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size`
as its dimension.
text_encoder_2 ([`CLIPTextModelWithProjection`]):
[CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
specifically the
[laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
variant.
text_encoder_3 ([`T5EncoderModel`]):
Frozen text-encoder. Stable Diffusion 3 uses
[T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the
[t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
tokenizer (`CLIPTokenizer`):
Tokenizer of class
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
tokenizer_2 (`CLIPTokenizer`):
Second Tokenizer of class
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
tokenizer_3 (`T5TokenizerFast`):
Tokenizer of class
[T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
"""
# 定义一个模型的 CPU 卸载顺序，指定各组件之间的依赖关系
model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->transformer->vae"
# 定义可选组件的空列表，可能在子类中进行扩展
_optional_components = []
# 定义回调张量输入的名称列表，这些输入将在处理过程中被回调
_callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"]
# 初始化类的构造函数
def __init__(
self,
transformer: SD3Transformer2DModel, # 输入的 2D 转换器模型
scheduler: FlowMatchEulerDiscreteScheduler, # 调度器，用于控制模型的生成过程
vae: AutoencoderKL, # 自编码器，用于生成图像
text_encoder: CLIPTextModelWithProjection, # 文本编码器，将文本转为向量
tokenizer: CLIPTokenizer, # 文本分词器，将文本拆分为标记
text_encoder_2: CLIPTextModelWithProjection, # 第二个文本编码器
tokenizer_2: CLIPTokenizer, # 第二个文本分词器
text_encoder_3: T5EncoderModel, # 第三个文本编码器
tokenizer_3: T5TokenizerFast, # 第三个文本分词器
pag_applied_layers: Union[str, List[str]] = "blocks.1", # 默认应用于第一个变换层
):
# 调用父类的构造函数
super().__init__()
# 注册各个模块，使其可以在模型中使用
self.register_modules(
vae=vae, # 注册自编码器
text_encoder=text_encoder, # 注册文本编码器
text_encoder_2=text_encoder_2, # 注册第二个文本编码器
text_encoder_3=text_encoder_3, # 注册第三个文本编码器
tokenizer=tokenizer, # 注册文本分词器
tokenizer_2=tokenizer_2, # 注册第二个文本分词器
tokenizer_3=tokenizer_3, # 注册第三个文本分词器
transformer=transformer, # 注册转换器
scheduler=scheduler, # 注册调度器
)
# 根据 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.tokenizer_max_length = (
self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
)
# 获取转换器的默认样本大小
self.default_sample_size = (
self.transformer.config.sample_size
if hasattr(self, "transformer") and self.transformer is not None
else 128
)
# 设置应用于 PAG 的层以及注意力处理器
self.set_pag_applied_layers(
pag_applied_layers, pag_attn_processors=(PAGCFGJointAttnProcessor2_0(), PAGJointAttnProcessor2_0())
)
# 从其他模块复制的函数，用于获取 T5 提示嵌入
def _get_t5_prompt_embeds(
self,
prompt: Union[str, List[str]] = None, # 输入的提示，字符串或字符串列表
num_images_per_prompt: int = 1, # 每个提示生成的图像数量
max_sequence_length: int = 256, # 最大序列长度
device: Optional[torch.device] = None, # 指定设备（如 CPU 或 GPU）
dtype: Optional[torch.dtype] = None, # 指定数据类型
):
# 如果没有指定设备，则使用实例的执行设备
device = device or self._execution_device
# 如果没有指定数据类型，则使用文本编码器的数据类型
dtype = dtype or self.text_encoder.dtype
# 如果提示词是字符串，则将其放入列表中，否则保持原样
prompt = [prompt] if isinstance(prompt, str) else prompt
# 获取提示词的批处理大小
batch_size = len(prompt)
# 如果文本编码器未初始化，则返回全零的张量
if self.text_encoder_3 is None:
return torch.zeros(
(
# 张量的第一维为批处理大小乘以每个提示生成的图像数量
batch_size * num_images_per_prompt,
# 张量的第二维为最大令牌长度
self.tokenizer_max_length,
# 张量的第三维为变换器的联合注意力维度
self.transformer.config.joint_attention_dim,
),
# 指定设备
device=device,
# 指定数据类型
dtype=dtype,
)
# 使用 tokenizer_3 编码提示词，返回张量形式的输入
text_inputs = self.tokenizer_3(
prompt,
# 填充到最大长度
padding="max_length",
# 设置最大序列长度
max_length=max_sequence_length,
# 截断超出部分
truncation=True,
# 添加特殊令牌
add_special_tokens=True,
# 返回 PyTorch 张量
return_tensors="pt",
)
# 获取编码后的输入 ID
text_input_ids = text_inputs.input_ids
# 获取未截断的输入 ID
untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
# 检查未截断的 ID 是否比截断的 ID 更长且不相等
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
# 解码被截断的部分并记录警告
removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
logger.warning(
"The following part of your input was truncated because `max_sequence_length` is set to "
f" {max_sequence_length} tokens: {removed_text}"
)
# 将输入 ID 转移到指定设备并通过文本编码器获得提示嵌入
prompt_embeds = self.text_encoder_3(text_input_ids.to(device))[0]
# 获取文本编码器的数据类型
dtype = self.text_encoder_3.dtype
# 将提示嵌入转换为指定的数据类型和设备
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
# 解包提示嵌入的形状信息
_, seq_len, _ = prompt_embeds.shape
# 为每个生成的图像复制文本嵌入和注意力掩码，使用适合 mps 的方法
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
# 重塑张量以符合(batch_size * num_images_per_prompt, seq_len, -1)的形状
prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
# 返回处理后的提示嵌入
return prompt_embeds
# 从 diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds 复制
def _get_clip_prompt_embeds(
self,
# 输入的提示词，可以是字符串或字符串列表
prompt: Union[str, List[str]],
# 每个提示生成的图像数量，默认为 1
num_images_per_prompt: int = 1,
# 可选的设备参数
device: Optional[torch.device] = None,
# 可选的跳过参数
clip_skip: Optional[int] = None,
# 使用的 clip 模型索引，默认为 0
clip_model_index: int = 0,
):
# 如果未指定设备，则使用执行环境中的默认设备
device = device or self._execution_device
# 初始化 CLIP 模型的分词器和文本编码器列表
clip_tokenizers = [self.tokenizer, self.tokenizer_2]
clip_text_encoders = [self.text_encoder, self.text_encoder_2]
# 根据索引选择当前使用的分词器和文本编码器
tokenizer = clip_tokenizers[clip_model_index]
text_encoder = clip_text_encoders[clip_model_index]
# 将提示文本转换为列表形式（如果输入为字符串则转为单元素列表），确定批处理大小
prompt = [prompt] if isinstance(prompt, str) else prompt
batch_size = len(prompt)
# 使用指定的分词器对提示文本进行编码，设置最大长度和填充方式，返回 PyTorch 张量
text_inputs = tokenizer(
prompt,
padding="max_length",
max_length=self.tokenizer_max_length,
truncation=True,
return_tensors="pt",
)
# 获取输入文本的 ID
text_input_ids = text_inputs.input_ids
# 获取未截断的文本 ID，并进行比较，如果存在截断则记录警告信息
untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
removed_text = tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
logger.warning(
"The following part of your input was truncated because CLIP can only handle sequences up to"
f" {self.tokenizer_max_length} tokens: {removed_text}"
)
# 使用文本编码器对输入文本进行编码，输出包含隐藏状态
prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
pooled_prompt_embeds = prompt_embeds[0]
# 根据 clip_skip 参数选择合适的隐藏状态
if clip_skip is None:
prompt_embeds = prompt_embeds.hidden_states[-2]
else:
prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
# 将编码结果转换为指定的数据类型和设备类型
prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
# 获取文本编码的序列长度
_, seq_len, _ = prompt_embeds.shape
# 使用 MPS 友好的方法复制文本嵌入以生成每个提示的生成结果
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
# 复制池化后的文本嵌入以生成每个提示的生成结果
pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1)
pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
# 返回生成的提示文本嵌入结果
return prompt_embeds, pooled_prompt_embeds
# 从 diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt 复制过来的
# 定义编码提示的函数，接受多个参数以支持不同的提示输入
def encode_prompt(
# 第一个提示，可以是字符串或字符串列表
self,
prompt: Union[str, List[str]],
# 第二个提示，可以是字符串或字符串列表
prompt_2: Union[str, List[str]],
# 第三个提示，可以是字符串或字符串列表
prompt_3: Union[str, List[str]],
# 可选参数，指定设备（如 GPU 或 CPU）
device: Optional[torch.device] = None,
# 每个提示生成的图像数量，默认是 1
num_images_per_prompt: int = 1,
# 是否进行无分类器引导，默认是 True
do_classifier_free_guidance: bool = True,
# 可选的负面提示，可以是字符串或字符串列表
negative_prompt: Optional[Union[str, List[str]]] = None,
# 可选的第二个负面提示，可以是字符串或字符串列表
negative_prompt_2: Optional[Union[str, List[str]]] = None,
# 可选的第三个负面提示，可以是字符串或字符串列表
negative_prompt_3: Optional[Union[str, List[str]]] = None,
# 可选的提示嵌入，类型为浮点张量
prompt_embeds: Optional[torch.FloatTensor] = None,
# 可选的负面提示嵌入，类型为浮点张量
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
# 可选的池化提示嵌入，类型为浮点张量
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
# 可选的负面池化提示嵌入，类型为浮点张量
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
# 可选参数，指定跳过的 CLIP 层数
clip_skip: Optional[int] = None,
# 最大序列长度，默认是 256
max_sequence_length: int = 256,
# 可选的 LORA 缩放因子
lora_scale: Optional[float] = None,
# 从 diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.check_inputs 复制而来
# 定义输入检查的函数，确保输入的有效性
def check_inputs(
self,
# 第一个提示
prompt,
# 第二个提示
prompt_2,
# 第三个提示
prompt_3,
# 图像高度
height,
# 图像宽度
width,
# 可选的负面提示
negative_prompt=None,
# 可选的第二个负面提示
negative_prompt_2=None,
# 可选的第三个负面提示
negative_prompt_3=None,
# 可选的提示嵌入
prompt_embeds=None,
# 可选的负面提示嵌入
negative_prompt_embeds=None,
# 可选的池化提示嵌入
pooled_prompt_embeds=None,
# 可选的负面池化提示嵌入
negative_pooled_prompt_embeds=None,
# 可选的回调，用于步骤结束时的张量输入
callback_on_step_end_tensor_inputs=None,
# 可选的最大序列长度
max_sequence_length=None,
# 从 diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_latents 复制而来
# 定义准备潜在张量的函数，生成潜在张量
def prepare_latents(
self,
# 批次大小
batch_size,
# 潜在张量的通道数
num_channels_latents,
# 图像高度
height,
# 图像宽度
width,
# 数据类型
dtype,
# 设备类型
device,
# 随机数生成器
generator,
# 可选的潜在张量
latents=None,
):
# 如果给定了潜在张量，则将其转移到指定设备和数据类型
if latents is not None:
return latents.to(device=device, dtype=dtype)
# 计算潜在张量的形状
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."
)
# 生成随机潜在张量
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# 返回生成的潜在张量
return latents
# 定义一个只读属性，用于获取引导缩放因子
@property
def guidance_scale(self):
return self._guidance_scale
# 定义一个只读属性，用于获取 CLIP 跳过的层数
@property
def clip_skip(self):
return self._clip_skip
# 定义一个只读属性，检查是否进行无分类器引导
# 依据 Imagen 论文中的公式定义引导缩放
# guidance_scale = 1 表示没有进行无分类器引导
@property
def do_classifier_free_guidance(self):
return self._guidance_scale > 1
# 定义一个只读属性
# 定义获取联合注意力参数的方法
def joint_attention_kwargs(self):
# 返回类属性 _joint_attention_kwargs 的值
return self._joint_attention_kwargs
# 定义 num_timesteps 属性
@property
def num_timesteps(self):
# 返回类属性 _num_timesteps 的值
return self._num_timesteps
# 定义 interrupt 属性
@property
def interrupt(self):
# 返回类属性 _interrupt 的值
return self._interrupt
# 禁用梯度计算并替换示例文档字符串
@torch.no_grad()
@replace_example_docstring(EXAMPLE_DOC_STRING)
# 定义类的可调用方法
def __call__(
# 定义输入提示，默认为 None
prompt: Union[str, List[str]] = None,
# 定义第二个提示，默认为 None
prompt_2: Optional[Union[str, List[str]]] = None,
# 定义第三个提示，默认为 None
prompt_3: Optional[Union[str, List[str]]] = None,
# 定义图像高度，默认为 None
height: Optional[int] = None,
# 定义图像宽度，默认为 None
width: Optional[int] = None,
# 定义推理步骤数，默认为 28
num_inference_steps: int = 28,
# 定义时间步列表，默认为 None
timesteps: List[int] = None,
# 定义引导比例，默认为 7.0
guidance_scale: float = 7.0,
# 定义负提示，默认为 None
negative_prompt: Optional[Union[str, List[str]]] = None,
# 定义第二个负提示，默认为 None
negative_prompt_2: Optional[Union[str, List[str]]] = None,
# 定义第三个负提示，默认为 None
negative_prompt_3: Optional[Union[str, List[str]]] = None,
# 定义每个提示生成的图像数量，默认为 1
num_images_per_prompt: Optional[int] = 1,
# 定义生成器，默认为 None
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 定义潜在变量，默认为 None
latents: Optional[torch.FloatTensor] = None,
# 定义提示嵌入，默认为 None
prompt_embeds: Optional[torch.FloatTensor] = None,
# 定义负提示嵌入，默认为 None
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
# 定义池化提示嵌入，默认为 None
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
# 定义负池化提示嵌入，默认为 None
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
# 定义输出类型，默认为 "pil"
output_type: Optional[str] = "pil",
# 定义是否返回字典，默认为 True
return_dict: bool = True,
# 定义联合注意力参数，默认为 None
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
# 定义跳过剪辑的数量，默认为 None
clip_skip: Optional[int] = None,
# 定义步骤结束时的回调，默认为 None
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 定义步骤结束时张量输入的回调，默认为 ["latents"]
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# 定义最大序列长度，默认为 256
max_sequence_length: int = 256,
# 定义页面缩放比例，默认为 3.0
pag_scale: float = 3.0,
# 定义自适应页面缩放比例，默认为 0.0
pag_adaptive_scale: float = 0.0,
.\diffusers\pipelines\pag\pipeline_pag_sd_animatediff.py
# 版权声明，标明版权归 HuggingFace 团队所有
#
# 根据 Apache 许可证第 2.0 版授权（“许可证”）；
# 除非遵守许可证，否则您不得使用此文件。
# 您可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，软件按“原样”分发，
# 不提供任何形式的保证或条件，无论是明示还是暗示。
# 请参见许可证以获取管理权限和
# 限制的具体条款。
import inspect # 导入 inspect 模块，用于获取活跃对象的信息
from typing import Any, Callable, Dict, List, Optional, Union # 导入类型提示，用于增强代码的可读性和可维护性
import torch # 导入 PyTorch 库，用于深度学习和张量计算
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection # 导入 Transformers 库中的 CLIP 相关模型和处理器
from ...image_processor import PipelineImageInput # 从自定义模块导入图像处理输入类
from ...loaders import IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin # 导入加载器混合类，用于处理不同模型加载
from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel, UNetMotionModel # 导入不同模型类
from ...models.lora import adjust_lora_scale_text_encoder # 导入函数，用于调整文本编码器的 LoRA 比例
from ...models.unets.unet_motion_model import MotionAdapter # 导入运动适配器类，用于处理动态模型
from ...schedulers import KarrasDiffusionSchedulers # 导入 Karras 扩散调度器类
from ...utils import ( # 导入多个工具函数和常量
USE_PEFT_BACKEND, # 用于指示是否使用 PEFT 后端的常量
logging, # 导入日志模块
replace_example_docstring, # 导入替换示例文档字符串的函数
scale_lora_layers, # 导入用于缩放 LoRA 层的函数
unscale_lora_layers, # 导入用于取消缩放 LoRA 层的函数
)
from ...utils.torch_utils import randn_tensor # 从工具模块导入生成随机张量的函数
from ...video_processor import VideoProcessor # 导入视频处理器类
from ..animatediff.pipeline_output import AnimateDiffPipelineOutput # 导入动画扩散管道输出类
from ..free_init_utils import FreeInitMixin # 导入自由初始化混合类
from ..free_noise_utils import AnimateDiffFreeNoiseMixin # 导入动画扩散自由噪声混合类
from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin # 导入扩散管道和稳定扩散混合类
from .pag_utils import PAGMixin # 导入 PAG 混合类
logger = logging.get_logger(__name__) # 创建一个记录器，用于日志记录，使用模块的名称
EXAMPLE_DOC_STRING = """ # 定义示例文档字符串的开始
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
```py # 示例文档字符串的结束
``` # 示例文档字符串的结束
Examples:
```py
# 导入 PyTorch 库
>>> import torch
# 从 diffusers 库导入相关的类
>>> from diffusers import AnimateDiffPAGPipeline, MotionAdapter, DDIMScheduler
# 从 diffusers.utils 导入 GIF 导出工具
>>> from diffusers.utils import export_to_gif
# 定义模型的 ID，用于加载预训练模型
>>> model_id = "SG161222/Realistic_Vision_V5.1_noVAE"
# 定义运动适配器的 ID，用于加载相应的适配器
>>> motion_adapter_id = "guoyww/animatediff-motion-adapter-v1-5-2"
# 从预训练模型加载运动适配器
>>> motion_adapter = MotionAdapter.from_pretrained(motion_adapter_id)
# 从预训练模型加载调度器，设置调度参数
>>> scheduler = DDIMScheduler.from_pretrained(
... model_id, subfolder="scheduler", beta_schedule="linear", steps_offset=1, clip_sample=False
... )
# 从预训练模型加载动画差异管道，并将其移动到 CUDA 设备
>>> pipe = AnimateDiffPAGPipeline.from_pretrained(
... model_id,
... motion_adapter=motion_adapter,
... scheduler=scheduler,
... pag_applied_layers=["mid"],
... torch_dtype=torch.float16,
... ).to("cuda")
# 生成视频，设置提示词和参数
>>> video = pipe(
... prompt="car, futuristic cityscape with neon lights, street, no human",
... negative_prompt="low quality, bad quality",
... num_inference_steps=25,
... guidance_scale=6.0,
... pag_scale=3.0,
... generator=torch.Generator().manual_seed(42),
... ).frames[0] # 获取生成的第一帧
# 导出生成的视频为 GIF 格式
>>> export_to_gif(video, "animatediff_pag.gif")
```
# 定义 AnimateDiffPAGPipeline 类，继承自多个基类以实现文本到视频的生成
class AnimateDiffPAGPipeline(
# 继承自 DiffusionPipeline 基类
DiffusionPipeline,
# 继承自 StableDiffusionMixin 以便于稳定扩散相关功能
StableDiffusionMixin,
# 继承自 TextualInversionLoaderMixin 用于加载文本反演嵌入
TextualInversionLoaderMixin,
# 继承自 IPAdapterMixin 用于加载 IP 适配器
IPAdapterMixin,
# 继承自 StableDiffusionLoraLoaderMixin 用于处理 LoRA 权重
StableDiffusionLoraLoaderMixin,
# 继承自 FreeInitMixin 用于初始化功能
FreeInitMixin,
# 继承自 AnimateDiffFreeNoiseMixin 以处理动画差异的噪声
AnimateDiffFreeNoiseMixin,
# 继承自 PAGMixin 用于应用 Perturbed Attention Guidance
PAGMixin,
):
r"""
文本到视频生成的管道，使用
[AnimateDiff](https://huggingface.co/docs/diffusers/en/api/pipelines/animatediff) 和 [Perturbed Attention
Guidance](https://huggingface.co/docs/diffusers/en/using-diffusers/pag)。
该模型继承自 [`DiffusionPipeline`]。请查阅超类文档，了解所有管道实现的通用方法
（下载、保存、在特定设备上运行等）。
该管道还继承了以下加载方法：
- [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] 用于加载文本反演嵌入
- [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] 用于加载 LoRA 权重
- [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] 用于保存 LoRA 权重
- [`~loaders.IPAdapterMixin.load_ip_adapter`] 用于加载 IP 适配器
参数：
vae ([`AutoencoderKL`]):
变分自编码器（VAE）模型，用于编码和解码图像的潜在表示。
text_encoder ([`CLIPTextModel`]):
冻结的文本编码器（[clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)）。
tokenizer (`CLIPTokenizer`):
[`~transformers.CLIPTokenizer`] 用于文本标记化。
unet ([`UNet2DConditionModel`]):
[`UNet2DConditionModel`] 用于创建 UNetMotionModel 来去噪编码的视频潜在。
motion_adapter ([`MotionAdapter`]):
[`MotionAdapter`]，与 `unet` 结合使用以去噪编码的视频潜在。
scheduler ([`SchedulerMixin`]):
与 `unet` 结合使用以去噪编码的图像潜在的调度器，可以是
[`DDIMScheduler`], [`LMSDiscreteScheduler`] 或 [`PNDMScheduler`]。
"""
# 定义模型在 CPU 上的卸载顺序
model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
# 定义可选组件列表
_optional_components = ["feature_extractor", "image_encoder", "motion_adapter"]
# 定义需要回调的张量输入
_callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
# 初始化方法，定义构造函数参数
def __init__(
# VAE 模型，用于图像的编码和解码
vae: AutoencoderKL,
# 文本编码器，用于处理输入文本
text_encoder: CLIPTextModel,
# 文本标记化器
tokenizer: CLIPTokenizer,
# UNet 模型，处理条件模型或运动模型
unet: Union[UNet2DConditionModel, UNetMotionModel],
# 动作适配器，用于去噪处理
motion_adapter: MotionAdapter,
# 调度器，用于模型的调度控制
scheduler: KarrasDiffusionSchedulers,
# 可选的特征提取器，默认为 None
feature_extractor: CLIPImageProcessor = None,
# 可选的图像编码器，默认为 None
image_encoder: CLIPVisionModelWithProjection = None,
# 应用的 PAG 层，可以是字符串或字符串列表
pag_applied_layers: Union[str, List[str]] = "mid_block.*attn1", # ["mid"], ["down_blocks.1"]
):
# 调用父类的初始化方法
super().__init__()
# 检查给定的 unet 是否为 UNet2DConditionModel 类型
if isinstance(unet, UNet2DConditionModel):
# 从 UNet2DConditionModel 创建 UNetMotionModel 实例，并传入 motion_adapter
unet = UNetMotionModel.from_unet2d(unet, motion_adapter)
# 注册多个模块，便于后续使用
self.register_modules(
# 注册变换自编码器模块
vae=vae,
# 注册文本编码器模块
text_encoder=text_encoder,
# 注册分词器模块
tokenizer=tokenizer,
# 注册 UNet 模块
unet=unet,
# 注册运动适配器模块
motion_adapter=motion_adapter,
# 注册调度器模块
scheduler=scheduler,
# 注册特征提取器模块
feature_extractor=feature_extractor,
# 注册图像编码器模块
image_encoder=image_encoder,
)
# 计算 VAE 的缩放因子，基于 VAE 配置中的块输出通道数量
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
# 初始化视频处理器，不进行缩放，并设置 VAE 缩放因子
self.video_processor = VideoProcessor(do_resize=False, vae_scale_factor=self.vae_scale_factor)
# 设置应用于 PAG 的层
self.set_pag_applied_layers(pag_applied_layers)
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt 复制而来，num_images_per_prompt -> num_videos_per_prompt
def encode_prompt(
self,
# 输入的提示文本
prompt,
# 指定设备（CPU 或 GPU）
device,
# 每个提示生成的图像数量
num_images_per_prompt,
# 是否进行分类器自由引导
do_classifier_free_guidance,
# 可选的负面提示文本
negative_prompt=None,
# 可选的提示嵌入张量
prompt_embeds: Optional[torch.Tensor] = None,
# 可选的负面提示嵌入张量
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 可选的 Lora 缩放因子
lora_scale: Optional[float] = None,
# 可选的剪辑跳过参数
clip_skip: Optional[int] = None,
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image 复制而来
def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
# 获取图像编码器参数的数据类型
dtype = next(self.image_encoder.parameters()).dtype
# 如果输入的图像不是张量，则进行特征提取
if not isinstance(image, torch.Tensor):
# 使用特征提取器将图像转换为张量，并返回其像素值
image = self.feature_extractor(image, return_tensors="pt").pixel_values
# 将图像张量移动到指定设备，并设置数据类型
image = image.to(device=device, dtype=dtype)
# 如果需要输出隐藏状态
if output_hidden_states:
# 获取图像编码器的隐藏状态的倒数第二个输出
image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
# 根据每个提示的图像数量重复隐藏状态
image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
# 获取无条件图像的隐藏状态
uncond_image_enc_hidden_states = self.image_encoder(
torch.zeros_like(image), output_hidden_states=True
).hidden_states[-2]
# 同样重复无条件隐藏状态
uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
num_images_per_prompt, dim=0
)
# 返回图像和无条件图像的隐藏状态
return image_enc_hidden_states, uncond_image_enc_hidden_states
else:
# 获取图像编码器的图像嵌入
image_embeds = self.image_encoder(image).image_embeds
# 根据每个提示的图像数量重复图像嵌入
image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
# 创建与图像嵌入形状相同的全零无条件图像嵌入
uncond_image_embeds = torch.zeros_like(image_embeds)
# 返回图像嵌入和无条件图像嵌入
return image_embeds, uncond_image_embeds
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds 复制而来
def prepare_ip_adapter_image_embeds(
# 输入的适配器图像
self, ip_adapter_image,
# 输入的适配器图像嵌入
ip_adapter_image_embeds,
# 指定设备（CPU 或 GPU）
device,
# 每个提示生成的图像数量
num_images_per_prompt,
# 是否进行分类器自由引导
do_classifier_free_guidance
# 处理图像嵌入和分类器自由引导的逻辑
):
# 初始化图像嵌入列表
image_embeds = []
# 如果启用分类器自由引导，初始化负图像嵌入列表
if do_classifier_free_guidance:
negative_image_embeds = []
# 如果输入适配器图像嵌入为空
if ip_adapter_image_embeds is None:
# 如果输入适配器图像不是列表，则将其转换为列表
if not isinstance(ip_adapter_image, list):
ip_adapter_image = [ip_adapter_image]
# 检查输入适配器图像的数量是否与 IP 适配器的数量一致
if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers):
# 抛出值错误，说明输入不匹配
raise ValueError(
f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
)
# 遍历输入适配器图像和对应的图像投影层
for single_ip_adapter_image, image_proj_layer in zip(
ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers
):
# 确定输出隐藏状态的标志
output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
# 编码单个图像，获取图像嵌入和负图像嵌入
single_image_embeds, single_negative_image_embeds = self.encode_image(
single_ip_adapter_image, device, 1, output_hidden_state
)
# 将单个图像嵌入添加到列表中
image_embeds.append(single_image_embeds[None, :])
# 如果启用分类器自由引导，将负图像嵌入添加到列表中
if do_classifier_free_guidance:
negative_image_embeds.append(single_negative_image_embeds[None, :])
else:
# 遍历已有的图像嵌入
for single_image_embeds in ip_adapter_image_embeds:
# 如果启用分类器自由引导，将嵌入分成负图像和图像嵌入
if do_classifier_free_guidance:
single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
negative_image_embeds.append(single_negative_image_embeds)
# 将图像嵌入添加到列表中
image_embeds.append(single_image_embeds)
# 初始化适配器图像嵌入列表
ip_adapter_image_embeds = []
# 遍历图像嵌入列表，并根据每个提示的图像数量进行重复
for i, single_image_embeds in enumerate(image_embeds):
single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
# 如果启用分类器自由引导，将负图像嵌入重复并连接
if do_classifier_free_guidance:
single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0)
single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], dim=0)
# 将图像嵌入移至指定设备
single_image_embeds = single_image_embeds.to(device=device)
# 将处理后的嵌入添加到适配器图像嵌入列表中
ip_adapter_image_embeds.append(single_image_embeds)
# 返回最终的适配器图像嵌入列表
return ip_adapter_image_embeds
# 从 diffusers.pipelines.animatediff.pipeline_animatediff.AnimateDiffPipeline.decode_latents 复制的代码
# 解码潜在表示，返回解码后的视频张量
def decode_latents(self, latents, decode_chunk_size: int = 16):
# 根据配置的缩放因子对潜在表示进行缩放
latents = 1 / self.vae.config.scaling_factor * latents
# 获取潜在表示的批次大小、通道数、帧数、高度和宽度
batch_size, channels, num_frames, height, width = latents.shape
# 重新排列和调整潜在表示的形状，方便后续处理
latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width)
video = [] # 初始化视频列表以存储解码结果
# 按照解码块大小遍历潜在表示
for i in range(0, latents.shape[0], decode_chunk_size):
# 获取当前块的潜在表示
batch_latents = latents[i : i + decode_chunk_size]
# 解码当前块的潜在表示，并提取样本
batch_latents = self.vae.decode(batch_latents).sample
# 将解码后的块添加到视频列表
video.append(batch_latents)
# 将所有解码块连接成一个张量
video = torch.cat(video)
# 重新调整视频张量的形状，便于后续处理
video = video[None, :].reshape((batch_size, num_frames, -1) + video.shape[2:]).permute(0, 2, 1, 3, 4)
# 将视频张量转换为 float32 类型，以确保兼容性
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 论文中的 η，范围应在 [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
# 从 diffusers.pipelines.pia.pipeline_pia.PIAPipeline.check_inputs 复制
def check_inputs(
self,
prompt,
height,
width,
negative_prompt=None,
prompt_embeds=None,
negative_prompt_embeds=None,
ip_adapter_image=None,
ip_adapter_image_embeds=None,
callback_on_step_end_tensor_inputs=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}.")
# 检查回调输入是否存在且不全在已注册的回调输入中
if callback_on_step_end_tensor_inputs is not None and not all(
k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
):
# 抛出异常，如果回调输入不在已注册的输入中
raise ValueError(
f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
)
# 检查是否同时提供了 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."
)
# 检查是否两个参数都未提供
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}."
)
# 检查是否同时提供了 ip_adapter_image 和 ip_adapter_image_embeds
if ip_adapter_image is not None and ip_adapter_image_embeds is not None:
# 抛出异常，提示只能提供其中一个
raise ValueError(
"Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."
)
# 检查 ip_adapter_image_embeds 是否存在
if ip_adapter_image_embeds is not None:
# 检查其类型是否为列表
if not isinstance(ip_adapter_image_embeds, list):
# 抛出异常，显示实际类型
raise ValueError(
f"`ip_adapter_image_embeds` has to be of type `list` but is {type(ip_adapter_image_embeds)}"
)
# 检查列表中第一个元素的维度是否为 3D 或 4D
elif ip_adapter_image_embeds[0].ndim not in [3, 4]:
# 抛出异常，显示实际维度
raise ValueError(
f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
)
# 从 diffusers.pipelines.animatediff.pipeline_animatediff.AnimateDiffPipeline.prepare_latents 复制的代码
# 准备潜在变量的方法，接收多个参数以配置潜在变量的生成
def prepare_latents(
self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None
):
# 如果启用了 FreeNoise，按照 [FreeNoise](https://arxiv.org/abs/2310.15169) 的公式 (7) 生成潜在变量
if self.free_noise_enabled:
# 调用 _prepare_latents_free_noise 方法生成潜在变量
latents = self._prepare_latents_free_noise(
batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents
)
# 检查生成器的类型和数量是否与请求的批量大小匹配
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."
)
# 定义潜在变量的形状，依据批量大小和其他参数计算
shape = (
batch_size,
num_channels_latents,
num_frames,
height // self.vae_scale_factor,
width // self.vae_scale_factor,
)
# 如果未提供潜在变量，则生成随机的潜在变量
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 clip_skip(self):
return self._clip_skip
# 定义一个属性，判断是否执行无分类器引导
# `guidance_scale` 类似于 Imagen 论文中公式 (2) 的引导权重 `w`
# `guidance_scale = 1` 表示不进行分类器自由引导
@property
def do_classifier_free_guidance(self):
return self._guidance_scale > 1
# 定义一个属性，返回交叉注意力的关键字参数
@property
def cross_attention_kwargs(self):
return self._cross_attention_kwargs
# 定义一个属性，返回时间步数的数量
@property
def num_timesteps(self):
return self._num_timesteps
# 禁用梯度计算，以减少内存消耗和提高计算效率
@torch.no_grad()
# 替换示例文档字符串为预定义的文档字符串
@replace_example_docstring(EXAMPLE_DOC_STRING)
# 定义可调用对象的特殊方法，用于执行某些操作
def __call__(
# 接收的提示信息，可以是字符串或字符串列表，默认为 None
self,
prompt: Optional[Union[str, List[str]]] = None,
# 要生成的帧数，默认为 16
num_frames: Optional[int] = 16,
# 输出图像的高度，默认为 None
height: Optional[int] = None,
# 输出图像的宽度，默认为 None
width: Optional[int] = None,
# 推理步骤的数量，默认为 50
num_inference_steps: int = 50,
# 指导比例，默认为 7.5，影响生成图像的质量
guidance_scale: float = 7.5,
# 负向提示信息，可以是字符串或字符串列表，默认为 None
negative_prompt: Optional[Union[str, List[str]]] = None,
# 每个提示生成的视频数量，默认为 1
num_videos_per_prompt: Optional[int] = 1,
# 控制噪声的强度，默认为 0.0
eta: float = 0.0,
# 随机数生成器，可以是单个或多个 torch.Generator，默认为 None
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 生成的潜在变量，可以是 torch.Tensor，默认为 None
latents: Optional[torch.Tensor] = None,
# 提示的嵌入表示，可以是 torch.Tensor，默认为 None
prompt_embeds: Optional[torch.Tensor] = None,
# 负向提示的嵌入表示，可以是 torch.Tensor，默认为 None
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 图像输入的适配器图像，默认为 None
ip_adapter_image: Optional[PipelineImageInput] = None,
# 适配器图像的嵌入表示，可以是 torch.Tensor 列表，默认为 None
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
# 输出类型，默认为 "pil"，表示返回 PIL 图像
output_type: Optional[str] = "pil",
# 是否返回字典格式，默认为 True
return_dict: bool = True,
# 交叉注意力的额外参数，默认为 None
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
# 跳过的 CLIP 层数，默认为 None
clip_skip: Optional[int] = None,
# 每步结束时的回调函数，默认为 None
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 每步结束时回调时传入的张量输入的名称，默认为 ["latents"]
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# 解码时的块大小，默认为 16
decode_chunk_size: int = 16,
# PAG 的缩放比例，默认为 3.0
pag_scale: float = 3.0,
# PAG 的自适应缩放，默认为 0.0
pag_adaptive_scale: float = 0.0,
