Buckets:

mishig's picture
|
download
raw
14.6 kB

キャッシュ

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 は、連続するノイズ除去ステップ間で、スペクトル進化認識型(SEA)指標を比較します。累積された指標の変化が閾値を下回る場合、計算コストの高い Transformer ブロックスタックをスキップし、キャッシュされた残差からその出力を予測します。指標は、生成出力を構成するビジュアル潜在変数から構築します。画像から動画や動画から動画への生成のように、出力軌跡の一部を構成するクリーンな条件付けフレームは含めます。生成を条件付けますが出力の一部ではない独立したビジュアルヒントは除外します。テキスト条件付けはビジュアル潜在変数ではないため除外します。

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

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

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

他の video transformer は、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には、条件付きと条件なしのブランチが連結された状態が含まれます。キャリブレーション処理はこの点を自動的に考慮し、結合した挙動を表す単一の比率の配列を生成します。この生成された配列は、修正なしで直接使用できます。

Xet Storage Details

Size:
14.6 kB
·
Xet hash:
0194952b1fb6d7b48c289fa09b4294297fa5ef880e1637efc252fdc437be10fe

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.