diffusers/main/fr/optimization/cache.md

Mise en cache

Le cache réutilise les sorties des couches intermédiaires entre les étapes de débruitage pour accélérer l'inférence. Il utilise plus de mémoire et ne nécessite pas d'entraînement. Activez une méthode avec une configuration sur un transformer qui prend en charge le cache.

Choisir une méthode de cache

Choisissez une méthode en fonction de la quantité de configuration que vous allez définir et de la précision requise. Chaque méthode équilibre la vitesse, la mémoire et la fidélité des sorties par rapport à une exécution non mise en cache de manière différente. Comparez les sorties avec et sans cache sur vos propres prompts avant de vous fier à une méthode.

Méthode	À utiliser lorsque	Notes
Cache KV de texte	Image NucleusMoE, réutilisation des projections K/V du texte entre les étapes	Actuellement, les blocs d’images NucleusMoE sont interceptés
SeaCache	Génération de vidéos Cosmos 3	Les paramètres ne sont souvent pas transférables d'un modèle à l'autre
FirstBlockCache	Voulez-vous un seul réglage principal de vitesse/qualité sur un transformateur enregistré	Un threshold contrôle la fréquence à laquelle les étapes sont ignorées
MagCache	Avez-vous des rapports d'amplitude pour votre point de contrôle et votre planificateur, ou allez-vous d'abord effectuer une calibration	Les ratios sont spécifiques au point de contrôle et au planificateur
TaylorSeer	Souhaitez-vous prédire des activations ultérieures à partir d'étapes antérieures	Un max_order plus élevé utilise plus de mémoire
PAB	Vidéo, prêt à ajuster la réutilisation de l'attention (intervalles de blocs et de pas de temps à sauter par type d'attention)	Les plages à ignorer nécessitent un réglage spécifique à chaque modèle
FasterCache	Comme PAB, avec en plus le saut optionnel de la branche CFG	Expérimental

Diffusion de l'attention pyramidale

Pyramid Attention Broadcast (PAB) approxime l'attention sur les étapes de débruitage en réutilisant les sorties d'attention pour certains blocs et certaines étapes temporelles, plutôt que de tout recalculer à chaque étape. La configuration distingue les types d'attention (spatiale, temporelle, croisée) lorsque le modèle les possède. Tous les modèles vidéo n'exposent pas les trois types ; définissez uniquement les plages correspondant aux blocs dont vous disposez.

Chaque type utilise un *_attention_block_skip_range (fréquence de recalcul par rapport à la réutilisation au sein de la fenêtre) et un *_attention_timestep_skip_range (étapes de débruitage pouvant être ignorées). Vous devez passer current_timestep_callback afin que l’hook puisse lire l’étape actuelle du pipeline. Des sauts plus larges ou plus agressifs signifient généralement plus de vitesse et un risque de qualité accru.

Passez un PyramidAttentionBroadcastConfig pour l’activer.

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 met en cache et réutilise les caractéristiques d’attention de manière similaire à PAB. Il peut également ignorer la branche inconditionnelle sous guidage sans classifieur et l’estimer à partir de la branche conditionnelle lorsque les latents successifs sont suffisamment redondants.

Passez un FasterCacheConfig pour l’activer. Comme pour PAB, définissez *_attention_block_skip_range et *_attention_timestep_skip_range pour les types d’attention dont vous disposez, ainsi que les options de saut de branche CFG si vous les souhaitez.

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 compare les indicateurs Spectral Evolution Aware (SEA) entre les étapes successives de débruitage. Lorsque la variation accumulée reste sous un seuil, il passe outre la pile de blocs transformer et prédit la sortie à partir des résidus mis en cache. La méthode est approximative et conçue pour la génération de vidéos.

Les adaptateurs intégrés à SeaCache incluent :

• Cosmos 3 est l'intégration principale, optimisée et testée.
• Wan T2V utilise le chemin générique de blocs répétés comme démonstration de la manière de fournir les latents visuels bruts à SeaCache. Les mêmes paramètres de cache ne sont peut-être pas transférables à Wan ou à d'autres variantes de Wan.

Activez SeaCache sur le transformer. La boucle de débruitage de Cosmos 3 associe l’étape du planificateur, la valeur sigma et le nombre d’étapes à chaque cache_context, de sorte qu’aucun paramètre 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))

SeaCache peut modifier les sorties. Appelez pipe.transformer.disable_cache() lorsque vous avez besoin que chaque étape exécute le transformeur complet. L’appel d’activation identique fonctionne avec Cosmos3OmniPipeline, Cosmos3OmniModularPipeline et Cosmos3DistilledModularPipeline.

