diffusers/main/id/optimization/cache.md

Penyimpanan sementara

Caching menggunakan kembali output lapisan perantara di antara langkah-langkah denoising untuk mempercepat inferensi. Metode ini menggunakan lebih banyak memori dan tidak memerlukan pelatihan. Aktifkan metode dengan konfigurasi pada transformer yang mendukung caching.

Pilih metode cache

Pilih metode berdasarkan seberapa banyak konfigurasi yang akan Anda atur dan tingkat kecocokan yang Anda butuhkan. Setiap metode menyeimbangkan kecepatan, memori, dan seberapa dekat keluaran dengan eksekusi tanpa cache secara berbeda. Bandingkan keluaran dengan dan tanpa cache pada prompt Anda sendiri sebelum mengandalkan suatu metode.

Metode	Gunakan saat	Catatan
KV Cache Teks	Gambar NucleusMoE, penggunaan ulang proyeksi K/V teks di antara langkah	Saat ini hook blok gambar NucleusMoE
SeaCache	Generasi video Cosmos 3	Pengaturan sering kali tidak dapat dipindahkan antar model
FirstBlockCache	Mau satu pengaturan utama kecepatan/kualitas pada transformer yang terdaftar	Satu threshold mengontrol seberapa sering langkah dilewati
MagCache	Apakah Anda memiliki rasio magnitudo untuk checkpoint dan scheduler Anda, atau akan melakukan kalibrasi terlebih dahulu	Rasio bersifat spesifik untuk checkpoint dan scheduler
TaylorSeer	Ingin memprediksi aktivasi di langkah kemudian dari langkah sebelumnya	max_order yang lebih tinggi menggunakan lebih banyak memori
PAB	Video, bersedia menyesuaikan penggunaan ulang perhatian (rentang blok dan lewati langkah waktu per jenis perhatian)	Rentang lewati perlu penyetelan per model
FasterCache	Seperti PAB, ditambah opsi untuk melewati cabang CFG	Eksperimental

Pyramid Attention Broadcast

Pyramid Attention Broadcast (PAB) memperkirakan perhatian di antara langkah-langkah denoising dengan menggunakan kembali output perhatian untuk beberapa blok dan timestep, alih-alih menghitung ulang setiap langkah. Config memisahkan jenis-jenis perhatian (spasial, temporal, silang) ketika model memilikinya. Tidak semua model video mengekspos ketiga jenis tersebut, dan hanya atur rentang yang sesuai dengan blok yang Anda miliki.

Setiap jenis menggunakan *_attention_block_skip_range (seberapa sering merekomputasi vs. menggunakan ulang di dalam jendela) dan *_attention_timestep_skip_range (timestep denoising mana yang boleh dilewati). Anda harus meneruskan current_timestep_callback agar hook dapat membaca timestep saat ini dari pipeline. Skip yang lebih lebar atau lebih agresif biasanya berarti lebih banyak kecepatan dan lebih banyak risiko kualitas.

Lewatkan PyramidAttentionBroadcastConfig untuk mengaktifkannya.

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 meng-cache dan menggunakan ulang fitur perhatian yang serupa dengan PAB. Ia juga dapat melewati cabang non-kondisional di bawah classifier-free guidance dan memperkirakan cabang tersebut dari cabang kondisional ketika latens yang berurutan cukup redundan.

Lewatkan FasterCacheConfig untuk mengaktifkannya. Seperti PAB, atur *_attention_block_skip_range dan *_attention_timestep_skip_range untuk jenis perhatian yang Anda miliki, serta opsi skip cabang CFG jika Anda menginginkannya.

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 membandingkan indikator Spectral Evolution Aware (SEA) antara langkah denoising yang berurutan. Ketika perubahan kumulatif tetap di bawah ambang batas, metode ini melewati tumpukan blok transformer dan memprediksi keluaran dari residual yang di-cache. Metode ini bersifat aproksimasi dan dirancang untuk generasi video.

Adaptor bawaan untuk SeaCache meliputi:

• Cosmos 3 adalah integrasi utama yang dioptimalkan dan diuji.
• Wan T2V menggunakan jalur blok berulang generik sebagai demo untuk cara menyediakan latens visual mentah ke SeaCache. Parameter cache yang sama mungkin tidak dapat diterapkan ke Wan atau varian Wan lainnya.

Aktifkan SeaCache pada transformer. Loop denoising Cosmos 3 menempelkan langkah scheduler, sigma, dan jumlah langkah ke setiap cache_context, sehingga tidak diperlukan parameter tambahan.

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 dapat mengubah output. Panggil pipe.transformer.disable_cache() ketika Anda membutuhkan setiap langkah untuk menjalankan transformer penuh. Panggilan enable yang sama berfungsi dengan Cosmos3OmniPipeline, Cosmos3OmniModularPipeline, dan Cosmos3DistilledModularPipeline.

