Caché

El almacenamiento en caché acelera la inferencia al guardar y reutilizar las salidas intermedias de diferentes capas, como las capas de atención y las capas de avance, en lugar de realizar toda la computación en cada paso de inferencia. Mejora significativamente la velocidad de generación a costa de más memoria y no requiere entrenamiento adicional.

Esta guía te muestra cómo usar los métodos de caché admitidos en Diffusers.

Pyramid Attention Broadcast

Pyramid Attention Broadcast (PAB) se basa en la observación de que las salidas de atención no difieren tanto entre los pasos de tiempo sucesivos del proceso de generación. Las diferencias de atención son menores en las capas de atención cruzada y, en general, se almacenan en caché durante un rango de pasos de tiempo más amplio. A esto le siguen las capas de atención temporal y espacial.

¡No todos los modelos de vídeo tienen tres tipos de atención (cruzada, temporal y espacial)!

PAB puede combinarse con otras técnicas, como el paralelismo de secuencias y el paralelismo de guía sin clasificador (paralelismo de datos), para la generación de video en tiempo casi real.

Configura y pasa un PyramidAttentionBroadcastConfig al transformer de un pipeline para habilitarlo. El parámetro spatial_attention_block_skip_range controla la frecuencia con la que se omiten los cálculos de atención en los bloques de atención espacial, y spatial_attention_timestep_skip_range es el rango de pasos temporales a omitir. Ten cuidado al elegir un rango apropiado, ya que un intervalo más pequeño puede provocar velocidades de inferencia más lentas y un intervalo más grande puede resultar en una menor calidad de generación.

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 almacena y reutiliza características de atención de manera similar a PAB, ya que las diferencias en la salida son pequeñas para cada paso de tiempo sucesivo.

Este método también puede optar por omitir la predicción de la rama incondicional, al usar la guía sin clasificador para el muestreo (común en la mayoría de los modelos base), y estimarla a partir de la predicción de la rama condicional si existe una redundancia significativa en las salidas latentes predichas entre los pasos de tiempo sucesivos.

Configura y pasa un FasterCacheConfig al transformer de un pipeline para habilitarlo.

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 compara los indicadores Spectral-Evolution-Aware (SEA) entre los pasos de eliminación de ruido sucesivos. Cuando el cambio acumulado del indicador permanece por debajo de un umbral, omite la pila de bloques de transformador costosa y predice su salida a partir de residuos en caché. Construya el indicador a partir de los latentes visuales que forman la salida generada. Incluya los marcos de condicionamiento limpios cuando formen parte de esa trayectoria de salida, como en la generación de imagen a video y de video a video. Excluya las pistas visuales separadas que condicionan la generación pero no forman parte de la salida. El condicionamiento de texto se excluye porque no es un latente visual.

Los paquetes de transferencia de Cosmos 3 tratan las pistas de control como secuencias visuales independientes, por lo que su adaptador las excluye del indicador. Las ramas de Control-CFG comparan la misma trayectoria de salida mientras conservan sus propios residuos en caché. Las pistas de control siguen condicionando al transformer.

La implementación proporciona adaptadores integrados para los siguientes modelos:

• Cosmos 3 es la integración principal optimizada y con pruebas de rendimiento. Almacena en caché la pila completa del decodificador a través de un límite de normalización posterior.
• Wan T2V utiliza la ruta de bloques repetidos genéricos en modo eager. Esta integración demuestra cómo otro transformador de video de un solo flujo puede proporcionar latentes visuales crudos a SeaCache; no es una afirmación de que los mismos parámetros de caché sean óptimos para Wan ni de que se admitan otras variantes de Wan.

