diffusers/main/ja/api/schedulers/tcd.md

TCDScheduler

Jianbin Zheng、Minghui Hu、Zhongyi Fan、Chaoyue Wang、Changxing Ding、Dacheng Tao、Tat-Jen Cham による Trajectory Consistency Distillation は、少数のステップで良好なサンプルを生成できる Strategic Stochastic Sampling（アルゴリズム 4）を導入しました。Consistency Models のマルチステップ スケジューラ（アルゴリズム 1）の高度な反復処理として位置づけられるこの手法は、軌道一貫性関数に特化した Strategic Stochastic Sampling です。

論文の抄録は以下の通りです：

Latent Consistency Model (LCM) は Consistency Model を潜在空間に拡張し、guided consistency distillation 技術を活用することで、テキストから画像を生成する処理の高速化において優れた性能を実現しています。しかし、LCM は明瞭さと細部の精緻さを兼ね備えた画像の生成に苦戦する傾向があることが観察されました。この制約に対処するため、まずその根本的な原因を掘り下げて明らかにしました。調査の結果、主な問題点は3つの異なる領域におけるエラーに起因することが判明しました。そこで、trajectory consistency function と strategic stochastic sampling を包含する Trajectory Consistency Distillation (TCD) を提案します。trajectory consistency function は、自己整合性の境界条件の範囲を広げることにより蒸留エラーを低減し、TCD に Probability Flow ODE の軌跡全体を正確に追跡する能力を付与します。また、strategic stochastic sampling は、多ステップの consistency sampling に固有の蓄積エラーを回避するために特化して設計されており、TCD モデルを補完するように細部まで調整されています。実験により、TCD は低 NFE において画像品質を大幅に向上させるだけでなく、高 NFE においても教師モデルと比較してより詳細な結果を生成することが示されました。

元のコードベースは jabir-zheng/TCD で確認できます。

TCDScheduler[[diffusers.TCDScheduler]]

diffusers.TCDScheduler[[diffusers.TCDScheduler]]

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)

ソース

パラメータ:

num_train_timesteps (int、デフォルトは1000) : モデルを学習するための拡散ステップ数。

beta_start (float、デフォルトは0.00085) : 推論の初期beta値。

beta_end (float、デフォルトは0.012) : 最終的なbeta値。

beta_schedule (str、デフォルトは"scaled_linear") : ベータスケジュール。モデルのステップ処理において、ベータの範囲からベータのシーケンスへのマッピングです。linear、scaled_linear、またはsquaredcos_cap_v2から選択します。

trained_betas (np.ndarray または list[float]、任意) : 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) : 各拡散ステップでは、そのステップと直前のステップのアルファの積の値が使用されます。最終ステップには直前のアルファが存在しません。このオプションがTrueの場合、直前のアルファの積は1に固定され、それ以外の場合はステップ0のアルファの値が使用されます。

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, 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 for more information.

timestep_scaling (float、デフォルトは 10.0) : 一貫性モデルの境界条件 c_skip と c_out を計算する際に、タイムステップに掛けられる係数です。この値を大きくすると、近似誤差は小さくなります（ただし、デフォルトの 10.0 における近似誤差はすでにかなり小さいです）。

rescale_betas_zero_snr (bool、デフォルトはFalse) : ベータを再スケーリングして、終端SNRをゼロにするかどうか。これにより、モデルは中間の明るさに限定されず、非常に明るく暗いサンプルを生成できるようになります。--offset_noiseとゆるく関連しています。

TCDScheduler は、論文 Trajectory Consistency Distillation で提案された Strategic Stochastic Sampling を取り入れ、元の Multistep Consistency Sampling を拡張して、制約のない軌跡の走査を可能にしています。

