Mise en cache

Le cache accélère l’inférence en stockant et en réutilisant les sorties intermédiaires de différentes couches, telles que les couches d’attention et les couches feedforward, plutôt qu’en effectuant l’intégralité du calcul à chaque étape d’inférence. Il améliore considérablement la vitesse de génération au prix d’une consommation mémoire accrue et ne nécessite pas d’entraînement supplémentaire.

Ce guide vous montre comment utiliser les méthodes de mise en cache prises en charge dans Diffusers.

Diffusion de l'attention pyramidale

Pyramid Attention Broadcast (PAB) repose sur l’observation que les sorties d’attention ne diffèrent guère entre les pas de temps successifs du processus de génération. Les différences d’attention sont les plus faibles dans les couches d’attention croisée et sont généralement mises en cache sur une plage de pas de temps plus large. Cela est suivi par les couches d’attention temporelle et d’attention spatiale.

Tous les modèles vidéo ne disposent pas de trois types d'attention (croisée, temporelle et spatiale) !

PAB peut être combiné avec d'autres techniques telles que le parallélisme de séquence et le parallélisme de guidage sans classifieur (parallélisme de données) pour une génération de vidéo quasi en temps réel.

Configurez et transmettez un PyramidAttentionBroadcastConfig au transformer d’un pipeline pour l’activer. Le paramètre spatial_attention_block_skip_range contrôle la fréquence à laquelle les calculs d’attention sont ignorés dans les blocs d’attention spatiale, et spatial_attention_timestep_skip_range définit la plage de pas de temps à ignorer. Veillez à choisir une plage appropriée, car un intervalle plus petit peut entraîner une vitesse d’inférence plus lente, tandis qu’un intervalle plus grand peut réduire la qualité de la génération.

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 met en cache et réutilise les caractéristiques d’attention de manière similaire à PAB, car les différences de sortie sont faibles à chaque pas de temps successif.

Cette méthode peut également choisir d’ignorer la prédiction de branche inconditionnelle, lors de l’utilisation du guidage sans classifieur pour l’échantillonnage (courant dans la plupart des modèles de base), et l’estimer à partir de la prédiction de branche conditionnelle s’il existe une redondance significative dans les sorties latentes prédites entre les pas de temps successifs.

Configurez et transmettez un FasterCacheConfig au transformer d’un pipeline pour l’activer.

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 compare les indicateurs Spectral-Evolution-Aware (SEA) entre les étapes successives de débruitage. Lorsque la variation cumulée de l'indicateur reste inférieure à un seuil, il contourne la pile coûteuse de blocs transformer et prédit sa sortie à partir des résidus mis en cache. Construisez l'indicateur à partir des latents visuels qui forment la sortie générée. Incluez les images de conditionnement propres lorsqu'elles font partie de cette trajectoire de sortie, comme dans la génération d'image vers vidéo et de vidéo vers vidéo. Excluez les indices visuels distincts qui conditionnent la génération mais ne font pas partie de la sortie. Le conditionnement textuel est exclu car il ne s'agit pas d'un latent visuel.

Les packs de transfert de Cosmos 3 gèrent les indices de contrôle comme des séquences visuelles distinctes, de sorte que son adaptateur les exclut de l'indicateur. Les branches Control-CFG comparent la même trajectoire de sortie tout en conservant leurs propres résidus mis en cache. Les indices de contrôle conditionnent toujours le transformeur.

L'implémentation fournit des adaptateurs intégrés pour les modèles suivants :

• Cosmos 3 est l'intégration principale, optimisée et benchmarkée. Elle met en cache la pile complète du décodeur jusqu'à la frontière de normalisation postérieure.
• Wan T2V utilise le chemin générique de blocs répétés en mode eager. Cette intégration montre comment un autre transformateur vidéo à flux unique peut fournir des latents visuels bruts à SeaCache ; il ne s’agit pas d’affirmer que les mêmes paramètres de cache sont optimaux pour Wan ni que d’autres variantes de Wan sont prises en charge.

D’autres transformers vidéo peuvent s’intégrer au chemin générique s’ils utilisent CacheMixin, exposent une liste de blocs répétés reconnue et enregistrent la disposition des entrées/sorties des blocs dans TransformerBlockRegistry. Le pipeline doit entrer dans un cache_context pour chaque appel au transformer, attacher step_index, sigma et num_inference_steps, et utiliser des noms de contexte distincts pour les trajectoires indépendantes telles que le guidage conditionnel et non conditionnel. Transmettre un raw_vision_callback qui renvoie les latents visuels formant la sortie générée lorsqu’aucun adaptateur intégré n’est disponible. Valider la qualité de la sortie et ajuster les paramètres de cache pour chaque modèle et planificateur ; le support et les résultats des benchmarks ne se transmettent pas automatiquement depuis Cosmos 3.