Pour utiliser SeaCache avec un autre transformeur vidéo, suivez Ajouter la mise en cache à un nouveau modèle. SeaCache nécessite également step_index, sigma et num_inference_steps dans chaque cache_context, ainsi qu'un raw_vision_callback lorsqu'aucun adaptateur intégré n'existe. Ajustez les paramètres en fonction du modèle et du planificateur.

FirstBlockCache

FirstBlockCacheConfig 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 passe outre les couches ultérieures coûteuses et réutilise la sortie précédente.

Activez-le via enable_cache afin que disable_cache et is_cache_enabled restent synchronisés. La valeur par défaut de threshold est 0.05. Une valeur plus élevée, comme 0.2, saute plus souvent pour un gain de vitesse supplémentaire, mais la qualité de la génération peut en souffrir.

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

TaylorSeer Cache

TaylorSeer Cache accélère l'inférence des modèles de diffusion grâce à des développements en série de Taylor sur les étapes de débruitage. Il prédit les activations des étapes ultérieures à partir des étapes antérieures et réutilise ces prédictions pendant plusieurs étapes, ce qui réduit le travail complet effectué par le transformer.

• 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 : Des ordres de Taylor plus élevés peuvent être plus précis mais utilisent plus de mémoire. Conservez cette valeur à 1 sauf si vous avez une raison de la modifier.

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 saute des blocs de transformeur en fonction de l'amplitude de la mise à jour résiduelle. Les amplitudes de mise à jour décroissent de manière prévisible au fil du débruitage, et MagCache suit un budget d'erreur à partir de rapports d'amplitude précalculés (mag_ratios) pour déterminer quand la réutilisation est sûre. Ces rapports sont spécifiques au point de contrôle et au planificateur. Les rapports issus d'un grand nombre d'étapes peuvent être interpolés vers un nombre d'étapes inférieur.

Les FLUX_MAG_RATIOS inclus 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.

MagCache suit deux étapes :

1. Calibration : exécutez une inférence avec calibrate=True. Le hook mesure les amplitudes des résidus et affiche les ratios calculés.
2. Inférence : désactivez le cache de calibration, puis transmettez ces ratios à MagCacheConfig pour accélérer le processus.

Le guidage sans classifieur peut affecter l’étalonnage. Les pipelines qui utilisent le CFG réel avec des contextes séquentiels, comme Flux lorsque true_cfg_scale > 1, entrent séparément dans cache_context("cond") et cache_context("uncond"). L’étalonnage peut afficher un tableau par contexte, mais vous devriez utiliser le tableau conditionnel dans la plupart des cas. Les pipelines qui appliquent le CFG par lot en concaténant les entrées conditionnelles et non conditionnelles (par exemple, CogVideoX) produisent un tableau unique et conjoint que vous pouvez utiliser directement.

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]

Cache KV de texte

TextKVCacheConfig calcule les projections de clés et de valeurs du texte une seule fois et les réutilise à chaque étape de débruitage, au lieu de les recalculer. apply_text_kv_cache() ne cible actuellement que NucleusMoEImageTransformerBlock. Activez-le avec 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]

Ajouter la mise en cache à un nouveau modèle

Les méthodes de cache s’attachent à un modèle via des hooks, et ces hooks dépendent de la classe du modèle, de ses blocs transformer et de la boucle de débruitage du pipeline. Configurez les trois éléments avant d’appeler enable_cache sur un nouveau modèle.

1. Hériter de CacheMixin dans la classe transformer. Cela ajoute enable_cache, disable_cache et cache_context.

2. Enregistrez la disposition de la sortie du bloc de transformer dans _register_transformer_blocks_metadata dans hooks/_helpers.py. Les méthodes qui ignorent ou réutilisent des blocs entiers, telles que FirstBlockCache, MagCache et SeaCache, l'utilisent pour localiser les états cachés dans les sorties du bloc.

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

3. Enveloppez chaque appel de transformer dans la boucle de débruitage du pipeline avec cache_context. Utilisez des noms distincts, tels que "cond" et "uncond", lorsque le classifier-free guidance exécute les branches conditionnelle et non conditionnelle en tant qu'appels séparés, afin que chaque branche conserve son propre état de cache.

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

PAB et FasterCache ont également besoin de current_timestep_callback afin que les hooks puissent lire le timestep actuel. Vérifiez la qualité de la sortie sur le nouveau modèle, car les réglages de cache optimisés pour un modèle ne se transmettent souvent pas à un autre.