diffusers-源码解析-三十六-
diffusers 源码解析（三十六）
.\diffusers\pipelines\latte\__init__.py
# 导入类型检查模块
from typing import TYPE_CHECKING
# 从相对路径导入多个实用工具函数和常量
from ...utils import (
DIFFUSERS_SLOW_IMPORT, # 用于判断是否为慢导入
OptionalDependencyNotAvailable, # 表示可选依赖不可用的异常
_LazyModule, # 懒加载模块的类
get_objects_from_module, # 从模块中获取对象的函数
is_torch_available, # 检查 PyTorch 是否可用的函数
is_transformers_available, # 检查 Transformers 是否可用的函数
)
# 存放空对象的字典
_dummy_objects = {}
# 存放导入结构的字典
_import_structure = {}
# 尝试检查是否可以使用 transformers 和 torch
try:
if not (is_transformers_available() and is_torch_available()): # 检查两个库的可用性
raise OptionalDependencyNotAvailable() # 如果不可用，抛出异常
except OptionalDependencyNotAvailable: # 捕获可选依赖不可用的异常
from ...utils import dummy_torch_and_transformers_objects # noqa F403 # 导入虚拟对象以避免错误
# 更新 _dummy_objects 字典，填充虚拟对象
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
else:
# 如果可用，将 "pipeline_latte" 加入导入结构
_import_structure["pipeline_latte"] = ["LattePipeline"]
# 检查是否进行类型检查或慢导入
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
try:
if not (is_transformers_available() and is_torch_available()): # 再次检查库的可用性
raise OptionalDependencyNotAvailable() # 如果不可用，抛出异常
except OptionalDependencyNotAvailable: # 捕获可选依赖不可用的异常
from ...utils.dummy_torch_and_transformers_objects import * # 导入虚拟对象以避免错误
else:
# 如果可用，导入 LattePipeline
from .pipeline_latte import LattePipeline
else:
import sys # 导入系统模块
# 用懒加载模块替换当前模块
sys.modules[__name__] = _LazyModule(
__name__, # 当前模块名称
globals()["__file__"], # 当前模块文件路径
_import_structure, # 导入结构
module_spec=__spec__, # 模块规范
)
# 将 _dummy_objects 中的虚拟对象添加到当前模块中
for name, value in _dummy_objects.items():
setattr(sys.modules[__name__], name, value) # 设置模块属性
.\diffusers\pipelines\ledits_pp\pipeline_leditspp_stable_diffusion.py
# 导入 inspect 模块，用于获取对象的信息
import inspect
# 导入 math 模块，提供数学函数
import math
# 从 itertools 导入 repeat，用于生成重复元素的迭代器
from itertools import repeat
# 从 typing 导入常用类型，用于类型注解
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
# 导入 torch 库，用于深度学习
import torch
# 从 torch.nn.functional 导入 F，提供各种神经网络功能
import torch.nn.functional as F
# 从 packaging 导入 version，用于版本比较
from packaging import version
# 从 transformers 导入 CLIP 相关组件，用于图像和文本处理
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
# 从上级目录导入 FrozenDict，用于不可变字典
from ...configuration_utils import FrozenDict
# 从上级目录导入图像处理相关类
from ...image_processor import PipelineImageInput, VaeImageProcessor
# 从上级目录导入加载器混合类
from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin
# 从上级目录导入模型
from ...models import AutoencoderKL, UNet2DConditionModel
# 从注意力处理器模块导入 Attention 和 AttnProcessor
from ...models.attention_processor import Attention, AttnProcessor
# 从 lora 模块导入调整 LORA 比例的函数
from ...models.lora import adjust_lora_scale_text_encoder
# 从稳定扩散管道导入安全检查器
from ...pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
# 从调度器导入 DDIM 和 DPMSolver 多步调度器
from ...schedulers import DDIMScheduler, DPMSolverMultistepScheduler
# 从工具模块导入各种工具函数
from ...utils import (
USE_PEFT_BACKEND, # 用于是否使用 PEFT 后端的标志
deprecate, # 用于标记弃用的函数
logging, # 用于日志记录
replace_example_docstring, # 用于替换示例文档字符串的函数
scale_lora_layers, # 用于缩放 LORA 层的函数
unscale_lora_layers, # 用于反缩放 LORA 层的函数
)
# 从 torch_utils 导入随机张量生成函数
from ...utils.torch_utils import randn_tensor
# 从管道工具模块导入扩散管道类
from ..pipeline_utils import DiffusionPipeline
# 从管道输出模块导入管道输出类
from .pipeline_output import LEditsPPDiffusionPipelineOutput, LEditsPPInversionPipelineOutput
# 创建日志记录器，使用当前模块名称
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，提供使用示例
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> import PIL # 导入 PIL 库，用于图像处理
>>> import requests # 导入 requests 库，用于 HTTP 请求
>>> import torch # 导入 torch 库，用于深度学习
>>> from io import BytesIO # 从 io 导入 BytesIO，用于字节流处理
>>> from diffusers import LEditsPPPipelineStableDiffusion # 从 diffusers 导入 LEditsPPPipelineStableDiffusion 类
>>> from diffusers.utils import load_image # 从 diffusers.utils 导入 load_image 函数
>>> pipe = LEditsPPPipelineStableDiffusion.from_pretrained( # 从预训练模型创建管道
... "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16 # 指定模型路径和数据类型
... )
>>> pipe = pipe.to("cuda") # 将管道移动到 GPU
>>> img_url = "https://www.aiml.informatik.tu-darmstadt.de/people/mbrack/cherry_blossom.png" # 图像的 URL
>>> image = load_image(img_url).convert("RGB") # 加载图像并转换为 RGB 格式
>>> _ = pipe.invert(image=image, num_inversion_steps=50, skip=0.1) # 对图像进行反演处理
>>> edited_image = pipe( # 使用管道进行图像编辑
... editing_prompt=["cherry blossom"], edit_guidance_scale=10.0, edit_threshold=0.75 # 设置编辑提示和参数
... ).images[0] # 获取编辑后的图像
```py
"""
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionAttendAndExcitePipeline 修改的类
class LeditsAttentionStore:
@staticmethod
# 静态方法，获取一个空的注意力存储
def get_empty_store():
# 返回一个字典，包含不同层次的注意力存储列表
return {"down_cross": [], "mid_cross": [], "up_cross": [], "down_self": [], "mid_self": [], "up_self": []}
# 定义一个可调用的方法，接收注意力权重和其他参数
def __call__(self, attn, is_cross: bool, place_in_unet: str, editing_prompts, PnP=False):
# 注意力权重的形状为：批大小 * 头大小, 序列长度查询, 序列长度键
if attn.shape[1] <= self.max_size: # 检查序列长度是否小于或等于最大大小
bs = 1 + int(PnP) + editing_prompts # 计算批次大小
skip = 2 if PnP else 1 # 确定跳过的步骤：如果是 PnP，跳过 2，否则跳过 1
# 将注意力权重拆分为指定批大小的张量，并调整维度顺序
attn = torch.stack(attn.split(self.batch_size)).permute(1, 0, 2, 3)
source_batch_size = int(attn.shape[1] // bs) # 计算源批次大小
# 调用 forward 方法，传入处理后的注意力权重
self.forward(attn[:, skip * source_batch_size :], is_cross, place_in_unet)
# 定义 forward 方法，用于处理注意力权重
def forward(self, attn, is_cross: bool, place_in_unet: str):
# 生成键值，用于存储不同位置的注意力
key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
# 将当前注意力权重添加到对应的存储列表中
self.step_store[key].append(attn)
# 定义在步骤之间的操作，可以选择存储当前步骤
def between_steps(self, store_step=True):
if store_step: # 如果需要存储当前步骤
if self.average: # 如果启用平均
if len(self.attention_store) == 0: # 如果注意力存储为空
self.attention_store = self.step_store # 直接赋值
else:
# 对现有注意力存储进行更新
for key in self.attention_store:
for i in range(len(self.attention_store[key])):
self.attention_store[key][i] += self.step_store[key][i]
else: # 如果不启用平均
if len(self.attention_store) == 0: # 如果注意力存储为空
self.attention_store = [self.step_store] # 作为第一个存储项
else:
self.attention_store.append(self.step_store) # 添加新的存储项
self.cur_step += 1 # 更新当前步骤计数
# 重置当前步骤的存储
self.step_store = self.get_empty_store()
# 定义获取特定步骤注意力的方法
def get_attention(self, step: int):
if self.average: # 如果启用平均
# 计算平均注意力，按当前步骤进行归一化
attention = {
key: [item / self.cur_step for item in self.attention_store[key]] for key in self.attention_store
}
else: # 否则，确保提供了步骤
assert step is not None # 断言步骤不为 None
attention = self.attention_store[step] # 获取指定步骤的注意力
return attention # 返回注意力
# 定义聚合注意力的方法，处理多个参数
def aggregate_attention(
self, attention_maps, prompts, res: Union[int, Tuple[int]], from_where: List[str], is_cross: bool, select: int
):
out = [[] for x in range(self.batch_size)] # 初始化输出列表
if isinstance(res, int): # 如果分辨率为整数
num_pixels = res**2 # 计算像素数量
resolution = (res, res) # 设置分辨率
else: # 如果分辨率为元组
num_pixels = res[0] * res[1] # 计算像素数量
resolution = res[:2] # 设置分辨率
# 遍历指定来源，提取注意力图
for location in from_where:
for bs_item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
for batch, item in enumerate(bs_item):
if item.shape[1] == num_pixels: # 检查当前项的形状是否匹配
# 重塑注意力图并选择指定项
cross_maps = item.reshape(len(prompts), -1, *resolution, item.shape[-1])[select]
out[batch].append(cross_maps) # 将提取的映射添加到输出
# 合并输出，计算每个批次的平均值
out = torch.stack([torch.cat(x, dim=0) for x in out]) # 在第一个维度上堆叠
out = out.sum(1) / out.shape[1] # 在头维度上取平均
return out # 返回聚合后的输出
# 初始化方法，用于创建类的实例，接受多个参数
def __init__(self, average: bool, batch_size=1, max_resolution=16, max_size: int = None):
# 获取一个空的存储结构，用于保存步骤数据
self.step_store = self.get_empty_store()
# 初始化注意力存储列表
self.attention_store = []
# 当前步骤计数器初始化为 0
self.cur_step = 0
# 设置是否计算平均值的标志
self.average = average
# 设置批处理大小，默认值为 1
self.batch_size = batch_size
# 如果没有指定最大大小，则计算为最大分辨率的平方
if max_size is None:
self.max_size = max_resolution**2
# 如果指定了最大大小且没有指定最大分辨率，则使用指定的最大大小
elif max_size is not None and max_resolution is None:
self.max_size = max_size
# 如果同时指定了最大分辨率和最大大小，则抛出错误
else:
raise ValueError("Only allowed to set one of max_resolution or max_size")
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionAttendAndExcitePipeline.GaussianSmoothing 修改而来
class LeditsGaussianSmoothing:
# 初始化函数，接收设备参数
def __init__(self, device):
# 定义高斯核的大小
kernel_size = [3, 3]
# 定义高斯核的标准差
sigma = [0.5, 0.5]
# 高斯核是每个维度高斯函数的乘积
kernel = 1
# 创建网格以便于计算高斯核
meshgrids = torch.meshgrid([torch.arange(size, dtype=torch.float32) for size in kernel_size])
# 遍历每个维度的大小、标准差和网格
for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
# 计算高斯核的均值
mean = (size - 1) / 2
# 计算高斯核值
kernel *= 1 / (std * math.sqrt(2 * math.pi)) * torch.exp(-(((mgrid - mean) / (2 * std)) ** 2))
# 确保高斯核的值之和为1
kernel = kernel / torch.sum(kernel)
# 将高斯核重塑为深度卷积权重
kernel = kernel.view(1, 1, *kernel.size())
# 重复高斯核以适应卷积层的形状
kernel = kernel.repeat(1, *[1] * (kernel.dim() - 1))
# 将权重移动到指定设备上
self.weight = kernel.to(device)
# 调用函数，用于应用高斯滤波
def __call__(self, input):
"""
参数:
对输入应用高斯滤波。
input (torch.Tensor): 需要应用高斯滤波的输入。
返回:
filtered (torch.Tensor): 经过滤波的输出。
"""
# 使用卷积操作应用高斯滤波
return F.conv2d(input, weight=self.weight.to(input.dtype))
class LEDITSCrossAttnProcessor:
# 初始化函数，接收多个参数以设置注意力处理器
def __init__(self, attention_store, place_in_unet, pnp, editing_prompts):
# 存储注意力的变量
self.attnstore = attention_store
# 设置在 UNet 中的位置
self.place_in_unet = place_in_unet
# 存储编辑提示
self.editing_prompts = editing_prompts
# 存储 PnP 相关的变量
self.pnp = pnp
# 调用函数，处理注意力和隐藏状态
def __call__(
self,
attn: Attention,
hidden_states,
encoder_hidden_states,
attention_mask=None,
temb=None,
# 处理输入，确定批量大小和序列长度
):
batch_size, sequence_length, _ = (
# 如果没有编码器隐藏状态，则使用隐藏状态形状，否则使用编码器隐藏状态形状
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
)
# 准备注意力掩码，以便后续的注意力计算
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
# 将隐藏状态转换为查询向量
query = attn.to_q(hidden_states)
# 如果没有编码器隐藏状态，使用隐藏状态；否则，如果需要规范化，则规范化编码器隐藏状态
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
elif attn.norm_cross:
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
# 将编码器隐藏状态转换为键向量
key = attn.to_k(encoder_hidden_states)
# 将编码器隐藏状态转换为值向量
value = attn.to_v(encoder_hidden_states)
# 将查询向量转换为批处理维度
query = attn.head_to_batch_dim(query)
# 将键向量转换为批处理维度
key = attn.head_to_batch_dim(key)
# 将值向量转换为批处理维度
value = attn.head_to_batch_dim(value)
# 计算注意力得分
attention_probs = attn.get_attention_scores(query, key, attention_mask)
# 存储注意力得分，标记为交叉注意力
self.attnstore(
attention_probs,
is_cross=True,
place_in_unet=self.place_in_unet,
editing_prompts=self.editing_prompts,
PnP=self.pnp,
)
# 通过值向量和注意力得分计算新的隐藏状态
hidden_states = torch.bmm(attention_probs, value)
# 将隐藏状态转换回头部维度
hidden_states = attn.batch_to_head_dim(hidden_states)
# 进行线性投影
hidden_states = attn.to_out[0](hidden_states)
# 进行 dropout 操作
hidden_states = attn.to_out[1](hidden_states)
# 重新缩放输出
hidden_states = hidden_states / attn.rescale_output_factor
# 返回最终的隐藏状态
return hidden_states
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion 导入的 rescale_noise_cfg 函数
def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
"""
根据 `guidance_rescale` 对 `noise_cfg` 进行重新缩放。基于 [Common Diffusion Noise Schedules and
Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf) 的发现。见第 3.4 节
"""
# 计算文本噪声预测的标准差，保持维度
std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
# 计算配置噪声的标准差，保持维度
std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
# 根据标准差调整指导结果（修正过度曝光）
noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
# 按照指导缩放因子混合原始结果，以避免图像“平淡”
noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
# 返回调整后的噪声配置
return noise_cfg
# 定义 LEditsPPPipelineStableDiffusion 类，继承多个混合类
class LEditsPPPipelineStableDiffusion(
DiffusionPipeline, TextualInversionLoaderMixin, StableDiffusionLoraLoaderMixin, IPAdapterMixin, FromSingleFileMixin
):
"""
使用 LEDits++ 和 Stable Diffusion 进行文本图像编辑的管道。
此模型继承自 [`DiffusionPipeline`]，并建立在 [`StableDiffusionPipeline`] 之上。查看超类
文档以了解所有管道实现的通用方法（下载、保存、在特定设备上运行等）。
# 文档字符串，定义类的参数及其描述
Args:
vae ([`AutoencoderKL`]):
# 变分自编码器模型，用于将图像编码和解码为潜在表示
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
text_encoder ([`~transformers.CLIPTextModel`]):
# 冻结的文本编码器，Stable Diffusion 使用 CLIP 的文本部分
Frozen text-encoder. Stable Diffusion uses the text portion of
# 指定的 CLIP 模型变体
[CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
# 指定的 CLIP-ViT 大模型变体
[clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
tokenizer ([`~transformers.CLIPTokenizer`]):
# CLIPTokenizer 类的分词器，用于文本处理
Tokenizer of class
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
unet ([`UNet2DConditionModel`]):
# 条件 U-Net 架构，用于对编码的图像潜在表示进行去噪
Conditional U-Net architecture to denoise the encoded image latents.
scheduler ([`DPMSolverMultistepScheduler`] or [`DDIMScheduler`]):
# 调度器，用于与 `unet` 结合去噪编码的图像潜在表示
A scheduler to be used in combination with `unet` to denoise the encoded image latens. Can be one of
# 允许的调度器类型，若传入其他类型则默认为 DPMSolverMultistepScheduler
[`DPMSolverMultistepScheduler`] or [`DDIMScheduler`]. If any other scheduler is passed it will
automatically be set to [`DPMSolverMultistepScheduler`].
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/CompVis/stable-diffusion-v1-4) for details.
feature_extractor ([`~transformers.CLIPImageProcessor`]):
# 从生成的图像中提取特征，以作为 `safety_checker` 的输入
Model that extracts features from generated images to be used as inputs for the `safety_checker`.
"""
# 定义模型在 CPU 卸载时的顺序
model_cpu_offload_seq = "text_encoder->unet->vae"
# 定义在 CPU 卸载时排除的组件
_exclude_from_cpu_offload = ["safety_checker"]
# 定义回调张量输入
_callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
# 定义可选组件
_optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
# 初始化方法，定义类的构造函数
def __init__(
self,
# 变分自编码器模型实例
vae: AutoencoderKL,
# 文本编码器实例
text_encoder: CLIPTextModel,
# 分词器实例
tokenizer: CLIPTokenizer,
# U-Net 模型实例
unet: UNet2DConditionModel,
# 调度器实例
scheduler: Union[DDIMScheduler, DPMSolverMultistepScheduler],
# 安全检查器实例
safety_checker: StableDiffusionSafetyChecker,
# 特征提取器实例
feature_extractor: CLIPImageProcessor,
# 指示是否需要安全检查器的布尔值，默认为 True
requires_safety_checker: bool = True,
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker 复制的内容
# 运行安全检查器，对输入图像进行安全性检测
def run_safety_checker(self, image, device, dtype):
# 检查安全检查器是否已初始化
if self.safety_checker is None:
# 若未初始化，NSFW 概念标志设置为 None
has_nsfw_concept = None
else:
# 检查输入图像是否为张量
if torch.is_tensor(image):
# 如果是张量，则处理图像为 PIL 格式
feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
else:
# 如果不是张量，则将 NumPy 数组转换为 PIL 图像
feature_extractor_input = self.image_processor.numpy_to_pil(image)
# 提取图像特征，返回张量并移动到指定设备
safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
# 执行安全检查，返回处理后的图像和 NSFW 概念标志
image, has_nsfw_concept = self.safety_checker(
images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
)
# 返回处理后的图像和 NSFW 概念标志
return image, has_nsfw_concept
# 从 StableDiffusionPipeline 复制的解码潜变量方法
def decode_latents(self, latents):
# 提示解码方法已弃用，未来版本将移除
deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
# 根据 VAE 配置缩放潜变量
latents = 1 / self.vae.config.scaling_factor * latents
# 解码潜变量，获取图像
image = self.vae.decode(latents, return_dict=False)[0]
# 将图像数据归一化并限制在 [0, 1] 范围内
image = (image / 2 + 0.5).clamp(0, 1)
# 将图像转换为 float32 格式，以兼容 bfloat16，且不会造成显著开销
image = image.cpu().permute(0, 2, 3, 1).float().numpy()
# 返回处理后的图像
return image
# 从 StableDiffusionPipeline 复制的准备额外步骤参数的方法
def prepare_extra_step_kwargs(self, eta, generator=None):
# 准备调度器步骤的额外参数，不同调度器的参数签名可能不同
# eta (η) 仅在 DDIMScheduler 中使用，其他调度器会忽略
# eta 对应于 DDIM 论文中的 η，范围应在 [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
# 从 StableDiffusionPipeline 复制的输入检查方法
def check_inputs(
self,
# 定义检查的输入参数，负提示和编辑提示嵌入等
negative_prompt=None,
editing_prompt_embeddings=None,
negative_prompt_embeds=None,
callback_on_step_end_tensor_inputs=None,
):
# 检查回调输入是否为空，并确保所有输入都在回调张量输入列表中
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]}"
)
# 检查负提示和负提示嵌入是否同时存在
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."
)
# 检查编辑提示嵌入和负提示嵌入是否同时存在
if editing_prompt_embeddings is not None and negative_prompt_embeds is not None:
# 检查两者形状是否相同
if editing_prompt_embeddings.shape != negative_prompt_embeds.shape:
# 如果形状不匹配，抛出值错误
raise ValueError(
"`editing_prompt_embeddings` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got: `editing_prompt_embeddings` {editing_prompt_embeddings.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents 修改而来
def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, latents):
# 计算预期的张量形状
# shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
# 如果输入的 latents 形状与预期形状不匹配，抛出值错误
# if latents.shape != shape:
# raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
# 将输入的 latents 张量移动到指定设备
latents = latents.to(device)
# 根据调度器所需的标准差缩放初始噪声
latents = latents * self.scheduler.init_noise_sigma
# 返回处理后的 latents
return latents
def prepare_unet(self, attention_store, PnP: bool = False):
# 创建一个空字典用于存储注意力处理器
attn_procs = {}
# 遍历 UNet 中的注意力处理器键
for name in self.unet.attn_processors.keys():
# 根据名称前缀确定其在 UNet 中的位置
if name.startswith("mid_block"):
place_in_unet = "mid"
elif name.startswith("up_blocks"):
place_in_unet = "up"
elif name.startswith("down_blocks"):
place_in_unet = "down"
else:
continue
# 根据名称选择合适的注意力处理器
if "attn2" in name and place_in_unet != "mid":
attn_procs[name] = LEDITSCrossAttnProcessor(
attention_store=attention_store,
place_in_unet=place_in_unet,
pnp=PnP,
editing_prompts=self.enabled_editing_prompts,
)
else:
attn_procs[name] = AttnProcessor()
# 将注意力处理器设置到 UNet 中
self.unet.set_attn_processor(attn_procs)
# 定义一个编码提示的方法，接收多个参数
def encode_prompt(
self, # 当前对象的引用
device, # 设备类型，例如 CPU 或 GPU
num_images_per_prompt, # 每个提示生成的图像数量
enable_edit_guidance, # 是否启用编辑引导
negative_prompt=None, # 可选的负面提示
editing_prompt=None, # 可选的编辑提示
negative_prompt_embeds: Optional[torch.Tensor] = None, # 可选的负面提示嵌入
editing_prompt_embeds: Optional[torch.Tensor] = None, # 可选的编辑提示嵌入
lora_scale: Optional[float] = None, # 可选的 Lora 缩放因子
clip_skip: Optional[int] = None, # 可选的跳过的剪辑层数
# 获取指导重缩放的属性
@property
def guidance_rescale(self):
return self._guidance_rescale # 返回指导重缩放的值
# 获取剪辑跳过的属性
@property
def clip_skip(self):
return self._clip_skip # 返回剪辑跳过的值
# 获取交叉注意力参数的属性
@property
def cross_attention_kwargs(self):
return self._cross_attention_kwargs # 返回交叉注意力参数
# 装饰器，指示不计算梯度
@torch.no_grad()
# 替换示例文档字符串的装饰器
@replace_example_docstring(EXAMPLE_DOC_STRING)
# 定义可调用的方法，接收多个可选参数
def __call__(
self, # 方法本身的引用
negative_prompt: Optional[Union[str, List[str]]] = None, # 可选的负面提示
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, # 可选的生成器
output_type: Optional[str] = "pil", # 输出类型，默认是 PIL
return_dict: bool = True, # 是否返回字典格式的输出
editing_prompt: Optional[Union[str, List[str]]] = None, # 可选的编辑提示
editing_prompt_embeds: Optional[torch.Tensor] = None, # 可选的编辑提示嵌入
negative_prompt_embeds: Optional[torch.Tensor] = None, # 可选的负面提示嵌入
reverse_editing_direction: Optional[Union[bool, List[bool]]] = False, # 是否反向编辑方向
edit_guidance_scale: Optional[Union[float, List[float]]] = 5, # 编辑引导缩放因子
edit_warmup_steps: Optional[Union[int, List[int]]] = 0, # 编辑热身步数
edit_cooldown_steps: Optional[Union[int, List[int]]] = None, # 编辑冷却步数
edit_threshold: Optional[Union[float, List[float]]] = 0.9, # 编辑阈值
user_mask: Optional[torch.Tensor] = None, # 可选的用户掩码
sem_guidance: Optional[List[torch.Tensor]] = None, # 可选的语义引导
use_cross_attn_mask: bool = False, # 是否使用交叉注意力掩码
use_intersect_mask: bool = True, # 是否使用交集掩码
attn_store_steps: Optional[List[int]] = [], # 存储注意力的步骤列表
store_averaged_over_steps: bool = True, # 是否在步骤上存储平均值
cross_attention_kwargs: Optional[Dict[str, Any]] = None, # 可选的交叉注意力参数
guidance_rescale: float = 0.0, # 指导重缩放的默认值
clip_skip: Optional[int] = None, # 可选的剪辑跳过层数
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, # 可选的步骤结束回调
callback_on_step_end_tensor_inputs: List[str] = ["latents"], # 默认的步骤结束回调输入
**kwargs, # 接收其他关键字参数
# 装饰器，指示不计算梯度
@torch.no_grad()
# 定义反转的方法，接收多个参数
def invert(
self, # 方法本身的引用
image: PipelineImageInput, # 输入图像
source_prompt: str = "", # 源提示字符串
source_guidance_scale: float = 3.5, # 源指导缩放因子
num_inversion_steps: int = 30, # 反转步骤数量
skip: float = 0.15, # 跳过的比例
generator: Optional[torch.Generator] = None, # 可选的生成器
cross_attention_kwargs: Optional[Dict[str, Any]] = None, # 可选的交叉注意力参数
clip_skip: Optional[int] = None, # 可选的剪辑跳过层数
height: Optional[int] = None, # 可选的图像高度
width: Optional[int] = None, # 可选的图像宽度
resize_mode: Optional[str] = "default", # 图像调整大小的模式
crops_coords: Optional[Tuple[int, int, int, int]] = None, # 可选的裁剪坐标
# 装饰器，指示不计算梯度
@torch.no_grad()
# 定义一个编码图像的函数，接收多个参数以处理图像
def encode_image(self, image, dtype=None, height=None, width=None, resize_mode="default", crops_coords=None):
# 使用图像处理器预处理输入图像，调整其高度、宽度和裁剪坐标
image = self.image_processor.preprocess(
image=image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords
)
# 使用图像处理器后处理图像，输出类型为 PIL 图像
resized = self.image_processor.postprocess(image=image, output_type="pil")
# 检查输入图像的最大维度是否超过默认分辨率的 1.5 倍
if max(image.shape[-2:]) > self.vae.config["sample_size"] * 1.5:
# 记录警告信息，提示输入图像分辨率过高可能导致输出图像有严重伪影
logger.warning(
"Your input images far exceed the default resolution of the underlying diffusion model. "
"The output images may contain severe artifacts! "
"Consider down-sampling the input using the `height` and `width` parameters"
)
# 将图像转换为指定的数据类型
image = image.to(dtype)
# 使用 VAE 编码器对图像进行编码，获取潜在分布的模式
x0 = self.vae.encode(image.to(self.device)).latent_dist.mode()
# 将编码结果转换为指定的数据类型
x0 = x0.to(dtype)
# 将编码结果乘以缩放因子
x0 = self.vae.config.scaling_factor * x0
# 返回编码后的结果和处理后的图像
return x0, resized
# 计算 DDIM 噪声
def compute_noise_ddim(scheduler, prev_latents, latents, timestep, noise_pred, eta):
# 1. 获取前一个时间步的值（t-1）
prev_timestep = timestep - scheduler.config.num_train_timesteps // scheduler.num_inference_steps
# 2. 计算 alphas 和 betas
alpha_prod_t = scheduler.alphas_cumprod[timestep] # 当前时间步的累积 alpha
alpha_prod_t_prev = (
scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else scheduler.final_alpha_cumprod
) # 前一个时间步的累积 alpha，若为负则使用最终 alpha
beta_prod_t = 1 - alpha_prod_t # 当前时间步的 beta
# 3. 从预测噪声计算预测的原始样本，也称为“预测的 x_0”
pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5)
# 4. 裁剪“预测的 x_0”
if scheduler.config.clip_sample:
pred_original_sample = torch.clamp(pred_original_sample, -1, 1) # 将值限制在 -1 到 1 之间
# 5. 计算方差：“sigma_t(η)” -> 参见公式 (16)
variance = scheduler._get_variance(timestep, prev_timestep) # 获取方差
std_dev_t = eta * variance ** (0.5) # 计算标准差
# 6. 计算指向 x_t 的方向，参见公式 (12)
pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * noise_pred
# 修改以返回更新后的 xtm1（避免误差累积）
mu_xt = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction # 计算 mu_xt
if variance > 0.0:
noise = (prev_latents - mu_xt) / (variance ** (0.5) * eta) # 计算噪声
else:
noise = torch.tensor([0.0]).to(latents.device) # 方差为零时噪声设为零
return noise, mu_xt + (eta * variance**0.5) * noise # 返回噪声和更新后的样本
# 计算 SDE DPM PP 二阶噪声
def compute_noise_sde_dpm_pp_2nd(scheduler, prev_latents, latents, timestep, noise_pred, eta):
def first_order_update(model_output, sample): # 定义一阶更新
sigma_t, sigma_s = scheduler.sigmas[scheduler.step_index + 1], scheduler.sigmas[scheduler.step_index] # 获取当前和前一个 sigma
alpha_t, sigma_t = scheduler._sigma_to_alpha_sigma_t(sigma_t) # 将 sigma 转换为 alpha
alpha_s, sigma_s = scheduler._sigma_to_alpha_sigma_t(sigma_s) # 将前一个 sigma 转换为 alpha
lambda_t = torch.log(alpha_t) - torch.log(sigma_t) # 计算 lambda_t
lambda_s = torch.log(alpha_s) - torch.log(sigma_s) # 计算 lambda_s
h = lambda_t - lambda_s # 计算 h
mu_xt = (sigma_t / sigma_s * torch.exp(-h)) * sample + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output # 计算 mu_xt
mu_xt = scheduler.dpm_solver_first_order_update( # 更新 mu_xt
model_output=model_output, sample=sample, noise=torch.zeros_like(sample) # 将噪声设为零
)
sigma = sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) # 计算 sigma
if sigma > 0.0:
noise = (prev_latents - mu_xt) / sigma # 计算噪声
else:
noise = torch.tensor([0.0]).to(sample.device) # 方差为零时噪声设为零
prev_sample = mu_xt + sigma * noise # 计算前一个样本
return noise, prev_sample # 返回噪声和前一个样本
# 定义二阶更新函数，接受模型输出列表和样本
def second_order_update(model_output_list, sample): # timestep_list, prev_timestep, sample):
# 获取当前和前后时刻的 sigma 值
sigma_t, sigma_s0, sigma_s1 = (
scheduler.sigmas[scheduler.step_index + 1], # 当前时刻的 sigma
scheduler.sigmas[scheduler.step_index], # 上一时刻的 sigma
scheduler.sigmas[scheduler.step_index - 1], # 更早时刻的 sigma
)
# 将 sigma 转换为 alpha 和 sigma_t
alpha_t, sigma_t = scheduler._sigma_to_alpha_sigma_t(sigma_t)
alpha_s0, sigma_s0 = scheduler._sigma_to_alpha_sigma_t(sigma_s0)
alpha_s1, sigma_s1 = scheduler._sigma_to_alpha_sigma_t(sigma_s1)
# 计算 lambda 值
lambda_t = torch.log(alpha_t) - torch.log(sigma_t) # 当前时刻的 lambda
lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) # 上一时刻的 lambda
lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) # 更早时刻的 lambda
# 获取最后两个模型输出
m0, m1 = model_output_list[-1], model_output_list[-2]
# 计算 h 和 r0 值
h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 # h 和 h_0
r0 = h_0 / h # r0 的计算
D0, D1 = m0, (1.0 / r0) * (m0 - m1) # D0 和 D1 的计算
# 计算 mu_xt
mu_xt = (
(sigma_t / sigma_s0 * torch.exp(-h)) * sample # 根据样本和 sigma 计算
+ (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 # 加上 D0 的贡献
+ 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 # 加上 D1 的一半贡献
)
# 计算 sigma
sigma = sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) # 最终的 sigma 计算
# 根据 sigma 计算噪声
if sigma > 0.0:
noise = (prev_latents - mu_xt) / sigma # 正常情况下计算噪声
else:
noise = torch.tensor([0.0]).to(sample.device) # sigma 为零时，噪声为零
# 计算前一个样本
prev_sample = mu_xt + sigma * noise # 最终样本的计算
return noise, prev_sample # 返回噪声和前一个样本
# 初始化 step_index，如果未定义的话
if scheduler.step_index is None:
scheduler._init_step_index(timestep) # 初始化步进索引
# 将模型输出转换为合适格式
model_output = scheduler.convert_model_output(model_output=noise_pred, sample=latents)
# 更新模型输出列表
for i in range(scheduler.config.solver_order - 1):
scheduler.model_outputs[i] = scheduler.model_outputs[i + 1] # 向前移动输出
scheduler.model_outputs[-1] = model_output # 保存当前模型输出
# 根据 lower_order_nums 选择更新方法
if scheduler.lower_order_nums < 1:
noise, prev_sample = first_order_update(model_output, latents) # 一阶更新
else:
noise, prev_sample = second_order_update(scheduler.model_outputs, latents) # 二阶更新
# 更新 lower_order_nums
if scheduler.lower_order_nums < scheduler.config.solver_order:
scheduler.lower_order_nums += 1 # 增加 lower_order_nums
# 完成后增加步进索引
scheduler._step_index += 1 # 增加步进索引
return noise, prev_sample # 返回噪声和前一个样本
# 定义一个计算噪声的函数，接收调度器和可变参数
def compute_noise(scheduler, *args):
# 检查调度器是否为 DDIMScheduler 类型
if isinstance(scheduler, DDIMScheduler):
# 调用对应的函数计算 DDIM 噪声
return compute_noise_ddim(scheduler, *args)
# 检查调度器是否为 DPMSolverMultistepScheduler，并验证其配置
elif (
isinstance(scheduler, DPMSolverMultistepScheduler)
and scheduler.config.algorithm_type == "sde-dpmsolver++"
and scheduler.config.solver_order == 2
):
# 调用对应的函数计算 SDE-DPM 噪声（2阶）
return compute_noise_sde_dpm_pp_2nd(scheduler, *args)
else:
# 如果不支持的调度器类型，抛出未实现的错误
raise NotImplementedError
.\diffusers\pipelines\ledits_pp\pipeline_leditspp_stable_diffusion_xl.py
# 版权声明，表示此文件的版权信息及归属
# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# 许可条款，声明该文件根据 Apache License 2.0 进行许可
# Licensed under the Apache License, Version 2.0 (the "License");
# 该文件的使用需遵循许可证的规定
# you may not use this file except in compliance with the License.
# 可以在以下网址获取许可证
# you may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非在适用法律或书面协议中另有规定，软件按“原样”提供
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 不提供任何形式的明示或暗示的担保
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 请参阅许可证以了解有关权限和限制的具体信息
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect # 导入 inspect 模块，用于获取对象的内部信息
import math # 导入 math 模块，提供数学函数
from typing import Any, Callable, Dict, List, Optional, Tuple, Union # 从 typing 导入类型注解，用于类型提示
import torch # 导入 PyTorch 库，用于深度学习
import torch.nn.functional as F # 导入 PyTorch 的功能性神经网络模块，通常用于激活函数等
from transformers import ( # 从 transformers 库导入多个类和函数
CLIPImageProcessor, # 导入 CLIP 图像处理器类
CLIPTextModel, # 导入 CLIP 文本模型类
CLIPTextModelWithProjection, # 导入带有投影的 CLIP 文本模型类
CLIPTokenizer, # 导入 CLIP 分词器类
CLIPVisionModelWithProjection, # 导入带有投影的 CLIP 视觉模型类
)
from ...image_processor import PipelineImageInput, VaeImageProcessor # 从相对路径导入图像处理相关类
from ...loaders import ( # 从相对路径导入加载器相关类
FromSingleFileMixin, # 单文件混合加载器
IPAdapterMixin, # IP 适配器混合类
StableDiffusionXLLoraLoaderMixin, # Stable Diffusion XL LoRA 加载器混合类
TextualInversionLoaderMixin, # 文本反演加载器混合类
)
from ...models import AutoencoderKL, UNet2DConditionModel # 从相对路径导入模型类
from ...models.attention_processor import ( # 从相对路径导入注意力处理器相关类
Attention, # 注意力机制类
AttnProcessor, # 注意力处理器基类
AttnProcessor2_0, # 注意力处理器版本 2.0
XFormersAttnProcessor, # XFormers 注意力处理器
)
from ...models.lora import adjust_lora_scale_text_encoder # 从相对路径导入调整 LoRA 比例的函数
from ...schedulers import DDIMScheduler, DPMSolverMultistepScheduler # 从相对路径导入调度器类
from ...utils import ( # 从相对路径导入实用工具函数
USE_PEFT_BACKEND, # 用于 PEFT 后端的常量
is_invisible_watermark_available, # 检查是否可用隐形水印的函数
is_torch_xla_available, # 检查是否可用 Torch XLA 的函数
logging, # 日志记录模块
replace_example_docstring, # 替换示例文档字符串的函数
scale_lora_layers, # 缩放 LoRA 层的函数
unscale_lora_layers, # 反缩放 LoRA 层的函数
)
from ...utils.torch_utils import randn_tensor # 从相对路径导入生成随机张量的函数
from ..pipeline_utils import DiffusionPipeline # 从相对路径导入扩散管道类
from .pipeline_output import LEditsPPDiffusionPipelineOutput, LEditsPPInversionPipelineOutput # 导入管道输出相关类
# 如果隐形水印可用，则导入相应的水印处理类
if is_invisible_watermark_available():
from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker # 导入 Stable Diffusion XL 水印类
# 检查是否可用 Torch XLA，若可用则导入其核心模块
if is_torch_xla_available():
import torch_xla.core.xla_model as xm # 导入 Torch XLA 核心模型模块
XLA_AVAILABLE = True # 设置标志，表示 XLA 可用
else:
XLA_AVAILABLE = False # 设置标志，表示 XLA 不可用
logger = logging.get_logger(__name__) # 获取当前模块的日志记录器，便于后续日志记录使用
EXAMPLE_DOC_STRING = """ # 定义示例文档字符串的多行字符串，用于文档或示例说明
# 示例代码，展示如何使用 LE edits PP 管道进行图像编辑
Examples:
```py
# 导入所需的库
>>> import torch # 导入 PyTorch 库
>>> import PIL # 导入 Python Imaging Library (PIL)
>>> import requests # 导入请求库以获取网络资源
>>> from io import BytesIO # 导入 BytesIO 用于字节流处理
# 从 diffusers 库导入 LEditsPPPipelineStableDiffusionXL 类
>>> from diffusers import LEditsPPPipelineStableDiffusionXL
# 创建一个 LEditsPPPipelineStableDiffusionXL 的实例，使用预训练模型
>>> pipe = LEditsPPPipelineStableDiffusionXL.from_pretrained(
... "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
... )
# 将管道移动到 CUDA 设备以加速计算
>>> pipe = pipe.to("cuda")
# 定义一个函数下载并返回图像
>>> def download_image(url):
... # 发送 GET 请求获取图像
... response = requests.get(url)
... # 将响应内容转换为 RGB 格式的图像
... return PIL.Image.open(BytesIO(response.content)).convert("RGB")
# 图像的 URL
>>> img_url = "https://www.aiml.informatik.tu-darmstadt.de/people/mbrack/tennis.jpg"
# 下载图像并存储
>>> image = download_image(img_url)
# 使用管道进行图像反转操作，指定反转步骤和跳过比例
>>> _ = pipe.invert(image=image, num_inversion_steps=50, skip=0.2)
# 使用管道进行图像编辑，定义编辑提示和参数
>>> edited_image = pipe(
... editing_prompt=["tennis ball", "tomato"],
... reverse_editing_direction=[True, False],
... edit_guidance_scale=[5.0, 10.0],
... edit_threshold=[0.9, 0.85],
... ).images[0] # 获取编辑后的第一张图像
文档字符串，通常用于描述类或方法的功能
"""
从 diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.LeditsAttentionStore 复制的类
class LeditsAttentionStore:
# 静态方法，返回一个空的注意力存储结构
@staticmethod
def get_empty_store():
return {"down_cross": [], "mid_cross": [], "up_cross": [], "down_self": [], "mid_self": [], "up_self": []}
# 使对象可调用的方法，处理注意力矩阵
def __call__(self, attn, is_cross: bool, place_in_unet: str, editing_prompts, PnP=False):
# attn.shape = batch_size * head_size, seq_len query, seq_len_key
# 如果注意力矩阵的第二维小于等于最大大小
if attn.shape[1] <= self.max_size:
# 计算批次大小，考虑 PnP 和编辑提示
bs = 1 + int(PnP) + editing_prompts
skip = 2 if PnP else 1 # 跳过 PnP 和无条件
# 将注意力矩阵分割并重新排列维度
attn = torch.stack(attn.split(self.batch_size)).permute(1, 0, 2, 3)
# 计算源批次大小
source_batch_size = int(attn.shape[1] // bs)
# 调用前向传播方法
self.forward(attn[:, skip * source_batch_size :], is_cross, place_in_unet)
# 前向传播方法，存储注意力
def forward(self, attn, is_cross: bool, place_in_unet: str):
# 创建键，用于存储当前层的注意力
key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
# 将当前注意力添加到步骤存储中
self.step_store[key].append(attn)
# 在步骤之间调用的方法，决定是否存储步骤
def between_steps(self, store_step=True):
if store_step:
# 如果需要平均处理
if self.average:
# 如果注意力存储为空，初始化
if len(self.attention_store) == 0:
self.attention_store = self.step_store
else:
# 将步骤存储中的注意力与已有的注意力相加
for key in self.attention_store:
for i in range(len(self.attention_store[key])):
self.attention_store[key][i] += self.step_store[key][i]
else:
# 如果不平均处理，初始化或追加步骤存储
if len(self.attention_store) == 0:
self.attention_store = [self.step_store]
else:
self.attention_store.append(self.step_store)
# 当前步骤计数加一
self.cur_step += 1
# 重置步骤存储为一个空的注意力存储
self.step_store = self.get_empty_store()
# 获取特定步骤的注意力数据
def get_attention(self, step: int):
# 如果需要平均，计算平均注意力
if self.average:
attention = {
key: [item / self.cur_step for item in self.attention_store[key]] for key in self.attention_store
}
else:
# 断言步骤不为空
assert step is not None
# 获取指定步骤的注意力
attention = self.attention_store[step]
return attention
# 聚合注意力的方法，处理多个输入
def aggregate_attention(
self, attention_maps, prompts, res: Union[int, Tuple[int]], from_where: List[str], is_cross: bool, select: int
):
# 初始化一个列表，包含 self.batch_size 个空列表，用于存储输出
out = [[] for x in range(self.batch_size)]
# 检查 res 是否为整数
if isinstance(res, int):
# 计算每个像素的数量，假设为 res 的平方
num_pixels = res**2
# 设置分辨率为 (res, res)
resolution = (res, res)
else:
# 计算像素数量为 res 的两个维度的乘积
num_pixels = res[0] * res[1]
# 设置分辨率为 res 的前两个值
resolution = res[:2]
# 遍历 from_where 列表
for location in from_where:
# 遍历与当前位置相关的注意力图
for bs_item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
# 枚举当前批次和项目
for batch, item in enumerate(bs_item):
# 检查当前项目的第二维是否等于 num_pixels
if item.shape[1] == num_pixels:
# 将项目重塑为指定形状，并选择相关的项目
cross_maps = item.reshape(len(prompts), -1, *resolution, item.shape[-1])[select]
# 将重塑的项目添加到输出的相应批次中
out[batch].append(cross_maps)
# 将每个批次的输出合并成一个张量
out = torch.stack([torch.cat(x, dim=0) for x in out])
# 对每个头进行平均
out = out.sum(1) / out.shape[1]
# 返回最终的输出
return out
def __init__(self, average: bool, batch_size=1, max_resolution=16, max_size: int = None):
# 初始化一个空的存储，用于保存步骤信息
self.step_store = self.get_empty_store()
# 初始化一个空的注意力存储列表
self.attention_store = []
# 当前步骤初始化为 0
self.cur_step = 0
# 设置平均标志
self.average = average
# 设置批次大小
self.batch_size = batch_size
# 如果 max_size 为空，计算最大尺寸为 max_resolution 的平方
if max_size is None:
self.max_size = max_resolution**2
# 如果 max_size 不为空且 max_resolution 为空，设置最大尺寸为 max_size
elif max_size is not None and max_resolution is None:
self.max_size = max_size
# 如果两个都被设置，则抛出错误
else:
raise ValueError("Only allowed to set one of max_resolution or max_size")
从 diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion 导入的 LeditsGaussianSmoothing 类
class LeditsGaussianSmoothing:
# 初始化函数，接受设备参数
def init(self, device):
# 定义高斯核的大小
kernel_size = [3, 3]
# 定义高斯核的标准差
sigma = [0.5, 0.5]
# 高斯核是每个维度的高斯函数的乘积
kernel = 1
# 创建网格，用于生成高斯核
meshgrids = torch.meshgrid([torch.arange(size, dtype=torch.float32) for size in kernel_size])
# 遍历每个维度的大小、标准差和网格
for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
# 计算均值
mean = (size - 1) / 2
# 更新高斯核值
kernel *= 1 / (std * math.sqrt(2 * math.pi)) * torch.exp(-(((mgrid - mean) / (2 * std)) ** 2))
# 确保高斯核的所有值之和等于 1
kernel = kernel / torch.sum(kernel)
# 将高斯核重塑为深度可分离卷积的权重
kernel = kernel.view(1, 1, *kernel.size())
# 重复高斯核以适应多通道输入
kernel = kernel.repeat(1, *[1] * (kernel.dim() - 1))
# 将权重转移到指定设备上
self.weight = kernel.to(device)
# 调用函数，用于应用高斯滤波
def __call__(self, input):
"""
参数:
对输入应用高斯滤波。
input (torch.Tensor): 要应用高斯滤波的输入。
返回:
filtered (torch.Tensor): 滤波后的输出。
"""
# 使用卷积函数对输入应用高斯滤波
return F.conv2d(input, weight=self.weight.to(input.dtype))
从 diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion 导入的 LEDITSCrossAttnProcessor 类
class LEDITSCrossAttnProcessor:
# 初始化函数，接受注意力存储、在 UNet 中的位置、PnP 和编辑提示
def init(self, attention_store, place_in_unet, pnp, editing_prompts):
# 设置注意力存储
self.attnstore = attention_store
# 设置在 UNet 中的位置
self.place_in_unet = place_in_unet
# 设置编辑提示
self.editing_prompts = editing_prompts
# 设置 PnP 参数
self.pnp = pnp
# 调用函数，用于处理注意力和隐藏状态
def __call__(
self,
attn: Attention,
hidden_states,
encoder_hidden_states,
attention_mask=None,
temb=None,
):
# 获取批次大小、序列长度和特征维度，如果没有编码器隐藏状态，则使用隐藏状态的形状
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
)
# 准备注意力掩码，调整为适合序列长度和批次大小
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
# 将隐藏状态转换为查询向量
query = attn.to_q(hidden_states)
# 如果没有编码器隐藏状态，使用当前隐藏状态；如果存在且需要归一化，则归一化编码器隐藏状态
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
elif attn.norm_cross:
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
# 将编码器隐藏状态转换为键向量
key = attn.to_k(encoder_hidden_states)
# 将编码器隐藏状态转换为值向量
value = attn.to_v(encoder_hidden_states)
# 将查询向量调整为批次维度
query = attn.head_to_batch_dim(query)
# 将键向量调整为批次维度
key = attn.head_to_batch_dim(key)
# 将值向量调整为批次维度
value = attn.head_to_batch_dim(value)
# 计算注意力得分
attention_probs = attn.get_attention_scores(query, key, attention_mask)
# 存储注意力得分到内部存储中，标记为跨注意力
self.attnstore(
attention_probs,
is_cross=True,
place_in_unet=self.place_in_unet,
editing_prompts=self.editing_prompts,
PnP=self.pnp,
)
# 使用注意力得分和值向量进行批次矩阵乘法以获取新的隐藏状态
hidden_states = torch.bmm(attention_probs, value)
# 将隐藏状态从批次维度转换回头部维度
hidden_states = attn.batch_to_head_dim(hidden_states)
# 通过线性投影转换隐藏状态
hidden_states = attn.to_out[0](hidden_states)
# 应用 dropout 操作
hidden_states = attn.to_out[1](hidden_states)
# 将隐藏状态按输出缩放因子进行归一化
hidden_states = hidden_states / attn.rescale_output_factor
# 返回最终的隐藏状态
return hidden_states
定义一个名为 LEditsPPPipelineStableDiffusionXL 的类，继承多个混入类
class LEditsPPPipelineStableDiffusionXL(
# 继承自 DiffusionPipeline，提供扩散模型功能
DiffusionPipeline,
# 继承自 FromSingleFileMixin，允许从单一文件加载数据
FromSingleFileMixin,
# 继承自 StableDiffusionXLLoraLoaderMixin，支持加载 LoRA 权重
StableDiffusionXLLoraLoaderMixin,
# 继承自 TextualInversionLoaderMixin，支持文本反演功能
TextualInversionLoaderMixin,
# 继承自 IPAdapterMixin，提供图像处理适配功能
IPAdapterMixin,
):
"""
使用 LEDits++ 和 Stable Diffusion XL 进行文本图像编辑的管道。
此模型继承自 [`DiffusionPipeline`] 并基于 [`StableDiffusionXLPipeline`]。查看超类文档以获取
所有管道实现的通用方法（下载、保存、在特定设备上运行等）。
此外，管道还继承了以下加载方法：
- *LoRA*: [`LEditsPPPipelineStableDiffusionXL.load_lora_weights`]
- *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`]
以及以下保存方法：
- *LoRA*: [`loaders.StableDiffusionXLPipeline.save_lora_weights`]
```
# 参数说明
Args:
vae ([`AutoencoderKL`]):
# 变分自编码器模型，用于将图像编码为潜在表示并解码。
text_encoder ([`~transformers.CLIPTextModel`]):
# 冻结的文本编码器，Stable Diffusion XL 使用 CLIP 的文本部分。
# 具体使用 [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) 变体。
text_encoder_2 ([`~transformers.CLIPTextModelWithProjection`]):
# 第二个冻结文本编码器，Stable Diffusion XL 使用 CLIP 的文本和池化部分。
# 具体使用 [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k) 变体。
tokenizer ([`~transformers.CLIPTokenizer`]):
# CLIPTokenizer 类的分词器。
tokenizer_2 ([`~transformers.CLIPTokenizer`]):
# 第二个 CLIPTokenizer 类的分词器。
unet ([`UNet2DConditionModel`]):
# 条件 U-Net 架构，用于对编码的图像潜在表示进行去噪。
scheduler ([`DPMSolverMultistepScheduler`] or [`DDIMScheduler`]):
# 与 `unet` 结合使用的调度器，用于去噪编码的图像潜在表示。
# 可以是 [`DPMSolverMultistepScheduler`] 或 [`DDIMScheduler`]，如传入其他调度器，将默认设置为 [`DPMSolverMultistepScheduler`]。
force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`):
# 是否强制将负提示嵌入始终设置为 0。另见 `stabilityai/stable-diffusion-xl-base-1-0` 的配置。
add_watermarker (`bool`, *optional*):
# 是否使用 [invisible_watermark 库](https://github.com/ShieldMnt/invisible-watermark/) 对输出图像进行水印处理。
# 如果未定义，且包已安装，则默认设置为 True，否则不使用水印。
"""
# 定义模型的 CPU 离线加载顺序
model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae"
# 定义可选组件列表
_optional_components = [
# 分词器、第二分词器、文本编码器等可选组件
"tokenizer",
"tokenizer_2",
"text_encoder",
"text_encoder_2",
"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, # 变分自编码器模型
text_encoder: CLIPTextModel, # 文本编码器模型
text_encoder_2: CLIPTextModelWithProjection, # 第二文本编码器模型，带投影
tokenizer: CLIPTokenizer, # 文本分词器
tokenizer_2: CLIPTokenizer, # 第二文本分词器
unet: UNet2DConditionModel, # UNet2D条件模型
scheduler: Union[DPMSolverMultistepScheduler, DDIMScheduler], # 调度器，可以是多步DPMSolver或DDIM调度器
image_encoder: CLIPVisionModelWithProjection = None, # 图像编码器，带投影，可选
feature_extractor: CLIPImageProcessor = None, # 特征提取器，可选
force_zeros_for_empty_prompt: bool = True, # 是否强制将空提示处理为零
add_watermarker: Optional[bool] = None, # 是否添加水印，可选
):
# 调用父类的初始化方法
super().__init__()
# 注册多个模块，使其可供后续使用
self.register_modules(
vae=vae,
text_encoder=text_encoder,
text_encoder_2=text_encoder_2,
tokenizer=tokenizer,
tokenizer_2=tokenizer_2,
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)
# 初始化图像处理器，使用计算的缩放因子
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
# 检查调度器类型并初始化为DPMSolverMultistepScheduler（如果必要）
if not isinstance(scheduler, DDIMScheduler) and not isinstance(scheduler, DPMSolverMultistepScheduler):
self.scheduler = DPMSolverMultistepScheduler.from_config(
scheduler.config, algorithm_type="sde-dpmsolver++", solver_order=2
)
# 记录警告，说明调度器已更改
logger.warning(
"This pipeline only supports DDIMScheduler and DPMSolverMultistepScheduler. "
"The scheduler has been changed to DPMSolverMultistepScheduler."
)
# 设置默认样本大小
self.default_sample_size = self.unet.config.sample_size
# 如果add_watermarker为None，则根据可用性确定是否添加水印
add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available()
# 根据是否添加水印初始化水印对象
if add_watermarker:
self.watermark = StableDiffusionXLWatermarker()
else:
self.watermark = None
# 初始化反转步骤为None
self.inversion_steps = None
# 编码提示方法，用于处理输入提示
def encode_prompt(
self,
device: Optional[torch.device] = None, # 指定计算设备，可选
num_images_per_prompt: int = 1, # 每个提示生成的图像数量
negative_prompt: Optional[str] = None, # 负提示，可选
negative_prompt_2: Optional[str] = None, # 第二个负提示，可选
negative_prompt_embeds: Optional[torch.Tensor] = None, # 负提示的嵌入表示，可选
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, # 负提示的池化嵌入表示，可选
lora_scale: Optional[float] = None, # LoRA缩放因子，可选
clip_skip: Optional[int] = None, # 跳过的CLIP层数量，可选
enable_edit_guidance: bool = True, # 是否启用编辑指导
editing_prompt: Optional[str] = None, # 编辑提示，可选
editing_prompt_embeds: Optional[torch.Tensor] = None, # 编辑提示的嵌入表示，可选
editing_pooled_prompt_embeds: Optional[torch.Tensor] = None, # 编辑提示的池化嵌入表示，可选
# 从diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs复制的内容
# 为调度器步骤准备额外的关键字参数，因不同调度器的签名不同
def prepare_extra_step_kwargs(self, eta, generator=None):
# eta (η) 仅用于 DDIMScheduler，其他调度器将忽略该参数
# eta 对应于 DDIM 论文中的 η，范围应在 [0, 1] 之间
# 检查调度器步骤是否接受 eta 参数
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 初始化额外步骤关键字参数字典
extra_step_kwargs = {}
# 如果接受 eta，则将其添加到额外步骤参数中
if accepts_eta:
extra_step_kwargs["eta"] = eta
# 检查调度器步骤是否接受 generator 参数
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 如果接受 generator，则将其添加到额外步骤参数中
if accepts_generator:
extra_step_kwargs["generator"] = generator
# 返回准备好的额外步骤参数
return extra_step_kwargs
# 检查输入参数的有效性
def check_inputs(
self,
negative_prompt=None,
negative_prompt_2=None,
negative_prompt_embeds=None,
negative_pooled_prompt_embeds=None,
):
# 如果同时提供 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."
)
# 如果同时提供 negative_prompt_2 和 negative_prompt_embeds，抛出错误
elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
# 如果提供 negative_prompt_embeds 但未提供 negative_pooled_prompt_embeds，抛出错误
if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
raise ValueError(
"If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
)
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents 修改而来
def prepare_latents(self, device, latents):
# 将潜在变量移动到指定设备
latents = latents.to(device)
# 根据调度器所需的标准差缩放初始噪声
latents = latents * self.scheduler.init_noise_sigma
# 返回处理后的潜在变量
return latents
# 获取添加时间标识的辅助方法
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.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae 复制的代码
def upcast_vae(self):
# 获取 VAE 的数据类型
dtype = self.vae.dtype
# 将 VAE 转换为 float32 数据类型
self.vae.to(dtype=torch.float32)
# 检查 VAE 解码器中的注意力处理器是否使用 Torch 2.0 或 XFormers
use_torch_2_0_or_xformers = isinstance(
self.vae.decoder.mid_block.attentions[0].processor,
(
AttnProcessor2_0,
XFormersAttnProcessor,
),
)
# 如果使用 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 和嵌入维度，默认嵌入维度为 512，数据类型默认为 float32
self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32
# 函数返回嵌入向量，类型为 torch.Tensor
) -> 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):
# 要生成的嵌入维度
Dimension of the embeddings to generate.
dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
# 生成的嵌入数据类型
Data type of the generated embeddings.
Returns:
`torch.Tensor`: # 返回嵌入向量，形状为 (len(w), embedding_dim)
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 转换为目标数据类型并计算最终的嵌入
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
# 此属性定义了与方程 (2) 中引导权重 w 类似的引导尺度
# 引导尺度 = 1 表示没有进行分类器自由引导。
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# corresponds to doing no classifier free guidance.
@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
# 从指定管道复制的内容
# Copied from diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.LEditsPPPipelineStableDiffusion.prepare_unet
# 准备 UNet 模型，设置注意力处理器
def prepare_unet(self, attention_store, PnP: bool = False):
# 初始化一个空字典用于存储注意力处理器
attn_procs = {}
# 遍历 UNet 的注意力处理器的键
for name in self.unet.attn_processors.keys():
# 如果名字以 "mid_block" 开头，则设置位置为 "mid"
if name.startswith("mid_block"):
place_in_unet = "mid"
# 如果名字以 "up_blocks" 开头，则设置位置为 "up"
elif name.startswith("up_blocks"):
place_in_unet = "up"
# 如果名字以 "down_blocks" 开头，则设置位置为 "down"
elif name.startswith("down_blocks"):
place_in_unet = "down"
# 如果名字不符合以上条件，则跳过当前循环
else:
continue
# 如果名字包含 "attn2" 且位置不是 "mid"
if "attn2" in name and place_in_unet != "mid":
# 创建 LEDITSCrossAttnProcessor 实例并加入字典
attn_procs[name] = LEDITSCrossAttnProcessor(
attention_store=attention_store, # 传入注意力存储
place_in_unet=place_in_unet, # 传入在 UNet 中的位置
pnp=PnP, # 传入 PnP 标志
editing_prompts=self.enabled_editing_prompts, # 传入启用的编辑提示
)
# 否则，创建默认的 AttnProcessor 实例
else:
attn_procs[name] = AttnProcessor()
# 设置 UNet 的注意力处理器
self.unet.set_attn_processor(attn_procs)
# 在不计算梯度的情况下调用该函数
@torch.no_grad()
# 用示例文档字符串替换文档字符串
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
# 可选的去噪结束时间
denoising_end: Optional[float] = None,
# 可选的负提示，可以是字符串或字符串列表
negative_prompt: Optional[Union[str, List[str]]] = None,
# 第二个可选的负提示
negative_prompt_2: Optional[Union[str, List[str]]] = None,
# 可选的负提示嵌入
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 可选的负池化提示嵌入
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
# 可选的输入适配器图像
ip_adapter_image: Optional[PipelineImageInput] = None,
# 输出类型的可选参数，默认为 "pil"
output_type: Optional[str] = "pil",
# 是否返回字典的可选参数，默认为 True
return_dict: bool = True,
# 可选的交叉注意力参数
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
# 引导缩放因子的可选参数，默认为 0.0
guidance_rescale: float = 0.0,
# 左上角裁剪坐标的可选参数，默认为 (0, 0)
crops_coords_top_left: Tuple[int, int] = (0, 0),
# 可选的目标尺寸
target_size: Optional[Tuple[int, int]] = None,
# 可选的编辑提示
editing_prompt: Optional[Union[str, List[str]]] = None,
# 可选的编辑提示嵌入
editing_prompt_embeddings: Optional[torch.Tensor] = None,
# 可选的编辑池化提示嵌入
editing_pooled_prompt_embeds: Optional[torch.Tensor] = None,
# 可选的反向编辑方向，默认为 False
reverse_editing_direction: Optional[Union[bool, List[bool]]] = False,
# 编辑引导缩放因子的可选参数，默认为 5
edit_guidance_scale: Optional[Union[float, List[float]]] = 5,
# 编辑热身步骤的可选参数，默认为 0
edit_warmup_steps: Optional[Union[int, List[int]]] = 0,
# 编辑冷却步骤的可选参数
edit_cooldown_steps: Optional[Union[int, List[int]]] = None,
# 编辑阈值的可选参数，默认为 0.9
edit_threshold: Optional[Union[float, List[float]]] = 0.9,
# 可选的语义引导张量列表
sem_guidance: Optional[List[torch.Tensor]] = None,
# 使用交叉注意力掩膜的可选参数，默认为 False
use_cross_attn_mask: bool = False,
# 使用交集掩膜的可选参数，默认为 False
use_intersect_mask: bool = False,
# 用户掩膜的可选参数
user_mask: Optional[torch.Tensor] = None,
# 注意力存储步骤的可选列表，默认为空列表
attn_store_steps: Optional[List[int]] = [],
# 是否在步骤间平均存储的可选参数，默认为 True
store_averaged_over_steps: bool = True,
# 可选的跳过剪辑的参数
clip_skip: Optional[int] = None,
# 可选的步骤结束回调函数
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
# 步骤结束时的张量输入回调参数，默认为 ["latents"]
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# 其他可选参数
**kwargs,
# 在不计算梯度的情况下调用该函数
@torch.no_grad()
# 从 diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.LEditsPPPipelineStableDiffusion.encode_image 修改而来
# 定义一个方法用于编码图像，接受多个可选参数
def encode_image(self, image, dtype=None, height=None, width=None, resize_mode="default", crops_coords=None):
# 预处理图像，调整大小和裁剪，返回处理后的图像
image = self.image_processor.preprocess(
image=image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords
)
# 后处理图像，转换为 PIL 格式
resized = self.image_processor.postprocess(image=image, output_type="pil")
# 检查图像的最大尺寸是否超过配置的采样大小的 1.5 倍
if max(image.shape[-2:]) > self.vae.config["sample_size"] * 1.5:
# 记录警告信息，提示用户输入图像分辨率过高
logger.warning(
"Your input images far exceed the default resolution of the underlying diffusion model. "
"The output images may contain severe artifacts! "
"Consider down-sampling the input using the `height` and `width` parameters"
)
# 将图像转换为指定的设备和数据类型
image = image.to(self.device, dtype=dtype)
# 检查是否需要将数据类型上调
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
# 如果需要上调数据类型
if needs_upcasting:
# 将图像转换为浮点数
image = image.float()
# 上调 VAE 模型的数据类型
self.upcast_vae()
# 使用 VAE 编码图像，获取潜在分布的模式
x0 = self.vae.encode(image).latent_dist.mode()
# 将模式转换为指定的数据类型
x0 = x0.to(dtype)
# 如果需要，转换回 fp16 数据类型
if needs_upcasting:
self.vae.to(dtype=torch.float16)
# 根据配置的缩放因子调整潜在向量
x0 = self.vae.config.scaling_factor * x0
# 返回潜在向量和调整后的图像
return x0, resized
# 装饰器，表示该方法在执行时不计算梯度
@torch.no_grad()
# 定义一个方法用于反转图像，接受多个可选参数
def invert(
self,
image: PipelineImageInput,
source_prompt: str = "",
source_guidance_scale=3.5,
negative_prompt: str = None,
negative_prompt_2: str = None,
num_inversion_steps: int = 50,
skip: float = 0.15,
generator: Optional[torch.Generator] = None,
crops_coords_top_left: Tuple[int, int] = (0, 0),
num_zero_noise_steps: int = 3,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
从 diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.rescale_noise_cfg 复制
def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
"""
根据 guidance_rescale 重新缩放 noise_cfg。基于 Common Diffusion Noise Schedules and
Sample Steps are Flawed 的研究结果。见第 3.4 节
"""
# 计算 noise_pred_text 的标准差，沿着指定维度，保持维度
std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
# 计算 noise_cfg 的标准差，沿着指定维度，保持维度
std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
# 通过标准差的比值重新缩放指导结果（修正过度曝光）
noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
# 用指导结果与原始结果混合，通过 guidance_rescale 因子避免“普通”图像
noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
# 返回重新缩放后的噪声配置
return noise_cfg
从 diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.compute_noise_ddim 复制
def compute_noise_ddim(scheduler, prev_latents, latents, timestep, noise_pred, eta):
# 1. 获取前一步值（t-1）
prev_timestep = timestep - scheduler.config.num_train_timesteps // scheduler.num_inference_steps
# 2. 计算 alphas 和 betas
# 当前时间步的累积 alpha 值
alpha_prod_t = scheduler.alphas_cumprod[timestep]
# 前一步的累积 alpha 值，若为负则使用最终的 alpha 值
alpha_prod_t_prev = (
scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else scheduler.final_alpha_cumprod
)
# 当前时间步的 beta 值
beta_prod_t = 1 - alpha_prod_t
# 3. 从预测噪声计算预测的原始样本，称为 "predicted x_0"
pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5)
# 4. 限制 "predicted x_0" 的范围
if scheduler.config.clip_sample:
# 将预测样本限制在 -1 和 1 之间
pred_original_sample = torch.clamp(pred_original_sample, -1, 1)
# 5. 计算方差：σ_t(η)，见公式 (16)
# σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
variance = scheduler._get_variance(timestep, prev_timestep)
# 计算标准差
std_dev_t = eta * variance ** (0.5)
# 6. 计算指向 x_t 的方向，见公式 (12)
pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * noise_pred
# 修改以返回更新后的 xtm1，以避免误差累积
mu_xt = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
# 如果方差大于 0，计算噪声
if variance > 0.0:
noise = (prev_latents - mu_xt) / (variance ** (0.5) * eta)
else:
# 否则噪声设置为零
noise = torch.tensor([0.0]).to(latents.device)
# 返回计算出的噪声和更新后的 mu_xt
return noise, mu_xt + (eta * variance**0.5) * noise
从 diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.compute_noise_sde_dpm_pp_2nd 复制
def compute_noise_sde_dpm_pp_2nd(scheduler, prev_latents, latents, timestep, noise_pred, eta):
# 定义一阶更新函数，输入模型输出和样本
def first_order_update(model_output, sample): # timestep, prev_timestep, sample):
# 获取当前和前一个步骤的 sigma 值
sigma_t, sigma_s = scheduler.sigmas[scheduler.step_index + 1], scheduler.sigmas[scheduler.step_index]
# 将 sigma 转换为 alpha 和 sigma_t
alpha_t, sigma_t = scheduler._sigma_to_alpha_sigma_t(sigma_t)
alpha_s, sigma_s = scheduler._sigma_to_alpha_sigma_t(sigma_s)
# 计算 lambda_t 和 lambda_s
lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
lambda_s = torch.log(alpha_s) - torch.log(sigma_s)
# 计算 h 值
h = lambda_t - lambda_s
# 计算 mu_xt，结合样本和模型输出
mu_xt = (sigma_t / sigma_s * torch.exp(-h)) * sample + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output
# 使用调度器的 DPM 解决方案更新 mu_xt
mu_xt = scheduler.dpm_solver_first_order_update(
model_output=model_output, sample=sample, noise=torch.zeros_like(sample)
)
# 计算 sigma
sigma = sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h))
# 如果 sigma 大于 0，计算噪声
if sigma > 0.0:
noise = (prev_latents - mu_xt) / sigma
else:
# 否则设定噪声为 0
noise = torch.tensor([0.0]).to(sample.device)
# 计算前一个样本
prev_sample = mu_xt + sigma * noise
# 返回噪声和前一个样本
return noise, prev_sample
# 定义二阶更新函数，输入模型输出列表和样本
def second_order_update(model_output_list, sample): # timestep_list, prev_timestep, sample):
# 获取当前和前两个步骤的 sigma 值
sigma_t, sigma_s0, sigma_s1 = (
scheduler.sigmas[scheduler.step_index + 1],
scheduler.sigmas[scheduler.step_index],
scheduler.sigmas[scheduler.step_index - 1],
)
# 将 sigma 转换为 alpha 和 sigma_t
alpha_t, sigma_t = scheduler._sigma_to_alpha_sigma_t(sigma_t)
alpha_s0, sigma_s0 = scheduler._sigma_to_alpha_sigma_t(sigma_s0)
alpha_s1, sigma_s1 = scheduler._sigma_to_alpha_sigma_t(sigma_s1)
# 计算 lambda 值
lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
# 获取最后两个模型输出
m0, m1 = model_output_list[-1], model_output_list[-2]
# 计算 h 和 h_0
h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1
r0 = h_0 / h
# 设定 D0 和 D1
D0, D1 = m0, (1.0 / r0) * (m0 - m1)
# 计算 mu_xt
mu_xt = (
(sigma_t / sigma_s0 * torch.exp(-h)) * sample
+ (alpha_t * (1 - torch.exp(-2.0 * h))) * D0
+ 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1
)
# 计算 sigma
sigma = sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h))
# 如果 sigma 大于 0，计算噪声
if sigma > 0.0:
noise = (prev_latents - mu_xt) / sigma
else:
# 否则设定噪声为 0
noise = torch.tensor([0.0]).to(sample.device)
# 计算前一个样本
prev_sample = mu_xt + sigma * noise
# 返回噪声和前一个样本
return noise, prev_sample
# 如果调度器的步骤索引为 None，初始化步骤索引
if scheduler.step_index is None:
scheduler._init_step_index(timestep)
# 将模型输出转换为可用格式
model_output = scheduler.convert_model_output(model_output=noise_pred, sample=latents)
# 更新模型输出列表
for i in range(scheduler.config.solver_order - 1):
scheduler.model_outputs[i] = scheduler.model_outputs[i + 1]
scheduler.model_outputs[-1] = model_output
# 根据低阶数量决定使用一阶或二阶更新
if scheduler.lower_order_nums < 1:
noise, prev_sample = first_order_update(model_output, latents)
else:
noise, prev_sample = second_order_update(scheduler.model_outputs, latents)
# 如果当前调度器的低阶数量小于配置中的求解器阶数
if scheduler.lower_order_nums < scheduler.config.solver_order:
# 增加低阶数量的计数
scheduler.lower_order_nums += 1
# 完成后，将步骤索引增加一
scheduler._step_index += 1
# 返回噪声和之前的样本
return noise, prev_sample
从 diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion 复制的代码
def compute_noise(scheduler, *args):
# 检查调度器是否为 DDIMScheduler 实例
if isinstance(scheduler, DDIMScheduler):
# 调用 DDIM 调度器的噪声计算函数并返回结果
return compute_noise_ddim(scheduler, *args)
# 检查调度器是否为 DPMSolverMultistepScheduler 实例，且满足特定配置
elif (
isinstance(scheduler, DPMSolverMultistepScheduler)
and scheduler.config.algorithm_type == "sde-dpmsolver++"
and scheduler.config.solver_order == 2
):
# 调用 SDE DPM 的二阶噪声计算函数并返回结果
return compute_noise_sde_dpm_pp_2nd(scheduler, *args)
else:
# 如果不满足以上条件，抛出未实现的错误
raise NotImplementedError
# `.\diffusers\pipelines\ledits_pp\pipeline_output.py`
```py
# 从 dataclass 模块导入 dataclass 装饰器
from dataclasses import dataclass
# 从 typing 模块导入 List, Optional, 和 Union 类型提示
from typing import List, Optional, Union
# 导入 numpy 库并重命名为 np
import numpy as np
# 导入 PIL.Image 模块
import PIL.Image
# 从上级目录的 utils 模块导入 BaseOutput 类
from ...utils import BaseOutput
# 定义 LEditsPPDiffusionPipelineOutput 类，继承自 BaseOutput
@dataclass
class LEditsPPDiffusionPipelineOutput(BaseOutput):
"""
LEdits++ Diffusion 流水线的输出类。
参数：
images (`List[PIL.Image.Image]` 或 `np.ndarray`)
长度为 `batch_size` 的去噪 PIL 图像列表或形状为 `(batch_size, height, width,
num_channels)` 的 NumPy 数组。
nsfw_content_detected (`List[bool]`)
表示生成图像是否包含“不适合工作” (nsfw) 内容的列表，如果无法进行安全检查则为 `None`。
"""
# 图像属性，可以是 PIL 图像列表或 NumPy 数组
images: Union[List[PIL.Image.Image], np.ndarray]
# NSFW 内容检测结果，可能是布尔值列表或 None
nsfw_content_detected: Optional[List[bool]]
# 定义 LEditsPPInversionPipelineOutput 类，继承自 BaseOutput
@dataclass
class LEditsPPInversionPipelineOutput(BaseOutput):
"""
LEdits++ Diffusion 流水线的输出类。
参数：
input_images (`List[PIL.Image.Image]` 或 `np.ndarray`)
长度为 `batch_size` 的裁剪和调整大小后的输入图像列表，或形状为 `
(batch_size, height, width, num_channels)` 的 NumPy 数组。
vae_reconstruction_images (`List[PIL.Image.Image]` 或 `np.ndarray`)
所有输入图像的 VAE 重建列表，作为长度为 `batch_size` 的 PIL 图像列表，或形状为
` (batch_size, height, width, num_channels)` 的 NumPy 数组。
"""
# 输入图像属性，可以是 PIL 图像列表或 NumPy 数组
images: Union[List[PIL.Image.Image], np.ndarray]
# VAE 重建图像属性，可以是 PIL 图像列表或 NumPy 数组
vae_reconstruction_images: Union[List[PIL.Image.Image], np.ndarray]
.\diffusers\pipelines\ledits_pp\__init__.py
# 从 typing 模块导入 TYPE_CHECKING，用于类型检查的标志
from typing import TYPE_CHECKING
# 从 utils 模块导入多个功能和常量
from ...utils import (
DIFFUSERS_SLOW_IMPORT, # 导入用于标识慢导入的常量
OptionalDependencyNotAvailable, # 导入表示可选依赖不可用的异常
_LazyModule, # 导入懒加载模块的类
get_objects_from_module, # 导入从模块获取对象的函数
is_torch_available, # 导入检查 PyTorch 是否可用的函数
is_transformers_available, # 导入检查 Transformers 是否可用的函数
)
# 初始化一个空字典，用于存储虚拟对象
_dummy_objects = {}
# 初始化一个空字典，用于存储模块导入结构
_import_structure = {}
# 尝试检测可选依赖项是否可用
try:
# 检查 Transformers 和 PyTorch 是否都可用，如果不可用则抛出异常
if not (is_transformers_available() and is_torch_available()):
raise OptionalDependencyNotAvailable()
# 捕获可选依赖不可用的异常
except OptionalDependencyNotAvailable:
# 从 utils 模块导入虚拟对象的模块，忽略相关的警告
from ...utils import dummy_torch_and_transformers_objects # noqa F403
# 更新虚拟对象字典，将虚拟对象加入到字典中
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
# 如果没有异常，执行以下代码
else:
# 将稳定扩散的管道添加到导入结构字典中
_import_structure["pipeline_leditspp_stable_diffusion"] = ["LEditsPPPipelineStableDiffusion"]
# 将 XL 版本的稳定扩散管道添加到导入结构字典中
_import_structure["pipeline_leditspp_stable_diffusion_xl"] = ["LEditsPPPipelineStableDiffusionXL"]
# 将管道输出的相关类添加到导入结构字典中
_import_structure["pipeline_output"] = ["LEditsPPDiffusionPipelineOutput", "LEditsPPDiffusionPipelineOutput"]
# 如果是类型检查或者使用了慢导入标志
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
# 尝试检测可选依赖项是否可用
try:
# 检查 Transformers 和 PyTorch 是否都可用，如果不可用则抛出异常
if not (is_transformers_available() and is_torch_available()):
raise OptionalDependencyNotAvailable()
# 捕获可选依赖不可用的异常
except OptionalDependencyNotAvailable:
# 从虚拟对象的模块导入所有内容
from ...utils.dummy_torch_and_transformers_objects import *
else:
# 从稳定扩散的管道模块导入相关类
from .pipeline_leditspp_stable_diffusion import (
LEditsPPDiffusionPipelineOutput, # 导入稳定扩散管道输出类
LEditsPPInversionPipelineOutput, # 导入逆向稳定扩散管道输出类
LEditsPPPipelineStableDiffusion, # 导入稳定扩散管道类
)
# 从 XL 版本的稳定扩散管道模块导入相关类
from .pipeline_leditspp_stable_diffusion_xl import LEditsPPPipelineStableDiffusionXL
# 如果不是类型检查且没有使用慢导入标志
else:
# 导入 sys 模块以便操作模块系统
import sys
# 使用懒加载模块创建当前模块的 lazy 实例，并将其赋值给当前模块
sys.modules[__name__] = _LazyModule(
__name__, # 当前模块的名称
globals()["__file__"], # 当前模块的文件路径
_import_structure, # 模块导入结构
module_spec=__spec__, # 模块的规格
)
# 将虚拟对象字典中的所有对象添加到当前模块中
for name, value in _dummy_objects.items():
setattr(sys.modules[__name__], name, value)