Untuk menggunakan SeaCache dengan video transformer lain, ikuti Menambahkan caching ke model baru. SeaCache juga membutuhkan step_index, sigma, dan num_inference_steps di setiap cache_context, serta raw_vision_callback jika tidak ada adapter bawaan. Sesuaikan parameter untuk setiap model dan scheduler.

FirstBlockCache

FirstBlockCacheConfig memeriksa seberapa besar perubahan pada lapisan awal denoiser dari satu timestep ke timestep berikutnya. Jika perubahannya kecil, model akan melewati lapisan-lapisan akhir yang mahal dan menggunakan kembali output sebelumnya.

Aktifkan melalui enable_cache agar disable_cache dan is_cache_enabled tetap sinkron. Nilai threshold bawaan adalah 0.05. Nilai yang lebih tinggi seperti 0.2 akan lebih sering dilewati untuk kecepatan ekstra, tetapi kualitas generasi mungkin menurun.

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))

Cache TaylorSeer

TaylorSeer Cache mempercepat inferensi difusi dengan ekspansi deret Taylor di antara langkah-langkah denoising. Metode ini memprediksi aktivasi langkah berikutnya dari langkah sebelumnya dan menggunakan ulang prediksi tersebut selama beberapa langkah sehingga transformer melakukan lebih sedikit kerja penuh.

• cache_interval: Jumlah langkah untuk menggunakan kembali output yang di-cache sebelum melakukan forward pass penuh
• disable_cache_before_step: Langkah-langkah awal yang menggunakan komputasi penuh untuk mengumpulkan data untuk aproksimasi
• max_order: Orde Taylor yang lebih tinggi dapat lebih akurat tetapi menggunakan lebih banyak memori. Biarkan ini tetap 1 kecuali Anda memiliki alasan untuk mengubahnya.

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 melewati blok transformer berdasarkan magnitudo pembaruan residual. Magnitudo pembaruan meluruh secara prediktif selama denoising, dan MagCache melacak anggaran error dari rasio magnitudo yang dihitung sebelumnya (mag_ratios) untuk menentukan kapan reuse aman. Rasio-rasio tersebut spesifik untuk checkpoint dan scheduler. Rasio dari jumlah langkah yang tinggi dapat diinterpolasi ke jumlah langkah yang lebih sedikit.

FLUX_MAG_RATIOS yang disertakan diukur pada FLUX.1. Model lain, termasuk Flux2 dan Flux2 Klein, memerlukan kalibrasi tersendiri.

MagCache mengikuti dua langkah:

1. Kalibrasi: Jalankan inferensi sekali dengan calibrate=True. Hook mengukur magnitudo residual dan mencetak rasio yang dihitung.
2. Inferensi: Nonaktifkan cache kalibrasi, lalu teruskan rasio tersebut ke MagCacheConfig untuk akselerasi.

Classifier-free guidance dapat memengaruhi kalibrasi. Pipeline yang menggunakan true CFG dengan konteks berurutan, seperti Flux ketika true_cfg_scale > 1, masuk ke cache_context("cond") dan cache_context("uncond") secara terpisah. Kalibrasi mungkin mencetak satu array per konteks, tetapi Anda sebaiknya menggunakan array kondisional dalam kebanyakan kasus. Pipeline yang melakukan batching CFG dengan menggabungkan input kondisional dan non-kondisional (misalnya, CogVideoX) menghasilkan satu array gabungan yang dapat Anda gunakan langsung.

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]

Teks KV Cache

TextKVCacheConfig menghitung proyeksi kunci dan nilai teks sekali dan memakainya kembali di setiap langkah denoising, alih-alih menghitungnya ulang. apply_text_kv_cache() saat ini hanya mengaitkan NucleusMoEImageTransformerBlock. Aktifkan dengan 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]

Menambahkan caching ke model baru

Metode cache menempel pada model melalui hook, dan hook tersebut bergantung pada kelas model, blok transformer-nya, dan loop denoising pipeline. Siapkan ketiga komponen tersebut sebelum memanggil enable_cache pada model baru.

1. Warisi dari CacheMixin di kelas transformer. Ini menambahkan enable_cache, disable_cache, dan cache_context.

2. Daftarkan tata letak keluaran blok transformer di _register_transformer_blocks_metadata dalam hooks/_helpers.py. Metode yang melewati atau menggunakan kembali seluruh blok, seperti FirstBlockCache, MagCache, dan SeaCache, menggunakannya untuk menemukan hidden states dalam keluaran blok.

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

3. Bungkus setiap panggilan transformer dalam loop denoising pipeline dengan cache_context. Gunakan nama yang terpisah, seperti "cond" dan "uncond", ketika classifier-free guidance menjalankan cabang kondisional dan non-kondisional sebagai panggilan terpisah, sehingga setiap cabang mempertahankan status cache-nya sendiri.

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

PAB dan FasterCache juga membutuhkan current_timestep_callback agar hook dapat membaca timestep saat ini. Periksa kualitas output pada model baru, karena pengaturan cache yang dioptimalkan untuk satu model sering kali tidak dapat diterapkan ke model lain.