Caché

El almacenamiento en caché reutiliza las salidas de las capas intermedias entre los pasos de eliminación de ruido para acelerar la inferencia. Utiliza más memoria y no requiere entrenamiento. Habilite un método con una configuración en un transformer que admita el almacenamiento en caché.

Elija un método de caché

Elija un método según la cantidad de configuración que establecerá y el ajuste que necesita. Cada método equilibra de manera diferente la velocidad, la memoria y la similitud de las salidas con la ejecución sin caché. Compare las salidas con y sin caché en sus propios prompts antes de confiar en un método.

Pyramid Attention Broadcast

Pyramid Attention Broadcast (PAB) aproxima la atención a lo largo de los pasos de desruido reutilizando las salidas de atención para algunos bloques y pasos de tiempo, en lugar de recalcularlos en cada paso. Config separa los tipos de atención (espacial, temporal, cruzada) cuando el modelo los tiene. No todos los modelos de vídeo exponen los tres, y solo se deben establecer los rangos que coincidan con los bloques disponibles.

Cada tipo utiliza un *_attention_block_skip_range (con qué frecuencia recalcular o reutilizar dentro de la ventana) y un *_attention_timestep_skip_range (qué pasos de desruido pueden omitirse). Debe pasar current_timestep_callback para que el hook pueda leer el paso actual del pipeline. Omisiones más amplias o agresivas suelen implicar más velocidad y más riesgo de pérdida de calidad.

Pase un PyramidAttentionBroadcastConfig para habilitarlo.

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 almacena y reutiliza características de atención de manera similar a PAB. También puede omitir la rama incondicional bajo la guía libre de clasificador y estimarla a partir de la rama condicional cuando los latentes sucesivos son lo suficientemente redundantes.

Pase una FasterCacheConfig para habilitarla. Al igual que con PAB, configure *_attention_block_skip_range y *_attention_timestep_skip_range para los tipos de atención que tenga, además de las opciones de omisión de la rama CFG cuando las desee.

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 compara los indicadores de evolución espectral consciente (SEA) entre los pasos de eliminación de ruido sucesivos. Cuando el cambio acumulado se mantiene por debajo de un umbral, omite la pila de bloques del transformador y predice la salida a partir de los residuos en caché. El método es aproximado y está diseñado para la generación de video.

Los adaptadores integrados para SeaCache incluyen:

• Cosmos 3 es la integración principal optimizada y con pruebas de rendimiento.
• Wan T2V utiliza la ruta de bloques repetidos genéricos como demostración de cómo proporcionar los latentes visuales crudos a SeaCache. Es posible que los mismos parámetros de caché no sean transferibles a Wan u otras variantes de Wan.

Habilite SeaCache en el transformer. El bucle de desruido de Cosmos 3 adjunta el paso del programador, sigma y el recuento de pasos a cada cache_context, por lo que no se necesitan parámetros adicionales.

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 puede alterar las salidas. Llama a pipe.transformer.disable_cache() cuando necesites que cada paso ejecute el transformador completo. La misma llamada de activación funciona con Cosmos3OmniPipeline, Cosmos3OmniModularPipeline y Cosmos3DistilledModularPipeline.

Para usar SeaCache con otro video transformer, siga Add caching to a new model. SeaCache también necesita step_index, sigma y num_inference_steps en cada cache_context, y un raw_vision_callback cuando no exista un adaptador integrado. Ajuste los parámetros según el modelo y el programador.

FirstBlockCache

FirstBlockCacheConfig comprueba cuánto cambian las primeras capas del denoiser de un paso de tiempo al siguiente. Si el cambio es pequeño, el modelo omite las capas posteriores más costosas y reutiliza la salida anterior.

Actívalo mediante enable_cache para que disable_cache y is_cache_enabled se mantengan sincronizados. El valor predeterminado de threshold es 0.05. Un valor más alto, como 0.2, omite con más frecuencia para ganar velocidad, pero la calidad de la generación puede disminuir.

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 acelera la inferencia de difusión mediante expansiones en serie de Taylor a lo largo de los pasos de desruido. Predice las activaciones de pasos posteriores a partir de los anteriores y reutiliza esas predicciones durante varios pasos, de modo que el transformer realiza menos trabajo completo.

• cache_interval: Número de pasos para reutilizar las salidas en caché antes de realizar una pasada completa hacia adelante
• disable_cache_before_step: Pasos iniciales que utilizan cálculos completos para recopilar datos para las aproximaciones
• max_order: Los órdenes de Taylor más altos pueden ser más precisos, pero usan más memoria. Mantenga este valor en 1 a menos que tenga una razón para cambiarlo.

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 omite bloques del transformador según la magnitud de la actualización residual. Las magnitudes de actualización decaen de forma predecible a lo largo del desruido, y MagCache sigue un presupuesto de error a partir de relaciones de magnitudes precalculadas (mag_ratios) para decidir cuándo es seguro reutilizar. Esas relaciones son específicas del checkpoint y del programador. Las relaciones de un número alto de pasos pueden interpolarse hacia un número menor de pasos.

Las FLUX_MAG_RATIOS incluidas se midieron en FLUX.1. Otros modelos, incluidos Flux2 y Flux2 Klein, requieren su propia ejecución de calibración.

MagCache sigue dos pasos:

1. Calibración: Ejecuta la inferencia una vez con calibrate=True. El hook mide las magnitudes residuales e imprime las proporciones calculadas.
2. Inferencia: deshabilita la caché de calibración y luego pasa esas proporciones a MagCacheConfig para la aceleración.

La guía sin clasificador puede afectar la calibración. Las tuberías que usan CFG real con contextos secuenciales, como Flux cuando true_cfg_scale > 1, entran en cache_context("cond") y cache_context("uncond") por separado. La calibración puede imprimir un array por contexto, pero en la mayoría de los casos debes usar el array condicional. Las tuberías que agrupan CFG concatenando entradas condicionales y no condicionales (por ejemplo, CogVideoX) producen un array conjunto único que puedes usar directamente.

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]

Caché de KV de texto

TextKVCacheConfig calcula las proyecciones de clave y valor del texto una sola vez y las reutiliza en cada paso de desruido, en lugar de recalcularlas. apply_text_kv_cache() actualmente solo aplica a NucleusMoEImageTransformerBlock. Actívalo con 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]

Agregar caché a un nuevo modelo

Los métodos de caché se adjuntan a un modelo a través de hooks, y estos hooks dependen de la clase del modelo, sus bloques transformer y el bucle de desruido del pipeline. Configure los tres antes de llamar a enable_cache en un modelo nuevo.

1. Hereda de CacheMixin en la clase transformer. Esto añade enable_cache, disable_cache y cache_context.

2. Registre el diseño de salida del bloque del transformer en _register_transformer_blocks_metadata en hooks/_helpers.py. Los métodos que omiten o reutilizan bloques completos, como FirstBlockCache, MagCache y SeaCache, lo utilizan para localizar los estados ocultos en las salidas del bloque.

   Copiado

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

3. Envuelva cada llamada al transformer en el bucle de desruido del pipeline con cache_context. Utilice nombres separados, como "cond" y "uncond", cuando la guía libre de clasificador ejecute las ramas condicional y incondicional como llamadas separadas, para que cada rama conserve su propio estado de caché.

   Copiado

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

PAB y FasterCache también necesitan current_timestep_callback para que los hooks puedan leer el timestep actual. Comprueba la calidad de la salida en el nuevo modelo, ya que las configuraciones de caché ajustadas para un modelo a menudo no se transfieren a otro.