Diffusers documentation

キャッシュ

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

キャッシュ

Cachingは、推論の各ステップで計算全体を実行する代わりに、attention層やfeedforward層など異なる層の中間出力を保存して再利用することで、推論を高速化します。これにより生成速度が大幅に向上しますが、メモリ使用量が増加するトレードオフがあり、追加の学習は不要です。

このガイドでは、Diffusersでサポートされているキャッシュメソッドの使用方法について説明します。

Pyramid Attention Broadcast

Pyramid Attention Broadcast (PAB)は、生成プロセスの連続するタイムステップ間でアテンションの出力がそれほど大きくは変わらないという観察に基づいています。アテンションの差分はクロスアテンション層で最も小さく、一般的にはより長いタイムステップ範囲にわたってキャッシュされます。その後、時間アテンション層と空間アテンション層が続きます。

すべての動画モデルが3種類の注意機構(cross、temporal、spatial)を備えているわけではありません!

PABは、シーケンス並列化やclassifier-free guidance並列化(データ並列化)などの他の技術と組み合わせることで、ほぼリアルタイムの動画生成を実現できます。

パイプラインのtransformerにPyramidAttentionBroadcastConfigを設定して渡すことで、これを有効にできます。spatial_attention_block_skip_rangeは、空間注意ブロック内で注意計算をスキップする頻度を制御し、spatial_attention_timestep_skip_rangeはスキップするタイムステップの範囲を指定します。適切な範囲を選ぶよう注意してください。間隔が小さいと推論速度が遅くなる可能性があり、間隔が大きいと生成品質が低下する可能性があります。

import torch
from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig

pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
pipeline.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,
)
pipeline.transformer.enable_cache(config)

FasterCache

FasterCache は、各連続タイムステップにおける出力の差分が小さいため、PAB と同様にアテンション特徴量をキャッシュして再利用します。

このメソッドは、サンプリング時にclassifier-free guidanceを使用する場合(多くのベースモデルで一般的です)、無条件ブランチの予測をスキップし、連続するタイムステップ間の予測された潜在出力に顕著な冗長性がある場合は、条件付きブランチの予測から推定することを選択する場合があります。

有効化するには、FasterCacheConfigを設定して、パイプラインのtransformerに渡します。

import torch
from diffusers import CogVideoXPipeline, FasterCacheConfig

pipe line= CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
pipeline.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",
)
pipeline.transformer.enable_cache(config)

SeaCache

SeaCache は、連続するノイズ除去ステップ間で Spectral-Evolution-Aware (SEA) 指標を比較します。累積された指標の変化が閾値を下回る場合、計算コストの高い transformer block スタックをスキップし、キャッシュされた残差からその出力を予測します。指標は、生成された出力を構成するビジュアル潜在変数から構築します。画像から動画や動画から動画への生成のように、出力軌道の一部であるクリーンな条件付けフレームは含めます。生成を条件付けますが出力の一部ではない個別のビジュアルヒントは除外します。テキスト条件付けはビジュアル潜在変数ではないため除外されます。

Cosmos 3 Transfer パックは、コントロールヒントを独立したビジュアルシーケンスとして扱うため、そのアダプターはインジケーターからこれらを除外します。 Control-CFG 分岐は、同じ出力軌道を比較しながら、それぞれのカッシュされた残差を保持します。コントロールヒントは引き続きトランスフォーマーを条件付けます。

実装では、以下のモデル用のアダプターが組み込まれています:

  • Cosmos 3は、主に最適化およびベンチマークされた統合です。これは、post-normalization境界を介して、デコーダスタック全体をキャッシュします。
  • Wan T2Vは、eagerモードで汎用の繰り返しブロックパスを使用します。この統合は、別の単一ストリーム動画transformerがSeaCacheに生(raw)のビジョン潜在変数(latents)をどのように提供できるかを示すものであり、Wanに対して同じキャッシュパラメータが最適であるという主張や、他のWanバリアントがサポートされていることを意味するものではありません。

他の video transformers は、CacheMixin を使用し、認識可能な繰り返しブロックのリストを公開し、TransformerBlockRegistry にブロックの入出力レイアウトを登録している場合、汎用パスと統合できます。パイプラインは、すべての transformer 呼び出しに対して cache_context に入る必要があり、step_index、sigma、num_inference_steps を付与し、条件付きと無条件のガイダンスのような独立したトラジェクトリには別々のコンテキスト名を使用する必要があります。組み込みのアダプターが利用できない場合は、生成された出力を構成する視覚的潜在変数を返す raw_vision_callback を渡してください。出力品質を検証し、各モデルとスケジューラに対してキャッシュのパラメータを調整してください。サポートやベンチマークの結果は、Cosmos 3 から自動的に引き継がれるものではありません。

