# TCDScheduler

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

该论文的摘要如下：

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

原始代码库可在 [jabir-zheng/TCD](https://github.com/jabir-zheng/TCD) 中找到。

## TCDScheduler[[diffusers.TCDScheduler]]

#### diffusers.TCDScheduler[[diffusers.TCDScheduler]]

```python
diffusers.TCDScheduler(num_train_timesteps: int = 1000, beta_start: float = 0.00085, beta_end: float = 0.012, beta_schedule: str = 'scaled_linear', trained_betas: numpy.ndarray | list[float] | None = None, original_inference_steps: int = 50, clip_sample: bool = False, clip_sample_range: float = 1.0, set_alpha_to_one: bool = True, steps_offset: int = 0, prediction_type: str = 'epsilon', thresholding: bool = False, dynamic_thresholding_ratio: float = 0.995, sample_max_value: float = 1.0, timestep_spacing: str = 'leading', timestep_scaling: float = 10.0, rescale_betas_zero_snr: bool = False)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_tcd.py#L141)

**参数：**

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](https://huggingface.co/papers/2210.02303) 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`, defaults to `"leading"`) : 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.

timestep_scaling (`float`，默认值为 `10.0`) : 在计算一致性模型边界条件 `c_skip` 和 `c_out` 时，时间步长将被乘以的系数。增大该值会降低近似误差（尽管在默认值 `10.0` 下，近似误差已经相当小）。

rescale_betas_zero_snr (`bool`，默认为 `False`)：是否将 betas 重新缩放，使其具有零终端 SNR。这使模型能够生成非常明亮和非常暗的样本，而不是将其限制在中等亮度的样本上。与 [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506) 大致相关。

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

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

该模型继承自 [SchedulerMixin](/docs/diffusers/main/zh/api/schedulers/overview#diffusers.SchedulerMixin) 和 [ConfigMixin](/docs/diffusers/main/zh/api/configuration#diffusers.ConfigMixin)。[~ConfigMixin](/docs/diffusers/main/zh/api/configuration#diffusers.ConfigMixin) 负责存储通过调度器的 `__init__` 函数传入的所有配置属性，例如 `num_train_timesteps`。可以通过 `scheduler.config.num_train_timesteps` 访问这些属性。[SchedulerMixin](/docs/diffusers/main/zh/api/schedulers/overview#diffusers.SchedulerMixin) 通过 [SchedulerMixin.save_pretrained()](/docs/diffusers/main/zh/api/schedulers/overview#diffusers.SchedulerMixin.save_pretrained) 和 [from_pretrained()](/docs/diffusers/main/zh/api/schedulers/overview#diffusers.SchedulerMixin.from_pretrained) 函数提供通用的加载和保存功能。

#### add_noise[[diffusers.TCDScheduler.add_noise]]

```python
add_noise(original_samples: Tensor, noise: Tensor, timesteps: IntTensor)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_tcd.py#L702)

**参数：**

original_samples (`torch.Tensor`) : 将添加噪声的原始样本。

noise (`torch.Tensor`) : 要添加到样本中的噪声。

timesteps (`torch.IntTensor`)：指示每个样本噪声水平的时间步。

**返回：** `torch.Tensor`

带噪声的样本。

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

#### get_velocity[[diffusers.TCDScheduler.get_velocity]]

```python
get_velocity(sample: Tensor, noise: Tensor, timesteps: IntTensor)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_tcd.py#L745)

**参数：**

sample (`torch.Tensor`)：输入样本。

noise (`torch.Tensor`) : 噪声张量。

timesteps (`torch.IntTensor`)：用于速度计算的 timesteps。

**返回：** `torch.Tensor`

计算得到的速度。

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

#### index_for_timestep[[diffusers.TCDScheduler.index_for_timestep]]

```python
index_for_timestep(timestep: typing.Union[float, torch.Tensor], schedule_timesteps: typing.Optional[torch.Tensor] = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_tcd.py#L263)

**参数：**

timestep (`float` 或 `torch.Tensor`)：要在调度器中查找的时间步值。

schedule_timesteps (`torch.Tensor`, *可选*) : 要搜索的时间步调度。如果为 `None`，则使用 `self.timesteps`。

**返回值：** `int`

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

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

#### previous_timestep[[diffusers.TCDScheduler.previous_timestep]]

```python
previous_timestep(timestep: int)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_tcd.py#L790)

**参数：**

timestep (`int`)：当前时间步。

**返回：** `int` 或 `torch.Tensor`

上一个时间步。

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

#### scale_model_input[[diffusers.TCDScheduler.scale_model_input]]

```python
scale_model_input(sample: Tensor, timestep: typing.Union[int, torch.Tensor, NoneType] = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_tcd.py#L342)

**参数：**

sample (`torch.Tensor`)：输入样本。

timestep (`int` 或 `torch.Tensor`，*可选*)：扩散链中的当前时间步。

**返回：** `torch.Tensor`

经过缩放的输入样本。

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

#### set_begin_index[[diffusers.TCDScheduler.set_begin_index]]

```python
set_begin_index(begin_index: int = 0)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_tcd.py#L332)

**参数：**

begin_index (`int`，默认为 `0`)：调度器的起始索引。

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

#### set_timesteps[[diffusers.TCDScheduler.set_timesteps]]

```python
set_timesteps(num_inference_steps: int | None = None, device: typing.Union[str, torch.device, NoneType] = None, original_inference_steps: int | None = None, timesteps: list[int] | None = None, strength: float = 1.0)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_tcd.py#L432)

**参数：**

num_inference_steps (`int`, *可选*) : 使用预训练模型生成样本时所用的扩散步数。如果使用了该参数，`timesteps` 必须为 `None`。

device (`str` 或 `torch.device`，*可选*)：时间步应移动到的设备。如果为 `None`，则不移动时间步。

original_inference_steps (`int`, *optional*) : 原始推理步数，用于生成线性间隔的时间步调度（这与标准的 `diffusers` 实现不同）。随后，我们将从这个调度中按索引均匀间隔地选取 `num_inference_steps` 个时间步，并将其作为最终的时间步调度。如果未设置，则默认为 `original_inference_steps` 属性。

timesteps (`list[int]`, *可选*) : 用于支持任意时间步间隔的自定义时间步。如果为 `None`，则使用训练/蒸馏时间步调度上时间步之间等间隔的默认时间步间隔策略。如果传入了 `timesteps`，则 `num_inference_steps` 必须为 `None`。

strength (`float`，默认为 `1.0`)：用于确定在使用 img2img、inpaint 等进行推理时所用的时间步数。

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

#### 步骤[[diffusers.TCDScheduler.step]]

```python
step(model_output: Tensor, timestep: typing.Union[int, torch.Tensor], sample: Tensor, eta: float = 0.3, generator: typing.Optional[torch.Generator] = None, return_dict: bool = True)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_tcd.py#L594)

**参数：**

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](/docs/diffusers/main/zh/api/schedulers/tcd#diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput) 或 `tuple`。

**返回：** [TCDSchedulerOutput](/docs/diffusers/main/zh/api/schedulers/tcd#diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput) 或 `tuple[torch.Tensor, torch.Tensor]`

If return_dict is `True`, [TCDSchedulerOutput](/docs/diffusers/main/zh/api/schedulers/tcd#diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput) is returned, otherwise a tuple is returned containing the previous sample and the predicted noised sample.

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

## TCDSchedulerOutput[[diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput]]

#### diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput[[diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput]]

```python
diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput(prev_sample: Tensor, pred_noised_sample: typing.Optional[torch.Tensor] = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_tcd.py#L35)

**参数：**

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` 函数输出的输出类。