Cosmos 3

SeaCache est désactivé par défaut. Activez-le sur le transformer ; la boucle de débruitage de Cosmos 3 associe l'étape active du planificateur, la valeur sigma et le nombre d'étapes à chaque appel cache_context, de sorte qu'aucun câblage supplémentaire n'est nécessaire :

from diffusers import Cosmos3OmniPipeline, SeaCacheConfig

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

Cette API au niveau du modèle fonctionne avec Cosmos3OmniPipeline, Cosmos3OmniModularPipeline et Cosmos3DistilledModularPipeline. SeaCache est une optimisation approximative et peut modifier les sorties générées. Appelez pipe.transformer.disable_cache() lorsque vous avez besoin que chaque étape de débruitage exécute le transformeur complet.

FirstBlockCache

FirstBlock Cache mesure dans quelle mesure les premières couches du débruiteur évoluent d'un pas de temps à l'autre. Si la variation est faible, le modèle saute les couches ultérieures coûteuses et réutilise la sortie précédente.

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 Cache

TaylorSeer Cache accélère l’inférence des modèles de diffusion en utilisant des développements en série de Taylor pour approximer et mettre en cache les activations intermédiaires entre les étapes de débruitage. La méthode prédit les sorties futures à partir des calculs passés et les réutilise à des intervalles spécifiés afin de réduire les calculs redondants.

Ce mécanisme de mise en cache offre d'excellents résultats avec un surcoût mémoire minimal. Pour une analyse détaillée des performances, consultez nos résultats ici.

Pour activer TaylorSeer Cache, créez un TaylorSeerCacheConfig et passez-le au transformer de votre pipeline :

• cache_interval : Nombre d'étapes pour réutiliser les sorties mises en cache avant d'effectuer un passage avant complet
• disable_cache_before_step : Étapes initiales qui utilisent des calculs complets pour collecter les données nécessaires aux approximations
• max_order : précision de l'approximation (en théorie, des valeurs plus élevées améliorent la qualité mais augmentent l'utilisation de la mémoire, nous recommandons de le définir sur 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 accélère l'inférence en sautant des blocs de transformeurs en fonction de l'amplitude de la mise à jour résiduelle. Il observe que l'amplitude des mises à jour (Output - Input) diminue de manière prévisible au fil du processus de diffusion. En accumulant un « budget d'erreur » basé sur des rapports d'amplitude précalculés, il décide dynamiquement quand sauter le calcul et réutiliser le résidu précédent.

MagCache s’appuie sur les ratios d’amplitude (mag_ratios), qui décrivent cette courbe de décroissance. Ces ratios sont spécifiques au point de contrôle du modèle et à l’ordonnanceur. Les FLUX_MAG_RATIOS fournis ont été mesurés sur FLUX.1 ; d’autres modèles, y compris Flux2 et Flux2 Klein, nécessitent leur propre exécution de calibration.

Pour utiliser MagCache, vous suivez généralement un processus en deux étapes : Calibration et Inference.

1. Calibration : exécutez une inférence avec calibrate=True. L’hook mesurera les amplitudes des résidus et affichera les ratios calculés dans la console.
2. Inférence : Transmettez ces ratios à MagCacheConfig pour activer l’accélération.

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 représente la courbe de décroissance intrinsèque de l’amplitude du modèle. Les ratios calibrés pour un nombre élevé d’étapes (par exemple, 50) peuvent être réutilisés pour un nombre d’étapes inférieur (par exemple, 20). L’implémentation utilise l’interpolation pour adapter la courbe au nombre actuel d’étapes d’inférence.

Pour les pipelines qui exécutent le Classifier-Free Guidance de manière séquentielle (comme Kandinsky 5.0), le journal de calibration peut afficher deux tableaux : un pour le passage Conditionnel et un pour le passage Inconditionnel. Dans la plupart des cas, vous devez utiliser le premier tableau (Conditionnel).

Pour les pipelines qui exécutent le Classifier-Free Guidance de manière batched (comme SDXL ou Flux), les hidden_states traités par le modèle contiennent à la fois les branches conditionnelles et non conditionnelles concaténées. Le processus de calibration en tient automatiquement compte, produisant un tableau unique de ratios qui représente le comportement conjoint. Vous pouvez utiliser ce tableau résultant directement, sans modification.