diffusers 源码解析（六十）
# 版权所有 2024 Katherine Crowson 和 The HuggingFace 团队。所有权利保留。
#
# 根据 Apache 许可证，版本 2.0（“许可证”）进行许可；
# 除非遵循该许可证，否则您不得使用此文件。
# 您可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面协议另有规定，分发在许可证下的软件均按“原样”提供，
# 不附带任何形式的担保或条件，无论是明示还是暗示。
# 请参见许可证以获取特定语言的权限和
# 限制条款。
import math # 导入数学模块以进行数学运算
from dataclasses import dataclass # 从数据类模块导入dataclass装饰器
from typing import List, Optional, Tuple, Union # 导入类型注解，用于类型检查和文档
import numpy as np # 导入NumPy库以进行数值计算
import torch # 导入PyTorch库以进行深度学习操作
from ..configuration_utils import ConfigMixin, register_to_config # 从配置工具导入混合配置和注册功能
from ..utils import BaseOutput, logging # 从实用工具导入基本输出类和日志记录功能
from ..utils.torch_utils import randn_tensor # 从PyTorch工具导入生成随机张量的功能
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin # 从调度工具导入调度器类
logger = logging.get_logger(__name__) # 获取当前模块的日志记录器实例，名称为当前模块名
@dataclass
# 从 diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput 复制，DDPM 转换为 EulerDiscrete
class EulerDiscreteSchedulerOutput(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 # 可选的预测去噪样本，默认为None
# 从 diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar 复制
def betas_for_alpha_bar(
num_diffusion_timesteps, # 扩散时间步的数量
max_beta=0.999, # 使用的最大beta值，防止出现奇点
alpha_transform_type="cosine", # alpha_bar的噪声调度类型，默认为"cosine"
):
"""
创建一个beta调度，离散化给定的alpha_t_bar函数，该函数定义了
（1-beta）随时间的累积乘积，从 t = [0,1]。
包含一个alpha_bar函数，该函数接受参数t并将其转换为
该部分扩散过程的（1-beta）的累积乘积。
参数：
num_diffusion_timesteps (`int`): 要生成的betas数量。
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_fn 函数，使用余弦函数进行变换
def alpha_bar_fn(t):
# 计算余弦值的平方，生成平滑过渡的 alpha 值
return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
# 检查指定的 alpha 变换类型是否为 "exp"
elif alpha_transform_type == "exp":
# 定义 alpha_bar_fn 函数，使用指数函数进行变换
def alpha_bar_fn(t):
# 计算指数衰减值，返回随着时间衰减的 alpha 值
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 列表转换为浮点数张量并返回
return torch.tensor(betas, dtype=torch.float32)
# 从 diffusers.schedulers.scheduling_ddim 导入 rescale_zero_terminal_snr 方法
def rescale_zero_terminal_snr(betas):
"""
根据文献 https://arxiv.org/pdf/2305.08891.pdf (算法 1) 对 betas 进行重新缩放，使终端 SNR 为零
参数:
betas (`torch.Tensor`):
用于初始化调度器的 betas。
返回:
`torch.Tensor`: 重新缩放后的 betas，使终端 SNR 为零
"""
# 将 betas 转换为 alphas_bar_sqrt
alphas = 1.0 - betas # 计算 alphas
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_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 = alphas_bar[1:] / alphas_bar[:-1] # 还原累积乘积
alphas = torch.cat([alphas_bar[0:1], alphas]) # 拼接第一个值
betas = 1 - alphas # 计算 betas
return betas # 返回重新缩放后的 betas
class EulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
"""
Euler 调度器。
此模型继承自 [`SchedulerMixin`] 和 [`ConfigMixin`]。请查看超类文档以获取库为所有调度器实现的通用方法，如加载和保存。
# 参数说明部分
Args:
num_train_timesteps (`int`, defaults to 1000): # 训练模型的扩散步骤数量，默认为1000。
The number of diffusion steps to train the model.
beta_start (`float`, defaults to 0.0001): # 推理的起始 beta 值，默认为0.0001。
The starting `beta` value of inference.
beta_end (`float`, defaults to 0.02): # 最终 beta 值，默认为0.02。
The final `beta` value.
beta_schedule (`str`, defaults to `"linear"`): # beta 调度方式，默认为“线性”，可选“线性”或“缩放线性”。
The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
`linear` or `scaled_linear`.
trained_betas (`np.ndarray`, *optional*): # 可选，直接传入 beta 数组以绕过 beta_start 和 beta_end。
Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
prediction_type (`str`, defaults to `epsilon`, *optional*): # 可选，调度函数的预测类型，默认为 `epsilon`。
Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
`sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
Video](https://imagen.research.google/video/paper.pdf) paper).
interpolation_type(`str`, defaults to `"linear"`, *optional*): # 可选，插值类型用于计算调度去噪步骤的中间 sigma，默认为“线性”。
The interpolation type to compute intermediate sigmas for the scheduler denoising steps. Should be one of
`"linear"` or `"log_linear"`.
use_karras_sigmas (`bool`, *optional*, defaults to `False`): # 可选，是否在采样过程中使用 Karras sigma，默认为 False。
Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
the sigmas are determined according to a sequence of noise levels {σi}.
timestep_spacing (`str`, defaults to `"linspace"`): # 时间步缩放方式，默认为“线性空间”。
The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
steps_offset (`int`, defaults to 0): # 推理步骤的偏移量，默认为0，某些模型可能需要。
An offset added to the inference steps, as required by some model families.
rescale_betas_zero_snr (`bool`, defaults to `False`): # 可选，是否将 beta 重新缩放为零终端 SNR，默认为 False。
Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
dark samples instead of limiting it to samples with medium brightness. Loosely related to
[`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
final_sigmas_type (`str`, defaults to `"zero"`): # 最终 sigma 值在采样过程中的噪声调度，默认为“零”。
The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
"""
# 兼容的调度器列表，提取 KarrasDiffusionSchedulers 中的名称
_compatibles = [e.name for e in KarrasDiffusionSchedulers]
# 设定顺序，通常用于调度器的优先级
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的调度方式，默认为线性
beta_schedule: str = "linear",
# 训练的beta值，默认为None
trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
# 预测类型，默认为"epsilon"
prediction_type: str = "epsilon",
# 插值类型，默认为线性
interpolation_type: str = "linear",
# 是否使用Karras的sigma，默认为False
use_karras_sigmas: Optional[bool] = False,
# sigma的最小值，默认为None
sigma_min: Optional[float] = None,
# sigma的最大值，默认为None
sigma_max: Optional[float] = None,
# 时间步间距类型，默认为"linspace"
timestep_spacing: str = "linspace",
# 时间步类型，默认为"discrete"，可选"discrete"或"continuous"
timestep_type: str = "discrete", # can be "discrete" or "continuous"
# 步骤偏移量，默认为0
steps_offset: int = 0,
# 是否重新缩放beta以实现零SNR，默认为False
rescale_betas_zero_snr: bool = False,
# 最终sigma的类型，默认为"zero"，可选"zero"或"sigma_min"
final_sigmas_type: str = "zero", # can be "zero" or "sigma_min"
):
# 如果提供了训练的beta值，将其转换为张量
if trained_betas is not None:
self.betas = torch.tensor(trained_betas, dtype=torch.float32)
# 如果beta调度方式为线性，则生成线性beta值
elif beta_schedule == "linear":
self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
# 如果beta调度方式为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调度方式为squaredcos_cap_v2，则使用Glide余弦调度生成beta值
elif beta_schedule == "squaredcos_cap_v2":
self.betas = betas_for_alpha_bar(num_train_timesteps)
# 如果不支持的调度方式，抛出错误
else:
raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
# 如果需要重新缩放beta以实现零SNR，则执行缩放
if rescale_betas_zero_snr:
self.betas = rescale_zero_terminal_snr(self.betas)
# 计算alpha值，等于1减去beta值
self.alphas = 1.0 - self.betas
# 计算alpha的累积乘积
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
# 如果需要重新缩放beta以实现零SNR，调整最后一个alpha值以避免inf
if rescale_betas_zero_snr:
# 近乎0但不为0，以避免第一个sigma为inf
self.alphas_cumprod[-1] = 2**-24
# 计算sigma值，通过公式得到并反转
sigmas = (((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5).flip(0)
# 创建时间步数组，从0到num_train_timesteps - 1
timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy()
# 将时间步转换为张量
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
# 可设值，初始化为None
self.num_inference_steps = None
# TODO: 支持所有预测类型和时间步类型的完整EDM缩放
# 如果时间步类型为连续且预测类型为v_prediction，则计算时间步
if timestep_type == "continuous" and prediction_type == "v_prediction":
self.timesteps = torch.Tensor([0.25 * sigma.log() for sigma in sigmas])
# 否则直接使用预定义的时间步
else:
self.timesteps = timesteps
# 将sigma与零张量连接以形成最终的sigma
self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
# 标记输入缩放是否被调用
self.is_scale_input_called = False
# 设置是否使用Karras sigma的标志
self.use_karras_sigmas = use_karras_sigmas
# 初始化步骤索引和开始索引为None
self._step_index = None
self._begin_index = None
# 将sigma移至CPU以减少CPU和GPU之间的通信
self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication
# 设为属性
@property
# 初始化噪声标准差
def init_noise_sigma(self):
# 获取初始噪声分布的标准差
max_sigma = max(self.sigmas) if isinstance(self.sigmas, list) else self.sigmas.max()
# 如果时间步长间隔配置为 "linspace" 或 "trailing"，返回最大标准差
if self.config.timestep_spacing in ["linspace", "trailing"]:
return max_sigma
# 否则返回 (max_sigma**2 + 1) 的平方根
return (max_sigma**2 + 1) ** 0.5
# 当前时间步的索引计数器属性
@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
# 缩放模型输入以确保与调度器的互换性
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
"""
确保与需要根据当前时间步缩放去噪模型输入的调度器互换性。通过 `(sigma**2 + 1) ** 0.5` 缩放去噪模型输入以匹配 Euler 算法。
参数:
sample (`torch.Tensor`):
输入样本。
timestep (`int`, *可选*):
扩散链中的当前时间步。
返回:
`torch.Tensor`:
缩放后的输入样本。
"""
# 如果步索引为空，初始化步索引
if self.step_index is None:
self._init_step_index(timestep)
# 获取当前时间步的标准差
sigma = self.sigmas[self.step_index]
# 按照公式缩放样本输入
sample = sample / ((sigma**2 + 1) ** 0.5)
# 标记输入缩放已被调用
self.is_scale_input_called = True
# 返回缩放后的样本
return sample
# 设置时间步的函数
def set_timesteps(
self,
num_inference_steps: int = None,
device: Union[str, torch.device] = None,
timesteps: Optional[List[int]] = None,
sigmas: Optional[List[float]] = None,
# 将 sigma 转换为时间步
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)
# 高索引为低索引加1
high_idx = low_idx + 1
# 获取低和高的对数 sigma
low = log_sigmas[low_idx]
high = log_sigmas[high_idx]
# 插值计算 sigma
w = (low - log_sigma) / (low - high)
# 限制插值权重在0和1之间
w = np.clip(w, 0, 1)
# 将插值转换为时间范围
t = (1 - w) * low_idx + w * high_idx
# 重新塑形为与 sigma 形状相同
t = t.reshape(sigma.shape)
# 返回转换后的时间步
return t
# 从指定的 GitHub 链接复制的代码
def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
"""构造 Karras 等人 (2022) 的噪声调度。"""
# 确保其他调度器复制此函数时不会出现问题的 hack
# TODO: 将此逻辑添加到其他调度器中
if hasattr(self.config, "sigma_min"):
# 从配置中获取 sigma_min 值
sigma_min = self.config.sigma_min
else:
# 如果没有，则设置为 None
sigma_min = None
if hasattr(self.config, "sigma_max"):
# 从配置中获取 sigma_max 值
sigma_max = self.config.sigma_max
else:
# 如果没有，则设置为 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 # 在论文中使用的值
# 创建一个线性空间，从 0 到 1，共 num_inference_steps 个点
ramp = np.linspace(0, 1, num_inference_steps)
# 计算 sigma_min 和 sigma_max 的逆 rho 次方
min_inv_rho = sigma_min ** (1 / rho)
max_inv_rho = sigma_max ** (1 / rho)
# 计算新的 sigma 值
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
# 返回计算得到的 sigma 值
return sigmas
def index_for_timestep(self, timestep, schedule_timesteps=None):
# 如果没有提供 schedule_timesteps，则使用类中的 timesteps
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
# 查找与 timestep 匹配的索引
indices = (schedule_timesteps == timestep).nonzero()
# 第一个 `step` 的 sigma 索引总是第二个索引
# 如果只有一个，则使用最后一个
# 这样可以确保在去噪调度中间开始时不跳过 sigma
pos = 1 if len(indices) > 1 else 0
# 返回所需的索引值
return indices[pos].item()
def _init_step_index(self, timestep):
# 如果 begin_index 为 None
if self.begin_index is None:
# 如果 timestep 是张量，则将其移动到 timesteps 的设备
if isinstance(timestep, torch.Tensor):
timestep = timestep.to(self.timesteps.device)
# 根据 timestep 初始化步骤索引
self._step_index = self.index_for_timestep(timestep)
else:
# 否则，使用 _begin_index
self._step_index = self._begin_index
def step(
self,
model_output: torch.Tensor,
timestep: Union[float, torch.Tensor],
sample: torch.Tensor,
s_churn: float = 0.0,
s_tmin: float = 0.0,
s_tmax: float = float("inf"),
s_noise: float = 1.0,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
def add_noise(
self,
original_samples: torch.Tensor,
noise: torch.Tensor,
timesteps: torch.Tensor,
) -> torch.Tensor:
# 确保 sigmas 和 timesteps 与 original_samples 具有相同的设备和数据类型
sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
# 检查设备类型是否为 "mps" 且 timesteps 是否为浮点数
if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
# mps 不支持 float64，因此将 timesteps 转换为 float32
schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32)
timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
else:
# 将 schedule_timesteps 转换为与 original_samples 相同的设备类型
schedule_timesteps = self.timesteps.to(original_samples.device)
# 将 timesteps 转换为与 original_samples 相同的设备类型
timesteps = timesteps.to(original_samples.device)
# 当 scheduler 用于训练时，self.begin_index 为 None，或者管道未实现 set_begin_index
if self.begin_index is None:
# 根据 timesteps 和 schedule_timesteps 计算步长索引
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]
# 根据步长索引从 sigmas 中提取相应的 sigma 值并展平
sigma = sigmas[step_indices].flatten()
# 如果 sigma 的形状小于 original_samples 的形状，则增加维度
while len(sigma.shape) < len(original_samples.shape):
sigma = sigma.unsqueeze(-1)
# 计算带噪声的样本，原始样本加上噪声乘以 sigma
noisy_samples = original_samples + noise * sigma
# 返回带噪声的样本
return noisy_samples
# 计算给定样本、噪声和时间步长的速度
def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
# 检查 timesteps 是否为整数类型
if (
isinstance(timesteps, int)
or isinstance(timesteps, torch.IntTensor)
or isinstance(timesteps, torch.LongTensor)
):
# 抛出错误，提示不支持整数作为时间步
raise ValueError(
(
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
" `EulerDiscreteScheduler.get_velocity()` is not supported. Make sure to pass"
" one of the `scheduler.timesteps` as a timestep."
),
)
# 如果使用 MPS 设备且时间步是浮点数类型
if sample.device.type == "mps" and torch.is_floating_point(timesteps):
# MPS 不支持 float64 类型，转换为 float32
schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
timesteps = timesteps.to(sample.device, dtype=torch.float32)
else:
# 将时间步转换为当前设备的默认类型
schedule_timesteps = self.timesteps.to(sample.device)
timesteps = timesteps.to(sample.device)
# 获取每个时间步对应的索引
step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps]
# 将 alpha 累积乘积转换到样本的设备
alphas_cumprod = self.alphas_cumprod.to(sample)
# 计算 sqrt(alpha) 的积
sqrt_alpha_prod = alphas_cumprod[step_indices] ** 0.5
# 将结果展平为一维
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
# 扩展维度，直到与样本形状匹配
while len(sqrt_alpha_prod.shape) < len(sample.shape):
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
# 计算 sqrt(1 - alpha) 的积
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[step_indices]) ** 0.5
# 将结果展平为一维
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
# 扩展维度，直到与样本形状匹配
while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
# 计算速度，结合噪声和样本
velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
# 返回计算得到的速度
return velocity
# 返回训练时间步的数量
def __len__(self):
return self.config.num_train_timesteps
# 版权所有 2024 Katherine Crowson 和 HuggingFace 团队。保留所有权利。
#
# 根据 Apache 许可证，版本 2.0（“许可证”）许可；
# 除非遵循该许可证，否则您不得使用此文件。
# 您可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，软件在许可证下分发是按“原样”基础进行的，
# 不提供任何形式的明示或暗示的保证或条件。
# 有关许可证的特定权限和限制，请参见许可证。
# 从 dataclasses 模块导入 dataclass 装饰器
from dataclasses import dataclass
# 从 typing 模块导入可选类型、元组和联合类型
from typing import Optional, Tuple, Union
# 导入 flax 库
import flax
# 导入 jax 的 numpy 模块
import jax.numpy as jnp
# 从配置工具中导入 ConfigMixin 和 register_to_config
from ..configuration_utils import ConfigMixin, register_to_config
# 从调度工具中导入相关类和函数
from .scheduling_utils_flax import (
CommonSchedulerState, # 导入通用调度器状态
FlaxKarrasDiffusionSchedulers, # 导入 Karras 扩散调度器
FlaxSchedulerMixin, # 导入调度器混合类
FlaxSchedulerOutput, # 导入调度器输出类
broadcast_to_shape_from_left, # 导入从左侧广播到形状的函数
)
# 定义一个调度器状态的数据类，使用 flax 的结构化数据类装饰器
@flax.struct.dataclass
class EulerDiscreteSchedulerState:
common: CommonSchedulerState # 包含通用调度器状态的属性
# 可设置的值
init_noise_sigma: jnp.ndarray # 初始噪声的标准差
timesteps: jnp.ndarray # 时间步长的数组
sigmas: jnp.ndarray # sigma 值的数组
num_inference_steps: Optional[int] = None # 推理步骤的数量，可选
# 类方法，用于创建 EulerDiscreteSchedulerState 实例
@classmethod
def create(
cls, common: CommonSchedulerState, init_noise_sigma: jnp.ndarray, timesteps: jnp.ndarray, sigmas: jnp.ndarray
):
# 返回一个新的实例，使用提供的参数初始化
return cls(common=common, init_noise_sigma=init_noise_sigma, timesteps=timesteps, sigmas=sigmas)
# 定义一个输出类，继承自 FlaxSchedulerOutput
@dataclass
class FlaxEulerDiscreteSchedulerOutput(FlaxSchedulerOutput):
state: EulerDiscreteSchedulerState # 包含调度器状态的属性
# 定义一个调度器类，继承自 FlaxSchedulerMixin 和 ConfigMixin
class FlaxEulerDiscreteScheduler(FlaxSchedulerMixin, ConfigMixin):
"""
Euler 调度器（算法 2），参考 Karras 等人（2022） https://arxiv.org/abs/2206.00364。基于 Katherine Crowson 的原始
k-diffusion 实现：
https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L51
[`~ConfigMixin`] 处理在调度器的 `__init__` 函数中传递的所有配置属性的存储，
例如 `num_train_timesteps`。可以通过 `scheduler.config.num_train_timesteps` 访问它们。
[`SchedulerMixin`] 提供通过 [`SchedulerMixin.save_pretrained`] 和
[`~SchedulerMixin.from_pretrained`] 函数进行通用加载和保存的功能。
"""
# 文档字符串，描述类初始化方法的参数
Args:
num_train_timesteps (`int`): 模型训练所用的扩散步骤数量
beta_start (`float`): 推理开始时的 `beta` 值
beta_end (`float`): 最终的 `beta` 值
beta_schedule (`str`):
beta 调度，表示从 beta 范围到模型步进的 beta 序列的映射。可选值为
`linear` 或 `scaled_linear`
trained_betas (`jnp.ndarray`, optional):
直接传递 beta 数组给构造函数以绕过 `beta_start`、`beta_end` 等选项
prediction_type (`str`, default `epsilon`, optional):
调度函数的预测类型，选项包括 `epsilon`（预测扩散过程的噪声）、`sample`（直接预测噪声样本）或 `v_prediction`（见第 2.4 节 https://imagen.research.google/video/paper.pdf）
dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`):
用于参数和计算的 `dtype`
"""
# 获取所有兼容的 FlaxKarrasDiffusionSchedulers 名称
_compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers]
# 声明 dtype 变量，类型为 jnp.dtype
dtype: jnp.dtype
# 属性，指示该类是否有状态
@property
def has_state(self):
# 返回 True，表示该调度器具有状态
return True
# 注册到配置的初始化方法
@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 调度，默认为 "linear"
beta_schedule: str = "linear",
# 可选参数，直接传递 beta 数组
trained_betas: Optional[jnp.ndarray] = None,
# 设置预测类型，默认为 "epsilon"
prediction_type: str = "epsilon",
# 设置时间步间距，默认为 "linspace"
timestep_spacing: str = "linspace",
# 设置数据类型，默认为 jnp.float32
dtype: jnp.dtype = jnp.float32,
):
# 将 dtype 赋值给实例变量
self.dtype = dtype
# 创建状态的方法，接受可选的公共状态参数
def create_state(self, common: Optional[CommonSchedulerState] = None) -> EulerDiscreteSchedulerState:
# 如果没有传入公共状态，则创建一个新的
if common is None:
common = CommonSchedulerState.create(self)
# 生成时间步的数组，逆序排列
timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1]
# 计算每个时间步的标准差
sigmas = ((1 - common.alphas_cumprod) / common.alphas_cumprod) ** 0.5
# 通过插值调整 sigmas 的值
sigmas = jnp.interp(timesteps, jnp.arange(0, len(sigmas)), sigmas)
# 在 sigmas 后附加一个 0.0 值
sigmas = jnp.concatenate([sigmas, jnp.array([0.0], dtype=self.dtype)])
# 标准化初始噪声分布的标准差
if self.config.timestep_spacing in ["linspace", "trailing"]:
# 若时间步间距为 linspace 或 trailing，取最大 sigmas 值
init_noise_sigma = sigmas.max()
else:
# 否则计算初始化噪声标准差
init_noise_sigma = (sigmas.max() ** 2 + 1) ** 0.5
# 返回 EulerDiscreteSchedulerState 的实例，包含公共状态、初始噪声标准差、时间步和 sigmas
return EulerDiscreteSchedulerState.create(
common=common,
init_noise_sigma=init_noise_sigma,
timesteps=timesteps,
sigmas=sigmas,
)
# 定义一个方法，缩放去噪模型输入，以匹配欧拉算法
def scale_model_input(self, state: EulerDiscreteSchedulerState, sample: jnp.ndarray, timestep: int) -> jnp.ndarray:
"""
缩放去噪模型输入，计算方式为 `(sigma**2 + 1) ** 0.5`，以匹配欧拉算法。
参数:
state (`EulerDiscreteSchedulerState`):
`FlaxEulerDiscreteScheduler` 状态数据类实例。
sample (`jnp.ndarray`):
当前正在通过扩散过程创建的样本实例。
timestep (`int`):
扩散链中的当前离散时间步。
返回:
`jnp.ndarray`: 缩放后的输入样本
"""
# 获取当前时间步对应的索引
(step_index,) = jnp.where(state.timesteps == timestep, size=1)
# 提取索引的第一个元素
step_index = step_index[0]
# 根据索引获取 sigma 值
sigma = state.sigmas[step_index]
# 将样本数据缩放
sample = sample / ((sigma**2 + 1) ** 0.5)
# 返回缩放后的样本
return sample
# 定义一个方法，设置扩散链中使用的时间步
def set_timesteps(
self, state: EulerDiscreteSchedulerState, num_inference_steps: int, shape: Tuple = ()
) -> EulerDiscreteSchedulerState:
"""
设置扩散链中使用的时间步。支持在推理之前运行的功能。
参数:
state (`EulerDiscreteSchedulerState`):
`FlaxEulerDiscreteScheduler` 状态数据类实例。
num_inference_steps (`int`):
生成样本时使用的扩散步骤数。
"""
# 根据配置的时间步间隔类型生成时间步
if self.config.timestep_spacing == "linspace":
# 生成线性间隔的时间步
timesteps = jnp.linspace(self.config.num_train_timesteps - 1, 0, num_inference_steps, dtype=self.dtype)
elif self.config.timestep_spacing == "leading":
# 计算步骤比率，并生成时间步
step_ratio = self.config.num_train_timesteps // num_inference_steps
timesteps = (jnp.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(float)
timesteps += 1
else:
# 抛出异常，时间步间隔类型无效
raise ValueError(
f"timestep_spacing must be one of ['linspace', 'leading'], got {self.config.timestep_spacing}"
)
# 计算 sigma 值
sigmas = ((1 - state.common.alphas_cumprod) / state.common.alphas_cumprod) ** 0.5
# 在时间步上插值 sigma 值
sigmas = jnp.interp(timesteps, jnp.arange(0, len(sigmas)), sigmas)
# 在 sigma 后面附加 0.0
sigmas = jnp.concatenate([sigmas, jnp.array([0.0], dtype=self.dtype)])
# 标准差初始化噪声分布
if self.config.timestep_spacing in ["linspace", "trailing"]:
# 对于线性或尾随间隔，初始化噪声的 sigma 为 sigma 的最大值
init_noise_sigma = sigmas.max()
else:
# 否则，计算初始化噪声的 sigma
init_noise_sigma = (sigmas.max() ** 2 + 1) ** 0.5
# 替换状态中的时间步和 sigma 信息
return state.replace(
timesteps=timesteps,
sigmas=sigmas,
num_inference_steps=num_inference_steps,
init_noise_sigma=init_noise_sigma,
)
# 定义一个步骤方法，处理模型输出和当前状态
def step(
self,
state: EulerDiscreteSchedulerState,
model_output: jnp.ndarray,
timestep: int,
sample: jnp.ndarray,
return_dict: bool = True,
) -> Union[FlaxEulerDiscreteSchedulerOutput, Tuple]:
"""
通过逆向 SDE 预测上一个时间步的样本。核心函数用于从学习到的模型输出（通常是预测的噪声）传播扩散过程。
Args:
state (`EulerDiscreteSchedulerState`):
`FlaxEulerDiscreteScheduler` 状态数据类实例。
model_output (`jnp.ndarray`): 来自学习扩散模型的直接输出。
timestep (`int`): 当前扩散链中的离散时间步。
sample (`jnp.ndarray`):
当前通过扩散过程生成的样本实例。
order: 多步推理的系数。
return_dict (`bool`): 返回元组而非 FlaxEulerDiscreteScheduler 类的选项。
Returns:
[`FlaxEulerDiscreteScheduler`] 或 `tuple`: 如果 `return_dict` 为 True，则返回 [`FlaxEulerDiscreteScheduler`]，
否则返回一个 `tuple`。返回元组时，第一个元素是样本张量。
"""
# 检查推理步骤数量是否为 None，如果是，则抛出错误
if state.num_inference_steps is None:
raise ValueError(
"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
)
# 获取当前时间步在 timesteps 中的索引
(step_index,) = jnp.where(state.timesteps == timestep, size=1)
step_index = step_index[0] # 提取索引值
# 获取当前时间步对应的 sigma 值
sigma = state.sigmas[step_index]
# 1. 从 sigma 缩放的预测噪声计算预测的原始样本 (x_0)
if self.config.prediction_type == "epsilon":
# 使用 epsilon 进行预测
pred_original_sample = sample - sigma * model_output
elif self.config.prediction_type == "v_prediction":
# 使用 v_prediction 进行预测，计算方法涉及线性组合
pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
else:
# 如果提供的 prediction_type 无效，抛出错误
raise ValueError(
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`"
)
# 2. 转换为 ODE 导数
derivative = (sample - pred_original_sample) / sigma
# 计算 dt，表示当前 sigma 与下一个 sigma 的差值
dt = state.sigmas[step_index + 1] - sigma
# 计算前一个样本的值
prev_sample = sample + derivative * dt
# 如果不返回字典，则返回元组
if not return_dict:
return (prev_sample, state)
# 否则返回 FlaxEulerDiscreteSchedulerOutput 对象
return FlaxEulerDiscreteSchedulerOutput(prev_sample=prev_sample, state=state)
def add_noise(
self,
state: EulerDiscreteSchedulerState,
original_samples: jnp.ndarray,
noise: jnp.ndarray,
timesteps: jnp.ndarray,
) -> jnp.ndarray:
# 获取当前时间步对应的 sigma 值，并展平为一维
sigma = state.sigmas[timesteps].flatten()
# 将 sigma 广播到与噪声形状相同
sigma = broadcast_to_shape_from_left(sigma, noise.shape)
# 生成加噪声的样本
noisy_samples = original_samples + noise * sigma
# 返回加噪声的样本
return noisy_samples
def __len__(self):
# 返回训练时间步的数量
return self.config.num_train_timesteps
# 版权声明，说明文件归 Stability AI、Katherine Crowson 和 HuggingFace 团队所有
#
# 根据 Apache 许可证 2.0 版进行许可，用户必须遵守该许可证才能使用本文件
# 用户可在以下网址获取许可证的副本
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，软件以“原样”方式分发，不提供任何形式的明示或暗示的担保或条件
# 具体权限和限制请参阅许可证
import math # 导入数学模块以进行数学计算
from dataclasses import dataclass # 从 dataclasses 模块导入 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 ..utils import BaseOutput, logging # 从工具模块导入基础输出类和日志记录工具
from .scheduling_utils import SchedulerMixin # 从调度工具导入调度混合类
logger = logging.get_logger(__name__) # 创建一个日志记录器，用于记录本模块的日志信息
@dataclass
class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
"""
调度器 `step` 函数输出的输出类。
参数：
prev_sample (`torch.FloatTensor`，形状为 `(batch_size, num_channels, height, width)` 的图像):
计算的上一个时间步的样本 `(x_{t-1})`。 `prev_sample` 应作为下一次模型输入在去噪循环中使用。
"""
prev_sample: torch.FloatTensor # 上一个时间步的样本，使用 PyTorch 的浮点张量
class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
"""
欧拉调度器。
此模型继承自 [`SchedulerMixin`] 和 [`ConfigMixin`]。有关库为所有调度器实现的通用方法的文档，请查看父类文档。
参数：
num_train_timesteps (`int`，默认为 1000):
用于训练模型的扩散步骤数量。
timestep_spacing (`str`，默认为 `"linspace"`):
时间步的缩放方式。有关更多信息，请参阅 [Common Diffusion Noise Schedules and
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) 表 2。
shift (`float`，默认为 1.0):
时间步调度的偏移值。
"""
_compatibles = [] # 兼容性列表，初始化为空
order = 1 # 调度器的顺序设置为 1
@register_to_config
def __init__(
self,
num_train_timesteps: int = 1000, # 设置训练时间步数量，默认值为 1000
shift: float = 1.0, # 设置时间步偏移值，默认值为 1.0
use_dynamic_shifting=False, # 是否使用动态偏移，默认值为 False
base_shift: Optional[float] = 0.5, # 基础偏移值，默认值为 0.5
max_shift: Optional[float] = 1.15, # 最大偏移值，默认值为 1.15
base_image_seq_len: Optional[int] = 256, # 基础图像序列长度，默认值为 256
max_image_seq_len: Optional[int] = 4096, # 最大图像序列长度，默认值为 4096
):
# 创建一个从1到num_train_timesteps的等间隔数组，并反转其顺序
timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
# 将NumPy数组转换为PyTorch张量，并指定数据类型为float32
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
# 计算sigmas，作为timesteps与num_train_timesteps的比例
sigmas = timesteps / num_train_timesteps
if not use_dynamic_shifting:
# 如果不使用动态偏移，则根据shift调整sigmas的值
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
# 计算最终的timesteps，乘以num_train_timesteps以获得真实的时间步长
self.timesteps = sigmas * num_train_timesteps
# 初始化步骤索引和开始索引为None
self._step_index = None
self._begin_index = None
# 将sigmas移动到CPU，以减少CPU/GPU之间的通信
self.sigmas = sigmas.to("cpu")
# 获取sigmas中的最小值
self.sigma_min = self.sigmas[-1].item()
# 获取sigmas中的最大值
self.sigma_max = self.sigmas[0].item()
@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):
"""
设置调度器的开始索引。该函数应在推理之前从管道运行。
Args:
begin_index (`int`):
调度器的开始索引。
"""
# 设置开始索引
self._begin_index = begin_index
def scale_noise(
self,
sample: torch.FloatTensor,
timestep: Union[float, torch.FloatTensor],
noise: Optional[torch.FloatTensor] = None,
) -> torch.FloatTensor:
"""
前向过程在流匹配中
参数:
sample (`torch.FloatTensor`):
输入样本。
timestep (`int`, *可选*):
当前扩散链中的时间步。
返回:
`torch.FloatTensor`:
缩放后的输入样本。
"""
# 确保 sigmas 和 timesteps 的设备和数据类型与原始样本相同
sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
# 检查设备类型是否为 MPS，且时间步是否为浮点数
if sample.device.type == "mps" and torch.is_floating_point(timestep):
# MPS 不支持 float64 类型
schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) # 将时间步转换为 float32
timestep = timestep.to(sample.device, dtype=torch.float32) # 将时间步转换为 float32
else:
# 将时间步转换到样本的设备上，保持原始数据类型
schedule_timesteps = self.timesteps.to(sample.device)
timestep = timestep.to(sample.device) # 将时间步转移到样本的设备上
# 当 scheduler 用于训练时，self.begin_index 为 None，或 pipeline 没有实现 set_begin_index
if self.begin_index is None:
# 根据当前时间步索引计算步骤索引
step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
elif self.step_index is not None:
# 在第一次去噪步骤后调用 add_noise（用于修补）
step_indices = [self.step_index] * timestep.shape[0] # 重复当前步骤索引
else:
# 在第一次去噪步骤之前调用 add_noise 以创建初始潜变量（图像到图像）
step_indices = [self.begin_index] * timestep.shape[0] # 重复初始步骤索引
# 获取对应步骤的 sigma 值并展平
sigma = sigmas[step_indices].flatten()
# 扩展 sigma 的维度以匹配样本的形状
while len(sigma.shape) < len(sample.shape):
sigma = sigma.unsqueeze(-1)
# 通过 sigma 和噪声对样本进行加权组合
sample = sigma * noise + (1.0 - sigma) * sample
# 返回处理后的样本
return sample
# 将 sigma 转换为时间步 t
def _sigma_to_t(self, sigma):
return sigma * self.config.num_train_timesteps # 计算对应的时间步
# 根据给定的 mu、sigma 和时间张量 t 进行时间偏移计算
def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) # 计算时间偏移的值
# 设置时间步的函数
def set_timesteps(
self,
num_inference_steps: int = None, # 推理步骤数量（可选）
device: Union[str, torch.device] = None, # 设备类型（可选）
sigmas: Optional[List[float]] = None, # sigma 值列表（可选）
mu: Optional[float] = None, # mu 值（可选）
):
"""
设置用于扩散链的离散时间步（在推理前运行）。
参数：
num_inference_steps (`int`):
用于生成样本的扩散步骤数量，使用预训练模型时。
device (`str` 或 `torch.device`, *可选*):
时间步要移动到的设备。如果为 `None`，则时间步不被移动。
"""
# 如果使用动态移动并且 mu 为 None，则抛出值错误
if self.config.use_dynamic_shifting and mu is None:
raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`")
# 如果 sigmas 为 None，则初始化 num_inference_steps 和计算时间步
if sigmas is None:
self.num_inference_steps = num_inference_steps
# 生成从 sigma_max 到 sigma_min 的均匀时间步
timesteps = np.linspace(
self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
)
# 归一化时间步以计算 sigmas
sigmas = timesteps / self.config.num_train_timesteps
# 如果使用动态移动，调用时间移动函数处理 sigmas
if self.config.use_dynamic_shifting:
sigmas = self.time_shift(mu, 1.0, sigmas)
else:
# 使用配置中的偏移量调整 sigmas
sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
# 将 sigmas 转换为 PyTorch 张量，并设置数据类型和设备
sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
# 计算时间步与训练时间步的乘积
timesteps = sigmas * self.config.num_train_timesteps
# 将计算出的时间步移动到指定设备
self.timesteps = timesteps.to(device=device)
# 将 sigmas 与一个零张量拼接，扩展维度
self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
# 初始化步骤索引和开始索引为 None
self._step_index = None
self._begin_index = None
def index_for_timestep(self, timestep, schedule_timesteps=None):
# 如果没有提供计划时间步，则使用当前的时间步
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
# 找到与给定时间步相等的索引
indices = (schedule_timesteps == timestep).nonzero()
# 对于第一次“步骤”，所采用的 sigma 索引始终是第二个索引
# 如果只有一个索引，则使用最后一个索引
pos = 1 if len(indices) > 1 else 0
# 返回所需索引的值
return indices[pos].item()
def _init_step_index(self, timestep):
# 如果开始索引为 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
def step(
self,
model_output: torch.FloatTensor, # 模型输出的浮点张量
timestep: Union[float, torch.FloatTensor], # 当前的时间步，可以是浮点数或张量
sample: torch.FloatTensor, # 当前样本的浮点张量
s_churn: float = 0.0, # 额外的旋涡参数，默认值为 0.0
s_tmin: float = 0.0, # 最小时间步限制，默认值为 0.0
s_tmax: float = float("inf"), # 最大时间步限制，默认值为无穷大
s_noise: float = 1.0, # 噪声强度，默认值为 1.0
generator: Optional[torch.Generator] = None, # 随机数生成器，默认值为 None
return_dict: bool = True, # 是否返回字典格式的结果，默认值为 True
) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]:
"""
通过反向 SDE 从上一个时间步预测样本。该函数从学习模型输出（通常是预测的噪声）传播扩散过程。
参数：
model_output (`torch.FloatTensor`):
来自学习扩散模型的直接输出。
timestep (`float`):
当前扩散链中的离散时间步。
sample (`torch.FloatTensor`):
当前通过扩散过程创建的样本实例。
s_churn (`float`):
s_tmin (`float`):
s_tmax (`float`):
s_noise (`float`, defaults to 1.0):
添加到样本的噪声缩放因子。
generator (`torch.Generator`, *optional*):
随机数生成器。
return_dict (`bool`):
是否返回 [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] 或
元组。
返回：
[`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] 或 `tuple`:
如果 return_dict 为 `True`，则返回 [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`]，
否则返回一个元组，元组的第一个元素是样本张量。
"""
# 检查 timestep 是否为整数类型
if (
isinstance(timestep, int)
or isinstance(timestep, torch.IntTensor)
or isinstance(timestep, torch.LongTensor)
):
# 如果 timestep 是整数类型，抛出异常，提示不支持此用法
raise ValueError(
(
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
" `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
" one of the `scheduler.timesteps` as a timestep."
),
)
# 如果 step_index 为空，初始化 step_index
if self.step_index is None:
self._init_step_index(timestep)
# 为了避免计算 prev_sample 时的精度问题，将样本上升为 float32 类型
sample = sample.to(torch.float32)
# 获取当前和下一个时间步的 sigma 值
sigma = self.sigmas[self.step_index]
sigma_next = self.sigmas[self.step_index + 1]
# 计算上一个样本，依据当前样本和模型输出
prev_sample = sample + (sigma_next - sigma) * model_output
# 将样本转换回与模型兼容的数据类型
prev_sample = prev_sample.to(model_output.dtype)
# 完成后将步骤索引增加一
self._step_index += 1
# 如果不需要返回字典，则返回包含 prev_sample 的元组
if not return_dict:
return (prev_sample,)
# 返回 FlowMatchEulerDiscreteSchedulerOutput 对象，包含 prev_sample
return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
# 返回训练时间步的数量
def __len__(self):
return self.config.num_train_timesteps
# Copyright 2024 Stability AI, Katherine Crowson 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.
# 你可以在遵循许可证的情况下使用此文件。
# 你可以在以下网址获取许可证的副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面同意，否则根据许可证分发的软件是以“原样”基础提供的，
# 不提供任何形式的担保或条件，无论是明示或暗示的。
# 有关许可证的具体权限和限制，请参阅许可证。
from dataclasses import dataclass # 导入数据类装饰器，用于创建简单的类
from typing import Optional, Tuple, Union # 导入类型提示，Optional表示可选类型，Tuple和Union用于表示元组和联合类型
import numpy as np # 导入NumPy库，通常用于数组操作
import torch # 导入PyTorch库，通常用于深度学习
from ..configuration_utils import ConfigMixin, register_to_config # 从上级模块导入配置混合类和注册配置的装饰器
from ..utils import BaseOutput, logging # 从上级模块导入基础输出类和日志工具
from ..utils.torch_utils import randn_tensor # 从上级模块导入用于生成随机张量的函数
from .scheduling_utils import SchedulerMixin # 从当前模块导入调度混合类
logger = logging.get_logger(__name__) # 获取当前模块的日志记录器，命名为模块名称
@dataclass # 使用数据类装饰器，自动生成初始化和其他方法
class FlowMatchHeunDiscreteSchedulerOutput(BaseOutput): # 定义FlowMatchHeunDiscreteSchedulerOutput类，继承自BaseOutput
"""
Output class for the scheduler's `step` function output.
调度器`step`函数的输出类。
Args:
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
denoising loop.
"""
prev_sample: torch.FloatTensor # 定义一个属性prev_sample，类型为torch.FloatTensor，表示前一步的样本
class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin): # 定义FlowMatchHeunDiscreteScheduler类，继承自SchedulerMixin和ConfigMixin
"""
Heun scheduler.
这是一个Heun调度器。
This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
methods the library implements for all schedulers such as loading and saving.
该模型继承自[`SchedulerMixin`]和[`ConfigMixin`]。请查阅超类文档以获取库为所有调度器实现的通用方法，例如加载和保存。
Args:
num_train_timesteps (`int`, defaults to 1000):
The number of diffusion steps to train the model.
训练模型的扩散步数。
timestep_spacing (`str`, defaults to `"linspace"`):
The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
时间步的缩放方式。有关更多信息，请参阅[常见扩散噪声调度和采样步骤存在缺陷]的表2。
shift (`float`, defaults to 1.0):
The shift value for the timestep schedule.
时间步调度的位移值。
"""
_compatibles = [] # 定义一个类属性，用于存储兼容性信息，初始为空列表
order = 2 # 定义一个类属性，表示调度器的顺序，初始值为2
@register_to_config # 使用装饰器将此方法注册到配置中
def __init__( # 定义初始化方法
self, # 引用自身
num_train_timesteps: int = 1000, # 定义num_train_timesteps参数，默认为1000，表示训练步数
shift: float = 1.0, # 定义shift参数，默认为1.0，表示时间步调度的位移值
):
# 创建一个线性空间，范围从 1 到 num_train_timesteps，生成 num_train_timesteps 个点，并反转顺序
timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
# 将 numpy 数组转换为 PyTorch 的张量，并指定数据类型为 float32
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
# 计算 sigmas，sigmas 为 timesteps 相对于 num_train_timesteps 的比例
sigmas = timesteps / num_train_timesteps
# 根据 shift 调整 sigmas 的值，进行非线性缩放
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
# 将调整后的 sigmas 乘以 num_train_timesteps，保存为 self.timesteps
self.timesteps = sigmas * num_train_timesteps
# 初始化步索引为 None，表示当前未定义
self._step_index = None
# 初始化起始索引为 None，表示当前未定义
self._begin_index = None
# 将 sigmas 移动到 CPU，避免过多的 CPU/GPU 通信
self.sigmas = sigmas.to("cpu")
# 记录 sigmas 的最小值，将其转为标量
self.sigma_min = self.sigmas[-1].item()
# 记录 sigmas 的最大值，将其转为标量
self.sigma_max = self.sigmas[0].item()
@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
def scale_noise(
self,
sample: torch.FloatTensor,
timestep: Union[float, torch.FloatTensor],
noise: Optional[torch.FloatTensor] = None,
) -> torch.FloatTensor:
"""
流匹配中的前向过程
参数:
sample (`torch.FloatTensor`):
输入样本。
timestep (`int`, *可选*):
扩散链中的当前时间步。
返回:
`torch.FloatTensor`:
缩放后的输入样本。
"""
# 如果当前步索引为 None，则初始化步索引
if self.step_index is None:
self._init_step_index(timestep)
# 获取当前步索引对应的 sigma 值
sigma = self.sigmas[self.step_index]
# 使用 sigma 对噪声和样本进行加权组合，得到新的样本
sample = sigma * noise + (1.0 - sigma) * sample
# 返回处理后的样本
return sample
def _sigma_to_t(self, sigma):
# 将 sigma 转换为与训练时间步数相关的值
return sigma * self.config.num_train_timesteps
# 设置离散时间步长，用于扩散链（在推断之前运行）
def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
"""
设置用于扩散链的离散时间步长（在推断之前运行）。
参数：
num_inference_steps (`int`):
生成样本时使用的扩散步骤数量。
device (`str` 或 `torch.device`, *可选*):
时间步长应该移动到的设备。如果为 `None`，则时间步长不移动。
"""
# 将输入的推断步骤数量保存到实例变量中
self.num_inference_steps = num_inference_steps
# 生成从最大 sigma 到最小 sigma 的均匀时间步长
timesteps = np.linspace(
self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
)
# 将时间步长归一化到训练时间步数
sigmas = timesteps / self.config.num_train_timesteps
# 根据配置的偏移量调整 sigma 值
sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
# 将 numpy 数组转换为张量，并移动到指定设备
sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
# 计算最终的时间步长
timesteps = sigmas * self.config.num_train_timesteps
# 扩展时间步长，使其重复交替
timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2)])
# 将时间步长保存到实例变量中，并移动到指定设备
self.timesteps = timesteps.to(device=device)
# 为 sigma 添加零值的张量
sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
# 扩展 sigma，使其重复交替
self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])
# 清空导数和时间增量
self.prev_derivative = None
self.dt = None
# 初始化步长和起始索引
self._step_index = None
self._begin_index = None
# 根据时间步长获取对应的索引
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()
# 初始化步长索引
def _init_step_index(self, timestep):
# 如果起始索引为 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 state_in_first_order(self):
# 如果 dt 为 None，返回 True
return self.dt is None
# 执行一步更新
def step(
self,
model_output: torch.FloatTensor,
timestep: Union[float, torch.FloatTensor],
sample: torch.FloatTensor,
s_churn: float = 0.0,
s_tmin: float = 0.0,
s_tmax: float = float("inf"),
s_noise: float = 1.0,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
):
# 返回训练时间步数的数量
def __len__(self):
return self.config.num_train_timesteps
# 版权所有 2024 Katherine Crowson，HuggingFace团队和hlky。保留所有权利。
#
# 根据Apache许可证第2.0版（“许可证”）许可；
# 除非遵守许可证，否则您不得使用此文件。
# 您可以在以下网址获得许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有规定，按照许可证分发的软件是按“原样”基础分发，
# 不提供任何形式的明示或暗示的担保或条件。
# 有关许可证具体条款，请参阅许可证。
import math # 导入数学库以进行数学计算
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 .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput # 从调度工具中导入调度相关类
# 从diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar复制的函数
def betas_for_alpha_bar(
num_diffusion_timesteps, # 定义扩散时间步数
max_beta=0.999, # 定义最大beta值，默认值为0.999
alpha_transform_type="cosine", # 定义alpha变换类型，默认值为“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`): 调度器用于调整模型输出的betas
"""
if alpha_transform_type == "cosine": # 检查alpha变换类型是否为“cosine”
def alpha_bar_fn(t): # 定义alpha_bar函数，接受参数t
return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 # 计算cosine变换
elif alpha_transform_type == "exp": # 检查alpha变换类型是否为“exp”
def alpha_bar_fn(t): # 定义alpha_bar函数，接受参数t
return math.exp(t * -12.0) # 计算指数变换
else: # 如果变换类型不支持
raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") # 抛出错误
betas = [] # 初始化空列表以存储beta值
for i in range(num_diffusion_timesteps): # 遍历每个扩散时间步
t1 = i / num_diffusion_timesteps # 计算当前时间步t1
t2 = (i + 1) / num_diffusion_timesteps # 计算下一个时间步t2
betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) # 计算beta值并添加到列表中
return torch.tensor(betas, dtype=torch.float32) # 返回转换为PyTorch张量的beta列表
class HeunDiscreteScheduler(SchedulerMixin, ConfigMixin): # 定义Heun离散调度器类，继承自SchedulerMixin和ConfigMixin
"""
具有Heun步骤的离散beta调度器。
该模型继承自[`SchedulerMixin`]和[`ConfigMixin`]。查看超类文档以获取库为所有调度器实现的通用
方法，例如加载和保存。
# 参数定义部分，描述每个参数的作用和默认值
Args:
num_train_timesteps (`int`, defaults to 1000): # 模型训练的扩散步数
The number of diffusion steps to train the model.
beta_start (`float`, defaults to 0.0001): # 推理的起始 beta 值
The starting `beta` value of inference.
beta_end (`float`, defaults to 0.02): # 推理的最终 beta 值
The final `beta` value.
beta_schedule (`str`, defaults to `"linear"`): # beta 调度策略，映射 beta 范围到一系列 beta
The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
`linear` or `scaled_linear`.
trained_betas (`np.ndarray`, *optional*): # 直接传递 beta 数组以绕过 beta_start 和 beta_end
Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
prediction_type (`str`, defaults to `epsilon`, *optional*): # 调度函数的预测类型
Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
`sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
Video](https://imagen.research.google/video/paper.pdf) paper).
clip_sample (`bool`, defaults to `True`): # 为了数值稳定性裁剪预测样本
Clip the predicted sample for numerical stability.
clip_sample_range (`float`, defaults to 1.0): # 样本裁剪的最大幅度，仅在 clip_sample=True 时有效
The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
use_karras_sigmas (`bool`, *optional*, defaults to `False`): # 是否在采样过程中使用 Karras sigma
Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
the sigmas are determined according to a sequence of noise levels {σi}.
timestep_spacing (`str`, defaults to `"linspace"`): # 时间步的缩放方式
The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
steps_offset (`int`, defaults to 0): # 添加到推理步数的偏移量
An offset added to the inference steps, as required by some model families.
"""
# 创建与 KarrasDiffusionSchedulers 兼容的名称列表
_compatibles = [e.name for e in KarrasDiffusionSchedulers]
# 设置默认顺序为2
order = 2
# 注册到配置中的构造函数
@register_to_config
def __init__(
self,
num_train_timesteps: int = 1000, # 训练的扩散步数，默认1000
beta_start: float = 0.00085, # 合理的默认起始 beta 值
beta_end: float = 0.012, # 默认的最终 beta 值
beta_schedule: str = "linear", # 默认使用线性调度
trained_betas: Optional[Union[np.ndarray, List[float]]] = None, # 可选的 beta 数组
prediction_type: str = "epsilon", # 默认预测类型为 epsilon
use_karras_sigmas: Optional[bool] = False, # 默认不使用 Karras sigmas
clip_sample: Optional[bool] = False, # 默认不裁剪样本
clip_sample_range: float = 1.0, # 默认裁剪范围为1.0
timestep_spacing: str = "linspace", # 默认时间步缩放为线性
steps_offset: int = 0, # 默认步数偏移为0
):
# 检查是否提供了训练后的贝塔值
if trained_betas is not None:
# 将训练后的贝塔值转换为32位浮点张量
self.betas = torch.tensor(trained_betas, dtype=torch.float32)
# 检查贝塔调度方式是否为线性
elif beta_schedule == "linear":
# 生成从beta_start到beta_end的线性贝塔值，数量为num_train_timesteps
self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
# 检查贝塔调度方式是否为缩放线性
elif beta_schedule == "scaled_linear":
# 该调度方式特定于潜在扩散模型
# 生成从beta_start的平方根到beta_end的平方根的线性贝塔值，并平方
self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
# 检查贝塔调度方式是否为平方余弦（版本2）
elif beta_schedule == "squaredcos_cap_v2":
# Glide余弦调度
# 使用余弦调度生成贝塔值
self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="cosine")
# 检查贝塔调度方式是否为指数
elif beta_schedule == "exp":
# 使用指数调度生成贝塔值
self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="exp")
# 如果贝塔调度方式未实现，则抛出异常
else:
raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
# 计算每个时刻的alpha值
self.alphas = 1.0 - self.betas
# 计算alpha的累积乘积
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
# 设置所有时间步的值
self.set_timesteps(num_train_timesteps, None, num_train_timesteps)
# 指示是否使用Karras sigma值
self.use_karras_sigmas = use_karras_sigmas
# 初始化步索引和开始索引为None
self._step_index = None
self._begin_index = None
# 将sigma移动到CPU以避免过多的CPU/GPU通信
self.sigmas = self.sigmas.to("cpu")
# 从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()
@property
def init_noise_sigma(self):
# 返回初始噪声分布的标准差
# 如果时间步间距为线性或尾随，则返回sigma的最大值
if self.config.timestep_spacing in ["linspace", "trailing"]:
return self.sigmas.max()
# 否则返回sigma最大值的平方加1的平方根
return (self.sigmas.max() ** 2 + 1) ** 0.5
@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):
"""
设置调度器的开始索引。此函数应在推理前从管道运行。
参数:
begin_index (`int`):
调度器的开始索引。
"""
# 将开始索引赋值给实例变量
self._begin_index = begin_index
# 对模型输入进行缩放，确保与调度器的互换性
def scale_model_input(
self,
sample: torch.Tensor,
timestep: Union[float, torch.Tensor],
) -> torch.Tensor:
"""
确保与需要根据当前时间步缩放去噪模型输入的调度器互换性。
参数:
sample (`torch.Tensor`):
输入样本。
timestep (`int`, *可选*):
扩散链中的当前时间步。
返回:
`torch.Tensor`:
缩放后的输入样本。
"""
# 如果步骤索引为空，初始化步骤索引
if self.step_index is None:
self._init_step_index(timestep)
# 根据当前步骤索引获取sigma值
sigma = self.sigmas[self.step_index]
# 对输入样本进行缩放
sample = sample / ((sigma**2 + 1) ** 0.5)
# 返回缩放后的样本
return sample
# 设置时间步数，包含多个可选参数
def set_timesteps(
self,
num_inference_steps: Optional[int] = None,
device: Union[str, torch.device] = None,
num_train_timesteps: Optional[int] = None,
timesteps: Optional[List[int]] = None,
# 从 diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t 复制的
def _sigma_to_t(self, sigma, log_sigmas):
# 获取log 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
# 获取低和高的log sigma值
low = log_sigmas[low_idx]
high = log_sigmas[high_idx]
# 进行sigma的插值
w = (low - log_sigma) / (low - high)
# 限制w在0到1之间
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_euler_discrete.EulerDiscreteScheduler._convert_to_karras 复制的
# 将输入的 sigma 值转换为 Karras 的噪声调度
def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
"""构建 Karras et al. (2022) 的噪声调度。"""
# 确保其他调度器复制此函数时不会出错的临时修复
# TODO: 将此逻辑添加到其他调度器
if hasattr(self.config, "sigma_min"):
# 从配置中获取 sigma_min 值
sigma_min = self.config.sigma_min
else:
# 如果没有设置，则 sigma_min 为空
sigma_min = None
if hasattr(self.config, "sigma_max"):
# 从配置中获取 sigma_max 值
sigma_max = self.config.sigma_max
else:
# 如果没有设置，则 sigma_max 为空
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 # 论文中使用的常量值
# 创建从 0 到 1 的线性数组，长度为 num_inference_steps
ramp = np.linspace(0, 1, num_inference_steps)
# 计算 sigma_min 的倒数
min_inv_rho = sigma_min ** (1 / rho)
# 计算 sigma_max 的倒数
max_inv_rho = sigma_max ** (1 / rho)
# 根据倒数范围生成新的 sigma 值
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
# 返回生成的 sigma 值
return sigmas
# 判断是否在一阶状态下
@property
def state_in_first_order(self):
# 如果 dt 为 None，返回 True
return self.dt is None
# 从 diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index 复制而来
def _init_step_index(self, timestep):
# 如果 begin_index 为 None
if self.begin_index is None:
# 如果 timestep 是张量，则转换为与 timesteps 相同的设备
if isinstance(timestep, torch.Tensor):
timestep = timestep.to(self.timesteps.device)
# 根据当前时间步获取步骤索引
self._step_index = self.index_for_timestep(timestep)
else:
# 否则使用 _begin_index
self._step_index = self._begin_index
def step(
self,
model_output: Union[torch.Tensor, np.ndarray],
timestep: Union[float, torch.Tensor],
sample: Union[torch.Tensor, np.ndarray],
return_dict: bool = True,
# 从 diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise 复制而来
def add_noise(
self,
original_samples: torch.Tensor,
noise: torch.Tensor,
timesteps: torch.Tensor,
# 返回一个张量，可能用于后续计算
) -> torch.Tensor:
# 确保 sigmas 和 timesteps 与 original_samples 的设备和数据类型一致
sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
# 检查设备类型是否为 mps，且 timesteps 为浮点类型
if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
# mps 不支持 float64 类型
schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32)
timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
else:
# 将 timesteps 转换到 original_samples 的设备
schedule_timesteps = self.timesteps.to(original_samples.device)
timesteps = timesteps.to(original_samples.device)
# 当 scheduler 用于训练时，self.begin_index 为 None，或者管道未实现 set_begin_index
if self.begin_index is None:
# 根据 timesteps 获取对应的步进索引
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 的维度少于 original_samples，则在最后一维增加维度
while len(sigma.shape) < len(original_samples.shape):
sigma = sigma.unsqueeze(-1)
# 添加噪声到原始样本中，生成有噪声的样本
noisy_samples = original_samples + noise * sigma
# 返回有噪声的样本
return noisy_samples
# 返回训练时间步的数量
def __len__(self):
return self.config.num_train_timesteps
# 版权信息，声明版权所有者及相关许可信息
# Copyright 2024 Zhejiang University Team and The HuggingFace Team. All rights reserved.
#
# 根据 Apache License, Version 2.0 许可本文件；
# 仅在遵守许可证的情况下使用此文件。
# 可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面同意，软件
# 在许可证下按“原样”分发，不提供任何形式的保证或条件，
# 明示或暗示。
# 参见许可证以获取特定语言关于权限和
# 限制的信息。
# 导入数学库
import math
# 从类型提示库导入所需类型
from typing import List, Optional, Tuple, Union
# 导入 NumPy 库
import numpy as np
# 导入 PyTorch 库
import torch
# 从配置工具中导入 ConfigMixin 和注册到配置的装饰器
from ..configuration_utils import ConfigMixin, register_to_config
# 从调度工具中导入 SchedulerMixin 和 SchedulerOutput
from .scheduling_utils import SchedulerMixin, SchedulerOutput
class IPNDMScheduler(SchedulerMixin, ConfigMixin):
"""
一个四阶改进伪线性多步调度器。
该模型继承自 [`SchedulerMixin`] 和 [`ConfigMixin`]。查看父类文档以了解库为所有调度器实现的通用
方法，例如加载和保存。
参数：
num_train_timesteps (`int`, 默认为 1000):
用于训练模型的扩散步骤数量。
trained_betas (`np.ndarray`, *可选*):
直接传递一组 beta 数组到构造函数，以绕过 `beta_start` 和 `beta_end`。
"""
# 调度器的阶数设置为 1
order = 1
# 注册到配置的方法
@register_to_config
def __init__(
# 设置训练的扩散时间步数，默认为 1000
self, num_train_timesteps: int = 1000, trained_betas: Optional[Union[np.ndarray, List[float]]] = None
):
# 设置 `betas`，`alphas` 和 `timesteps`
self.set_timesteps(num_train_timesteps)
# 初始化噪声分布的标准差
self.init_noise_sigma = 1.0
# 当前仅支持 F-PNDM，即 Runge-Kutta 方法
# 有关算法的更多信息，请查看论文：https://arxiv.org/pdf/2202.09778.pdf
# 主要参考公式 (9)，(12)，(13) 和算法 2。
self.pndm_order = 4
# 运行时值的列表
self.ets = []
# 当前步骤索引初始化为 None
self._step_index = None
# 开始索引初始化为 None
self._begin_index = None
# 属性，获取当前时间步的索引
@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
# 设置用于扩散链的离散时间步（在推理前运行）
def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
"""
设置用于扩散链的离散时间步（在推理前运行）。
Args:
num_inference_steps (`int`):
生成样本时使用的扩散步骤数。
device (`str` or `torch.device`, *optional*):
时间步应移动到的设备。如果为 `None`，则不移动时间步。
"""
# 保存扩散步骤数到实例变量中
self.num_inference_steps = num_inference_steps
# 创建一个从 1 到 0 的线性空间，步数为 num_inference_steps + 1，去掉最后一个元素
steps = torch.linspace(1, 0, num_inference_steps + 1)[:-1]
# 在时间步的末尾添加一个 0.0
steps = torch.cat([steps, torch.tensor([0.0])])
# 如果训练的 beta 值存在，将其转换为张量
if self.config.trained_betas is not None:
self.betas = torch.tensor(self.config.trained_betas, dtype=torch.float32)
else:
# 否则根据步骤计算 beta 值
self.betas = torch.sin(steps * math.pi / 2) ** 2
# 计算 alpha 值，作为 beta 值的平方根
self.alphas = (1.0 - self.betas**2) ** 0.5
# 计算时间步的角度值并去掉最后一个元素
timesteps = (torch.atan2(self.betas, self.alphas) / math.pi * 2)[:-1]
# 将时间步移到指定设备上
self.timesteps = timesteps.to(device)
# 初始化空列表以存储 ets
self.ets = []
# 初始化步索引和开始索引为 None
self._step_index = None
self._begin_index = None
# 从 EulerDiscreteScheduler 复制的函数，获取给定时间步的索引
def index_for_timestep(self, timestep, schedule_timesteps=None):
# 如果没有提供时间步，则使用实例中的时间步
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
# 找到与给定时间步相等的索引
indices = (schedule_timesteps == timestep).nonzero()
# 确保第一个步骤的 sigma 索引总是第二个索引（或唯一一个时的最后一个索引）
pos = 1 if len(indices) > 1 else 0
# 返回找到的索引值
return indices[pos].item()
# 从 EulerDiscreteScheduler 复制的函数，初始化步索引
def _init_step_index(self, timestep):
# 如果开始索引为 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
# 步骤函数，接收模型输出和样本进行处理
def step(
self,
model_output: torch.Tensor,
timestep: Union[int, torch.Tensor],
sample: torch.Tensor,
return_dict: bool = True,
# 定义函数，返回调度器输出或元组
) -> Union[SchedulerOutput, Tuple]:
# 预测从上一个时间步生成的样本，通过反向 SDE 进行传播
"""
# 文档字符串，描述函数的目的和参数
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the linear multistep method. It performs one forward pass multiple times to approximate the solution.
Args:
model_output (`torch.Tensor`):
从学习的扩散模型直接输出的张量。
timestep (`int`):
当前扩散链中的离散时间步。
sample (`torch.Tensor`):
通过扩散过程创建的当前样本实例。
return_dict (`bool`):
是否返回一个 [`~schedulers.scheduling_utils.SchedulerOutput`] 或元组。
Returns:
[`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
如果 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)
# 当前时间步索引
timestep_index = self.step_index
# 前一个时间步索引
prev_timestep_index = self.step_index + 1
# 根据当前样本和模型输出计算 ets
ets = sample * self.betas[timestep_index] + model_output * self.alphas[timestep_index]
# 将计算得到的 ets 添加到列表中
self.ets.append(ets)
# 根据 ets 的长度进行不同的线性多步法计算
if len(self.ets) == 1:
ets = self.ets[-1]
elif len(self.ets) == 2:
ets = (3 * self.ets[-1] - self.ets[-2]) / 2
elif len(self.ets) == 3:
ets = (23 * self.ets[-1] - 16 * self.ets[-2] + 5 * self.ets[-3]) / 12
else:
ets = (1 / 24) * (55 * self.ets[-1] - 59 * self.ets[-2] + 37 * self.ets[-3] - 9 * self.ets[-4])
# 获取前一个样本
prev_sample = self._get_prev_sample(sample, timestep_index, prev_timestep_index, ets)
# 完成后步索引加一
self._step_index += 1
# 根据 return_dict 决定返回格式
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:
# 文档字符串，描述函数的目的和参数
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
Args:
sample (`torch.Tensor`):
输入样本。
Returns:
`torch.Tensor`:
缩放后的输入样本。
"""
# 返回未缩放的输入样本
return sample
# 定义一个私有方法，用于获取前一个样本
def _get_prev_sample(self, sample, timestep_index, prev_timestep_index, ets):
# 获取当前时间步的 alpha 值
alpha = self.alphas[timestep_index]
# 获取当前时间步的 sigma 值
sigma = self.betas[timestep_index]
# 获取前一个时间步的 alpha 值
next_alpha = self.alphas[prev_timestep_index]
# 获取前一个时间步的 sigma 值
next_sigma = self.betas[prev_timestep_index]
# 根据当前样本、sigma 和 ets 计算预测值
pred = (sample - sigma * ets) / max(alpha, 1e-8)
# 通过预测值、前一个时间步的 alpha 和 sigma 计算前一个样本
prev_sample = next_alpha * pred + ets * next_sigma
# 返回计算得到的前一个样本
return prev_sample
# 定义一个方法，用于获取训练时间步的数量
def __len__(self):
# 返回配置中训练时间步的数量
return self.config.num_train_timesteps
# 版权所有 2024 NVIDIA 和 HuggingFace 团队。保留所有权利。
#
# 根据 Apache 许可证第 2.0 版（“许可证”）进行许可；
# 除非遵循该许可证，否则不得使用此文件。
# 你可以在以下网址获得许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面同意，软件在“按原样”基础上分发，
# 不提供任何形式的明示或暗示的担保或条件。
# 有关许可证下特定权限和限制，请参见许可证。
# 从 dataclasses 模块导入 dataclass 装饰器，用于创建数据类
from dataclasses import dataclass
# 从 typing 模块导入 Optional、Tuple 和 Union 类型提示
from typing import Optional, Tuple, Union
# 导入 flax 库
import flax
# 导入 jax 库
import jax
# 导入 jax.numpy 模块，并命名为 jnp
import jax.numpy as jnp
# 从 jax 库导入 random 模块
from jax import random
# 从配置工具模块导入 ConfigMixin 和 register_to_config
from ..configuration_utils import ConfigMixin, register_to_config
# 从工具模块导入 BaseOutput 类
from ..utils import BaseOutput
# 从调度工具模块导入 FlaxSchedulerMixin 类
from .scheduling_utils_flax import FlaxSchedulerMixin
# 使用 flax 库的 struct.dataclass 装饰器定义 KarrasVeSchedulerState 数据类
@flax.struct.dataclass
class KarrasVeSchedulerState:
# 可设置的值
num_inference_steps: Optional[int] = None # 推理步骤数量，默认为 None
timesteps: Optional[jnp.ndarray] = None # 时间步数组，默认为 None
schedule: Optional[jnp.ndarray] = None # 调度数组，表示 sigma(t_i)，默认为 None
# 类方法，用于创建 KarrasVeSchedulerState 的实例
@classmethod
def create(cls):
return cls() # 返回类的实例
# 使用 dataclass 装饰器定义 FlaxKarrasVeOutput 数据类
@dataclass
class FlaxKarrasVeOutput(BaseOutput):
"""
调度器的步骤函数输出的输出类。
参数:
prev_sample (`jnp.ndarray`，形状为 `(batch_size, num_channels, height, width)`，用于图像):
计算的样本 (x_{t-1})，表示前一个时间步的样本。`prev_sample` 应作为下一个模型输入
在去噪循环中使用。
derivative (`jnp.ndarray`，形状为 `(batch_size, num_channels, height, width)`，用于图像):
预测原始图像样本 (x_0) 的导数。
state (`KarrasVeSchedulerState`): `FlaxKarrasVeScheduler` 的状态数据类。
"""
prev_sample: jnp.ndarray # 前一个时间步的样本
derivative: jnp.ndarray # 预测原始样本的导数
state: KarrasVeSchedulerState # 调度器状态
# 定义 FlaxKarrasVeScheduler 类，继承自 FlaxSchedulerMixin 和 ConfigMixin
class FlaxKarrasVeScheduler(FlaxSchedulerMixin, ConfigMixin):
"""
从 Karras 等人 [1] 定制的变异扩展 (VE) 模型进行随机采样。请参考 [1] 的算法 2 和表 1 中的 VE 列。
[1] Karras, Tero 等人。 "阐明基于扩散的生成模型的设计空间。"
https://arxiv.org/abs/2206.00364 [2] Song, Yang 等人。 "通过随机微分方程进行基于得分的生成建模。"
https://arxiv.org/abs/2011.13456
[`~ConfigMixin`] 处理存储在调度器 `__init__` 函数中传入的所有配置属性，
如 `num_train_timesteps`。它们可以通过 `scheduler.config.num_train_timesteps` 访问。
[`SchedulerMixin`] 通过 [`SchedulerMixin.save_pretrained`] 和
[`~SchedulerMixin.from_pretrained`] 函数提供一般的加载和保存功能。
有关参数的更多详细信息，请参见原始论文的附录 E：“阐明设计空间的扩散模型
# 引用论文的标题和链接，说明相关模型的研究
Diffusion-Based Generative Models." https://arxiv.org/abs/2206.00364. # 说明用于查找特定模型的最佳参数的表格位置
The grid search values used to find the
# 定义参数说明
optimal {s_noise, s_churn, s_min, s_max} for a specific model are described in Table 5 of the paper.
# 参数 sigma_min：最小噪声强度
Args:
sigma_min (`float`): minimum noise magnitude
# 参数 sigma_max：最大噪声强度
sigma_max (`float`): maximum noise magnitude
# 参数 s_noise：在采样过程中抵消细节丢失的额外噪声量，合理范围为 [1.000, 1.011]
s_noise (`float`): the amount of additional noise to counteract loss of detail during sampling.
A reasonable range is [1.000, 1.011].
# 参数 s_churn：控制整体随机性的参数，合理范围为 [0, 100]
s_churn (`float`): the parameter controlling the overall amount of stochasticity.
A reasonable range is [0, 100].
# 参数 s_min：开始添加噪声的 sigma 范围起始值，合理范围为 [0, 10]
s_min (`float`): the start value of the sigma range where we add noise (enable stochasticity).
A reasonable range is [0, 10].
# 参数 s_max：添加噪声的 sigma 范围结束值，合理范围为 [0.2, 80]
s_max (`float`): the end value of the sigma range where we add noise.
A reasonable range is [0.2, 80].
# 文档字符串结束
"""
# 属性 has_state：返回布尔值 True，表示存在状态
@property
def has_state(self):
return True
# 注册初始化方法，接收多个噪声参数，并设置默认值
@register_to_config
def __init__(
self,
sigma_min: float = 0.02, # 初始化时设置最小噪声强度的默认值
sigma_max: float = 100, # 初始化时设置最大噪声强度的默认值
s_noise: float = 1.007, # 初始化时设置额外噪声的默认值
s_churn: float = 80, # 初始化时设置整体随机性的默认值
s_min: float = 0.05, # 初始化时设置 sigma 范围起始值的默认值
s_max: float = 50, # 初始化时设置 sigma 范围结束值的默认值
):
# 方法体为空，表示初始化时无额外操作
pass
# 创建并返回 KarrasVeSchedulerState 状态实例
def create_state(self):
return KarrasVeSchedulerState.create()
# 设置扩散链使用的连续时间步，支持推理前运行
def set_timesteps(
self, state: KarrasVeSchedulerState, num_inference_steps: int, shape: Tuple = ()
) -> KarrasVeSchedulerState:
"""
# 定义时间步设置的文档字符串，包含参数说明
Sets the continuous timesteps used for the diffusion chain. Supporting function to be run before inference.
Args:
state (`KarrasVeSchedulerState`):
the `FlaxKarrasVeScheduler` state data class.
num_inference_steps (`int`):
the number of diffusion steps used when generating samples with a pre-trained model.
"""
# 生成从 num_inference_steps 到 0 的时间步数组
timesteps = jnp.arange(0, num_inference_steps)[::-1].copy()
# 根据公式计算时间步的调度
schedule = [
(
self.config.sigma_max**2
* (self.config.sigma_min**2 / self.config.sigma_max**2) ** (i / (num_inference_steps - 1))
)
for i in timesteps
]
# 替换状态，更新时间步数、调度和时间步数组
return state.replace(
num_inference_steps=num_inference_steps,
schedule=jnp.array(schedule, dtype=jnp.float32),
timesteps=timesteps,
)
# 添加噪声到输入样本的方法
def add_noise_to_input(
self,
state: KarrasVeSchedulerState, # 当前状态
sample: jnp.ndarray, # 输入样本数据
sigma: float, # 当前噪声强度
key: jax.Array, # 随机数生成的键值
) -> Tuple[jnp.ndarray, float]: # 定义函数返回一个元组，其中第一个元素是 jnp.ndarray 类型，第二个元素是 float 类型
"""
Explicit Langevin-like "churn" step of adding noise to the sample according to a factor gamma_i ≥ 0 to reach a
higher noise level sigma_hat = sigma_i + gamma_i*sigma_i. # 文档字符串，描述了函数的作用和参数
TODO Args: # TODO: 这里是待完成的参数说明部分
"""
if self.config.s_min <= sigma <= self.config.s_max: # 检查 sigma 是否在配置的最小和最大值之间
gamma = min(self.config.s_churn / state.num_inference_steps, 2**0.5 - 1) # 计算 gamma 值，限制其最大值
else: # 如果 sigma 不在范围内
gamma = 0 # 将 gamma 设置为 0
# sample eps ~ N(0, S_noise^2 * I) # 从正态分布中采样噪声 eps
key = random.split(key, num=1) # 将随机数生成器的 key 分裂，得到新的 key
eps = self.config.s_noise * random.normal(key=key, shape=sample.shape) # 根据配置的噪声标准差生成噪声
sigma_hat = sigma + gamma * sigma # 计算新的噪声水平 sigma_hat
sample_hat = sample + ((sigma_hat**2 - sigma**2) ** 0.5 * eps) # 更新样本，加上噪声
return sample_hat, sigma_hat # 返回更新后的样本和新的噪声水平
def step( # 定义 step 方法
self,
state: KarrasVeSchedulerState, # 当前状态对象，包含调度器的状态
model_output: jnp.ndarray, # 从学习到的扩散模型输出的张量
sigma_hat: float, # 当前的噪声水平
sigma_prev: float, # 上一个噪声水平
sample_hat: jnp.ndarray, # 更新后的样本
return_dict: bool = True, # 返回结果的类型，默认为 True
) -> Union[FlaxKarrasVeOutput, Tuple]: # 返回类型可以是 FlaxKarrasVeOutput 或元组
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise). # 文档字符串，描述了函数的功能
Args: # 参数说明部分
state (`KarrasVeSchedulerState`): the `FlaxKarrasVeScheduler` state data class. # 状态对象
model_output (`torch.Tensor` or `np.ndarray`): direct output from learned diffusion model. # 模型输出
sigma_hat (`float`): TODO # 待完成的参数说明
sigma_prev (`float`): TODO # 待完成的参数说明
sample_hat (`torch.Tensor` or `np.ndarray`): TODO # 待完成的参数说明
return_dict (`bool`): option for returning tuple rather than FlaxKarrasVeOutput class # 返回类型选项
Returns: # 返回值说明
[`~schedulers.scheduling_karras_ve_flax.FlaxKarrasVeOutput`] or `tuple`: Updated sample in the diffusion
chain and derivative. [`~schedulers.scheduling_karras_ve_flax.FlaxKarrasVeOutput`] if `return_dict` is
True, otherwise a `tuple`. When returning a tuple, the first element is the sample tensor. # 返回值的描述
"""
pred_original_sample = sample_hat + sigma_hat * model_output # 预测原始样本，结合当前样本和模型输出
derivative = (sample_hat - pred_original_sample) / sigma_hat # 计算样本与预测样本之间的导数
sample_prev = sample_hat + (sigma_prev - sigma_hat) * derivative # 计算上一个样本
if not return_dict: # 如果不返回字典
return (sample_prev, derivative, state) # 返回样本、导数和状态
return FlaxKarrasVeOutput(prev_sample=sample_prev, derivative=derivative, state=state) # 返回结果对象
def step_correct( # 定义 step_correct 方法
self,
state: KarrasVeSchedulerState, # 当前状态对象
model_output: jnp.ndarray, # 模型输出的张量
sigma_hat: float, # 当前噪声水平
sigma_prev: float, # 上一个噪声水平
sample_hat: jnp.ndarray, # 更新后的样本
sample_prev: jnp.ndarray, # 上一个样本
derivative: jnp.ndarray, # 样本的导数
return_dict: bool = True, # 返回类型选项
) -> Union[FlaxKarrasVeOutput, Tuple]:
"""
修正预测的样本，基于网络的输出 model_output。TODO 完成描述
参数:
state (`KarrasVeSchedulerState`): `FlaxKarrasVeScheduler` 的状态数据类。
model_output (`torch.Tensor` 或 `np.ndarray`): 从学习的扩散模型直接输出。
sigma_hat (`float`): TODO
sigma_prev (`float`): TODO
sample_hat (`torch.Tensor` 或 `np.ndarray`): TODO
sample_prev (`torch.Tensor` 或 `np.ndarray`): TODO
derivative (`torch.Tensor` 或 `np.ndarray`): TODO
return_dict (`bool`): 返回元组而不是 FlaxKarrasVeOutput 类的选项
返回:
prev_sample (TODO): 扩散链中更新的样本。derivative (TODO): TODO
"""
# 计算原始预测样本，基于前一个样本和模型输出的缩放
pred_original_sample = sample_prev + sigma_prev * model_output
# 计算修正后的导数，使用当前样本与预测样本的差值
derivative_corr = (sample_prev - pred_original_sample) / sigma_prev
# 更新前一个样本，考虑新的样本和导数的影响
sample_prev = sample_hat + (sigma_prev - sigma_hat) * (0.5 * derivative + 0.5 * derivative_corr)
# 检查是否返回字典形式的结果
if not return_dict:
# 返回样本、导数和状态的元组
return (sample_prev, derivative, state)
# 返回 FlaxKarrasVeOutput 对象，包含更新的样本、导数和状态
return FlaxKarrasVeOutput(prev_sample=sample_prev, derivative=derivative, state=state)
# 定义一个添加噪声的方法，尚未实现
def add_noise(self, state: KarrasVeSchedulerState, original_samples, noise, timesteps):
# 抛出未实现错误
raise NotImplementedError()
# 版权声明，声明文件的版权归Katherine Crowson、HuggingFace团队及hlky所有
#
# 根据Apache许可证第2.0版（“许可证”）进行许可；
# 除非遵循许可证，否则您不能使用此文件。
# 您可以在以下网址获取许可证副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面同意，否则根据许可证分发的软件是按“原样”基础提供的，
# 不提供任何形式的明示或暗示的担保或条件。
# 请参见许可证以获取有关权限和限制的具体规定。
# 导入数学库
import math
# 从类型提示模块导入List、Optional、Tuple和Union
from typing import List, Optional, Tuple, Union
# 导入numpy库并简化为np
import numpy as np
# 导入torch库
import torch
# 从配置工具模块导入ConfigMixin和register_to_config
from ..configuration_utils import ConfigMixin, register_to_config
# 从torch工具模块导入randn_tensor函数
from ..utils.torch_utils import randn_tensor
# 从调度工具模块导入KarrasDiffusionSchedulers、SchedulerMixin和SchedulerOutput
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值，默认为0.999
alpha_transform_type="cosine", # alpha变换类型，默认为“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_fn函数，接受参数t
def alpha_bar_fn(t):
# 返回经过余弦函数转换的值
return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
# 如果alpha变换类型为“exp”
elif alpha_transform_type == "exp":
# 定义alpha_bar_fn函数，接受参数t
def alpha_bar_fn(t):
# 返回指数衰减的值
return math.exp(t * -12.0)
# 如果alpha变换类型不受支持，抛出异常
else:
raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
betas = [] # 初始化一个空列表，用于存储beta值
# 遍历每个扩散时间步
for i in range(num_diffusion_timesteps):
t1 = i / num_diffusion_timesteps # 计算当前时间步的比例
t2 = (i + 1) / num_diffusion_timesteps # 计算下一个时间步的比例
# 计算beta值，并限制在max_beta范围内
betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
# 返回转换为torch张量的beta值，数据类型为float32
return torch.tensor(betas, dtype=torch.float32)
# 定义KDPM2AncestralDiscreteScheduler类，继承自SchedulerMixin和ConfigMixin
class KDPM2AncestralDiscreteScheduler(SchedulerMixin, ConfigMixin):
"""
KDPM2DiscreteScheduler与祖先采样，灵感来自于DPMSolver2和[Elucidating
the Design Space of Diffusion-Based Generative Models](https://huggingface.co/papers/2206.00364)论文中的算法2。
此模型继承自[`SchedulerMixin`]和[`ConfigMixin`]。请查阅超类文档以获取通用
# 方法库为所有调度器实现的功能，如加载和保存。
# 参数说明：
# num_train_timesteps（`int`，默认为1000）：
# 训练模型的扩散步骤数。
# beta_start（`float`，默认为0.00085）：
# 推理的起始`beta`值。
# beta_end（`float`，默认为0.012）：
# 最终的`beta`值。
# beta_schedule（`str`，默认为`"linear"`）：
# beta计划，将beta范围映射到一系列用于模型步进的betas。可选值为
# `linear`或`scaled_linear`。
# trained_betas（`np.ndarray`，*可选*）：
# 直接传递betas数组到构造函数，以绕过`beta_start`和`beta_end`。
# use_karras_sigmas（`bool`，*可选*，默认为`False`）：
# 是否在采样过程中使用Karras sigmas作为噪声调度中的步长。如果为`True`，
# sigmas根据一系列噪声水平{σi}确定。
# prediction_type（`str`，默认为`epsilon`，*可选*）：
# 调度器函数的预测类型；可以是`epsilon`（预测扩散过程的噪声），
# `sample`（直接预测有噪声的样本）或`v_prediction`（参见[Imagen Video](https://imagen.research.google/video/paper.pdf)论文第2.4节）。
# timestep_spacing（`str`，默认为`"linspace"`）：
# 时间步的缩放方式。有关更多信息，请参考[Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891)的表2。
# steps_offset（`int`，默认为0）：
# 添加到推理步骤的偏移量，某些模型系列所需。
# 兼容的调度器列表，提取KarrasDiffusionSchedulers的名称
_compatibles = [e.name for e in KarrasDiffusionSchedulers]
# 设置阶数为2
order = 2
# 注册到配置中的初始化方法
@register_to_config
def __init__(
# 定义num_train_timesteps参数，默认为1000
num_train_timesteps: int = 1000,
# 定义beta_start参数，默认为0.00085，设置合理的默认值
beta_start: float = 0.00085, # sensible defaults
# 定义beta_end参数，默认为0.012
beta_end: float = 0.012,
# 定义beta_schedule参数，默认为"linear"
beta_schedule: str = "linear",
# 定义trained_betas参数，默认为None，类型为np.ndarray或float列表
trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
# 定义use_karras_sigmas参数，默认为False
use_karras_sigmas: Optional[bool] = False,
# 定义prediction_type参数，默认为"epsilon"
prediction_type: str = "epsilon",
# 定义timestep_spacing参数，默认为"linspace"
timestep_spacing: str = "linspace",
# 定义steps_offset参数，默认为0
steps_offset: int = 0,
):
# 检查是否提供了训练好的 beta 值
if trained_betas is not None:
# 如果提供了，使用这些 beta 值创建一个张量
self.betas = torch.tensor(trained_betas, dtype=torch.float32)
# 检查 beta 调度类型是否为线性
elif beta_schedule == "linear":
# 创建一个从 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_start 的平方根到 beta_end 的平方根的线性间隔，然后平方
self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
# 检查 beta 调度类型是否为平方余弦调度 v2
elif beta_schedule == "squaredcos_cap_v2":
# Glide 余弦调度
# 使用指定的函数为 alpha_bar 计算 beta 值
self.betas = betas_for_alpha_bar(num_train_timesteps)
else:
# 如果没有匹配的调度类型，抛出未实现的错误
raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
# 计算 alpha 值，alpha 等于 1 减去 beta
self.alphas = 1.0 - self.betas
# 计算 alpha 的累积乘积
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
# 设置所有时间步的值
self.set_timesteps(num_train_timesteps, None, num_train_timesteps)
# 初始化步骤索引为 None
self._step_index = None
# 初始化开始索引为 None
self._begin_index = None
# 将 sigmas 张量移动到 CPU，以减少 CPU/GPU 之间的通信
self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication
@property
def init_noise_sigma(self):
# 返回初始噪声分布的标准差
if self.config.timestep_spacing in ["linspace", "trailing"]:
# 如果时间步间隔为线性或拖尾，返回 sigmas 的最大值
return self.sigmas.max()
# 否则，计算并返回基于 sigmas 最大值的结果
return (self.sigmas.max() ** 2 + 1) ** 0.5
@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):
"""
设置调度器的开始索引。该函数应在推理之前从管道运行。
Args:
begin_index (`int`):
调度器的开始索引。
"""
# 设置开始索引
self._begin_index = begin_index
def scale_model_input(
self,
# 输入的样本张量
sample: torch.Tensor,
# 当前时间步，可以是浮点数或张量
timestep: Union[float, torch.Tensor],
) -> torch.Tensor:
"""
确保与需要根据当前时间步缩放去噪模型输入的调度器互换性。
参数:
sample (`torch.Tensor`):
输入样本。
timestep (`int`, *可选*):
当前扩散链中的时间步。
返回:
`torch.Tensor`:
缩放后的输入样本。
"""
# 如果 step_index 未初始化，调用初始化函数
if self.step_index is None:
self._init_step_index(timestep)
# 如果当前状态为一阶状态，获取当前步的 sigma
if self.state_in_first_order:
sigma = self.sigmas[self.step_index]
else:
# 否则，获取插值后的 sigma
sigma = self.sigmas_interpol[self.step_index - 1]
# 对样本进行缩放，缩放因子为 sqrt(sigma^2 + 1)
sample = sample / ((sigma**2 + 1) ** 0.5)
# 返回缩放后的样本
return sample
def set_timesteps(
self,
num_inference_steps: int,
device: Union[str, torch.device] = None,
num_train_timesteps: Optional[int] = None,
# 从 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 = t.reshape(sigma.shape)
# 返回时间值
return 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 et al. (2022) 的噪声调度。"""
# 确保其他复制此函数的调度器不会出错的 hack
# TODO: 将此逻辑添加到其他调度器中
if hasattr(self.config, "sigma_min"):
# 如果配置中有 sigma_min，使用其值
sigma_min = self.config.sigma_min
else:
# 否则设为 None
sigma_min = None
if hasattr(self.config, "sigma_max"):
# 如果配置中有 sigma_max，使用其值
sigma_max = self.config.sigma_max
else:
# 否则设为 None
sigma_max = None
# 如果 sigma_min 仍为 None，使用输入 sigma 的最后一个值
sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
# 如果 sigma_max 仍为 None，使用输入 sigma 的第一个值
sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
rho = 7.0 # 论文中使用的值
# 创建从 0 到 1 的线性空间
ramp = np.linspace(0, 1, num_inference_steps)
# 计算 sigma 的逆
min_inv_rho = sigma_min ** (1 / rho)
max_inv_rho = sigma_max ** (1 / rho)
# 计算新的 sigma 值
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
# 返回新的 sigma 值
return sigmas
@property
def state_in_first_order(self):
# 返回 sample 是否为 None，判断是否处于一阶状态
return self.sample is 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 索引始终是第二个索引
# 如果只有一个索引，则使用最后一个索引
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):
# 如果开始索引为空
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(
# 模型输出，类型可以是 torch.Tensor 或 np.ndarray
model_output: Union[torch.Tensor, np.ndarray],
# 当前时间步，可以是浮点数或张量
timestep: Union[float, torch.Tensor],
# 输入样本，可以是张量或数组
sample: Union[torch.Tensor, np.ndarray],
# 可选的随机生成器
generator: Optional[torch.Generator] = None,
# 是否返回字典格式的结果，默认为 True
return_dict: bool = True,
# 从 diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise 复制而来
def add_noise(
# 原始样本，类型为 torch.Tensor
original_samples: torch.Tensor,
# 添加的噪声，类型为 torch.Tensor
noise: torch.Tensor,
# 时间步，类型为 torch.Tensor
timesteps: torch.Tensor,
) -> torch.Tensor:
# 确保 sigmas 和 timesteps 与 original_samples 具有相同的设备和数据类型
sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
# 如果设备类型是 "mps" 且 timesteps 是浮点数
if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
# mps 不支持 float64 类型
schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32)
# 将 timesteps 转换为与 original_samples 相同的设备和 float32 类型
timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
else:
# 将 schedule_timesteps 转换为与 original_samples 相同的设备
schedule_timesteps = self.timesteps.to(original_samples.device)
# 将 timesteps 转换为与 original_samples 相同的设备
timesteps = timesteps.to(original_samples.device)
# 当 scheduler 用于训练时 self.begin_index 为 None，或管道未实现 set_begin_index
if self.begin_index is None:
# 根据 schedule_timesteps 计算每个 timesteps 的索引
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]
# 根据 step_indices 获取对应的 sigma 值，并展平为一维
sigma = sigmas[step_indices].flatten()
# 在 sigma 的维度小于 original_samples 的维度时，扩展 sigma 维度
while len(sigma.shape) < len(original_samples.shape):
sigma = sigma.unsqueeze(-1)
# 生成带噪声的样本，原始样本加上噪声和 sigma 的乘积
noisy_samples = original_samples + noise * sigma
# 返回生成的带噪声样本
return noisy_samples
# 返回训练时间步的数量
def __len__(self):
return self.config.num_train_timesteps
浙公网安备 33010602011771号