Diffusers 文件
TCD排程器
並獲得增強的文件體驗
開始使用
TCD排程器
Jianbin Zheng、Minghui Hu、Zhongyi Fan、Chaoyue Wang、Changxing Ding、Dacheng Tao 和 Tat-Jen Cham 撰寫的《Trajectory Consistency Distillation》(軌跡一致性蒸餾)引入了一種戰略隨機取樣(演算法4),能夠在少量步驟內生成高質量樣本。它區別於一致性模型中的多步排程器(演算法1),戰略隨機取樣是專門為軌跡一致性函式量身定製的。
論文摘要如下:
潛變數一致性模型(LCM)將一致性模型擴充套件到潛變數空間,並利用引導一致性蒸餾技術,在加速文字到影像合成方面取得了令人印象深刻的效能。然而,我們觀察到LCM在生成具有清晰度和複雜細節的影像方面存在困難。為了解決這一限制,我們首先深入探討並闡明瞭其根本原因。我們的研究發現,主要問題源於三個不同領域的錯誤。因此,我們引入了軌跡一致性蒸餾(TCD),其中包括軌跡一致性函式和戰略隨機取樣。軌跡一致性函式透過拓寬自一致性邊界條件的範圍來減少蒸餾誤差,並賦予TCD準確跟蹤機率流ODE整個軌跡的能力。此外,戰略隨機取樣是專門設計用於規避多步一致性取樣中固有的累積誤差,並與TCD模型相得益彰。實驗表明,TCD不僅在低NFE(噪聲評估次數)下顯著提升了影像質量,而且在高NFE下相比教師模型產生了更詳細的結果。
原始程式碼庫可在jabir-zheng/TCD找到。
TCDScheduler
class diffusers.TCDScheduler
< 來源 >( num_train_timesteps: int = 1000 beta_start: float = 0.00085 beta_end: float = 0.012 beta_schedule: str = 'scaled_linear' trained_betas: typing.Union[numpy.ndarray, typing.List[float], NoneType] = 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.0001) — 推理的起始beta
值。 - beta_end (
float
, 預設為0.02) — 最終beta
值。 - beta_schedule (
str
, 預設為"linear"
) — beta排程器,將beta範圍對映到模型步進的beta序列。可選擇linear
、scaled_linear
或squaredcos_cap_v2
。 - trained_betas (
np.ndarray
, 可選) — 直接向建構函式傳遞beta陣列,以繞過beta_start
和beta_end
。 - original_inference_steps (
int
, 可選, 預設為50) — 用於生成線性間隔時間步計劃的預設推理步數,我們最終將從該計劃中抽取num_inference_steps
個均勻間隔的時間步,以形成最終的時間步計劃。 - clip_sample (
bool
, 預設為True
) — 裁剪預測樣本以保持數值穩定性。 - 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
, 預設為epsilon
, 可選) — 排程函式的預測型別;可以是epsilon
(預測擴散過程的噪聲)、sample
(直接預測噪聲樣本)或v_prediction
(參見Imagen Video論文第2.4節)。 - 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"
) — 時間步的縮放方式。更多資訊請參考《Common Diffusion Noise Schedules and Sample Steps are Flawed》(常見擴散噪聲排程和取樣步長存在缺陷)的表2。 - timestep_scaling (
float
, 預設為10.0) — 計算一致性模型邊界條件c_skip
和c_out
時,時間步將乘以的因子。增加此值將減小近似誤差(儘管預設值10.0
下的近似誤差已經很小)。 - rescale_betas_zero_snr (
bool
, 預設為False
) — 是否重新縮放beta值,使其終端SNR為零。這使得模型能夠生成非常亮和非常暗的樣本,而不是將其限制在亮度適中的樣本。與--offset_noise
大致相關。
TCDScheduler
結合了論文Trajectory Consistency Distillation
(軌跡一致性蒸餾)中引入的戰略隨機取樣
,擴充套件了原始的多步一致性取樣,以實現無限制的軌跡遍歷。
此程式碼基於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()函式提供通用的載入和儲存功能。
scale_model_input
< 來源 >( sample: Tensor timestep: typing.Optional[int] = None ) → torch.Tensor
確保與需要根據當前時間步縮放去噪模型輸入的排程器互換使用。
設定排程器的起始索引。此函式應在推理之前從管道中執行。
set_timesteps
< 來源 >( num_inference_steps: typing.Optional[int] = None device: typing.Union[str, torch.device] = None original_inference_steps: typing.Optional[int] = None timesteps: typing.Optional[typing.List[int]] = None strength: 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: Tensor timestep: int sample: Tensor eta: float = 0.3 generator: typing.Optional[torch._C.Generator] = None return_dict: bool = True ) → ~schedulers.scheduling_utils.TCDSchedulerOutput
或 tuple
引數
- model_output (
torch.Tensor
) — 從學習到的擴散模型中直接輸出。 - timestep (
int
) — 擴散鏈中當前的離散時間步長。 - sample (
torch.Tensor
) — 擴散過程建立的當前樣本例項。 - eta (
float
) — 一個隨機引數(在論文中被稱為gamma
),用於控制每一步的隨機性。當eta = 0時,表示確定性取樣,而當eta = 1時,表示完全隨機取樣。 - generator (
torch.Generator
, 可選) — 一個隨機數生成器。 - return_dict (
bool
, 可選, 預設為True
) — 是否返回TCDSchedulerOutput或tuple
。
返回
~schedulers.scheduling_utils.TCDSchedulerOutput
或tuple
如果return_dict為True
,則返回TCDSchedulerOutput,否則返回一個元組,其中第一個元素是樣本張量。
透過逆轉 SDE 預測前一個時間步的樣本。此函式從學習到的模型輸出(通常是預測的噪聲)傳播擴散過程。
TCDSchedulerOutput
class diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput
< 來源 >( prev_sample: Tensor pred_noised_sample: typing.Optional[torch.Tensor] = None )
排程器 step
函式輸出的輸出類。