Diffusers 文件

TCD排程器

Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

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序列。可選擇linearscaled_linearsquaredcos_cap_v2
  • trained_betas (np.ndarray, 可選) — 直接向建構函式傳遞beta陣列,以繞過beta_startbeta_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_skipc_out時,時間步將乘以的因子。增加此值將減小近似誤差(儘管預設值10.0下的近似誤差已經很小)。
  • rescale_betas_zero_snr (bool, 預設為False) — 是否重新縮放beta值,使其終端SNR為零。這使得模型能夠生成非常亮和非常暗的樣本,而不是將其限制在亮度適中的樣本。與--offset_noise大致相關。

TCDScheduler結合了論文Trajectory Consistency Distillation(軌跡一致性蒸餾)中引入的戰略隨機取樣,擴充套件了原始的多步一致性取樣,以實現無限制的軌跡遍歷。

此程式碼基於TCD的官方倉庫(https://github.com/jabir-zheng/TCD)。

此模型繼承自SchedulerMixinConfigMixin~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

引數

  • sample (torch.Tensor) — 輸入樣本。
  • timestep (int, 可選) — 擴散鏈中的當前時間步。

返回

torch.Tensor

一個縮放後的輸入樣本。

確保與需要根據當前時間步縮放去噪模型輸入的排程器互換使用。

set_begin_index

< >

( begin_index: int = 0 )

引數

  • begin_index (int) — 排程器的起始索引。

設定排程器的起始索引。此函式應在推理之前從管道中執行。

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 (strtorch.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.TCDSchedulerOutputtuple

引數

  • model_output (torch.Tensor) — 從學習到的擴散模型中直接輸出。
  • timestep (int) — 擴散鏈中當前的離散時間步長。
  • sample (torch.Tensor) — 擴散過程建立的當前樣本例項。
  • eta (float) — 一個隨機引數(在論文中被稱為gamma),用於控制每一步的隨機性。當eta = 0時,表示確定性取樣,而當eta = 1時,表示完全隨機取樣。
  • generator (torch.Generator, 可選) — 一個隨機數生成器。
  • return_dict (bool, 可選, 預設為 True) — 是否返回TCDSchedulerOutputtuple

返回

~schedulers.scheduling_utils.TCDSchedulerOutputtuple

如果return_dict為True,則返回TCDSchedulerOutput,否則返回一個元組,其中第一個元素是樣本張量。

透過逆轉 SDE 預測前一個時間步的樣本。此函式從學習到的模型輸出(通常是預測的噪聲)傳播擴散過程。

TCDSchedulerOutput

class 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 函式輸出的輸出類。

< > 在 GitHub 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.