このコードは、TCDの公式リポジトリ(https://github.com/jabir-zheng/TCD)に基づいています。

このモデルは SchedulerMixin と ConfigMixin を継承しています。~ConfigMixin は、num_train_timesteps のような、スケジューラの __init__ 関数に渡されるすべての設定属性を保存する処理を担当します。これらは scheduler.config.num_train_timesteps を介してアクセスできます。SchedulerMixin は、SchedulerMixin.save_pretrained() および from_pretrained() 関数を通じて、一般的な読み込みおよび保存機能を提供します。

add_noise[[diffusers.TCDScheduler.add_noise]]

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

ソース

パラメータ:

original_samples (torch.Tensor) : ノイズが加えられる元のサンプル。

noise (torch.Tensor) : サンプルに追加するノイズ。

timesteps (torch.IntTensor) : 各サンプルのノイズレベルを示すタイムステップ。

戻り値: torch.Tensor

ノイズ付きサンプル。

各タイムステップにおけるノイズの大きさに従って、元のサンプルにノイズを追加します（これはフォワード拡散プロセスです）。

get_velocity[[diffusers.TCDScheduler.get_velocity]]

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

ソース

パラメータ:

sample (torch.Tensor) : 入力サンプル。

noise (torch.Tensor) : ノイズのテンソル。

timesteps (torch.IntTensor) : 速度計算用のタイムステップ。

戻り値: torch.Tensor

計算された速度。

速度の式に従って、サンプルとノイズから速度予測を計算します。

index_for_timestep[[diffusers.TCDScheduler.index_for_timestep]]

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

ソース

パラメータ:

timestep (float または torch.Tensor) : スケジュール内で検索するタイムステップの値。

schedule_timesteps (torch.Tensor, 任意) : 検索対象のタイムステップスケジュール。None の場合は self.timesteps を使用します。

戻り値: int

スケジュールにおけるタイムステップのインデックスです。最初のステップでは、スケジュールの途中から開始する場合（画像から画像への生成など）にシグマをスキップすることを避けるため、複数の一致が存在する場合は2番目のインデックスを返します。

与えられたタイムステップのタイムステップスケジュールにおけるインデックスを特定します。

previous_timestep[[diffusers.TCDScheduler.previous_timestep]]

previous_timestep(timestep: int)

ソース

パラメータ:

timestep (int) : 現在のタイムステップ。

戻り値: int または torch.Tensor

直前のタイムステップ。

拡散チェーンにおける直前のタイムステップを計算します。

scale_model_input[[diffusers.TCDScheduler.scale_model_input]]

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

ソース

パラメータ:

sample (torch.Tensor) : 入力サンプル。

timestep (int または torch.Tensor、任意) : 拡散チェーンにおける現在のタイムステップ。

戻り値: torch.Tensor

スケーリングされた入力サンプル。

現在のタイムステップに応じて、デノイジングモデルの入力をスケーリングする必要があるスケジューラとの相互運用性を保証します。

set_begin_index[[diffusers.TCDScheduler.set_begin_index]]

set_begin_index(begin_index: int = 0)

ソース

パラメータ:

begin_index (int、デフォルトは0) : スケジューラの開始インデックス。

スケジューラの開始インデックスを設定します。この関数は、推論の前にパイプラインから実行する必要があります。

set_timesteps[[diffusers.TCDScheduler.set_timesteps]]

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)

ソース

パラメータ:

num_inference_steps (int、任意) : 事前学習済みモデルでサンプルを生成する際に使用する拡散ステップ数。使用する場合、timesteps は None である必要があります。

device (str または torch.device、任意) : timesteps を移動先とするデバイスです。None の場合、timesteps は移動されません。

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]]

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)

ソース

パラメータ:

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]

True の場合、TCDSchedulerOutput が返され、それ以外の場合は、直前のサンプルと予測されたノイズ付きサンプルを含むタプルが返されます。

前回のタイムステップからのサンプルを、SDEを逆方向に適用することで予測します。この関数は、学習済みモデルの出力（最もよく使われるのは予測ノイズ）から拡散過程を伝播させます。

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

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

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

ソース

パラメータ:

prev_sample (torch.Tensor、画像の場合は (batch_size, num_channels, height, width) の形状) : 前のタイムステップの計算済みサンプル (x_{t-1})。prev_sample は、デノイジングループの次のモデル入力として使用する必要があります。

pred_noised_sample (torch.Tensor、画像の場合は形状 (batch_size, num_channels, height, width)、任意) : 現在のタイムステップにおけるモデルの出力に基づいて予測されるノイズ付きサンプル (x_{s})。

スケジューラのstep関数の出力用のクラス。