Diffusers documentation

TCDScheduler

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

TCDScheduler

Jianbin Zheng、Minghui Hu、Zhongyi Fan、Chaoyue Wang、Changxing Ding、Dacheng Tao 和 Tat-Jen Cham 提出的 Trajectory Consistency Distillation 引入了一种策略性随机采样(Algorithm 4),能够在少量步数内生成优质样本。作为 Consistency Models 中多步调度器(Algorithm 1)的高级迭代版本,策略性随机采样专门针对轨迹一致性函数进行了定制。

该论文的摘要如下:

潜在一致性模型(LCM)将一致性模型扩展至潜在空间,并利用引导一致性蒸馏技术,在加速文本到图像合成方面取得了卓越的性能。然而,我们观察到 LCM 在生成兼具清晰度和细节丰富度的图像时存在困难。为了解决这一局限性,我们首先深入探究并阐明了其根本原因。我们的研究指出,主要问题源于三个不同方面的误差。因此,我们提出了轨迹一致性蒸馏(TCD),它包含轨迹一致性函数和策略性随机采样。轨迹一致性函数通过扩展自一致性边界条件的范围来减少蒸馏误差,并赋予 TCD 精确追踪概率流 ODE 完整轨迹的能力。此外,策略性随机采样专门设计用于规避多步一致性采样中固有的累积误差,并经过精心定制以配合 TCD 模型。实验表明,TCD 不仅在低 NFE 下显著提升了图像质量,而且在高 NFE 下相比教师模型能产生更精细的结果。

原始代码库可在 jabir-zheng/TCD 中找到。

TCDScheduler

class diffusers.TCDScheduler

< >

( num_train_timesteps: int = 1000beta_start: float = 0.00085beta_end: float = 0.012beta_schedule: str = 'scaled_linear'trained_betas: numpy.ndarray | list[float] | None = Noneoriginal_inference_steps: int = 50clip_sample: bool = Falseclip_sample_range: float = 1.0set_alpha_to_one: bool = Truesteps_offset: int = 0prediction_type: str = 'epsilon'thresholding: bool = Falsedynamic_thresholding_ratio: float = 0.995sample_max_value: float = 1.0timestep_spacing: str = 'leading'timestep_scaling: float = 10.0rescale_betas_zero_snr: bool = False )

