CacheDiT

CacheDiT accélère les pipelines DiT en réutilisant les sorties des blocs de transformateur entre les étapes de débruitage. Il ne nécessite pas d’entraînement et prend en charge la plupart des pipelines DiT de Diffusers, y compris Flux, Qwen-Image, Wan et HunyuanVideo. Diffusers propose également un cache intégré qui ne requiert aucune dépendance supplémentaire.

Installez CacheDiT depuis PyPI.

pip install -U cache-dit

Appel de cache_dit.supported_pipelines() pour lister les familles de pipelines prises en charge par CacheDiT.

import cache_dit

cache_dit.supported_pipelines()

Activer la mise en cache

Appel cache_dit.enable_cache sur un pipeline pour le mettre en cache avec les paramètres par défaut, puis exécutez le pipeline comme d’habitude.

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 fonctionne également avec torch.compile. Compilez le transformer après avoir appelé enable_cache. Consultez la documentation compile pour les paramètres qui évitent la recompilation avec des formes d'entrée dynamiques.

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

Appel de cache_dit.summary après l'inférence pour consigner le nombre d'étapes mises en cache et les différences résiduelles entre les étapes.

stats = cache_dit.summary(pipeline)

Appel de cache_dit.disable_cache pour restaurer le pipeline d'origine.

cache_dit.disable_cache(pipeline)

Configurer le cache

DBCache (Dual Block Cache) calcule les n premiers blocs (Fn) à chaque étape. Lorsque leur sortie varie très peu par rapport à l'étape précédente, il réutilise la sortie mise en cache pour les blocs restants, et il peut recalculer les n derniers blocs (Bn) pour la corriger. Le calibrateur TaylorSeer prédit la sortie mise en cache à partir des étapes antérieures au lieu de la réutiliser telle quelle.

enable_cache est activé par défaut avec DBCache, les 8 premiers blocs étant toujours calculés (F8B0) et 8 étapes de mise en route non mises en cache. Pour privilégier la qualité au détriment de la vitesse, augmentez Fn_compute_blocks ou diminuez residual_diff_threshold (valeur par défaut 0.08). Pour une qualité optimale à des taux de cache élevés, ajoutez le calibrateur 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),
)

Pour les pipelines pris en charge, CacheDiT sait déjà si CFG est exécuté comme un passage avant distinct. Pour les autres modèles, définissez enable_separate_cfg=True dans DBCacheConfig si le modèle exécute les passages conditionnel et non conditionnel séparément, ou False s’il les fusionne ou n’utilise pas CFG.

Consultez la documentation sur la conception de DBCache pour comprendre le fonctionnement des blocs Fn et Bn, ainsi que les benchmarks de cache pour les chiffres relatifs à la vitesse et à la qualité.

Étapes suivantes

• Pour les pipelines non encore pris en charge par CacheDiT, consultez la documentation de BlockAdapter.
• CacheDiT prend également en charge le parallélisme de contexte et la quantification.