Buckets:
CacheDiT
CacheDiTは、デノイジングステップ間でトランスフォーマーブロックの出力を再利用することで、DiTパイプラインを高速化します。学習は不要であり、Flux、Qwen-Image、Wan、HunyuanVideoを含む、DiffusersのほとんどのDiTパイプラインに対応しています。また、Diffusersには追加の依存関係を必要としない組み込みのキャッシュ機能も用意されています。
PyPIからCacheDiTをインストールします。
pip install -U cache-dit
CacheDiTがサポートするパイプラインファミリーの一覧を取得するには、cache_dit.supported_pipelines()を呼び出します。
import cache_dit
cache_dit.supported_pipelines()
キャッシュの有効化
パイプラインに対してcache_dit.enable_cacheを呼び出してデフォルト設定でキャッシュし、その後、パイプラインを通常どおり実行します。
import torch
import cache_dit
from diffusers import FluxPipeline
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev", dtype=torch.bfloat16
).to("cuda")
cache_dit.enable_cache(pipeline)
image = pipeline(
"A cat holding a sign that says hello world", num_inference_steps=28
).images[0]
CacheDiT は torch.compile との併用も可能です。enable_cache を呼び出した後に transformer をコンパイルしてください。動的な入力形状での再コンパイルを避けるための設定については、compile のドキュメントを参照してください。
pipeline.transformer = torch.compile(pipeline.transformer)
推論後にcache_dit.summaryを呼び出し、キャッシュされたステップ数とステップ間の残差の違いをログに記録します。
stats = cache_dit.summary(pipeline)
cache_dit.disable_cache を呼び出して、元のパイプラインを復元します。
cache_dit.disable_cache(pipeline)
キャッシュを設定
DBCache (Dual Block Cache) は、各ステップで最初の n ブロック (Fn) を計算します。その出力が前回のステップからほとんど変化しない場合、残りのブロックにはキャッシュされた出力を再利用し、最後の n ブロック (Bn) を再計算して補正を行うことができます。TaylorSeer のキャリブレーターは、キャッシュされた出力をそのまま再利用するのではなく、以前のステップの出力から予測します。
enable_cache はデフォルトで DBCache を使用し、最初の 8 ブロックは常に計算されます(F8B0)、また 8 ステップのキャッシュされないウォームアップが行われます。速度と品質のトレードオフを行うには、Fn_compute_blocks を増やすか、residual_diff_threshold(デフォルト 0.08)を下げてください。高キャッシュレートで最良の品質を得るには、TaylorSeer カリブレーターを追加してください。
from cache_dit import DBCacheConfig, TaylorSeerCalibratorConfig
cache_dit.enable_cache(
pipeline,
cache_config=DBCacheConfig(
max_warmup_steps=8,
Fn_compute_blocks=8,
Bn_compute_blocks=0,
residual_diff_threshold=0.12,
),
calibrator_config=TaylorSeerCalibratorConfig(taylorseer_order=1),
)
サポートされているパイプラインでは、CacheDiT は CFG が独立したフォワードパスとして実行されるかどうかをすでに把握しています。その他のモデルについては、モデルが条件付きパスと非条件付きパスを別々に実行する場合は DBCacheConfig で enable_separate_cfg=True を設定し、それらを融合している場合や CFG を使用しない場合は False を設定してください。
Fn ブロックと Bn ブロックの動作については DBCache の設計 のドキュメントを、速度と品質の数値については キャッシュのベンチマーク を参照してください。
次のステップ
- CacheDiTがまだサポートしていないパイプラインについては、BlockAdapterのドキュメントを参照してください。
- CacheDiT は コンテキスト並列性 と 量子化 もサポートしています。
Xet Storage Details
- Size:
- 4.96 kB
- Xet hash:
- 0aedaf2cd99489e1879baae891ceac7d797b31bf978f80230d8c72675eea665f
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.