参数

  • num_train_timesteps (int,默认为 1000) — 用于训练模型的扩散步数。
  • beta_start (float,默认为 0.00085) — 推理的起始 beta 值。
  • beta_end (float,默认为 0.012) — 最终的 beta 值。
  • beta_schedule (str,默认为 "scaled_linear") — Beta 调度,将 beta 范围映射为用于模型步进的一系列 beta 值。可选 linear、scaled_linear 或 squaredcos_cap_v2。
  • trained_betas(np.ndarray 或 list[float],可选)— 直接传入 betas 数组到构造函数中,以绕过 beta_start 和 beta_end。
  • original_inference_steps (int,默认为 50) — 用于生成线性间隔时间步调度的默认推理步数,最终将从中均匀选取 num_inference_steps 个时间步,以构成最终的时间步调度。
  • clip_sample (bool,默认为 False) — 对预测样本进行裁剪,以确保数值稳定性。
  • clip_sample_range (float,默认为 1.0) — 样本裁剪的最大幅度。仅在 clip_sample=True 时有效。
  • set_alpha_to_one (bool,默认为 True) — 每个扩散步骤都会使用当前步骤和上一步骤的 alpha 乘积值。对于最后一步,不存在上一步的 alpha。当此选项为 True 时,上一步的 alpha 乘积固定为 1,否则使用步骤 0 的 alpha 值。
  • steps_offset(int,默认为 0)— 添加到推理步数中的偏移量,某些模型系列需要此参数。
  • prediction_type (str, defaults to "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 paper).
  • thresholding (bool,默认为 False) — 是否使用“动态阈值”方法。此方法不适用于 Stable Diffusion 等潜在空间扩散模型。
  • dynamic_thresholding_ratio (float, 默认值为 0.995) — 动态阈值化方法的比例。仅在 thresholding=True 时有效。
  • sample_max_value (float,默认为 1.0) — 动态阈值化的阈值。仅在 thresholding=True 时有效。
  • timestep_spacing (str,默认为 "leading") — 时间步的缩放方式。更多信息请参阅 常见扩散噪声调度和采样步骤存在缺陷 中的表 2。
  • timestep_scaling (float,默认为 10.0) — 在计算一致性模型边界条件 c_skip 和 c_out 时,时间步(timesteps)将被乘以的系数。增大该值会降低近似误差(尽管在默认值 10.0 下,近似误差已经相当小)。
  • rescale_betas_zero_snr (bool,默认为 False) — 是否重新缩放 betas 以实现零终端 SNR。这使模型能够生成非常明亮和非常暗的样本,而不是将其限制在中等亮度的样本上。与 --offset_noise 大致相关。

TCDScheduler 引入了论文 Trajectory Consistency Distillation 中提出的 Strategic Stochastic Sampling,将原有的 Multistep Consistency Sampling 扩展为支持无限制的轨迹遍历。

此代码基于 TCD 的官方仓库(https://github.com/jabir-zheng/TCD)。

该模型继承自 SchedulerMixin 和 ConfigMixin。~ConfigMixin 负责存储通过调度器的 __init__ 函数传入的所有配置属性,例如 num_train_timesteps。可以通过 scheduler.config.num_train_timesteps 访问这些属性。SchedulerMixin 通过 SchedulerMixin.save_pretrained() 和 from_pretrained() 函数提供通用的加载和保存功能。

add_noise

< >

( original_samples: Tensornoise: Tensortimesteps: IntTensor ) → torch.Tensor

参数

  • original_samples (torch.Tensor) — 将添加噪声的原始样本。
  • noise (torch.Tensor) — 要添加到样本中的噪声。
  • timesteps (torch.IntTensor) — 指示每个样本噪声水平的时间步。

返回值

torch.Tensor

带噪声的样本。

根据每个时间步的噪声幅度,向原始样本添加噪声(即前向扩散过程)。

get_velocity

< >

( sample: Tensornoise: Tensortimesteps: IntTensor ) → torch.Tensor

参数

  • sample (torch.Tensor) — 输入样本。
  • noise (torch.Tensor) — 噪声张量。
  • timesteps (torch.IntTensor) — 用于速度计算的 timesteps。

返回值

torch.Tensor

计算得到的速度。

根据速度公式,从样本和噪声计算速度预测。

index_for_timestep

< >

( timestep: typing.Union[float, torch.Tensor]schedule_timesteps: typing.Optional[torch.Tensor] = None ) → int

参数

  • timestep(float 或 torch.Tensor)— 要在调度器中查找的时间步值。
  • schedule_timesteps (torch.Tensor,可选) — 要搜索的时间步调度。如果为 None,则使用 self.timesteps。

返回值

int

时间步在调度器中的索引。对于第一个步骤,如果存在多个匹配项,则返回第二个索引,以避免在从调度器中间开始(例如,用于图像到图像任务)时跳过某个 sigma。

查找给定时间步在时间步调度中的索引。

previous_timestep

< >

( timestep: int ) → int or torch.Tensor

参数

  • timestep (int) — 当前时间步。

返回值

int 或 torch.Tensor

上一个时间步。

计算扩散链中的前一个时间步。

scale_model_input

< >

( sample: Tensortimestep: typing.Union[int, torch.Tensor, NoneType] = None ) → torch.Tensor

参数

  • sample (torch.Tensor) — 输入样本。
  • timestep (int 或 torch.Tensor,可选) — 扩散链中的当前时间步。

返回值

torch.Tensor

经过缩放的输入样本。

确保与那些需要根据当前时间步对去噪模型输入进行缩放的调度器之间的可互换性。

set_begin_index

< >

( begin_index: int = 0 )

参数

  • begin_index (int,默认为 0) — 调度器的起始索引。

设置调度器的起始索引。此函数应在推理之前从流水线中运行。

set_timesteps

< >

( num_inference_steps: int | None = Nonedevice: typing.Union[str, torch.device, NoneType] = Noneoriginal_inference_steps: int | None = Nonetimesteps: list[int] | None = Nonestrength: float = 1.0 )

参数

  • num_inference_steps (int, 可选) — 使用预训练模型生成样本时所用的扩散步数。如果使用了该参数,timesteps 必须为 None。
  • device (str 或 torch.device,可选) — 时间步应移动到的设备。如果为 None,则不移动时间步。
  • original_inference_steps (int,可选) — 原始推理步数,用于生成线性间隔的时间步调度(与标准的 diffusers 实现不同)。随后,我们将按索引均匀地从该调度中选取 num_inference_steps 个时间步,并将其作为最终的时间步调度。如果未设置,则默认为 original_inference_steps 属性。
  • timesteps (list[int], 可选) — 用于支持任意时间步间隔的自定义时间步。如果为 None,则使用训练/蒸馏时间步调度中时间步等间隔的默认时间步间隔策略。如果传入了 timesteps,则 num_inference_steps 必须为 None。
  • strength (float,默认为 1.0) — 在使用 img2img、inpaint 等进行推理时,用于确定所使用的步数。

设置扩散链所使用的离散时间步(需在推理之前运行)。

步骤

< >

( model_output: Tensortimestep: typing.Union[int, torch.Tensor]sample: Tensoreta: float = 0.3generator: typing.Optional[torch.Generator] = Nonereturn_dict: bool = True ) → TCDSchedulerOutput or tuple[torch.Tensor, torch.Tensor]

参数

  • model_output (torch.Tensor) — 学习到的扩散模型的直接输出。
  • timestep (int 或 torch.Tensor) — 扩散链中的当前离散时间步。
  • sample (torch.Tensor) — 扩散过程生成的样本的当前实例。
  • eta (float,默认为 0.3) — 一个随机参数(在论文中称为 gamma),用于控制每一步的随机性。当 eta = 0 时,表示确定性采样;而 eta = 1 则表示完全随机采样。
  • generator (torch.Generator,可选) — 随机数生成器。
  • return_dict (bool,默认为 True) — 是否返回 TCDSchedulerOutput 或 tuple。

返回值

TCDSchedulerOutput 或 tuple[torch.Tensor, torch.Tensor]

If return_dict is True, TCDSchedulerOutput is returned, otherwise a tuple is returned containing the previous sample and the predicted noised sample.

通过逆向 SDE 预测上一时间步的样本。此函数将扩散过程从学习到的模型输出(通常是预测的噪声)进行传播。

TCDSchedulerOutput

class diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput

< >

( prev_sample: Tensorpred_noised_sample: typing.Optional[torch.Tensor] = None )

参数

  • prev_sample(对于图像,形状为 (batch_size, num_channels, height, width) 的 torch.Tensor)— 前一时间步的已计算样本 (x_{t-1})。在去噪循环中,应将 prev_sample 用作下一个模型输入。
  • pred_noised_sample (torch.Tensor,形状为 (batch_size, num_channels, height, width)(针对图像),可选) — 基于当前时间步的模型输出所预测的含噪样本 (x_{s})。

调度器 step 函数输出的输出类。

Update on GitHub