diffusers 源码解析（六十二）
# 版权声明，声明版权所有者及许可协议信息
# Copyright 2024 ETH Zurich Computer Vision Lab and The HuggingFace Team. All rights reserved.
#
# 根据 Apache License, Version 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 math
# 导入数据类装饰器以创建数据类
from dataclasses import dataclass
# 导入可选类型、元组和联合类型的类型提示
from typing import Optional, Tuple, Union
# 导入 NumPy 以进行数组操作
import numpy as np
# 导入 PyTorch 库以进行深度学习操作
import torch
# 从配置工具导入配置混合类和注册到配置的功能
from ..configuration_utils import ConfigMixin, register_to_config
# 从工具模块导入基本输出类
from ..utils import BaseOutput
# 从 PyTorch 工具模块导入随机张量生成函数
from ..utils.torch_utils import randn_tensor
# 从调度工具模块导入调度混合类
from .scheduling_utils import SchedulerMixin
# 定义调度器输出的输出类，继承自基本输出类
@dataclass
class RePaintSchedulerOutput(BaseOutput):
"""
调度器步进函数输出的输出类。
参数:
prev_sample (`torch.Tensor` 形状为 `(batch_size, num_channels, height, width)` 的图像):
前一个时间步的计算样本 (x_{t-1})。 `prev_sample` 应作为下一个模型输入
在去噪循环中使用。
pred_original_sample (`torch.Tensor` 形状为 `(batch_size, num_channels, height, width)` 的图像):
基于当前时间步模型输出的预测去噪样本 (x_{0})。 `pred_original_sample` 可用于预览进展或指导。
"""
# 前一个时间步的样本
prev_sample: torch.Tensor
# 预测的去噪样本
pred_original_sample: torch.Tensor
# 从 diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar 复制的函数
def betas_for_alpha_bar(
num_diffusion_timesteps,
max_beta=0.999,
alpha_transform_type="cosine",
):
"""
创建一个 beta 调度程序，离散化给定的 alpha_t_bar 函数，该函数定义了时间 t = [0,1] 的 (1-beta) 的累积乘积。
包含一个 alpha_bar 函数，该函数接受参数 t 并将其转换为到当前扩散过程部分的 (1-beta) 的累积乘积。
参数:
num_diffusion_timesteps (`int`): 生成的 beta 数量。
max_beta (`float`): 使用的最大 beta；使用小于 1 的值以
防止奇异性。
alpha_transform_type (`str`, *可选*, 默认为 `cosine`): alpha_bar 的噪声调度类型。
从 `cosine` 或 `exp` 中选择
返回:
betas (`np.ndarray`): 调度程序用于模型输出步骤的 beta 值
"""
# 如果选择了余弦作为变换类型
if alpha_transform_type == "cosine":
# 定义 alpha_bar 函数，基于余弦计算
def alpha_bar_fn(t):
return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
# 如果选择了指数作为变换类型
elif alpha_transform_type == "exp":
# 定义 alpha_bar 函数，基于指数计算
def alpha_bar_fn(t):
return math.exp(t * -12.0)
# 如果 alpha_transform_type 不受支持，则引发一个值错误，包含不支持的类型信息
else:
raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
# 初始化一个空列表，用于存储 beta 值
betas = []
# 遍历每一个扩散时间步长
for i in range(num_diffusion_timesteps):
# 计算当前时间步的比例 t1
t1 = i / num_diffusion_timesteps
# 计算下一个时间步的比例 t2
t2 = (i + 1) / num_diffusion_timesteps
# 计算 beta 值，确保不超过 max_beta，并添加到列表中
betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
# 将 beta 列表转换为 PyTorch 张量，并指定数据类型为 float32
return torch.tensor(betas, dtype=torch.float32)
# 定义一个名为 RePaintScheduler 的类，继承自 SchedulerMixin 和 ConfigMixin
class RePaintScheduler(SchedulerMixin, ConfigMixin):
"""
`RePaintScheduler` 是用于在给定掩码内进行 DDPM 修复的调度器。
该模型继承自 [`SchedulerMixin`] 和 [`ConfigMixin`]。有关库为所有调度器实现的通用
方法（如加载和保存）的文档，请查看超类文档。
参数：
num_train_timesteps (`int`, defaults to 1000):
训练模型的扩散步骤数量。
beta_start (`float`, defaults to 0.0001):
推理的起始 `beta` 值。
beta_end (`float`, defaults to 0.02):
最终的 `beta` 值。
beta_schedule (`str`, defaults to `"linear"`):
beta 调度，一个将 beta 范围映射到模型步进序列的映射。可选值有
`linear`、`scaled_linear`、`squaredcos_cap_v2` 或 `sigmoid`。
eta (`float`):
在扩散步骤中添加噪声的噪声权重。如果值在 0.0 和 1.0 之间，对应于 DDIM 调度器；
如果值在 -0.0 和 1.0 之间，对应于 DDPM 调度器。
trained_betas (`np.ndarray`, *optional*):
直接将 beta 数组传递给构造函数，以绕过 `beta_start` 和 `beta_end`。
clip_sample (`bool`, defaults to `True`):
将预测样本裁剪在 -1 和 1 之间，以确保数值稳定性。
"""
# 定义调度器的顺序为 1
order = 1
# 注册到配置的方法，初始化调度器的参数
@register_to_config
def __init__(
# 训练步骤数量，默认为 1000
self,
num_train_timesteps: int = 1000,
# 起始 beta 值，默认为 0.0001
beta_start: float = 0.0001,
# 最终 beta 值，默认为 0.02
beta_end: float = 0.02,
# beta 调度类型，默认为 "linear"
beta_schedule: str = "linear",
# 噪声权重，默认为 0.0
eta: float = 0.0,
# 可选的 beta 数组，默认为 None
trained_betas: Optional[np.ndarray] = None,
# 是否裁剪样本，默认为 True
clip_sample: bool = True,
):
# 检查训练好的 beta 值是否为 None
if trained_betas is not None:
# 将训练好的 beta 值转换为 PyTorch 张量
self.betas = torch.from_numpy(trained_betas)
# 检查 beta 调度是否为线性
elif beta_schedule == "linear":
# 创建线性范围的 beta 值，从 beta_start 到 beta_end
self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
# 检查 beta 调度是否为缩放线性
elif beta_schedule == "scaled_linear":
# 这个调度特定于潜在扩散模型
# 创建 beta 值，先取平方根，再平方
self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
# 检查 beta 调度是否为 squaredcos_cap_v2
elif beta_schedule == "squaredcos_cap_v2":
# Glide 余弦调度
# 使用 alpha_bar 函数生成 beta 值
self.betas = betas_for_alpha_bar(num_train_timesteps)
# 检查 beta 调度是否为 sigmoid
elif beta_schedule == "sigmoid":
# GeoDiff sigmoid 调度
# 创建范围从 -6 到 6 的线性 beta 值
betas = torch.linspace(-6, 6, num_train_timesteps)
# 使用 sigmoid 函数缩放 beta 值，并映射到 [beta_start, beta_end]
self.betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start
else:
# 如果提供的 beta_schedule 未实现，抛出异常
raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
# 计算 alphas，等于 1 减去 betas
self.alphas = 1.0 - self.betas
# 计算 alphas 的累积乘积
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
# 创建一个值为 1.0 的张量
self.one = torch.tensor(1.0)
# 初始化最终的累积 alpha 值
self.final_alpha_cumprod = torch.tensor(1.0)
# 初始化噪声分布的标准差
self.init_noise_sigma = 1.0
# 可设置的值
# 推理步骤数初始化为 None
self.num_inference_steps = None
# 创建一个张量，表示从 0 到 num_train_timesteps 的反向数组
self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy())
# 初始化 eta 值
self.eta = eta
# 定义缩放模型输入的方法
def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
"""
确保与需要根据当前时间步缩放去噪模型输入的调度器的互换性。
参数:
sample (`torch.Tensor`):
输入样本。
timestep (`int`, *可选*):
扩散链中的当前时间步。
返回:
`torch.Tensor`:
缩放后的输入样本。
"""
# 返回原样本，未进行缩放
return sample
# 定义设置时间步的方法
def set_timesteps(
self,
num_inference_steps: int,
jump_length: int = 10,
jump_n_sample: int = 10,
device: Union[str, torch.device] = None,
# 定义用于设置离散时间步长的函数，在推理前运行
):
"""
设置扩散链的离散时间步（在推理之前运行）。
Args:
num_inference_steps (`int`):
用于生成样本的扩散步骤数量，如果使用，则`timesteps`必须为`None`。
jump_length (`int`, defaults to 10):
在时间上向前跳跃的步数，在进行一次跳跃时向后移动时间（在RePaint论文中表示为“j”）。请参阅论文中的图9和10。
jump_n_sample (`int`, defaults to 10):
对于所选时间样本，进行向前时间跳跃的次数。请参阅论文中的图9和10。
device (`str` or `torch.device`, *optional*):
应将时间步移动到的设备。如果为`None`，则时间步不移动。
"""
# 选择最小值，确保推理步骤不超过训练时间步
num_inference_steps = min(self.config.num_train_timesteps, num_inference_steps)
# 将推理步骤数量保存到实例变量
self.num_inference_steps = num_inference_steps
# 初始化时间步列表
timesteps = []
# 创建跳跃字典以记录时间跳跃的信息
jumps = {}
# 根据给定的跳跃长度，计算每个时间步的跳跃数量
for j in range(0, num_inference_steps - jump_length, jump_length):
jumps[j] = jump_n_sample - 1
# 初始化时间变量为推理步骤数量
t = num_inference_steps
# 从推理步骤开始，逐步向前计算时间步
while t >= 1:
# 递减时间步
t = t - 1
# 将当前时间步添加到列表
timesteps.append(t)
# 检查当前时间步是否需要跳跃
if jumps.get(t, 0) > 0:
jumps[t] = jumps[t] - 1
# 对于每个跳跃长度，向前增加时间步
for _ in range(jump_length):
t = t + 1
# 将跳跃后的时间步添加到列表
timesteps.append(t)
# 将时间步数组乘以缩放因子以适应训练时间步
timesteps = np.array(timesteps) * (self.config.num_train_timesteps // self.num_inference_steps)
# 将时间步转换为张量并移动到指定设备
self.timesteps = torch.from_numpy(timesteps).to(device)
# 定义获取方差的函数
def _get_variance(self, t):
# 计算前一个时间步
prev_timestep = t - self.config.num_train_timesteps // self.num_inference_steps
# 获取当前和前一个时间步的累积alpha值
alpha_prod_t = self.alphas_cumprod[t]
alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
# 计算当前和前一个时间步的beta值
beta_prod_t = 1 - alpha_prod_t
beta_prod_t_prev = 1 - alpha_prod_t_prev
# 对于t > 0，计算预测方差，具体公式见论文中的公式（6）和（7）
# 预测方差的计算公式，获取前一个样本
variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
# 返回计算得到的方差
return variance
# 定义步骤函数，用于执行模型推理
def step(
self,
model_output: torch.Tensor,
timestep: int,
sample: torch.Tensor,
original_image: torch.Tensor,
mask: torch.Tensor,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
# 定义撤销步骤的函数，接受样本、时间步和可选的生成器
def undo_step(self, sample, timestep, generator=None):
# 计算每个推断步骤的训练时间步数
n = self.config.num_train_timesteps // self.num_inference_steps
# 循环 n 次，进行每个时间步的处理
for i in range(n):
# 获取当前时间步对应的 beta 值
beta = self.betas[timestep + i]
# 如果设备类型是 MPS，则处理随机噪声
if sample.device.type == "mps":
# 在 MPS 上 randn 生成的随机数不具可复现性
noise = randn_tensor(sample.shape, dtype=sample.dtype, generator=generator)
# 将噪声移动到样本所在的设备
noise = noise.to(sample.device)
else:
# 在其他设备上生成随机噪声
noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
# 更新样本，按照公式进行噪声混合
sample = (1 - beta) ** 0.5 * sample + beta**0.5 * noise
# 返回更新后的样本
return sample
# 定义添加噪声的函数，接受原始样本、噪声和时间步，返回张量
def add_noise(
self,
original_samples: torch.Tensor,
noise: torch.Tensor,
timesteps: torch.IntTensor,
) -> torch.Tensor:
# 抛出未实现的错误，提示使用指定的方法进行训练
raise NotImplementedError("Use `DDPMScheduler.add_noise()` to train for sampling with RePaint.")
# 定义返回训练时间步数的函数
def __len__(self):
# 返回配置中定义的训练时间步数
return self.config.num_train_timesteps
# 版权所有 2024 Shuchen Xue 等人在中国科学院大学团队和 HuggingFace 团队。保留所有权利。
#
# 根据 Apache 许可证第 2.0 版（“许可证”）进行许可；
# 除非遵守许可证，否则您不得使用此文件。
# 您可以在以下地址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面同意，软件
# 根据许可证分发是以“原样”基础进行的，
# 不附带任何形式的明示或暗示的担保或条件。
# 请参见许可证以获取特定语言管理权限和
# 限制条款。
# 免责声明：请查看 https://arxiv.org/abs/2309.05019
# 该代码库是基于 https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py 的修改版本
# 导入数学库，用于数学运算
import math
# 从 typing 模块导入所需的类型注解
from typing import Callable, List, Optional, Tuple, Union
# 导入 numpy 库，用于数组和数学运算
import numpy as np
# 导入 torch 库，用于张量计算
import torch
# 从配置工具模块导入配置混合类和注册配置函数
from ..configuration_utils import ConfigMixin, register_to_config
# 从 utils 模块导入弃用警告工具
from ..utils import deprecate
# 从 torch_utils 模块导入生成随机张量的函数
from ..utils.torch_utils import randn_tensor
# 从调度工具模块导入调度器类和输出类
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
# 从 diffusers.schedulers.scheduling_ddpm 导入 betas_for_alpha_bar 函数
def betas_for_alpha_bar(
num_diffusion_timesteps, # 需要生成的 beta 数量
max_beta=0.999, # 最大 beta 值，使用小于 1 的值以防止奇点
alpha_transform_type="cosine", # alpha 变换类型，默认为余弦
):
"""
创建一个 beta 调度程序，该程序离散化给定的 alpha_t_bar 函数，该函数定义了
随时间推移 (1-beta) 的累积乘积，从 t = [0,1] 开始。
包含一个 alpha_bar 函数，该函数接受参数 t 并将其转换为 (1-beta) 的累积乘积
直到扩散过程的那部分。
参数：
num_diffusion_timesteps (`int`): 生成的 beta 数量。
max_beta (`float`): 要使用的最大 beta；使用小于 1 的值以
防止奇点。
alpha_transform_type (`str`, *可选*, 默认为 `cosine`): alpha_bar 的噪声调度类型。
从 `cosine` 或 `exp` 中选择
返回：
betas (`np.ndarray`): 调度程序用于更新模型输出的 betas
"""
# 如果 alpha 变换类型为余弦
if alpha_transform_type == "cosine":
# 定义 alpha_bar 函数，返回余弦平方值
def alpha_bar_fn(t):
return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
# 如果 alpha 变换类型为指数
elif alpha_transform_type == "exp":
# 定义 alpha_bar 函数，返回指数衰减值
def alpha_bar_fn(t):
return math.exp(t * -12.0)
# 如果 alpha 变换类型不支持，则引发错误
else:
raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
betas = [] # 初始化 betas 列表
# 遍历每个扩散时间步
for i in range(num_diffusion_timesteps):
t1 = i / num_diffusion_timesteps # 当前时间步归一化
t2 = (i + 1) / num_diffusion_timesteps # 下一个时间步归一化
# 计算 beta 值并添加到列表中，限制最大值
betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
# 返回以 torch 张量形式表示的 betas
return torch.tensor(betas, dtype=torch.float32)
# 定义一个新的类 SASolverScheduler，继承自 SchedulerMixin 和 ConfigMixin
class SASolverScheduler(SchedulerMixin, ConfigMixin):
"""
`SASolverScheduler` 是一个快速的专用高阶求解器，用于扩散 SDEs。
# 该模型继承自 [`SchedulerMixin`] 和 [`ConfigMixin`]，查看超类文档以了解库为所有调度器实现的通用方法，例如加载和保存。
_compatibles = [e.name for e in KarrasDiffusionSchedulers] # 创建一个包含 KarrasDiffusionSchedulers 中所有调度器名称的列表
order = 1 # 设置调度器的默认顺序为 1
@register_to_config # 将该初始化方法注册到配置中
def __init__( # 定义初始化方法
self, # 当前实例
num_train_timesteps: int = 1000, # 训练时的时间步数，默认为 1000
beta_start: float = 0.0001, # beta 参数的起始值，默认为 0.0001
beta_end: float = 0.02, # beta 参数的结束值，默认为 0.02
beta_schedule: str = "linear", # beta 调度类型，默认为线性
trained_betas: Optional[Union[np.ndarray, List[float]]] = None, # 可选的训练 beta 值，默认为 None
predictor_order: int = 2, # 预测器的阶数，默认为 2
corrector_order: int = 2, # 校正器的阶数，默认为 2
prediction_type: str = "epsilon", # 预测类型，默认为 "epsilon"
tau_func: Optional[Callable] = None, # 可选的 tau 函数，默认为 None
thresholding: bool = False, # 是否使用阈值处理，默认为 False
dynamic_thresholding_ratio: float = 0.995, # 动态阈值比例，默认为 0.995
sample_max_value: float = 1.0, # 采样的最大值，默认为 1.0
algorithm_type: str = "data_prediction", # 算法类型，默认为 "data_prediction"
lower_order_final: bool = True, # 是否使用较低阶数的最终步骤，默认为 True
use_karras_sigmas: Optional[bool] = False, # 是否使用 Karras 的 sigma 值，默认为 False
lambda_min_clipped: float = -float("inf"), # lambda 的最小裁剪值，默认为负无穷大
variance_type: Optional[str] = None, # 可选的方差类型，默认为 None
timestep_spacing: str = "linspace", # 时间步间隔类型，默认为线性间隔
steps_offset: int = 0, # 时间步偏移量，默认为 0
):
# 如果训练好的 beta 参数不为 None，则将其转换为浮点张量
if trained_betas is not None:
self.betas = torch.tensor(trained_betas, dtype=torch.float32)
# 如果 beta_schedule 为 "linear"，则生成线性间隔的 beta 值
elif beta_schedule == "linear":
self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
# 如果 beta_schedule 为 "scaled_linear"，则按特定方式生成 beta 值
elif beta_schedule == "scaled_linear":
# 此调度非常特定于潜在扩散模型
self.betas = (
torch.linspace(
beta_start**0.5,
beta_end**0.5,
num_train_timesteps,
dtype=torch.float32,
)
** 2 # 平方以得到最终的 beta 值
)
# 如果 beta_schedule 为 "squaredcos_cap_v2"，则使用 Glide 余弦调度生成 beta 值
elif beta_schedule == "squaredcos_cap_v2":
# Glide 余弦调度
self.betas = betas_for_alpha_bar(num_train_timesteps)
# 如果 beta_schedule 不匹配任何已知类型，则抛出未实现的错误
else:
raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
# 计算 alpha 值为 1 减去对应的 beta 值
self.alphas = 1.0 - self.betas
# 计算 alpha 值的累积乘积
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
# 当前仅支持 VP 类型的噪声调度
self.alpha_t = torch.sqrt(self.alphas_cumprod) # 计算当前时刻的 alpha 值的平方根
self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) # 计算当前时刻的噪声标准差
self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) # 计算 lambda 值
self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5 # 计算 sigma 值
# 初始化噪声分布的标准差
self.init_noise_sigma = 1.0
# 检查算法类型是否有效，若无效则抛出未实现的错误
if algorithm_type not in ["data_prediction", "noise_prediction"]:
raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}")
# 可设定的值
self.num_inference_steps = None # 推理步骤数初始化为 None
# 创建时间步数组，从 0 到 num_train_timesteps-1 的线性间隔
timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy()
self.timesteps = torch.from_numpy(timesteps) # 将时间步数组转换为张量
# 初始化时间步列表和模型输出列表
self.timestep_list = [None] * max(predictor_order, corrector_order - 1)
self.model_outputs = [None] * max(predictor_order, corrector_order - 1)
# 如果 tau_func 为 None，设置默认的 tau 函数
if tau_func is None:
self.tau_func = lambda t: 1 if t >= 200 and t <= 800 else 0
else:
self.tau_func = tau_func # 否则使用传入的 tau 函数
# 根据算法类型确定是否预测 x0
self.predict_x0 = algorithm_type == "data_prediction"
self.lower_order_nums = 0 # 较低阶数初始化为 0
self.last_sample = None # 上一个样本初始化为 None
self._step_index = None # 当前步骤索引初始化为 None
self._begin_index = None # 起始步骤索引初始化为 None
self.sigmas = self.sigmas.to("cpu") # 将 sigmas 移动到 CPU，以减少 CPU/GPU 之间的通信
@property
def step_index(self):
"""
当前时间步的索引计数器。每次调度器步骤后增加 1。
"""
return self._step_index # 返回当前步骤索引
@property
def begin_index(self):
"""
第一个时间步的索引。应通过 `set_begin_index` 方法从管道设置。
"""
return self._begin_index # 返回起始步骤索引
# 从 diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index 复制
# 设置调度器的起始索引，默认值为0
def set_begin_index(self, begin_index: int = 0):
"""
设置调度器的起始索引。此函数应在推理前从管道运行。
Args:
begin_index (`int`):
调度器的起始索引。
"""
# 将提供的起始索引赋值给内部变量
self._begin_index = begin_index
# 从 diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample 复制的代码
def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
"""
"动态阈值处理：在每个采样步骤中，我们将 s 设置为 xt0 中的某个百分位绝对像素值（在时间步 t 的 x_0 预测），
如果 s > 1，则将 xt0 限制在范围 [-s, s]，然后除以 s。动态阈值处理将饱和像素（接近 -1 和 1 的像素）推向内部，
从而在每一步中积极防止像素饱和。我们发现动态阈值处理显著改善了照片现实感以及图像文本对齐，尤其是在使用非常大的引导权重时。"
https://arxiv.org/abs/2205.11487
"""
# 获取样本的数值类型
dtype = sample.dtype
# 获取批大小、通道数和其他维度
batch_size, channels, *remaining_dims = sample.shape
# 如果样本数据类型不是 float32 或 float64，则将其转换为 float
if dtype not in (torch.float32, torch.float64):
sample = sample.float() # 为百分位数计算向上转换，并且对 CPU 半精度未实现限制
# 将样本展平以进行每幅图像的百分位数计算
sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
# 获取样本的绝对值
abs_sample = sample.abs() # "某个百分位绝对像素值"
# 计算绝对样本的指定百分位数
s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
# 限制 s 的范围，以避免其小于1或超过最大值
s = torch.clamp(
s, min=1, max=self.config.sample_max_value
) # 当限制到 min=1 时，相当于标准剪切到 [-1, 1]
# 扩展 s 以适应批处理
s = s.unsqueeze(1) # (batch_size, 1) 因为限制会在维度0上广播
# 限制样本并将其缩放到范围 [-s, s]
sample = torch.clamp(sample, -s, s) / s # "将 xt0 限制在范围 [-s, s] 并除以 s"
# 将样本形状恢复为原始形状
sample = sample.reshape(batch_size, channels, *remaining_dims)
# 将样本转换回原始数据类型
sample = sample.to(dtype)
# 返回处理后的样本
return sample
# 从 diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t 复制的代码
def _sigma_to_t(self, sigma, log_sigmas):
# 获取 sigma 的对数值
log_sigma = np.log(np.maximum(sigma, 1e-10))
# 获取分布
dists = log_sigma - log_sigmas[:, np.newaxis]
# 获取 sigma 的范围
low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)
high_idx = low_idx + 1
# 获取低和高对数 sigma 值
low = log_sigmas[low_idx]
high = log_sigmas[high_idx]
# 插值计算 sigma
w = (low - log_sigma) / (low - high)
w = np.clip(w, 0, 1)
# 将插值转换为时间范围
t = (1 - w) * low_idx + w * high_idx
# 将 t 的形状调整为 sigma 的形状
t = t.reshape(sigma.shape)
# 返回计算得到的 t 值
return t
# 从 diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t 拷贝的函数
def _sigma_to_alpha_sigma_t(self, sigma):
# 根据 sigma 计算 alpha_t，公式为 1 / √(sigma² + 1)
alpha_t = 1 / ((sigma**2 + 1) ** 0.5)
# 计算 sigma_t，公式为 sigma * alpha_t
sigma_t = sigma * alpha_t
# 返回 alpha_t 和 sigma_t
return alpha_t, sigma_t
# 从 diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras 拷贝的函数
def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
"""构建 Karras 等人 (2022) 的噪声调度。"""
# 确保其他拷贝此函数的调度器不会出现错误的黑客方法
# TODO: 将此逻辑添加到其他调度器中
if hasattr(self.config, "sigma_min"):
# 如果配置中有 sigma_min，则使用该值
sigma_min = self.config.sigma_min
else:
# 否则将 sigma_min 设为 None
sigma_min = None
if hasattr(self.config, "sigma_max"):
# 如果配置中有 sigma_max，则使用该值
sigma_max = self.config.sigma_max
else:
# 否则将 sigma_max 设为 None
sigma_max = None
# 如果 sigma_min 为 None，则取 in_sigmas 的最后一个值
sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
# 如果 sigma_max 为 None，则取 in_sigmas 的第一个值
sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
# 设置 rho 值为 7.0，参考文献中的用法
rho = 7.0
# 创建一个从 0 到 1 的线性 ramp
ramp = np.linspace(0, 1, num_inference_steps)
# 计算 min_inv_rho 和 max_inv_rho
min_inv_rho = sigma_min ** (1 / rho)
max_inv_rho = sigma_max ** (1 / rho)
# 计算 sigmas，使用 ramp 插值
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
# 返回计算得到的 sigmas
return sigmas
def convert_model_output(
self,
model_output: torch.Tensor,
*args,
sample: torch.Tensor = None,
**kwargs,
# 计算指数负函数在给定区间的积分
def get_coefficients_exponential_negative(self, order, interval_start, interval_end):
"""
计算从 interval_start 到 interval_end 的积分 exp(-x) * x^order dx
"""
# 确保 order 只支持 0, 1, 2 和 3
assert order in [0, 1, 2, 3], "order is only supported for 0, 1, 2 and 3"
# 如果 order 为 0，使用对应的公式计算并返回结果
if order == 0:
return torch.exp(-interval_end) * (torch.exp(interval_end - interval_start) - 1)
# 如果 order 为 1，使用对应的公式计算并返回结果
elif order == 1:
return torch.exp(-interval_end) * (
(interval_start + 1) * torch.exp(interval_end - interval_start) - (interval_end + 1)
)
# 如果 order 为 2，使用对应的公式计算并返回结果
elif order == 2:
return torch.exp(-interval_end) * (
(interval_start**2 + 2 * interval_start + 2) * torch.exp(interval_end - interval_start)
- (interval_end**2 + 2 * interval_end + 2)
)
# 如果 order 为 3，使用对应的公式计算并返回结果
elif order == 3:
return torch.exp(-interval_end) * (
(interval_start**3 + 3 * interval_start**2 + 6 * interval_start + 6)
* torch.exp(interval_end - interval_start)
- (interval_end**3 + 3 * interval_end**2 + 6 * interval_end + 6)
)
# 计算给定区间内与指数相关的积分系数
def get_coefficients_exponential_positive(self, order, interval_start, interval_end, tau):
"""
计算从 interval_start 到 interval_end 的积分
公式为 exp(x(1+tau^2)) * x^order dx
"""
# 确保 order 的值只在支持的范围内
assert order in [0, 1, 2, 3], "order is only supported for 0, 1, 2 and 3"
# 变量替换后计算结束区间
interval_end_cov = (1 + tau**2) * interval_end
# 变量替换后计算起始区间
interval_start_cov = (1 + tau**2) * interval_start
# 处理 order 为 0 的情况
if order == 0:
return (
# 计算并返回积分结果
torch.exp(interval_end_cov) * (1 - torch.exp(-(interval_end_cov - interval_start_cov))) / (1 + tau**2)
)
# 处理 order 为 1 的情况
elif order == 1:
return (
# 计算并返回积分结果
torch.exp(interval_end_cov)
* (
(interval_end_cov - 1)
- (interval_start_cov - 1) * torch.exp(-(interval_end_cov - interval_start_cov))
)
/ ((1 + tau**2) ** 2)
)
# 处理 order 为 2 的情况
elif order == 2:
return (
# 计算并返回积分结果
torch.exp(interval_end_cov)
* (
(interval_end_cov**2 - 2 * interval_end_cov + 2)
- (interval_start_cov**2 - 2 * interval_start_cov + 2)
* torch.exp(-(interval_end_cov - interval_start_cov))
)
/ ((1 + tau**2) ** 3)
)
# 处理 order 为 3 的情况
elif order == 3:
return (
# 计算并返回积分结果
torch.exp(interval_end_cov)
* (
(interval_end_cov**3 - 3 * interval_end_cov**2 + 6 * interval_end_cov - 6)
- (interval_start_cov**3 - 3 * interval_start_cov**2 + 6 * interval_start_cov - 6)
* torch.exp(-(interval_end_cov - interval_start_cov))
)
/ ((1 + tau**2) ** 4)
)
# 计算系数的函数
def get_coefficients_fn(self, order, interval_start, interval_end, lambda_list, tau):
# 确保 order 的值在支持的范围内
assert order in [1, 2, 3, 4]
# 确保 lambda_list 的长度等于 order
assert order == len(lambda_list), "the length of lambda list must be equal to the order"
# 初始化系数列表
coefficients = []
# 计算拉格朗日多项式系数
lagrange_coefficient = self.lagrange_polynomial_coefficient(order - 1, lambda_list)
# 遍历 order 的范围
for i in range(order):
# 初始化单个系数
coefficient = 0
# 遍历 order 的范围
for j in range(order):
# 根据预测标志选择不同的积分计算方法
if self.predict_x0:
# 计算正指数的系数
coefficient += lagrange_coefficient[i][j] * self.get_coefficients_exponential_positive(
order - 1 - j, interval_start, interval_end, tau
)
else:
# 计算负指数的系数
coefficient += lagrange_coefficient[i][j] * self.get_coefficients_exponential_negative(
order - 1 - j, interval_start, interval_end
)
# 将计算得到的系数添加到列表中
coefficients.append(coefficient)
# 确保系数列表的长度与 order 匹配
assert len(coefficients) == order, "the length of coefficients does not match the order"
# 返回计算得到的系数列表
return coefficients
# 定义随机亚当斯-巴什福斯更新函数
def stochastic_adams_bashforth_update(
self,
model_output: torch.Tensor, # 模型输出的张量
*args, # 额外的参数，使用可变参数
sample: torch.Tensor, # 输入样本的张量
noise: torch.Tensor, # 噪声的张量
order: int, # 更新的顺序
tau: torch.Tensor, # 时间步长的张量
**kwargs, # 额外的关键字参数
# 定义随机亚当斯-莫尔顿更新函数
def stochastic_adams_moulton_update(
self,
this_model_output: torch.Tensor, # 当前模型输出的张量
*args, # 额外的参数，使用可变参数
last_sample: torch.Tensor, # 上一个样本的张量
last_noise: torch.Tensor, # 上一个噪声的张量
this_sample: torch.Tensor, # 当前样本的张量
order: int, # 更新的顺序
tau: torch.Tensor, # 时间步长的张量
**kwargs, # 额外的关键字参数
# 从 diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep 复制的函数
def index_for_timestep(self, timestep, schedule_timesteps=None):
# 如果未提供调度时间步，使用默认时间步
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
# 查找与当前时间步相匹配的索引候选
index_candidates = (schedule_timesteps == timestep).nonzero()
# 如果没有找到匹配的候选索引
if len(index_candidates) == 0:
# 设置步索引为时间步列表的最后一个索引
step_index = len(self.timesteps) - 1
# 如果找到多个匹配的候选索引
# 第一个步骤的 sigma 索引总是取第二个索引（或只有一个时取最后一个索引）
# 这样可以确保在去噪声调度中间开始时不会意外跳过一个 sigma
elif len(index_candidates) > 1:
# 取第二个候选索引作为步索引
step_index = index_candidates[1].item()
else:
# 取第一个候选索引作为步索引
step_index = index_candidates[0].item()
# 返回步索引
return step_index
# 从 diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index 复制的函数
def _init_step_index(self, timestep):
"""
初始化调度器的 step_index 计数器。
"""
# 如果 begin_index 为空
if self.begin_index is None:
# 如果时间步是张量类型，将其转换为与时间步设备一致
if isinstance(timestep, torch.Tensor):
timestep = timestep.to(self.timesteps.device)
# 根据时间步获取步索引
self._step_index = self.index_for_timestep(timestep)
else:
# 使用预定义的开始索引
self._step_index = self._begin_index
# 定义步函数
def step(
self,
model_output: torch.Tensor, # 模型输出的张量
timestep: int, # 当前时间步
sample: torch.Tensor, # 输入样本的张量
generator=None, # 可选的生成器
return_dict: bool = True, # 是否返回字典
# 定义缩放模型输入的函数
def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
"""
确保与需要根据当前时间步缩放去噪声模型输入的调度器的可互换性。
参数：
sample (`torch.Tensor`):
输入样本。
返回：
`torch.Tensor`:
缩放后的输入样本。
"""
# 直接返回输入样本
return sample
# 从 diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise 复制的函数
def add_noise(
self,
original_samples: torch.Tensor, # 原始样本的张量
noise: torch.Tensor, # 噪声的张量
timesteps: torch.IntTensor, # 时间步的整数张量
) -> torch.Tensor:
# 确保 alphas_cumprod 和 timestep 与 original_samples 具有相同的设备和数据类型
# 将 self.alphas_cumprod 移动到目标设备，以避免后续 add_noise 调用时冗余的 CPU 到 GPU 数据移动
self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device)
# 将 alphas_cumprod 转换为 original_samples 的数据类型
alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype)
# 将 timesteps 移动到 original_samples 的设备
timesteps = timesteps.to(original_samples.device)
# 计算 sqrt_alpha_prod 为 alphas_cumprod 在 timesteps 索引下的平方根
sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
# 将 sqrt_alpha_prod 展平为一维张量
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
# 如果 sqrt_alpha_prod 的维度少于 original_samples 的维度，则在最后一维添加一个新的维度
while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
# 计算 sqrt_one_minus_alpha_prod 为 1 减去 alphas_cumprod 在 timesteps 索引下的平方根
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
# 将 sqrt_one_minus_alpha_prod 展平为一维张量
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
# 如果 sqrt_one_minus_alpha_prod 的维度少于 original_samples 的维度，则在最后一维添加一个新的维度
while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
# 计算加噪声的样本，将原始样本与噪声按加权组合
noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
# 返回加噪声后的样本
return noisy_samples
# 定义获取对象长度的方法
def __len__(self):
# 返回配置中定义的训练时间步数
return self.config.num_train_timesteps
# 版权所有 2024 Google Brain 和 HuggingFace 团队。保留所有权利。
#
# 根据 Apache 许可证第 2.0 版（“许可证”）许可；
# 除非遵守许可证，否则您不得使用此文件。
# 您可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有规定，否则根据许可证分发的软件按“现状”提供，
# 不提供任何形式的保证或条件，无论是明示或暗示的。
# 有关许可证具体条款的权限和限制，请参阅许可证。
# 免责声明：此文件受到 https://github.com/yang-song/score_sde_pytorch 的强烈影响
# 导入数学库，用于数学运算
import math
# 从 dataclasses 模块导入 dataclass，用于创建数据类
from dataclasses import dataclass
# 从 typing 模块导入 Optional、Tuple 和 Union，用于类型注释
from typing import Optional, Tuple, Union
# 导入 PyTorch 库，用于深度学习模型的构建和训练
import torch
# 导入配置工具类和注册配置装饰器
from ..configuration_utils import ConfigMixin, register_to_config
# 导入基本输出类
from ..utils import BaseOutput
# 导入生成随机张量的工具函数
from ..utils.torch_utils import randn_tensor
# 导入调度相关的工具类
from .scheduling_utils import SchedulerMixin, SchedulerOutput
# 定义一个数据类，表示调度器的输出
@dataclass
class SdeVeOutput(BaseOutput):
"""
调度器 `step` 函数输出的输出类。
参数：
prev_sample (`torch.Tensor`，形状为 `(batch_size, num_channels, height, width)`，用于图像):
上一时间步的计算样本 `(x_{t-1})`。`prev_sample` 应作为去噪循环中的下一个模型输入。
prev_sample_mean (`torch.Tensor`，形状为 `(batch_size, num_channels, height, width)`，用于图像):
在前一时间步平均的 `prev_sample`。
"""
# 定义上一时间步的样本张量
prev_sample: torch.Tensor
# 定义上一时间步样本的平均值张量
prev_sample_mean: torch.Tensor
# 定义一个类，表示方差爆炸的随机微分方程调度器
class ScoreSdeVeScheduler(SchedulerMixin, ConfigMixin):
"""
`ScoreSdeVeScheduler` 是一个方差爆炸的随机微分方程 (SDE) 调度器。
此模型继承自 [`SchedulerMixin`] 和 [`ConfigMixin`]。请查看超类文档以获取库为所有调度器实现的通用方法，例如加载和保存。
参数：
num_train_timesteps (`int`，默认值为 1000):
用于训练模型的扩散步骤数。
snr (`float`，默认值为 0.15):
权衡来自 `model_output` 样本（来自网络）到随机噪声的步骤系数。
sigma_min (`float`，默认值为 0.01):
采样过程中 sigma 序列的初始噪声尺度。最小 sigma 应该反映数据的分布。
sigma_max (`float`，默认值为 1348.0):
用于传入模型的连续时间步范围的最大值。
sampling_eps (`float`，默认值为 1e-5):
采样结束值，时间步逐渐从 1 减少到 epsilon。
correct_steps (`int`，默认值为 1):
对生成样本执行的校正步骤数量。
"""
# 定义调度器的顺序
order = 1
# 将配置注册到调度器
@register_to_config
# 初始化类的构造函数
def __init__(
# 训练时间步的数量，默认为2000
num_train_timesteps: int = 2000,
# 信噪比，默认为0.15
snr: float = 0.15,
# 最小噪声标准差，默认为0.01
sigma_min: float = 0.01,
# 最大噪声标准差，默认为1348.0
sigma_max: float = 1348.0,
# 采样的最小值，默认为1e-5
sampling_eps: float = 1e-5,
# 校正步数，默认为1
correct_steps: int = 1,
):
# 设置初始噪声分布的标准差为最大值
self.init_noise_sigma = sigma_max
# 可设置的时间步值，初始为None
self.timesteps = None
# 设置噪声标准差
self.set_sigmas(num_train_timesteps, sigma_min, sigma_max, sampling_eps)
# 输入样本和时间步的缩放函数
def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
"""
确保可以与调度器兼容，根据当前时间步缩放去噪模型的输入。
参数：
sample (`torch.Tensor`):
输入样本。
timestep (`int`, *可选*):
扩散链中的当前时间步。
返回：
`torch.Tensor`:
一个缩放后的输入样本。
"""
# 返回未缩放的样本
return sample
# 设置扩散链的连续时间步
def set_timesteps(
# 推断步骤的数量
self, num_inference_steps: int, sampling_eps: float = None, device: Union[str, torch.device] = None
):
"""
在推断前设置用于扩散链的连续时间步。
参数：
num_inference_steps (`int`):
生成样本时使用的扩散步骤数量。
sampling_eps (`float`, *可选*):
最终时间步值（覆盖调度器实例化时提供的值）。
device (`str` 或 `torch.device`, *可选*):
时间步要移动到的设备。如果为`None`，则不移动。
"""
# 设置采样最小值，如果未提供则使用配置中的值
sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps
# 创建等间隔的时间步，并将其放在指定设备上
self.timesteps = torch.linspace(1, sampling_eps, num_inference_steps, device=device)
# 设置噪声标准差的函数
def set_sigmas(
# 推断步骤的数量
self, num_inference_steps: int, sigma_min: float = None, sigma_max: float = None, sampling_eps: float = None
):
"""
设置扩散链中使用的噪声尺度（在推理之前运行）。sigmas 控制样本更新中 `drift` 和 `diffusion` 组件的权重。
参数：
num_inference_steps (`int`):
生成样本时使用的扩散步骤数量，使用预训练模型。
sigma_min (`float`, optional):
初始噪声尺度值（覆盖调度器实例化时给定的值）。
sigma_max (`float`, optional):
最终噪声尺度值（覆盖调度器实例化时给定的值）。
sampling_eps (`float`, optional):
最终时间步值（覆盖调度器实例化时给定的值）。
"""
# 如果 sigma_min 为 None，则使用配置中的 sigma_min
sigma_min = sigma_min if sigma_min is not None else self.config.sigma_min
# 如果 sigma_max 为 None，则使用配置中的 sigma_max
sigma_max = sigma_max if sigma_max is not None else self.config.sigma_max
# 如果 sampling_eps 为 None，则使用配置中的 sampling_eps
sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps
# 如果 timesteps 为 None，则根据推理步骤和采样 epsilon 设置时间步
if self.timesteps is None:
self.set_timesteps(num_inference_steps, sampling_eps)
# 计算 sigmas，依据 sigma_min 和 sigma_max 以及时间步
self.sigmas = sigma_min * (sigma_max / sigma_min) ** (self.timesteps / sampling_eps)
# 生成离散的 sigmas，通过指数函数计算从 sigma_min 到 sigma_max 的数值
self.discrete_sigmas = torch.exp(torch.linspace(math.log(sigma_min), math.log(sigma_max), num_inference_steps))
# 将 sigmas 转换为张量形式，依据时间步进行计算
self.sigmas = torch.tensor([sigma_min * (sigma_max / sigma_min) ** t for t in self.timesteps])
# 获取与当前时间步相邻的 sigma 值
def get_adjacent_sigma(self, timesteps, t):
# 如果当前时间步为 0，则返回与 t 相同形状的零张量，否则返回离散 sigma 值
return torch.where(
timesteps == 0,
torch.zeros_like(t.to(timesteps.device)),
self.discrete_sigmas[timesteps - 1].to(timesteps.device),
)
# 步骤预测，处理模型输出、当前时间步和样本
def step_pred(
self,
model_output: torch.Tensor,
timestep: int,
sample: torch.Tensor,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
# 步骤校正，处理模型输出和样本
def step_correct(
self,
model_output: torch.Tensor,
sample: torch.Tensor,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
) -> Union[SchedulerOutput, Tuple]:
"""
根据网络的 `model_output` 校正预测样本。通常在对前一个时间步的预测后重复运行此函数。
参数：
model_output (`torch.Tensor`):
来自学习扩散模型的直接输出。
sample (`torch.Tensor`):
通过扩散过程创建的当前样本实例。
generator (`torch.Generator`, *可选*):
随机数生成器。
return_dict (`bool`, *可选*, 默认值为 `True`):
是否返回 [`~schedulers.scheduling_sde_ve.SdeVeOutput`] 或 `tuple`。
返回：
[`~schedulers.scheduling_sde_ve.SdeVeOutput`] 或 `tuple`:
如果 `return_dict` 为 `True`，返回 [`~schedulers.scheduling_sde_ve.SdeVeOutput`]，否则返回一个元组
其第一个元素是样本张量。
"""
# 检查时间步是否被设置，如果没有则抛出错误
if self.timesteps is None:
raise ValueError(
"`self.timesteps` 未设置，需要在创建调度器后运行 'set_timesteps'"
)
# 对于小批量大小，论文建议用 sqrt(d) 替代 norm(z)，其中 d 是 z 的维度
# 为校正生成噪声
noise = randn_tensor(sample.shape, layout=sample.layout, generator=generator).to(sample.device)
# 从 model_output、噪声和信噪比计算步长
grad_norm = torch.norm(model_output.reshape(model_output.shape[0], -1), dim=-1).mean()
noise_norm = torch.norm(noise.reshape(noise.shape[0], -1), dim=-1).mean()
step_size = (self.config.snr * noise_norm / grad_norm) ** 2 * 2
step_size = step_size * torch.ones(sample.shape[0]).to(sample.device)
# self.repeat_scalar(step_size, sample.shape[0])
# 计算校正样本：model_output 项和噪声项
step_size = step_size.flatten()
while len(step_size.shape) < len(sample.shape):
step_size = step_size.unsqueeze(-1) # 扩展 step_size 以匹配样本的维度
prev_sample_mean = sample + step_size * model_output # 计算样本的平均值
prev_sample = prev_sample_mean + ((step_size * 2) ** 0.5) * noise # 加上噪声以得到最终样本
if not return_dict: # 如果不返回字典，则以元组形式返回
return (prev_sample,)
return SchedulerOutput(prev_sample=prev_sample) # 返回校正后的样本
# 返回一个张量，表示加噪样本
) -> torch.Tensor:
# 确保时间步和 sigma 与原始样本具有相同的设备和数据类型
timesteps = timesteps.to(original_samples.device)
# 获取与时间步对应的 sigma 值，并确保它们在相同的设备上
sigmas = self.discrete_sigmas.to(original_samples.device)[timesteps]
# 如果存在噪声，则将其按 sigma 进行缩放；否则生成与原始样本相同形状的随机噪声
noise = (
noise * sigmas[:, None, None, None]
if noise is not None
else torch.randn_like(original_samples) * sigmas[:, None, None, None]
)
# 将噪声添加到原始样本，生成加噪样本
noisy_samples = noise + original_samples
# 返回加噪样本
return noisy_samples
# 定义返回训练时间步数的函数
def __len__(self):
# 返回配置中的训练时间步数
return self.config.num_train_timesteps
# 版权声明，表示文件的所有权和使用限制
# Copyright 2024 Google Brain 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.
# 声明：此文件受到 https://github.com/yang-song/score_sde_pytorch 的强烈影响
# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch
# 导入所需模块
from dataclasses import dataclass
from typing import Optional, Tuple, Union
import flax # 导入 flax 用于构建可训练的模型
import jax # 导入 jax 进行高效的数值计算
import jax.numpy as jnp # 导入 jax 的 numpy 模块
from jax import random # 导入 jax 的随机数生成模块
# 导入配置工具和调度器相关的功能
from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import FlaxSchedulerMixin, FlaxSchedulerOutput, broadcast_to_shape_from_left
@flax.struct.dataclass
class ScoreSdeVeSchedulerState:
# 定义可设置的状态变量
timesteps: Optional[jnp.ndarray] = None # 时间步，默认为 None
discrete_sigmas: Optional[jnp.ndarray] = None # 离散 sigma 值，默认为 None
sigmas: Optional[jnp.ndarray] = None # sigma 值，默认为 None
@classmethod
def create(cls):
# 创建并返回类的实例
return cls()
@dataclass
class FlaxSdeVeOutput(FlaxSchedulerOutput):
"""
ScoreSdeVeScheduler 的步骤函数输出的输出类。
参数：
state (`ScoreSdeVeSchedulerState`):
prev_sample (`jnp.ndarray`，形状为 `(batch_size, num_channels, height, width)`，用于图像):
计算的前一时间步的样本 (x_{t-1})。`prev_sample` 应作为下一次模型输入
在去噪循环中使用。
prev_sample_mean (`jnp.ndarray`，形状为 `(batch_size, num_channels, height, width)`，用于图像):
平均的 `prev_sample`。与 `prev_sample` 相同，仅在之前的时间步上进行了均值平均。
"""
state: ScoreSdeVeSchedulerState # 调度器状态
prev_sample: jnp.ndarray # 前一时间步的样本
prev_sample_mean: Optional[jnp.ndarray] = None # 可选的前一样本均值
class FlaxScoreSdeVeScheduler(FlaxSchedulerMixin, ConfigMixin):
"""
方差爆炸随机微分方程 (SDE) 调度器。
有关更多信息，请参见原始论文：https://arxiv.org/abs/2011.13456
[`~ConfigMixin`] 处理在调度器的 `__init__` 函数中传递的所有配置属性的存储
例如 `num_train_timesteps`。可以通过 `scheduler.config.num_train_timesteps` 访问。
[`SchedulerMixin`] 提供通过 [`SchedulerMixin.save_pretrained`] 和
[`~SchedulerMixin.from_pretrained`] 函数进行通用加载和保存的功能。
"""
# 函数参数说明，描述每个参数的意义和用途
Args:
num_train_timesteps (`int`): 模型训练时使用的扩散步骤数量
snr (`float`):
权重系数，影响模型输出样本与随机噪声之间的关系
sigma_min (`float`):
采样过程中 sigma 序列的初始噪声规模，最小 sigma 应该与数据的分布相符
sigma_max (`float`): 用于传入模型的连续时间步的最大值
sampling_eps (`float`): 采样的结束值，时间步从 1 渐进减少到 epsilon
correct_steps (`int`): 对生成样本执行的纠正步骤数量
"""
# 定义一个属性，指示是否具有状态
@property
def has_state(self):
return True
# 初始化方法，设置模型的参数
@register_to_config
def __init__(
self,
num_train_timesteps: int = 2000, # 训练步骤数量，默认为2000
snr: float = 0.15, # 信噪比，默认为0.15
sigma_min: float = 0.01, # 最小噪声尺度，默认为0.01
sigma_max: float = 1348.0, # 最大噪声尺度，默认为1348.0
sampling_eps: float = 1e-5, # 采样结束值，默认为1e-5
correct_steps: int = 1, # 纠正步骤数量，默认为1
):
pass # 初始化方法为空，不执行任何操作
# 创建状态方法，用于初始化调度器状态
def create_state(self):
# 创建调度器状态实例
state = ScoreSdeVeSchedulerState.create()
# 设置 sigma 值并返回更新后的状态
return self.set_sigmas(
state,
self.config.num_train_timesteps, # 使用配置中的训练步骤数量
self.config.sigma_min, # 使用配置中的最小 sigma
self.config.sigma_max, # 使用配置中的最大 sigma
self.config.sampling_eps, # 使用配置中的采样结束值
)
# 设置时间步的方法，更新状态以支持推理过程
def set_timesteps(
self, state: ScoreSdeVeSchedulerState, num_inference_steps: int, shape: Tuple = (), sampling_eps: float = None
) -> ScoreSdeVeSchedulerState:
"""
设置扩散链中使用的连续时间步，推理前需运行的支持函数。
Args:
state (`ScoreSdeVeSchedulerState`): `FlaxScoreSdeVeScheduler` 的状态数据类实例
num_inference_steps (`int`):
生成样本时使用的扩散步骤数量
sampling_eps (`float`, optional):
最终时间步值（覆盖调度器实例化时给定的值）。
"""
# 如果提供了采样结束值，则使用它，否则使用配置中的值
sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps
# 生成从1到采样结束值的均匀分布的时间步
timesteps = jnp.linspace(1, sampling_eps, num_inference_steps)
# 返回更新后的状态，包括新设置的时间步
return state.replace(timesteps=timesteps)
# 设置 sigma 值的方法，用于调整状态
def set_sigmas(
self,
state: ScoreSdeVeSchedulerState, # 当前状态实例
num_inference_steps: int, # 推理步骤数量
sigma_min: float = None, # 可选的最小 sigma 值
sigma_max: float = None, # 可选的最大 sigma 值
sampling_eps: float = None, # 可选的采样结束值
) -> ScoreSdeVeSchedulerState:
"""
设置扩散链使用的噪声尺度。推理前运行的辅助函数。
sigmas 控制样本更新中的 `drift` 和 `diffusion` 组件的权重。
Args:
state (`ScoreSdeVeSchedulerState`): `FlaxScoreSdeVeScheduler` 的状态数据类实例。
num_inference_steps (`int`):
生成样本时使用的扩散步骤数量，基于预训练模型。
sigma_min (`float`, optional):
初始噪声尺度值（覆盖在 Scheduler 实例化时给定的值）。
sigma_max (`float`, optional):
最终噪声尺度值（覆盖在 Scheduler 实例化时给定的值）。
sampling_eps (`float`, optional):
最终时间步值（覆盖在 Scheduler 实例化时给定的值）。
"""
# 如果 sigma_min 为 None，则使用配置中的 sigma_min
sigma_min = sigma_min if sigma_min is not None else self.config.sigma_min
# 如果 sigma_max 为 None，则使用配置中的 sigma_max
sigma_max = sigma_max if sigma_max is not None else self.config.sigma_max
# 如果 sampling_eps 为 None，则使用配置中的 sampling_eps
sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps
# 如果 state 的时间步为 None，则设置时间步
if state.timesteps is None:
state = self.set_timesteps(state, num_inference_steps, sampling_eps)
# 计算离散 sigma 的指数值，生成 num_inference_steps 个点
discrete_sigmas = jnp.exp(jnp.linspace(jnp.log(sigma_min), jnp.log(sigma_max), num_inference_steps))
# 计算当前时间步对应的 sigmas 数组
sigmas = jnp.array([sigma_min * (sigma_max / sigma_min) ** t for t in state.timesteps])
# 用离散 sigma 和 sigmas 替换状态中的相应值
return state.replace(discrete_sigmas=discrete_sigmas, sigmas=sigmas)
def get_adjacent_sigma(self, state, timesteps, t):
# 返回与当前时间步相邻的 sigma 值，如果时间步为 0，则返回与之相同形状的零数组
return jnp.where(timesteps == 0, jnp.zeros_like(t), state.discrete_sigmas[timesteps - 1])
def step_pred(
self,
state: ScoreSdeVeSchedulerState,
model_output: jnp.ndarray,
timestep: int,
sample: jnp.ndarray,
key: jax.Array,
return_dict: bool = True,
) -> Union[FlaxSdeVeOutput, Tuple]:
"""
预测在前一个时间步的样本，通过反向 SDE 实现。核心功能是从学习的模型输出（通常是预测的噪声）传播扩散过程。
参数:
state (`ScoreSdeVeSchedulerState`): `FlaxScoreSdeVeScheduler` 状态数据类实例。
model_output (`jnp.ndarray`): 来自学习的扩散模型的直接输出。
timestep (`int`): 扩散链中的当前离散时间步。
sample (`jnp.ndarray`):
当前正在通过扩散过程创建的样本实例。
generator: 随机数生成器。
return_dict (`bool`): 选项，决定返回元组还是 `FlaxSdeVeOutput` 类。
返回:
[`FlaxSdeVeOutput`] 或 `tuple`: 如果 `return_dict` 为 True，则返回 [`FlaxSdeVeOutput`]，否则返回一个元组。当返回元组时，第一个元素是样本张量。
"""
# 检查状态的时间步是否设置，如果没有则抛出错误
if state.timesteps is None:
raise ValueError(
"`state.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"
)
# 将当前时间步扩展到与样本数量相同的形状
timestep = timestep * jnp.ones(
sample.shape[0],
)
# 计算离散时间步的索引
timesteps = (timestep * (len(state.timesteps) - 1)).long()
# 获取当前时间步对应的 sigma 值
sigma = state.discrete_sigmas[timesteps]
# 获取与当前时间步相邻的 sigma 值
adjacent_sigma = self.get_adjacent_sigma(state, timesteps, timestep)
# 初始化漂移项为与样本形状相同的零张量
drift = jnp.zeros_like(sample)
# 计算扩散项，表示噪声的变化
diffusion = (sigma**2 - adjacent_sigma**2) ** 0.5
# 根据论文中的公式 6，模型输出由网络建模为 grad_x log pt(x)
# 公式 47 显示了 SDE 模型与祖先采样方法的类比
diffusion = diffusion.flatten() # 将扩散项展平为一维
diffusion = broadcast_to_shape_from_left(diffusion, sample.shape) # 将扩散项广播到样本形状
drift = drift - diffusion**2 * model_output # 更新漂移项
# 根据扩散项采样噪声
key = random.split(key, num=1) # 分裂随机数生成器以获得新密钥
noise = random.normal(key=key, shape=sample.shape) # 生成与样本相同形状的随机噪声
prev_sample_mean = sample - drift # 计算前一个样本的均值，减去漂移项
# TODO 检查变量 diffusion 是否为噪声的正确缩放项
prev_sample = prev_sample_mean + diffusion * noise # 加上扩散场对噪声的影响
# 如果不返回字典，则返回一个元组，包含前一个样本、前一个样本均值和状态
if not return_dict:
return (prev_sample, prev_sample_mean, state)
# 返回 FlaxSdeVeOutput 对象，包含前一个样本、前一个样本均值和状态
return FlaxSdeVeOutput(prev_sample=prev_sample, prev_sample_mean=prev_sample_mean, state=state)
def step_correct(
self,
state: ScoreSdeVeSchedulerState,
model_output: jnp.ndarray,
sample: jnp.ndarray,
key: jax.Array,
return_dict: bool = True,
) -> Union[FlaxSdeVeOutput, Tuple]:
"""
根据网络输出的 model_output 校正预测样本，通常在每个时间步后反复执行。
参数：
state (`ScoreSdeVeSchedulerState`): `FlaxScoreSdeVeScheduler` 状态数据类实例。
model_output (`jnp.ndarray`): 从学习的扩散模型直接输出。
sample (`jnp.ndarray`):
当前由扩散过程创建的样本实例。
generator: 随机数生成器。
return_dict (`bool`): 选项，用于返回元组而不是 FlaxSdeVeOutput 类。
返回：
[`FlaxSdeVeOutput`] 或 `tuple`: 如果 `return_dict` 为 True 则返回 [`FlaxSdeVeOutput`]，否则返回 `tuple`。返回元组时，第一个元素是样本张量。
"""
# 检查 timesteps 是否被设置，如果未设置则抛出错误
if state.timesteps is None:
raise ValueError(
"`state.timesteps` 未设置，需要在创建调度器后运行 'set_timesteps'"
)
# 对于小批量大小，论文建议用 sqrt(d) 替代 norm(z)，其中 d 是 z 的维度
# 为校正生成噪声
key = random.split(key, num=1) # 将随机数生成器的 key 拆分
noise = random.normal(key=key, shape=sample.shape) # 生成与样本相同形状的噪声
# 从 model_output、噪声和 snr 计算步长
grad_norm = jnp.linalg.norm(model_output) # 计算 model_output 的范数
noise_norm = jnp.linalg.norm(noise) # 计算噪声的范数
step_size = (self.config.snr * noise_norm / grad_norm) ** 2 * 2 # 计算步长
step_size = step_size * jnp.ones(sample.shape[0]) # 扩展步长到样本批次大小
# 计算校正样本：model_output 项和噪声项
step_size = step_size.flatten() # 将步长展平
step_size = broadcast_to_shape_from_left(step_size, sample.shape) # 将步长广播到样本形状
prev_sample_mean = sample + step_size * model_output # 计算校正样本均值
prev_sample = prev_sample_mean + ((step_size * 2) ** 0.5) * noise # 添加噪声得到最终样本
# 根据 return_dict 的值决定返回格式
if not return_dict:
return (prev_sample, state) # 返回元组形式的样本和状态
return FlaxSdeVeOutput(prev_sample=prev_sample, state=state) # 返回 FlaxSdeVeOutput 对象
def __len__(self):
return self.config.num_train_timesteps # 返回训练时间步的数量
# 版权声明，说明此代码的所有权归斯坦福大学团队和HuggingFace团队所有
#
# 根据Apache许可证，版本2.0（“许可证”）进行许可；
# 除非遵循许可证，否则不得使用此文件。
# 可以在以下网址获取许可证的副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，软件
# 根据许可证分发是在“按原样”基础上进行的，
# 不提供任何形式的担保或条件，包括明示或暗示的。
# 请参见许可证以获取有关特定语言的权限和
# 限制的详细信息。
# 免责声明：此代码受到以下项目的强烈影响 https://github.com/pesser/pytorch_diffusion
# 和 https://github.com/hojonathanho/diffusion
import math # 导入数学模块以进行数学运算
from dataclasses import dataclass # 从数据类模块导入dataclass装饰器
from typing import List, Optional, Tuple, Union # 从typing模块导入类型提示
import numpy as np # 导入numpy库以进行数值计算
import torch # 导入PyTorch库以进行张量操作
from ..configuration_utils import ConfigMixin, register_to_config # 从配置工具导入配置混合类和注册函数
from ..schedulers.scheduling_utils import SchedulerMixin # 从调度工具导入调度混合类
from ..utils import BaseOutput, logging # 从工具模块导入基本输出类和日志记录工具
from ..utils.torch_utils import randn_tensor # 从PyTorch工具导入随机张量生成函数
logger = logging.get_logger(__name__) # 创建一个日志记录器以记录当前模块的日志，禁用pylint的命名警告
@dataclass # 将类标记为数据类
class TCDSchedulerOutput(BaseOutput): # 定义调度器输出类，继承基本输出类
"""
调度器的`step`函数输出的输出类。
参数：
prev_sample (`torch.Tensor`，形状为`(batch_size, num_channels, height, width)`的图像):
先前时间步的计算样本`(x_{t-1})`。`prev_sample`应作为去噪循环中的下一个模型输入使用。
pred_noised_sample (`torch.Tensor`，形状为`(batch_size, num_channels, height, width)`的图像):
基于当前时间步的模型输出的预测噪声样本`(x_{s})`。
"""
prev_sample: torch.Tensor # 定义先前样本属性，类型为torch张量
pred_noised_sample: Optional[torch.Tensor] = None # 定义可选的预测噪声样本属性，默认值为None
# 从diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar复制
def betas_for_alpha_bar( # 定义用于生成beta调度的函数
num_diffusion_timesteps, # 传播时间步数的参数
max_beta=0.999, # 最大beta的默认值
alpha_transform_type="cosine", # alpha变换类型的默认值为“余弦”
):
"""
创建一个beta调度，离散化给定的alpha_t_bar函数，该函数定义了时间上(1-beta)的累积乘积，从t = [0,1]。
包含一个alpha_bar函数，该函数接受一个参数t并将其转换为(1-beta)的累积乘积
到该传播过程的部分。
参数：
num_diffusion_timesteps (`int`): 要生成的beta数量。
max_beta (`float`): 要使用的最大beta；使用小于1的值以防止奇异性。
alpha_transform_type (`str`, *可选*, 默认值为`cosine`): alpha_bar的噪声调度类型。
从`cosine`或`exp`中选择
返回：
betas (`np.ndarray`): 调度程序用于模型输出的步进的beta。
"""
# 检查 alpha_transform_type 是否为 "cosine"
if alpha_transform_type == "cosine":
# 定义 alpha_bar_fn 函数，计算余弦平方值
def alpha_bar_fn(t):
# 根据公式计算余弦函数值的平方，调整时间参数 t
return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
# 检查 alpha_transform_type 是否为 "exp"
elif alpha_transform_type == "exp":
# 定义 alpha_bar_fn 函数，计算指数衰减值
def alpha_bar_fn(t):
# 计算指数函数的衰减值，使用负的时间参数乘以常数
return math.exp(t * -12.0)
# 如果 alpha_transform_type 不是已支持的类型，抛出错误
else:
raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
# 初始化空列表，用于存储 beta 值
betas = []
# 循环遍历每个扩散时间步
for i in range(num_diffusion_timesteps):
# 计算当前时间步的归一化值 t1
t1 = i / num_diffusion_timesteps
# 计算下一个时间步的归一化值 t2
t2 = (i + 1) / num_diffusion_timesteps
# 计算 beta 值，并将其添加到列表中，取最小值以限制
betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
# 将 beta 列表转换为张量并返回，数据类型为浮点32
return torch.tensor(betas, dtype=torch.float32)
# 从 diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr 复制而来
def rescale_zero_terminal_snr(betas: torch.Tensor) -> torch.Tensor:
"""
根据 https://arxiv.org/pdf/2305.08891.pdf (算法 1) 重新缩放 betas，使其具有零终端 SNR
参数:
betas (`torch.Tensor`):
调度器初始化时使用的 betas。
返回:
`torch.Tensor`: 具有零终端 SNR 的重新缩放 betas
"""
# 将 betas 转换为 alphas_bar_sqrt
alphas = 1.0 - betas # 计算 alphas，取 1 减去 betas
alphas_cumprod = torch.cumprod(alphas, dim=0) # 计算 alphas 的累积乘积
alphas_bar_sqrt = alphas_cumprod.sqrt() # 计算 alphas 的平方根
# 存储旧值。
alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() # 克隆第一个 alphas_bar_sqrt 的值
alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() # 克隆最后一个 alphas_bar_sqrt 的值
# 移动，使最后一个时间步为零。
alphas_bar_sqrt -= alphas_bar_sqrt_T # 从每个值中减去最后一个值
# 缩放，使第一个时间步恢复到旧值。
alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) # 根据旧值进行缩放
# 将 alphas_bar_sqrt 转换为 betas
alphas_bar = alphas_bar_sqrt**2 # 还原平方根
alphas = alphas_bar[1:] / alphas_bar[:-1] # 还原累积乘积
alphas = torch.cat([alphas_bar[0:1], alphas]) # 连接第一个值与后续值
betas = 1 - alphas # 计算 betas
return betas # 返回重新缩放的 betas
class TCDScheduler(SchedulerMixin, ConfigMixin):
"""
`TCDScheduler` 结合了论文 `Trajectory Consistency Distillation` 中提出的 `Strategic Stochastic Sampling`，
扩展了原始的多步一致性采样，允许无限制的轨迹遍历。
该代码基于 TCD 的官方仓库 (https://github.com/jabir-zheng/TCD)。
该模型继承自 [`SchedulerMixin`] 和 [`ConfigMixin`]。 [`~ConfigMixin`] 负责存储在调度器的 `__init__` 函数中传递的所有配置属性，
如 `num_train_timesteps`。可以通过 `scheduler.config.num_train_timesteps` 访问。[`SchedulerMixin`] 提供通用的加载和保存
功能，通过 [`SchedulerMixin.save_pretrained`] 和 [`~SchedulerMixin.from_pretrained`] 函数实现。
"""
order = 1 # 设置调度器的顺序为 1
@register_to_config
def __init__(
self,
num_train_timesteps: int = 1000, # 训练时间步的数量，默认为 1000
beta_start: float = 0.00085, # beta 的起始值，默认为 0.00085
beta_end: float = 0.012, # beta 的结束值，默认为 0.012
beta_schedule: str = "scaled_linear", # beta 的调度策略，默认为 "scaled_linear"
trained_betas: Optional[Union[np.ndarray, List[float]]] = None, # 已训练的 betas，默认为 None
original_inference_steps: int = 50, # 原始推断步骤的数量，默认为 50
clip_sample: bool = False, # 是否裁剪样本，默认为 False
clip_sample_range: float = 1.0, # 样本裁剪范围，默认为 1.0
set_alpha_to_one: bool = True, # 是否将 alpha 设置为 1，默认为 True
steps_offset: int = 0, # 步骤偏移量，默认为 0
prediction_type: str = "epsilon", # 预测类型，默认为 "epsilon"
thresholding: bool = False, # 是否使用阈值处理，默认为 False
dynamic_thresholding_ratio: float = 0.995, # 动态阈值比例，默认为 0.995
sample_max_value: float = 1.0, # 样本的最大值，默认为 1.0
timestep_spacing: str = "leading", # 时间步间隔类型，默认为 "leading"
timestep_scaling: float = 10.0, # 时间步缩放因子，默认为 10.0
rescale_betas_zero_snr: bool = False, # 是否重新缩放 betas 以实现零 SNR，默认为 False
):
# 检查已训练的贝塔值是否为 None
if trained_betas is not None:
# 将训练好的贝塔值转换为张量，数据类型为 float32
self.betas = torch.tensor(trained_betas, dtype=torch.float32)
# 如果贝塔调度方式为线性
elif beta_schedule == "linear":
# 创建从 beta_start 到 beta_end 的线性贝塔值
self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
# 如果贝塔调度方式为缩放线性
elif beta_schedule == "scaled_linear":
# 该调度方式特定于潜在扩散模型
self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
# 如果贝塔调度方式为 squaredcos_cap_v2
elif beta_schedule == "squaredcos_cap_v2":
# Glide 余弦调度
self.betas = betas_for_alpha_bar(num_train_timesteps)
# 如果以上都不符合，抛出未实现错误
else:
raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
# 针对零 SNR 进行重新缩放
if rescale_betas_zero_snr:
# 调整贝塔值以适应零终端 SNR
self.betas = rescale_zero_terminal_snr(self.betas)
# 计算 alphas，作为 1 减去贝塔值
self.alphas = 1.0 - self.betas
# 计算 alphas 的累积乘积
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
# 在每一步的 ddim 中查看之前的 alphas_cumprod
# 最后一步时没有前一个 alphas_cumprod，因为我们已经在 0 处
# `set_alpha_to_one` 决定是否将此参数设置为 1 或使用最后一个非前驱的 alpha
self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
# 初始化噪声分布的标准差
self.init_noise_sigma = 1.0
# 可设置的值
self.num_inference_steps = None
# 反转的时间步长数组
self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
self.custom_timesteps = False
# 初始化步长和起始索引
self._step_index = None
self._begin_index = None
# 从 diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep 复制的函数
def index_for_timestep(self, timestep, schedule_timesteps=None):
# 如果未提供调度时间步，则使用默认时间步
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
# 找到与当前时间步匹配的索引
indices = (schedule_timesteps == timestep).nonzero()
# 第一步的 sigma 索引总是第二个索引（如果只有一个则为最后一个）
# 确保在去噪调度中不意外跳过 sigma
pos = 1 if len(indices) > 1 else 0
# 返回所需索引的值
return indices[pos].item()
# 从 diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index 复制的函数
def _init_step_index(self, timestep):
# 如果 begin_index 为 None
if self.begin_index is None:
# 如果时间步是张量，转移到相应设备
if isinstance(timestep, torch.Tensor):
timestep = timestep.to(self.timesteps.device)
# 初始化步长索引
self._step_index = self.index_for_timestep(timestep)
else:
# 否则，使用预设的起始索引
self._step_index = self._begin_index
@property
# 定义一个方法，返回当前的步骤索引
def step_index(self):
return self._step_index
# 定义一个属性，返回开始索引
@property
def begin_index(self):
"""
返回第一个时间步的索引。应通过 `set_begin_index` 方法从管道设置。
"""
return self._begin_index
# 从 diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index 复制而来
# 定义一个方法，用于设置调度器的开始索引
def set_begin_index(self, begin_index: int = 0):
"""
设置调度器的开始索引。此函数应在推理之前从管道运行。
参数:
begin_index (`int`):
调度器的开始索引。
"""
self._begin_index = begin_index
# 定义一个方法，确保与需要根据当前时间步缩放去噪模型输入的调度器互操作
def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
"""
确保与需要根据当前时间步缩放去噪模型输入的调度器互操作。
参数:
sample (`torch.Tensor`):
输入样本。
timestep (`int`, *可选*):
扩散链中的当前时间步。
返回:
`torch.Tensor`:
一个缩放后的输入样本。
"""
return sample
# 从 diffusers.schedulers.scheduling_ddim.DDIMScheduler._get_variance 复制而来
# 定义一个方法，计算当前和前一个时间步的方差
def _get_variance(self, timestep, prev_timestep):
# 获取当前时间步的累积 alpha 值
alpha_prod_t = self.alphas_cumprod[timestep]
# 获取前一个时间步的累积 alpha 值，如果无效则使用最终累积值
alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
# 计算当前时间步的 beta 值
beta_prod_t = 1 - alpha_prod_t
# 计算前一个时间步的 beta 值
beta_prod_t_prev = 1 - alpha_prod_t_prev
# 根据公式计算方差
variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
return variance
# 从 diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample 复制而来
# 定义动态阈值采样的方法，输入为样本张量，返回处理后的张量
def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
"""
"动态阈值处理：在每个采样步骤中，我们将 s 设置为 xt0（在时间步 t 的 x_0 预测）中的某个百分位绝对像素值，
如果 s > 1，则将 xt0 阈值处理为范围 [-s, s]，然后除以 s。动态阈值处理将饱和像素（接近 -1 和 1 的像素）
向内推，这样可以在每一步主动防止像素饱和。我们发现动态阈值处理显著提高了照片真实感以及图像-文本对齐，
尤其是在使用非常大的引导权重时。"
https://arxiv.org/abs/2205.11487
"""
# 获取输入样本的数值类型
dtype = sample.dtype
# 获取样本的批次大小、通道数和剩余维度
batch_size, channels, *remaining_dims = sample.shape
# 如果数据类型不是 float32 或 float64，则转换为 float 类型
if dtype not in (torch.float32, torch.float64):
sample = sample.float() # 为了进行分位数计算而上升数据类型，并且 clamp 对 CPU half 类型未实现
# 将样本展平，以便对每个图像进行分位数计算
sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
# 计算样本的绝对值
abs_sample = sample.abs() # "某个百分位绝对像素值"
# 计算绝对样本的指定百分位
s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
# 将 s 限制在 [1, sample_max_value] 的范围内
s = torch.clamp(
s, min=1, max=self.config.sample_max_value
) # 当限制到 min=1 时，相当于标准剪切到 [-1, 1]
# 在维度 1 上扩展 s 的维度，以便广播
s = s.unsqueeze(1) # (batch_size, 1) 因为 clamp 会在维度 0 上广播
# 将样本限制在 [-s, s] 范围内，然后除以 s
sample = torch.clamp(sample, -s, s) / s # "我们将 xt0 阈值处理为 [-s, s] 的范围，然后除以 s"
# 重新调整样本的形状为原来的形状
sample = sample.reshape(batch_size, channels, *remaining_dims)
# 将样本转换回原来的数据类型
sample = sample.to(dtype)
# 返回处理后的样本
return sample
# 定义设置时间步的方法，接受多个参数用于配置
def set_timesteps(
self,
num_inference_steps: Optional[int] = None, # 推理步骤的数量
device: Union[str, torch.device] = None, # 设备类型（字符串或 torch 设备）
original_inference_steps: Optional[int] = None, # 原始推理步骤数量
timesteps: Optional[List[int]] = None, # 时间步列表
strength: float = 1.0, # 强度参数，默认为 1.0
# 定义步骤的方法，处理模型输出与样本
def step(
self,
model_output: torch.Tensor, # 模型输出的张量
timestep: int, # 当前时间步
sample: torch.Tensor, # 输入样本张量
eta: float = 0.3, # 噪声参数，默认为 0.3
generator: Optional[torch.Generator] = None, # 可选的随机数生成器
return_dict: bool = True, # 是否返回字典格式的结果，默认为 True
# 从 diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise 复制的内容
def add_noise(
self,
original_samples: torch.Tensor, # 原始样本的张量
noise: torch.Tensor, # 噪声的张量
timesteps: torch.IntTensor, # 时间步的整数张量
) -> torch.Tensor:
# 确保 alphas_cumprod 和 timestep 的设备和数据类型与 original_samples 相同
# 将 self.alphas_cumprod 移动到设备上，以避免后续 add_noise 调用时的冗余 CPU 到 GPU 数据移动
self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device)
# 将 alphas_cumprod 转换为与 original_samples 相同的数据类型
alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype)
# 将 timesteps 移动到 original_samples 的设备上
timesteps = timesteps.to(original_samples.device)
# 计算 sqrt_alpha_prod 为 alphas_cumprod 在 timesteps 位置的平方根
sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
# 将 sqrt_alpha_prod 扁平化
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
# 如果 sqrt_alpha_prod 的维度小于 original_samples，则增加维度
while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
# 计算 sqrt_one_minus_alpha_prod 为 (1 - alphas_cumprod 在 timesteps 位置) 的平方根
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
# 将 sqrt_one_minus_alpha_prod 扁平化
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
# 如果 sqrt_one_minus_alpha_prod 的维度小于 original_samples，则增加维度
while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
# 计算带噪声的样本 noisy_samples
noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
# 返回带噪声的样本
return noisy_samples
# 从 diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity 复制而来
def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor:
# 确保 alphas_cumprod 和 timestep 的设备和数据类型与 sample 相同
self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device)
# 将 alphas_cumprod 转换为与 sample 相同的数据类型
alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype)
# 将 timesteps 移动到 sample 的设备上
timesteps = timesteps.to(sample.device)
# 计算 sqrt_alpha_prod 为 alphas_cumprod 在 timesteps 位置的平方根
sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
# 将 sqrt_alpha_prod 扁平化
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
# 如果 sqrt_alpha_prod 的维度小于 sample，则增加维度
while len(sqrt_alpha_prod.shape) < len(sample.shape):
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
# 计算 sqrt_one_minus_alpha_prod 为 (1 - alphas_cumprod 在 timesteps 位置) 的平方根
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
# 将 sqrt_one_minus_alpha_prod 扁平化
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
# 如果 sqrt_one_minus_alpha_prod 的维度小于 sample，则增加维度
while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
# 计算速度 velocity
velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
# 返回速度
return velocity
# 返回训练时间步数的数量
def __len__(self):
return self.config.num_train_timesteps
# 从 diffusers.schedulers.scheduling_ddpm.DDPMScheduler.previous_timestep 复制而来
# 定义一个方法，用于获取给定时间步的前一个时间步
def previous_timestep(self, timestep):
# 检查是否存在自定义时间步
if self.custom_timesteps:
# 找到当前时间步在时间步数组中的索引
index = (self.timesteps == timestep).nonzero(as_tuple=True)[0][0]
# 如果当前时间步是最后一个，设置前一个时间步为 -1
if index == self.timesteps.shape[0] - 1:
prev_t = torch.tensor(-1)
# 否则，获取下一个时间步
else:
prev_t = self.timesteps[index + 1]
# 如果没有自定义时间步
else:
# 计算推理步骤的数量，优先使用实例属性，否则使用配置的训练时间步
num_inference_steps = (
self.num_inference_steps if self.num_inference_steps else self.config.num_train_timesteps
)
# 计算前一个时间步
prev_t = timestep - self.config.num_train_timesteps // num_inference_steps
# 返回前一个时间步
return prev_t
# 版权所有 2024 Kakao Brain 和 HuggingFace Team。保留所有权利。
#
# 根据 Apache 许可证第 2.0 版（“许可证”）授权；
# 除非遵循许可证，否则您不得使用此文件。
# 您可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面同意，软件在许可证下分发均为“按原样”提供，
# 不提供任何形式的保证或条件，无论是明示还是暗示。
# 请参见许可证以获取与权限和
# 限制相关的特定语言。
# 导入数学模块
import math
# 从 dataclasses 模块导入 dataclass 装饰器
from dataclasses import dataclass
# 导入类型提示相关的模块
from typing import Optional, Tuple, Union
# 导入 NumPy 库
import numpy as np
# 导入 PyTorch 库
import torch
# 从配置工具导入 ConfigMixin 和 register_to_config
from ..configuration_utils import ConfigMixin, register_to_config
# 从工具模块导入 BaseOutput
from ..utils import BaseOutput
# 从 torch_utils 导入 randn_tensor 函数
from ..utils.torch_utils import randn_tensor
# 从调度工具导入 SchedulerMixin
from .scheduling_utils import SchedulerMixin
@dataclass
# 从 diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput 复制的类，包含 DDPM 到 UnCLIP 的转换
class UnCLIPSchedulerOutput(BaseOutput):
"""
调度器 `step` 函数输出的输出类。
参数：
prev_sample (`torch.Tensor`，形状为 `(batch_size, num_channels, height, width)`，用于图像):
计算出的前一个时间步的样本 `(x_{t-1})`。 `prev_sample` 应在去噪循环中作为下一个模型输入使用。
pred_original_sample (`torch.Tensor`，形状为 `(batch_size, num_channels, height, width)`，用于图像):
基于当前时间步的模型输出的预测去噪样本 `(x_{0})`。
`pred_original_sample` 可用于预览进度或进行引导。
"""
# 上一个样本的张量
prev_sample: torch.Tensor
# 可选的预测去噪样本
pred_original_sample: Optional[torch.Tensor] = None
# 从 diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar 复制的函数
def betas_for_alpha_bar(
num_diffusion_timesteps,
max_beta=0.999,
alpha_transform_type="cosine",
):
"""
创建一个 beta 计划，该计划离散化给定的 alpha_t_bar 函数，该函数定义了 (1-beta) 随时间的累积乘积
从 t = [0,1] 开始。
包含一个 alpha_bar 函数，该函数接受参数 t，并将其转换为扩散过程至该部分的 (1-beta) 的累积乘积。
参数：
num_diffusion_timesteps (`int`): 生成的 beta 数量。
max_beta (`float`): 使用的最大 beta 值；使用小于 1 的值以
防止奇异性。
alpha_transform_type (`str`, *可选*, 默认值为 `cosine`): alpha_bar 的噪声调度类型。
从 `cosine` 或 `exp` 中选择。
返回：
betas (`np.ndarray`): 调度器用于更新模型输出的 betas
"""
# 检查 alpha 转换类型是否为 "cosine"
if alpha_transform_type == "cosine":
# 定义 alpha_bar 函数，接受参数 t
def alpha_bar_fn(t):
# 计算 t 的 alpha_bar 值，使用余弦函数
return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
# 检查 alpha_transform_type 是否为 "exp"
elif alpha_transform_type == "exp":
# 定义 alpha_bar_fn 函数，用于计算指数衰减
def alpha_bar_fn(t):
# 计算并返回指数衰减值
return math.exp(t * -12.0)
# 如果 alpha_transform_type 不是已知类型，抛出异常
else:
raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
# 初始化空列表 betas 用于存储每个时间步的 beta 值
betas = []
# 遍历每个扩散时间步
for i in range(num_diffusion_timesteps):
# 计算当前时间步的归一化值 t1
t1 = i / num_diffusion_timesteps
# 计算下一个时间步的归一化值 t2
t2 = (i + 1) / num_diffusion_timesteps
# 计算 beta 值并添加到 betas 列表，确保不超过 max_beta
betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
# 将 betas 转换为张量并返回
return torch.tensor(betas, dtype=torch.float32)
# 定义 UnCLIPScheduler 类，继承自 SchedulerMixin 和 ConfigMixin
class UnCLIPScheduler(SchedulerMixin, ConfigMixin):
"""
NOTE: do not use this scheduler. The DDPM scheduler has been updated to support the changes made here. This
scheduler will be removed and replaced with DDPM.
This is a modified DDPM Scheduler specifically for the karlo unCLIP model.
This scheduler has some minor variations in how it calculates the learned range variance and dynamically
re-calculates betas based off the timesteps it is skipping.
The scheduler also uses a slightly different step ratio when computing timesteps to use for inference.
See [`~DDPMScheduler`] for more information on DDPM scheduling
Args:
num_train_timesteps (`int`): number of diffusion steps used to train the model.
variance_type (`str`):
options to clip the variance used when adding noise to the denoised sample. Choose from `fixed_small_log`
or `learned_range`.
clip_sample (`bool`, default `True`):
option to clip predicted sample between `-clip_sample_range` and `clip_sample_range` for numerical
stability.
clip_sample_range (`float`, default `1.0`):
The range to clip the sample between. See `clip_sample`.
prediction_type (`str`, default `epsilon`, optional):
prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion process)
or `sample` (directly predicting the noisy sample`)
"""
# 注册构造函数到配置中
@register_to_config
def __init__(
# 定义训练时间步的数量，默认为 1000
self,
num_train_timesteps: int = 1000,
# 定义方差类型，默认为 'fixed_small_log'
variance_type: str = "fixed_small_log",
# 是否裁剪样本，默认为 True
clip_sample: bool = True,
# 样本裁剪范围，默认为 1.0
clip_sample_range: Optional[float] = 1.0,
# 预测类型，默认为 'epsilon'
prediction_type: str = "epsilon",
# β 调度方式，默认为 'squaredcos_cap_v2'
beta_schedule: str = "squaredcos_cap_v2",
):
# 检查 β 调度方式是否有效
if beta_schedule != "squaredcos_cap_v2":
raise ValueError("UnCLIPScheduler only supports `beta_schedule`: 'squaredcos_cap_v2'")
# 根据训练时间步计算 beta 值
self.betas = betas_for_alpha_bar(num_train_timesteps)
# 计算 alpha 值
self.alphas = 1.0 - self.betas
# 计算累积的 alpha 值
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
# 初始化一个值为 1.0 的张量
self.one = torch.tensor(1.0)
# 初始化噪声分布的标准差
self.init_noise_sigma = 1.0
# 可设置的推理步数
self.num_inference_steps = None
# 创建时间步张量，从 num_train_timesteps 到 0 递减
self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy())
# 设置方差类型
self.variance_type = variance_type
# 定义缩放模型输入的方法
def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
Args:
sample (`torch.Tensor`): input sample
timestep (`int`, optional): current timestep
Returns:
`torch.Tensor`: scaled input sample
"""
# 返回未处理的样本，未进行缩放
return sample
# 定义设置离散时间步长的方法，用于扩散链
def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
"""
设置用于扩散链的离散时间步长。此方法应在推理前运行。
注意：此调度器使用与其他扩散调度器略有不同的步长比例。
不同的步长比例是为了模仿原始的 karlo 实现，并不影响结果的质量或准确性。
参数:
num_inference_steps (`int`):
生成样本时使用的扩散步数。
"""
# 将输入的推理步骤数量保存到实例变量
self.num_inference_steps = num_inference_steps
# 计算步长比例
step_ratio = (self.config.num_train_timesteps - 1) / (self.num_inference_steps - 1)
# 生成时间步长数组，反向排列并转换为整数类型
timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
# 将时间步长转换为张量并移动到指定设备
self.timesteps = torch.from_numpy(timesteps).to(device)
# 定义获取方差的方法，计算当前和前一个时间步的方差
def _get_variance(self, t, prev_timestep=None, predicted_variance=None, variance_type=None):
# 如果没有提供前一个时间步，默认为当前时间步减一
if prev_timestep is None:
prev_timestep = t - 1
# 获取当前时间步和前一个时间步的累积 alpha 值
alpha_prod_t = self.alphas_cumprod[t]
alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.one
# 计算当前和前一个时间步的 beta 值
beta_prod_t = 1 - alpha_prod_t
beta_prod_t_prev = 1 - alpha_prod_t_prev
# 根据前后时间步计算 beta 值
if prev_timestep == t - 1:
beta = self.betas[t]
else:
beta = 1 - alpha_prod_t / alpha_prod_t_prev
# 对于 t > 0，计算预测方差
# x_{t-1} ~ N(pred_prev_sample, variance) == 将方差添加到预测样本中
variance = beta_prod_t_prev / beta_prod_t * beta
# 如果没有提供方差类型，则使用配置中的类型
if variance_type is None:
variance_type = self.config.variance_type
# 针对训练稳定性进行的一些特殊处理
if variance_type == "fixed_small_log":
# 将方差限制到最小值，然后计算对数
variance = torch.log(torch.clamp(variance, min=1e-20))
# 计算最终方差
variance = torch.exp(0.5 * variance)
elif variance_type == "learned_range":
# 注意与 DDPM 调度器的区别
min_log = variance.log()
max_log = beta.log()
# 计算方差的比例
frac = (predicted_variance + 1) / 2
# 计算最终方差
variance = frac * max_log + (1 - frac) * min_log
# 返回计算得到的方差
return variance
# 定义步骤方法，处理模型输出和样本
def step(
self,
model_output: torch.Tensor,
timestep: int,
sample: torch.Tensor,
prev_timestep: Optional[int] = None,
generator=None,
return_dict: bool = True,
# 从 diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise 复制的方法
def add_noise(
self,
original_samples: torch.Tensor,
noise: torch.Tensor,
timesteps: torch.IntTensor,
# 返回类型为 torch.Tensor
) -> torch.Tensor:
# 确保 alphas_cumprod 和 timestep 与 original_samples 具有相同的设备和数据类型
# 将 self.alphas_cumprod 移动到相应的设备，以避免后续 add_noise 调用中的冗余 CPU 到 GPU 数据移动
self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device)
# 将 alphas_cumprod 转换为 original_samples 的数据类型
alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype)
# 将 timesteps 移动到 original_samples 的设备
timesteps = timesteps.to(original_samples.device)
# 计算 alphas_cumprod 在 timesteps 位置的平方根
sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
# 将结果展平为一维
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
# 如果 sqrt_alpha_prod 的维度小于 original_samples 的维度，则添加新的维度
while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
# 计算 (1 - alphas_cumprod) 在 timesteps 位置的平方根
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
# 将结果展平为一维
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
# 如果 sqrt_one_minus_alpha_prod 的维度小于 original_samples 的维度，则添加新的维度
while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
# 计算带噪声的样本
noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
# 返回带噪声的样本
return noisy_samples
# 版权信息，指明作者及团队
# Copyright 2024 TSAIL Team and The HuggingFace Team. All rights reserved.
#
# 许可信息，指明使用此文件的条件
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# 说明如何获取许可证
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 说明法律责任和条件
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 查看许可证的具体权限和限制
# See the License for the specific language governing permissions and
# limitations under the License.
# 免责声明，指向相关信息链接
# DISCLAIMER: check https://arxiv.org/abs/2302.04867 and https://github.com/wl-zhao/UniPC for more info
# 此代码基于指定 GitHub 项目进行修改
# The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
# 导入数学模块
import math
# 从 typing 模块导入所需的类型
from typing import List, Optional, Tuple, Union
# 导入 NumPy 库
import numpy as np
# 导入 PyTorch 库
import torch
# 从配置工具导入相关混合类和配置注册函数
from ..configuration_utils import ConfigMixin, register_to_config
# 从工具模块导入弃用函数
from ..utils import deprecate
# 从调度工具导入所需的调度类和输出类
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
# 定义函数，生成 beta 调度
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
def betas_for_alpha_bar(
num_diffusion_timesteps,
max_beta=0.999,
alpha_transform_type="cosine",
):
"""
创建 beta 调度，以离散化给定的 alpha_t_bar 函数，该函数定义了
(1-beta) 的累积乘积，时间范围从 t = [0,1]。
包含一个 alpha_bar 函数，该函数接受参数 t，并将其转换为
在扩散过程中到该部分的 (1-beta) 的累积乘积。
参数:
num_diffusion_timesteps (`int`): 生成的 beta 数量。
max_beta (`float`): 使用的最大 beta 值；使用低于 1 的值以
防止奇点。
alpha_transform_type (`str`, *optional*, default to `cosine`): alpha_bar 的噪声调度类型。
可选择 `cosine` 或 `exp`
返回:
betas (`np.ndarray`): 调度器用于步骤模型输出的 betas
"""
# 根据 alpha_transform_type 确定 alpha_bar 函数
if alpha_transform_type == "cosine":
# 定义余弦变换的 alpha_bar 函数
def alpha_bar_fn(t):
return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
elif alpha_transform_type == "exp":
# 定义指数变换的 alpha_bar 函数
def alpha_bar_fn(t):
return math.exp(t * -12.0)
else:
# 抛出不支持的类型错误
raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
# 初始化 beta 列表
betas = []
# 遍历每个扩散时间步
for i in range(num_diffusion_timesteps):
# 计算当前时间步 t1
t1 = i / num_diffusion_timesteps
# 计算下一个时间步 t2
t2 = (i + 1) / num_diffusion_timesteps
# 计算并添加 beta 值，限制在 max_beta 之下
betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
# 返回以 tensor 格式的 beta 值
return torch.tensor(betas, dtype=torch.float32)
# 定义函数，重新缩放终止信噪比
# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr
def rescale_zero_terminal_snr(betas):
"""
# 重新缩放 betas，使其具有零终端信噪比，参考文献 https://arxiv.org/pdf/2305.08891.pdf (算法 1)
# 参数:
# betas (`torch.Tensor`):
# 用于初始化调度器的 betas。
# 返回:
# `torch.Tensor`: 具有零终端信噪比的重新缩放后的 betas
"""
# 将 betas 转换为 alphas_bar_sqrt
alphas = 1.0 - betas # 计算 alphas，表示每个时间步的 alpha 值
alphas_cumprod = torch.cumprod(alphas, dim=0) # 计算 alphas 的累积乘积
alphas_bar_sqrt = alphas_cumprod.sqrt() # 计算累积乘积的平方根
# 存储旧值。
alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() # 记录第一个时间步的平方根值
alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() # 记录最后一个时间步的平方根值
# 进行平移，使最后一个时间步为零。
alphas_bar_sqrt -= alphas_bar_sqrt_T # 从所有值中减去最后一个值
# 进行缩放，使第一个时间步恢复为旧值。
alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) # 缩放以恢复第一个值
# 将 alphas_bar_sqrt 转换回 betas
alphas_bar = alphas_bar_sqrt**2 # 还原平方根，得到 alphas_bar
alphas = alphas_bar[1:] / alphas_bar[:-1] # 还原累积乘积，计算每个时间步的 alpha
alphas = torch.cat([alphas_bar[0:1], alphas]) # 将第一个 alpha 添加到结果中
betas = 1 - alphas # 计算 betas，作为 1 减去 alphas
return betas # 返回重新缩放后的 betas
# 定义一个名为 UniPCMultistepScheduler 的类，继承自 SchedulerMixin 和 ConfigMixin
class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
"""
`UniPCMultistepScheduler` 是一个无训练的框架，旨在快速采样扩散模型。
该模型继承自 [`SchedulerMixin`] 和 [`ConfigMixin`]。有关库为所有调度程序实现的通用
方法（例如加载和保存）的文档，请查看超类文档。
"""
# 定义与 KarrasDiffusionSchedulers 兼容的名称列表
_compatibles = [e.name for e in KarrasDiffusionSchedulers]
# 设置调度程序的顺序为 1
order = 1
# 使用 @register_to_config 装饰器将该方法注册到配置
@register_to_config
def __init__(
# 设置训练时间步的数量，默认值为 1000
num_train_timesteps: int = 1000,
# 设置 beta 的起始值，默认值为 0.0001
beta_start: float = 0.0001,
# 设置 beta 的结束值，默认值为 0.02
beta_end: float = 0.02,
# 设置 beta 的调度方式，默认为线性
beta_schedule: str = "linear",
# 训练好的 beta 值，可选，默认为 None
trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
# 设置求解器的阶数，默认值为 2
solver_order: int = 2,
# 设置预测类型，默认为 epsilon
prediction_type: str = "epsilon",
# 是否启用阈值处理，默认为 False
thresholding: bool = False,
# 设置动态阈值处理比例，默认值为 0.995
dynamic_thresholding_ratio: float = 0.995,
# 设置样本的最大值，默认值为 1.0
sample_max_value: float = 1.0,
# 是否预测 x0，默认为 True
predict_x0: bool = True,
# 设置求解器类型，默认值为 "bh2"
solver_type: str = "bh2",
# 是否在最后一步降低阶数，默认为 True
lower_order_final: bool = True,
# 禁用校正器的步骤列表，默认为空
disable_corrector: List[int] = [],
# 设置求解器 p，默认为 None
solver_p: SchedulerMixin = None,
# 是否使用 Karras 的 sigma 值，可选，默认为 False
use_karras_sigmas: Optional[bool] = False,
# 设置时间步的间隔方式，默认为 "linspace"
timestep_spacing: str = "linspace",
# 设置时间步的偏移量，默认为 0
steps_offset: int = 0,
# 设置最终 sigma 的类型，可选，默认为 "zero"
final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min"
# 是否重新缩放 beta 以适应零 SNR，默认为 False
rescale_betas_zero_snr: bool = False,
):
# 检查是否提供了已训练的 beta 值
if trained_betas is not None:
# 将训练的 beta 值转换为 32 位浮点张量
self.betas = torch.tensor(trained_betas, dtype=torch.float32)
# 根据线性调度生成 beta 值
elif beta_schedule == "linear":
self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
# 根据缩放线性调度生成 beta 值，特定于潜在扩散模型
elif beta_schedule == "scaled_linear":
self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
# 使用平方余弦调度生成 beta 值
elif beta_schedule == "squaredcos_cap_v2":
# Glide 余弦调度
self.betas = betas_for_alpha_bar(num_train_timesteps)
else:
# 如果调度未实现，抛出异常
raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
# 如果需要重新缩放 beta 值，调用相应方法
if rescale_betas_zero_snr:
self.betas = rescale_zero_terminal_snr(self.betas)
# 计算 alphas 值
self.alphas = 1.0 - self.betas
# 计算 alphas 的累积乘积
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
# 如果需要重新缩放，设置 alphas_cumprod 的最后一个值
if rescale_betas_zero_snr:
# 设置接近于 0 的值，以避免第一个 sigma 为无穷大
self.alphas_cumprod[-1] = 2**-24
# 目前仅支持 VP 类型噪声调度
self.alpha_t = torch.sqrt(self.alphas_cumprod)
self.sigma_t = torch.sqrt(1 - self.alphas_cumprod)
self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t)
# 计算 sigmas 值
self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5
# 初始化噪声分布的标准差
self.init_noise_sigma = 1.0
# 检查求解器类型是否合法
if solver_type not in ["bh1", "bh2"]:
if solver_type in ["midpoint", "heun", "logrho"]:
# 如果求解器是某些类型，注册到配置中
self.register_to_config(solver_type="bh2")
else:
# 如果求解器未实现，抛出异常
raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}")
# 设置预测 x0 值
self.predict_x0 = predict_x0
# 可设置值的初始化
self.num_inference_steps = None
# 创建时间步的逆序数组
timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy()
# 将时间步转换为张量
self.timesteps = torch.from_numpy(timesteps)
# 初始化模型输出列表
self.model_outputs = [None] * solver_order
# 初始化时间步列表
self.timestep_list = [None] * solver_order
# 初始化较低阶数目
self.lower_order_nums = 0
# 设置校正器禁用标志
self.disable_corrector = disable_corrector
# 设置求解器参数
self.solver_p = solver_p
# 初始化最后样本
self.last_sample = None
# 初始化步索引
self._step_index = None
# 初始化开始索引
self._begin_index = None
# 将 sigmas 移动到 CPU，减少 CPU/GPU 通信
self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication
@property
def step_index(self):
"""
当前时间步的索引计数器。每次调度步骤后增加 1。
"""
return self._step_index
@property
def begin_index(self):
"""
第一个时间步的索引。应通过 `set_begin_index` 方法从管道设置。
"""
return self._begin_index
# 从 diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index 复制而来
def set_begin_index(self, begin_index: int = 0):
"""
设置调度器的起始索引。此函数应在推理之前从管道中运行。
参数：
begin_index (`int`):
调度器的起始索引。
"""
# 将给定的起始索引存储到对象的属性中
self._begin_index = begin_index
# 从 diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample 复制而来
def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
"""
“动态阈值化：在每个采样步骤中，我们将 s 设置为 xt0（时间步 t 处 x_0 的预测）中的某个百分位绝对像素值，
如果 s > 1，则将 xt0 阈值化到范围 [-s, s]，然后除以 s。动态阈值化推动饱和像素（那些接近 -1 和 1 的像素）向内移动，从而在每个步骤中积极防止像素饱和。我们发现，动态阈值化显著改善了照片真实感以及更好的图像-文本对齐，尤其是在使用非常大的引导权重时。”
https://arxiv.org/abs/2205.11487
"""
# 获取输入样本的类型
dtype = sample.dtype
# 获取样本的批量大小、通道数以及剩余维度
batch_size, channels, *remaining_dims = sample.shape
# 如果样本类型不是 float32 或 float64，则将其转换为 float 类型
if dtype not in (torch.float32, torch.float64):
sample = sample.float() # 为了进行分位数计算进行类型提升，且 clamp 在 CPU 的 half 类型上未实现
# 将样本展平，以便在每个图像上进行分位数计算
sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
# 计算样本的绝对值
abs_sample = sample.abs() # “某个百分位绝对像素值”
# 计算每个样本的动态阈值 s
s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
# 将 s 限制在 min=1 和 max=self.config.sample_max_value 之间
s = torch.clamp(
s, min=1, max=self.config.sample_max_value
) # 当限制为 min=1 时，相当于标准的 [-1, 1] 裁剪
# 将 s 的形状调整为 (batch_size, 1)，以便在维度 0 上进行广播
s = s.unsqueeze(1) # (batch_size, 1) 因为 clamp 会在维度 0 上广播
# 对样本进行阈值化，并将其除以 s
sample = torch.clamp(sample, -s, s) / s # “我们将 xt0 阈值化到范围 [-s, s]，然后除以 s”
# 将样本的形状恢复到原始形状
sample = sample.reshape(batch_size, channels, *remaining_dims)
# 将样本转换回原始的数据类型
sample = sample.to(dtype)
# 返回处理后的样本
return sample
# 从 diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t 复制而来
# 定义私有方法 _sigma_to_t，接受 sigma 和 log_sigmas 作为参数
def _sigma_to_t(self, sigma, log_sigmas):
# 计算 sigma 的对数值，确保 sigma 不小于 1e-10，以避免取对数时出现负无穷
log_sigma = np.log(np.maximum(sigma, 1e-10))
# 计算 log_sigma 与 log_sigmas 之间的距离
dists = log_sigma - log_sigmas[:, np.newaxis]
# 计算 sigmas 的范围，获取低索引和高索引
low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)
high_idx = low_idx + 1
# 根据低索引和高索引获取 log_sigmas 的对应值
low = log_sigmas[low_idx]
high = log_sigmas[high_idx]
# 进行插值计算，得到权重 w
w = (low - log_sigma) / (low - high)
# 将权重限制在 [0, 1] 范围内
w = np.clip(w, 0, 1)
# 将插值转换为时间范围 t
t = (1 - w) * low_idx + w * high_idx
# 调整 t 的形状以匹配 sigma 的形状
t = t.reshape(sigma.shape)
# 返回计算得到的时间范围 t
return t
# 定义私有方法 _sigma_to_alpha_sigma_t，接受 sigma 作为参数
# 从 DPMSolverMultistepScheduler 中复制而来
def _sigma_to_alpha_sigma_t(self, sigma):
# 计算 alpha_t，作为 sigma 的归一化因子
alpha_t = 1 / ((sigma**2 + 1) ** 0.5)
# 根据 alpha_t 计算 sigma_t
sigma_t = sigma * alpha_t
# 返回 alpha_t 和 sigma_t
return alpha_t, sigma_t
# 定义私有方法 _convert_to_karras，接受 in_sigmas 和 num_inference_steps 作为参数
# 从 EulerDiscreteScheduler 中复制而来
def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
"""构造 Karras 等人 (2022) 的噪声调度。"""
# 确保其他调度器在复制此函数时不会出现问题的黑客方式
# TODO: 将此逻辑添加到其他调度器
if hasattr(self.config, "sigma_min"):
# 如果配置中有 sigma_min，使用其值
sigma_min = self.config.sigma_min
else:
# 否则将 sigma_min 设置为 None
sigma_min = None
if hasattr(self.config, "sigma_max"):
# 如果配置中有 sigma_max，使用其值
sigma_max = self.config.sigma_max
else:
# 否则将 sigma_max 设置为 None
sigma_max = None
# 设置 sigma_min 和 sigma_max 的默认值
sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
# 定义 rho 为 7.0，这是论文中使用的值
rho = 7.0
# 创建从 0 到 1 的 ramp 数组，长度为 num_inference_steps
ramp = np.linspace(0, 1, num_inference_steps)
# 计算 sigma_min 和 sigma_max 的倒数
min_inv_rho = sigma_min ** (1 / rho)
max_inv_rho = sigma_max ** (1 / rho)
# 计算 sigmas，使用 ramp 进行插值
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
# 返回计算得到的 sigmas
return sigmas
# 定义 convert_model_output 方法，接受 model_output 和其他参数
def convert_model_output(
self,
model_output: torch.Tensor,
*args,
sample: torch.Tensor = None,
**kwargs,
# 该函数将模型输出转换为 UniPC 算法所需的相应类型
) -> torch.Tensor:
r"""
将模型输出转换为 UniPC 算法所需的对应类型。
Args:
model_output (`torch.Tensor`):
学习的扩散模型的直接输出。
timestep (`int`):
扩散链中的当前离散时间步。
sample (`torch.Tensor`):
由扩散过程创建的当前样本实例。
Returns:
`torch.Tensor`:
转换后的模型输出。
"""
# 获取时间步，优先从位置参数 args 获取，如果 args 为空，则从关键字参数 kwargs 获取
timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
# 如果样本为 None，则尝试从位置参数 args 获取第二个参数作为样本
if sample is None:
if len(args) > 1:
sample = args[1]
else:
# 如果样本仍然缺失，则抛出异常
raise ValueError("missing `sample` as a required keyward argument")
# 如果时间步不为空，警告用户该参数已被弃用
if timestep is not None:
deprecate(
"timesteps",
"1.0.0",
"Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
)
# 获取当前步骤索引对应的 sigma 值
sigma = self.sigmas[self.step_index]
# 将 sigma 转换为 alpha 和 sigma_t
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
# 如果预测 x0，则根据不同的预测类型计算 x0_pred
if self.predict_x0:
if self.config.prediction_type == "epsilon":
# 计算 x0_pred
x0_pred = (sample - sigma_t * model_output) / alpha_t
elif self.config.prediction_type == "sample":
# 直接使用模型输出作为 x0_pred
x0_pred = model_output
elif self.config.prediction_type == "v_prediction":
# 根据 v_prediction 计算 x0_pred
x0_pred = alpha_t * sample - sigma_t * model_output
else:
# 如果 prediction_type 不符合预期，则抛出异常
raise ValueError(
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
" `v_prediction` for the UniPCMultistepScheduler."
)
# 如果启用了阈值处理，则对 x0_pred 应用阈值
if self.config.thresholding:
x0_pred = self._threshold_sample(x0_pred)
# 返回计算后的 x0_pred
return x0_pred
else:
# 否则根据不同的预测类型返回相应的输出
if self.config.prediction_type == "epsilon":
return model_output
elif self.config.prediction_type == "sample":
# 根据公式计算 epsilon
epsilon = (sample - alpha_t * model_output) / sigma_t
return epsilon
elif self.config.prediction_type == "v_prediction":
# 根据公式计算 epsilon
epsilon = alpha_t * model_output + sigma_t * sample
return epsilon
else:
# 如果 prediction_type 不符合预期，则抛出异常
raise ValueError(
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
" `v_prediction` for the UniPCMultistepScheduler."
)
# 定义 multistep_uni_p_bh_update 函数
def multistep_uni_p_bh_update(
self,
model_output: torch.Tensor,
*args,
# 定义样本和顺序的可选参数
sample: torch.Tensor = None,
order: int = None,
**kwargs,
# 多步骤更新函数，进行模型输出的更新
def multistep_uni_c_bh_update(
self,
this_model_output: torch.Tensor, # 当前模型输出的张量
*args, # 可变参数，用于传递其他参数
last_sample: torch.Tensor = None, # 上一个样本，默认为 None
this_sample: torch.Tensor = None, # 当前样本，默认为 None
order: int = None, # 步骤顺序，默认为 None
**kwargs, # 关键字参数，用于传递其他参数
# 从 diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep 复制而来
def index_for_timestep(self, timestep, schedule_timesteps=None): # 定义时间步的索引获取函数
if schedule_timesteps is None: # 如果未提供调度时间步
schedule_timesteps = self.timesteps # 使用对象的时间步
index_candidates = (schedule_timesteps == timestep).nonzero() # 找出与当前时间步匹配的索引候选
if len(index_candidates) == 0: # 如果没有匹配的索引候选
step_index = len(self.timesteps) - 1 # 设定步索引为时间步的最后一个索引
# 对于第一个“步骤”所取的 sigma 索引
# 总是第二个索引（如果只有一个，则为最后一个）
# 这样可以确保我们不会意外跳过 sigma
# 如果我们从去噪调度的中间开始（例如图像到图像）
elif len(index_candidates) > 1: # 如果有多个匹配的候选
step_index = index_candidates[1].item() # 选择第二个索引
else: # 否则只有一个匹配
step_index = index_candidates[0].item() # 选择第一个索引
return step_index # 返回找到的步索引
# 从 diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index 复制而来
def _init_step_index(self, timestep): # 初始化调度器的步索引计数器
"""
Initialize the step_index counter for the scheduler.
"""
if self.begin_index is None: # 如果开始索引为空
if isinstance(timestep, torch.Tensor): # 如果时间步是张量
timestep = timestep.to(self.timesteps.device) # 将其转换到时间步的设备
self._step_index = self.index_for_timestep(timestep) # 获取并设置当前步索引
else: # 如果开始索引不为空
self._step_index = self._begin_index # 使用提供的开始索引
def step( # 定义步骤函数
self,
model_output: torch.Tensor, # 模型输出的张量
timestep: Union[int, torch.Tensor], # 当前时间步，可以是整数或张量
sample: torch.Tensor, # 当前样本的张量
return_dict: bool = True, # 是否返回字典形式的结果，默认为 True
# 定义函数，返回类型为调度器输出或元组
) -> Union[SchedulerOutput, Tuple]:
"""
通过反向 SDE 从前一个时间步预测样本。该函数使用多步 UniPC 传播样本。
参数：
model_output (`torch.Tensor`):
从学习的扩散模型直接输出的张量。
timestep (`int`):
当前在扩散链中的离散时间步。
sample (`torch.Tensor`):
通过扩散过程创建的当前样本实例。
return_dict (`bool`):
是否返回 [`~schedulers.scheduling_utils.SchedulerOutput`] 或元组。
返回：
[`~schedulers.scheduling_utils.SchedulerOutput`] 或元组:
如果 return_dict 为 `True`，则返回 [`~schedulers.scheduling_utils.SchedulerOutput`]，否则返回元组，
其中第一个元素是样本张量。
"""
# 检查推理步骤是否为 None，若是则引发错误
if self.num_inference_steps is None:
raise ValueError(
"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
)
# 如果步骤索引为 None，则初始化步骤索引
if self.step_index is None:
self._init_step_index(timestep)
# 检查是否使用修正器，条件是步骤索引大于 0，并且前一个步骤没有被禁用且上一个样本不为 None
use_corrector = (
self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None
)
# 转换模型输出以便后续使用
model_output_convert = self.convert_model_output(model_output, sample=sample)
# 如果使用修正器，更新样本
if use_corrector:
sample = self.multistep_uni_c_bh_update(
this_model_output=model_output_convert,
last_sample=self.last_sample,
this_sample=sample,
order=self.this_order,
)
# 更新模型输出和时间步列表
for i in range(self.config.solver_order - 1):
self.model_outputs[i] = self.model_outputs[i + 1]
self.timestep_list[i] = self.timestep_list[i + 1]
# 将当前模型输出和时间步存入最后的位置
self.model_outputs[-1] = model_output_convert
self.timestep_list[-1] = timestep
# 根据配置决定当前阶数
if self.config.lower_order_final:
this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index)
else:
this_order = self.config.solver_order
# 设置当前阶数并进行多步的热身
self.this_order = min(this_order, self.lower_order_nums + 1) # warmup for multistep
assert self.this_order > 0 # 确保当前阶数大于 0
# 更新最后的样本
self.last_sample = sample
# 更新前一个样本
prev_sample = self.multistep_uni_p_bh_update(
model_output=model_output, # 传递原始未转换的模型输出，以防使用 solver-p
sample=sample,
order=self.this_order,
)
# 如果低阶数量小于配置的解算器阶数，则增加低阶数量
if self.lower_order_nums < self.config.solver_order:
self.lower_order_nums += 1
# 完成后将步骤索引增加一
self._step_index += 1
# 如果不返回字典，则返回前一个样本作为元组
if not return_dict:
return (prev_sample,)
# 返回调度器输出对象，包含前一个样本
return SchedulerOutput(prev_sample=prev_sample)
# 定义一个方法，用于缩放模型输入，接受一个张量样本及可变参数
def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
"""
确保与需要根据当前时间步缩放去噪模型输入的调度器的互换性。
Args:
sample (`torch.Tensor`):
输入样本。
Returns:
`torch.Tensor`:
缩放后的输入样本。
"""
# 直接返回输入样本，不做任何处理
return sample
# 从 diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise 拷贝
def add_noise(
self,
original_samples: torch.Tensor, # 原始样本张量
noise: torch.Tensor, # 噪声张量
timesteps: torch.IntTensor, # 时间步张量
) -> torch.Tensor:
# 确保 sigmas 和 timesteps 的设备和数据类型与 original_samples 相同
sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
# 检查设备类型，如果是 MPS 并且时间步是浮点数
if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
# MPS 不支持 float64，因此将时间步转换为 float32
schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32)
timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
else:
# 将调度时间步转换为与原始样本相同的设备
schedule_timesteps = self.timesteps.to(original_samples.device)
timesteps = timesteps.to(original_samples.device)
# 如果 begin_index 为 None，表示调度器用于训练或管道未实现 set_begin_index
if self.begin_index is None:
# 计算每个时间步对应的步骤索引
step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps]
elif self.step_index is not None:
# add_noise 在第一次去噪步骤之后调用（用于修复）
step_indices = [self.step_index] * timesteps.shape[0]
else:
# add_noise 在第一次去噪步骤之前调用以创建初始潜在图像（img2img）
step_indices = [self.begin_index] * timesteps.shape[0]
# 根据步骤索引提取 sigma，并展平
sigma = sigmas[step_indices].flatten()
# 如果 sigma 的维度小于原始样本的维度，则增加维度
while len(sigma.shape) < len(original_samples.shape):
sigma = sigma.unsqueeze(-1)
# 将 sigma 转换为 alpha_t 和 sigma_t
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
# 根据 alpha_t 和 sigma_t 生成带噪声的样本
noisy_samples = alpha_t * original_samples + sigma_t * noise
# 返回带噪声的样本
return noisy_samples
# 返回训练时间步的数量
def __len__(self):
return self.config.num_train_timesteps
浙公网安备 33010602011771号