Diffusers documentation
TCDScheduler
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
class diffusers.TCDScheduler
< source >( 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") — ベータスケジュールです。モデルのステップ処理において、ベータの範囲からベータのシーケンスへの変換を定義します。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 beepsilon(predicts the noise of the diffusion process),sample(directly predicts the noisy sample) orv_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) — consistency model の境界条件であるc_skipおよびc_outを計算する際に、timestep に乗算される係数です。この値を大きくすると近似誤差は小さくなります（ただし、デフォルト値の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
< source >( original_samples: Tensornoise: Tensortimesteps: IntTensor ) → torch.Tensor
各タイムステップにおけるノイズの大きさに従って、元のサンプルにノイズを追加します（これはフォワード拡散プロセスです）。
get_velocity
< source >( sample: Tensornoise: Tensortimesteps: IntTensor ) → torch.Tensor
速度の式に従って、サンプルとノイズから速度予測を計算します。
index_for_timestep
< source >( timestep: typing.Union[float, torch.Tensor]schedule_timesteps: typing.Optional[torch.Tensor] = None ) → int
与えられたタイムステップのタイムステップスケジュールにおけるインデックスを特定します。
previous_timestep
< source >( timestep: int ) → int or torch.Tensor
拡散チェーンにおける直前のタイムステップを計算します。
scale_model_input
< source >( sample: Tensortimestep: typing.Union[int, torch.Tensor, NoneType] = None ) → torch.Tensor
現在のタイムステップに応じて、デノイジングモデルの入力をスケーリングする必要があるスケジューラーとの相互運用性を保証します。
set_begin_index
< source >( begin_index: int = 0 )
スケジューラの開始インデックスを設定します。この関数は、推論の前にパイプラインから実行する必要があります。
set_timesteps
< source >( 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、省略可) — timesteps を移動先とするデバイスです。Noneの場合、timesteps は移動されません。 - 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 などを使用する場合、推論に使用するタイムステップ数を決定するために使用されます。
拡散チェーンで使用される離散タイムステップを設定します（推論の実行前に実行する必要があります）。
ステップ
< source >( 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) —tupleを返すか、TCDSchedulerOutput を返すかを指定します。
戻り値
TCDSchedulerOutput または tuple[torch.Tensor, torch.Tensor]
True の場合、TCDSchedulerOutput が返され、それ以外の場合は、直前のサンプルと予測されたノイズ付きサンプルを含むタプルが返されます。
前のタイムステップのサンプルを、SDEを逆方向に適用することで予測します。この関数は、学習済みモデルの出力（多くの場合、予測されたノイズ）から拡散過程を伝播させます。
TCDSchedulerOutput
class diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput
< source >( prev_sample: Tensorpred_noised_sample: typing.Optional[torch.Tensor] = None )
スケジューラーのstep関数の出力用のクラス。