キャッシュ

Cachingは、デノイジングステップ間で中間層の出力を再利用して推論を高速化します。メモリ使用量が増加しますが、トレーニングは不要です。Cachingをサポートするtransformerで、設定を指定してメソッドを有効にします。

キャッシュメソッドを選択

設定する設定項目の量と、必要な精度に応じて、適切な方法を選択してください。各方法は、速度、メモリ使用量、およびキャッシュなしの実行との出力の一致度合いを、それぞれ異なるバランスで調整しています。ある方法に頼る前に、ご自身のプロンプトに対して、キャッシュありとキャッシュなしの出力を比較してください。

Pyramid Attention Broadcast

Pyramid Attention Broadcast (PAB) は、一部のブロックとタイムステップについて注意の出力を再利用し、すべてのステップで再計算しないことで、デノイジングステップ間の注意を近似します。モデルがそれらを持つ場合、Config は注意の種類（空間、時間、クロス）を分離します。すべての動画モデルが3種類すべてを公開するわけではなく、存在するブロックに一致する範囲のみを設定してください。

各タイプは、*_attention_block_skip_range（ウィンドウ内で再計算と再利用を切り替える頻度）と、*_attention_timestep_skip_range（スキップ可能なデノイジングタイムステップ）を使用します。フックがパイプラインの現在のタイムステップを読み取れるように、current_timestep_callback を渡す必要があります。スキップ範囲が広いほど、またはより積極的であるほど、速度は向上しますが、品質リスクも高くなります。

有効にするには、PyramidAttentionBroadcastConfig を渡します。

import torch
from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig

pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

config = PyramidAttentionBroadcastConfig(
    spatial_attention_block_skip_range=2,
    spatial_attention_timestep_skip_range=(100, 800),
    current_timestep_callback=lambda: pipe.current_timestep,
)
pipe.transformer.enable_cache(config)

FasterCache

FasterCacheは、PABと同様に、アテンション特徴量をキャッシュして再利用します。また、classifier-free guidanceにおける無条件ブランチをスキップし、連続する潜在変数が十分に冗長である場合、条件付きブランチから推定することもできます。

有効にするには、FasterCacheConfig を渡します。PABと同様に、使用しているアテンションの種類について*_attention_block_skip_rangeと*_attention_timestep_skip_rangeを設定し、必要に応じてCFGブランチのスキップオプションも設定します。

import torch
from diffusers import CogVideoXPipeline, FasterCacheConfig

pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

config = FasterCacheConfig(
    spatial_attention_block_skip_range=2,
    spatial_attention_timestep_skip_range=(-1, 681),
    current_timestep_callback=lambda: pipe.current_timestep,
    attention_weight_callback=lambda _: 0.3,
    unconditional_batch_skip_range=5,
    unconditional_batch_timestep_skip_range=(-1, 781),
    tensor_format="BFCHW",
)
pipe.transformer.enable_cache(config)

SeaCache

SeaCacheは、連続するノイズ除去ステップ間でSpectral Evolution Aware (SEA)指標を比較します。累積変化が閾値を下回る場合、transformerブロックスタックをスキップし、キャッシュされた残差から出力を予測します。この手法は近似であり、動画生成向けに設計されています。

SeaCacheの組み込みアダプターには、以下が含まれます:

• Cosmos 3は、主に最適化およびベンチマークされた統合です。
• Wan T2Vは、SeaCacheに生ビジョン潜在変数をどのように提供するかを示すデモとして、汎用的な繰り返しブロックパスを使用しています。同じキャッシュパラメータがWanやその他のWanバリアントにそのまま適用できるわけではありません。

transformer で SeaCache を有効にします。Cosmos 3 のデノイジングループは、各 cache_context にスケジューラーのステップ、シグマ、ステップ数を付与するため、追加のパラメータは不要です。

from diffusers import Cosmos3OmniPipeline, SeaCacheConfig

pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano")
pipe.transformer.enable_cache(SeaCacheConfig(threshold=0.2, max_consecutive_cached=2))

SeaCache は出力を変更する可能性があります。すべてのステップで完全な transformer を実行する必要がある場合は、pipe.transformer.disable_cache() を呼び出してください。同じ有効化の呼び出しは、Cosmos3OmniPipeline、Cosmos3OmniModularPipeline、および Cosmos3DistilledModularPipeline でも機能します。

別の動画 Transformer で SeaCache を使用するには、新しいモデルにキャッシュを追加するの手順に従ってください。SeaCache には、すべての cache_context で step_index、sigma、num_inference_steps が必要であり、組み込みのアダプターが存在しない場合は raw_vision_callback も必要です。モデルとスケジューラーごとにパラメータを調整してください。

FirstBlockCache

FirstBlockCacheConfig は、denoiser の初期レイヤーがタイムステップごとにどの程度変化するかを確認します。変化が小さい場合、モデルはコストの高い後続のレイヤーをスキップし、前回の出力を再利用します。

enable_cache を使用して有効化し、disable_cache と is_cache_enabled が同期した状態を維持してください。デフォルトの threshold は 0.05 です。0.2 のような高い値を設定すると、より頻繁にスキップされるため速度が向上しますが、生成品質が低下する可能性があります。

import torch
from diffusers import DiffusionPipeline, FirstBlockCacheConfig

pipe = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16
)
pipe.transformer.enable_cache(FirstBlockCacheConfig(threshold=0.2))

TaylorSeer キャッシュ

TaylorSeer Cacheは、デノイジングステップにわたるテイラー級数展開を用いて拡散モデルの推論を高速化します。初期ステップの活性化から後続ステップの活性化を予測し、その予測を複数ステップにわたって再利用することで、Transformerの完全な計算量を削減します。

