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 はコンテキスト並列性と量子化もサポートしています。

Update on GitHub