Cosmos 3

SeaCache はデフォルトで無効になっています。transformer でこれを有効にしてください。Cosmos 3 のデノイジングループは、各 cache_context 呼び出しにアクティブなスケジューラーのステップ、sigma、ステップ数を付与するため、追加の配線は不要です。

from diffusers import Cosmos3OmniPipeline, SeaCacheConfig

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

このモデルレベルのAPIは、Cosmos3OmniPipeline、Cosmos3OmniModularPipeline、および Cosmos3DistilledModularPipeline で動作します。SeaCache は近似による最適化であり、生成される出力が変わる可能性があります。すべてのデノイジングステップでトランスフォーマーを完全に実行する必要がある場合は、pipe.transformer.disable_cache() を呼び出してください。

FirstBlockCache

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

import torch
from diffusers import DiffusionPipeline
from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig

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

TaylorSeer キャッシュ

TaylorSeer Cache は、テイラー級数展開を用いてデノイジングステップ間の中間アクティベーションを近似・キャッシュすることで、拡散モデルの推論を高速化します。この手法は、過去の計算に基づいて将来の出力を予測し、指定された間隔でそれらを再利用することで、冗長な計算を削減します。

このキャッシュ機構は、追加のメモリオーバーヘッドを最小限に抑えつつ、優れた結果をもたらします。詳細なパフォーマンス分析については、こちらをご覧ください。

TaylorSeer Cache を有効にするには、TaylorSeerCacheConfig を作成し、パイプラインの transformer に渡します。

  • cache_interval: 完全なフォワードパスを実行する前に、キャッシュされた出力を再利用するステップ数
  • disable_cache_before_step: 近似のためのデータを収集するために、完全な計算を行う初期ステップ
  • max_order: 近似の精度(理論上、値を高くすると品質が向上しますが、メモリ使用量も増加します。推奨値は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ブロックの計算をスキップすることで推論を高速化します。拡散プロセスにおいて、更新量(Output - Input)の大きさが予測可能に減衰することを観察しています。事前計算された大きさの比率に基づいて「エラー予算」を蓄積し、計算をスキップして前回の残差を再利用するタイミングを動的に決定します。

MagCache は、この減衰曲線を記述する Magnitude Ratios(mag_ratios)に依存しています。これらの比率は、モデルのチェックポイントとスケジューラごとに固有です。バンドルされている FLUX_MAG_RATIOS は FLUX.1 で測定されたものであり、Flux2 や Flux2 Klein を含む他のモデルには、それぞれ独自のキャリブレーション実行が必要です。

MagCache を使用するには、通常、キャリブレーションと推論という2つのステップに従います。

  1. キャリブレーション: calibrate=True を指定して推論を一度実行します。フックが残差の大きさを測定し、計算された比率をコンソールに出力します。
  2. 推論: これらの比率をMagCacheConfigに渡すことで、高速化を有効にします。
import torch
from diffusers import FluxPipeline, MagCacheConfig

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)
# Logs will print something like: "MagCache Calibration Results: [1.0, 1.37, 0.97, 0.87]"

# 2. Inference Step
# Apply the specific ratios obtained from calibration for optimized speed.
# Note: For Flux models, you can also import defaults: 
# from diffusers.hooks.mag_cache import 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]

mag_ratiosは、モデルの固有のマグニチュード減衰曲線を表します。ステップ数が多い場合(例:50)に較正された比率は、ステップ数が少ない場合(例:20)にも再利用できます。実装では、補間を用いて曲線を現在の推論ステップ数にマッピングします。

Classifier-Free Guidance を逐次実行するパイプライン(Kandinsky 5.0 など)では、キャリブレーションログに 2 つの配列が出力される場合があります。1 つは Conditional パス用、もう 1 つは Unconditional パス用です。ほとんどの場合、最初の配列(Conditional)を使用してください。

Classifier-Free Guidance をバッチ方式で実行するパイプライン(SDXL や Flux など)では、モデルが処理するhidden_statesには、条件付きと条件なしのブランチが連結された状態が含まれます。キャリブレーション処理はこの点を自動的に考慮し、結合した挙動を表す単一の比率の配列を生成します。この生成された配列は、修正なしで直接使用できます。

Update on GitHub