• cache_interval: 完全なフォワードパスを実行する前に、キャッシュされた出力を再利用するステップ数
• disable_cache_before_step: 近似のためのデータを収集するために、完全な計算を行う初期ステップ
• max_order: Taylor の次数を高くすると精度が向上する場合がありますが、メモリ使用量が増えます。変更する理由がない限り、1 のままにしてください。

import torch
from diffusers import FluxPipeline, TaylorSeerCacheConfig

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

config = TaylorSeerCacheConfig(
    cache_interval=5,
    max_order=1,
    disable_cache_before_step=10,
    taylor_factors_dtype=torch.bfloat16,
)
pipe.transformer.enable_cache(config)

MagCache

MagCacheは、残差更新の大きさに基づいてtransformerブロックをスキップします。更新の大きさは、デノイジングの過程で予測可能な形で減衰し、MagCacheは事前計算された大きさの比率（mag_ratios）から誤差予算を追跡することで、再利用が安全なタイミングを判断します。これらの比率はチェックポイントとスケジューラに依存します。ステップ数が大きい場合の比率は、ステップ数が少ない場合に補間して適用できます。

バンドルされているFLUX_MAG_RATIOSはFLUX.1で測定されたものです。Flux2やFlux2 Kleinを含む他のモデルには、それぞれ独自のキャリブレーション実行が必要です。

MagCacheは2つのステップに従います:

1. キャリブレーション: calibrate=Trueを指定して推論を1回実行します。フックが残差の大きさを測定し、計算された比率を出力します。
2. 推論: キャリブレーションキャッシュを無効化し、その後、これらの比率をMagCacheConfigに渡して高速化します。

Classifier-free guidance はキャリブレーションに影響を与える可能性があります。Flux のように true_cfg_scale > 1 の場合に true CFG とシーケンシャルコンテキストを使用するパイプラインでは、cache_context("cond") と cache_context("uncond") がそれぞれ別々に呼び出されます。キャリブレーションではコンテキストごとに1つの配列が出力される場合がありますが、ほとんどの場合は条件付きの配列を使用してください。条件付き入力と非条件付き入力を連結して CFG をバッチ処理するパイプライン（CogVideoX など）では、そのまま直接使用できる単一の結合配列が生成されます。

import torch
from diffusers import FluxPipeline, MagCacheConfig
from diffusers.hooks.mag_cache import FLUX_MAG_RATIOS

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

# 1. Calibration Step
# Run full inference to measure model behavior.
calib_config = MagCacheConfig(calibrate=True, num_inference_steps=4)
pipe.transformer.enable_cache(calib_config)

# Run a prompt to trigger calibration
pipe("A cat playing chess", num_inference_steps=4)
# Prints: [MagCache] Calibration Complete. Copy these values to MagCacheConfig(mag_ratios=...):

# 2. Inference Step
# Disable calibration hooks before enabling MagCache for inference.
pipe.transformer.disable_cache()

# Apply ratios from calibration, or use the Flux defaults:
# mag_ratios=FLUX_MAG_RATIOS
mag_config = MagCacheConfig(
    mag_ratios=[1.0, 1.37, 0.97, 0.87],
    num_inference_steps=4
)

pipe.transformer.enable_cache(mag_config)

image = pipe("A cat playing chess", num_inference_steps=4).images[0]

テキスト KV キャッシュ

TextKVCacheConfig は、テキストのキーとバリューの射影を一度だけ計算し、再計算せずにすべてのデノイジングステップで再利用します。現在、apply_text_kv_cache() は NucleusMoEImageTransformerBlock に対してのみフックされています。enable_cache を有効にして有効化してください。

import torch
from diffusers import NucleusMoEImagePipeline, TextKVCacheConfig

pipe = NucleusMoEImagePipeline.from_pretrained(
    "NucleusAI/NucleusMoE-Image", dtype=torch.bfloat16
)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

pipe.transformer.enable_cache(TextKVCacheConfig())

image = pipe("A cat holding a sign that says hello world", num_inference_steps=50).images[0]

新しいモデルにキャッシュを追加する

Cacheメソッドはフックを介してモデルに接続され、これらのフックはモデルクラス、そのtransformerブロック、およびパイプラインのdenoisingループに依存します。enable_cacheを新しいモデルに呼び出す前に、これら3つすべてを設定してください。

1. transformer クラスで CacheMixin を継承します。これにより、enable_cache、disable_cache、cache_context が追加されます。

2. hooks/_helpers.py の _register_transformer_blocks_metadata に、transformer ブロックの出力レイアウトを登録します。FirstBlockCache、MagCache、SeaCache など、ブロック全体をスキップまたは再利用するメソッドは、ブロックの出力から隠れ状態を見つけるためにこれを使用します。

   コピーしました

   TransformerBlockRegistry.register(
       model_class=MyTransformerBlock,
       metadata=TransformerBlockMetadata(
           return_hidden_states_index=0,
           return_encoder_hidden_states_index=None,
       ),
   )

3. パイプラインのデノイジングループ内の各transformer呼び出しをcache_contextでラップします。classifier-free guidanceが条件付きと条件なしのブランチを個別の呼び出しとして実行する場合、"cond"や"uncond"のような別々の名前を使用し、各ブランチが独自のキャッシュ状態を保持できるようにします。

   コピーしました

   with self.transformer.cache_context("cond"):
       noise_pred = self.transformer(hidden_states=latents, timestep=timestep, ...)[0]

PABおよびFasterCacheも、フックが現在のタイムステップを読み取れるようにcurrent_timestep_callbackが必要です。新しいモデルでの出力品質を確認してください。1つのモデル用に調整されたキャッシュ設定は、別のモデルには適用されないことが多いためです。