Otros transformers de vídeo pueden integrarse con la ruta genérica cuando utilizan CacheMixin, exponen una lista de bloques repetidos reconocida y registran la disposición de entrada/salida de los bloques en TransformerBlockRegistry. El pipeline debe entrar en un cache_context para cada llamada al transformer, adjuntar step_index, sigma y num_inference_steps, y utilizar nombres de contexto separados para trayectorias independientes, como la guía condicional y la incondicional. Pasa un raw_vision_callback que devuelva los latentes visuales que forman la salida generada cuando no hay un adaptador integrado disponible. Valida la calidad de la salida y ajusta los parámetros de caché para cada modelo y programador; el soporte y los resultados de las pruebas de rendimiento no se transfieren automáticamente desde Cosmos 3.

Cosmos 3

SeaCache está desactivado de forma predeterminada. Actívalo en el transformer; el bucle de desruido de Cosmos 3 adjunta el paso activo del programador, sigma y el recuento de pasos a cada llamada de cache_context, por lo que no se necesita ninguna configuración adicional:

from diffusers import Cosmos3OmniPipeline, SeaCacheConfig

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

Esta API a nivel de modelo funciona con Cosmos3OmniPipeline, Cosmos3OmniModularPipeline y Cosmos3DistilledModularPipeline. SeaCache es una optimización aproximada y puede alterar las salidas generadas. Llama a pipe.transformer.disable_cache() cuando necesites que cada paso de desruido ejecute el transformador completo.

FirstBlockCache

FirstBlock Cache 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.

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 acelera la inferencia de difusión mediante el uso de expansiones en serie de Taylor para aproximar y almacenar en caché las activaciones intermedias entre los pasos de desruido. El método predice las salidas futuras basándose en cálculos anteriores y las reutiliza en intervalos especificados para reducir cálculos redundantes.

Este mecanismo de caché ofrece resultados sólidos con un sobrecoste de memoria adicional mínimo. Para un análisis detallado del rendimiento, consulte nuestros hallazgos aquí.

Para habilitar TaylorSeer Cache, cree una TaylorSeerCacheConfig y pásela al transformer de su pipeline:

• 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: Precisión de la aproximación (en teoría, valores más altos mejoran la calidad pero aumentan el uso de memoria, aunque recomendamos configurarlo en 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 acelera la inferencia omitiendo bloques de transformador en función de la magnitud de la actualización residual. Observa que la magnitud de las actualizaciones (Output - Input) decae de forma predecible a lo largo del proceso de difusión. Al acumular un “presupuesto de error” basado en relaciones de magnitud precalculadas, decide dinámicamente cuándo omitir el cálculo y reutilizar la residual anterior.

MagCache se basa en Magnitude Ratios (mag_ratios), que describen esta curva de decaimiento. Estas proporciones son específicas del punto de control del modelo y del programador. El FLUX_MAG_RATIOS incluido se midió en FLUX.1; otros modelos, incluidos Flux2 y Flux2 Klein, requieren su propia ejecución de calibración.

Para usar MagCache, normalmente se sigue un proceso de dos pasos: Calibración e Inferencia.

1. Calibración: Ejecuta la inferencia una vez con calibrate=True. El hook medirá las magnitudes residuales e imprimirá las proporciones calculadas en la consola.
2. Inferencia: Pasa estas proporciones a MagCacheConfig para habilitar la aceleración.

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 representa la curva de decaimiento de magnitud intrínseca del modelo. Las proporciones calibradas para un número alto de pasos (p. ej., 50) pueden reutilizarse para un número menor de pasos (p. ej., 20). La implementación utiliza interpolación para mapear la curva al número actual de pasos de inferencia.

Para los pipelines que ejecutan Classifier-Free Guidance de forma secuencial (como Kandinsky 5.0), el registro de calibración podría imprimir dos matrices: una para el paso Conditional y otra para el paso Unconditional. En la mayoría de los casos, debes usar la primera matriz (Conditional).

Para los pipelines que ejecutan Classifier-Free Guidance de manera batched (como SDXL o Flux), los hidden_states procesados por el modelo contienen tanto las ramas condicionales como las incondicionales concatenadas. El proceso de calibración tiene esto en cuenta automáticamente, produciendo un único array de ratios que representa el comportamiento conjunto. Puedes usar este array resultante directamente sin modificación.