CacheDiT

CacheDiT mempercepat pipeline DiT dengan menggunakan ulang output blok transformer di antara langkah-langkah denoising. Metode ini tidak memerlukan pelatihan dan mendukung sebagian besar pipeline DiDiffusers, termasuk Flux, Qwen-Image, Wan, dan HunyuanVideo. Diffusers juga memiliki penyimpanan cache bawaan yang tidak memerlukan dependensi tambahan.

Install CacheDiT dari PyPI.

pip install -U cache-dit

Panggil cache_dit.supported_pipelines() untuk menampilkan daftar keluarga pipeline yang didukung oleh CacheDiT.

import cache_dit

cache_dit.supported_pipelines()

Aktifkan penyangga

Panggil cache_dit.enable_cache pada pipeline untuk meng-cache-nya dengan pengaturan default, lalu jalankan pipeline seperti biasa.

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 juga berfungsi dengan torch.compile. Kompilasi transformer setelah Anda memanggil enable_cache. Lihat dokumentasi compile untuk pengaturan yang menghindari kompilasi ulang dengan bentuk input dinamis.

pipeline.transformer = torch.compile(pipeline.transformer)

Panggil cache_dit.summary setelah inferensi untuk mencatat berapa banyak langkah yang di-cache dan perbedaan residual antar langkah.

stats = cache_dit.summary(pipeline)

Panggil cache_dit.disable_cache untuk memulihkan pipeline asli.

cache_dit.disable_cache(pipeline)

Konfigurasi cache

DBCache (Dual Block Cache) menghitung n blok pertama (Fn) pada setiap langkah. Ketika output mereka nyaris tidak berubah dari langkah sebelumnya, DBCache menggunakan kembali output yang di-cache untuk blok-blok sisanya, dan dapat menghitung ulang n blok terakhir (Bn) untuk memperbaikinya. Kalibrator TaylorSeer memprediksi output yang di-cache dari langkah-langkah sebelumnya, alih-alih menggunakannya apa adanya.

enable_cache secara default menggunakan DBCache dengan 8 blok pertama yang selalu dihitung (F8B0) dan 8 langkah pemanasan tanpa cache. Untuk mengorbankan kecepatan demi kualitas, naikkan Fn_compute_blocks atau turunkan residual_diff_threshold (default 0.08). Untuk kualitas terbaik pada tingkat cache tinggi, tambahkan kalibrator 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),
)

Untuk pipeline yang didukung, CacheDiT sudah mengetahui apakah CFG dijalankan sebagai forward pass terpisah. Untuk model lain, atur enable_separate_cfg=True di DBCacheConfig jika model menjalankan jalur kondisional dan non-kondisional secara terpisah, atau False jika model menggabungkannya atau tidak menggunakan CFG.

Lihat dokumen desain DBCache untuk memahami cara kerja blok Fn dan Bn, serta benchmark cache untuk angka kecepatan dan kualitas.

Langkah selanjutnya

• Untuk pipeline yang belum didukung oleh CacheDiT, lihat dokumentasi BlockAdapter.
• CacheDiT juga mendukung paralelisme konteks dan kuantisasi.