Cache

O cache acelera a inferência armazenando e reutilizando saídas intermediárias de diferentes camadas, como as camadas de atenção e feedforward, em vez de executar todo o cálculo em cada etapa de inferência. Ele melhora significativamente a velocidade de geração à custa de mais memória e não requer treinamento adicional.

Este guia mostra como usar os métodos de cache suportados no Diffusers.

Pyramid Attention Broadcast

Pyramid Attention Broadcast (PAB) baseia-se na observação de que as saídas de atenção não diferem muito entre os timesteps sucessivos do processo de geração. As diferenças de atenção são menores nas camadas de atenção cruzada e, em geral, são armazenadas em cache por um intervalo de timesteps mais longo. Isso é seguido pelas camadas de atenção temporal e atenção espacial.

Nem todos os modelos de vídeo possuem os três tipos de atenção (cross, temporal e espacial)!

O PAB pode ser combinado com outras técnicas, como paralelismo de sequência e paralelismo de orientação sem classificador (paralelismo de dados), para geração de vídeo em tempo quase real.

Configure e passe um PyramidAttentionBroadcastConfig para o transformer de um pipeline para ativá-lo. O spatial_attention_block_skip_range controla a frequência com que os cálculos de atenção são ignorados nos blocos de atenção espacial, e o spatial_attention_timestep_skip_range é o intervalo de timesteps a serem ignorados. Tenha cuidado ao escolher um intervalo apropriado, pois um intervalo menor pode levar a velocidades de inferência mais lentas e um intervalo maior pode resultar em menor qualidade de geração.

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 faz o cache e reutiliza os recursos de atenção de forma semelhante ao PAB, pois as diferenças de saída são pequenas para cada passo de tempo sucessivo.

Este método também pode optar por pular a previsão de ramo incondicional, ao usar a orientação sem classificador para amostragem (comum na maioria dos modelos base), e estimá-la a partir da previsão de ramo condicional, caso haja redundância significativa nas saídas latentes previstas entre os passos de tempo sucessivos.

Configure e passe um FasterCacheConfig para o transformer de um pipeline para ativá-lo.

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 os indicadores Spectral-Evolution-Aware (SEA) entre as etapas sucessivas de remoção de ruído. Quando a mudança acumulada do indicador permanece abaixo de um limiar, ele ignora a pilha de blocos de transformer de alto custo e prevê sua saída a partir de resíduos em cache. Construa o indicador a partir dos latentes visuais que formam a saída gerada. Inclua os quadros de condicionamento limpos quando eles fazem parte dessa trajetória de saída, como na geração de imagem para vídeo e de vídeo para vídeo. Exclua dicas visuais separadas que condicionam a geração, mas que não fazem parte da saída. O condicionamento de texto é excluído porque não é um latente visual.

Os pacotes de transferência do Cosmos 3 tratam as dicas de controle como sequências visuais separadas, de modo que seu adaptador as exclui do indicador. Os ramos do Control-CFG comparam a mesma trajetória de saída, mas mantêm seus próprios resíduos em cache. As dicas de controle ainda condicionam o transformer.

A implementação fornece adaptadores integrados para os seguintes modelos:

• Cosmos 3 é a integração principal otimizada e benchmarkada. Ele faz o cache da pilha completa do decodificador por meio de uma fronteira de pós-normalização.
• Wan T2V usa o caminho de bloco repetido genérico em modo eager. Essa integração demonstra como outro transformer de vídeo de fluxo único pode fornecer latentes de visão brutos ao SeaCache; não é uma afirmação de que os mesmos parâmetros de cache sejam ótimos para o Wan ou de que outras variantes do Wan sejam suportadas.

Outros transformers de vídeo podem se integrar ao caminho genérico quando utilizam CacheMixin, expõem uma lista de blocos repetidos reconhecida e registram o layout de entrada/saída dos blocos em TransformerBlockRegistry. O pipeline deve entrar em um cache_context para cada chamada do transformer, anexar step_index, sigma e num_inference_steps, e usar nomes de contexto separados para trajetórias independentes, como o direcionamento condicional e incondicional. Passe um raw_vision_callback que retorne os latentes visuais que formam a saída gerada quando nenhum adaptador embutido estiver disponível. Valide a qualidade da saída e ajuste os parâmetros de cache para cada modelo e agendador; o suporte e os resultados de benchmark não são transferidos automaticamente do Cosmos 3.

Cosmos 3

SeaCache é desativado por padrão. Ative-o no transformer; o loop de desruidificação do Cosmos 3 anexa o passo ativo do agendador, sigma e a contagem de passos a cada chamada de cache_context, portanto, nenhum cabeamento extra é necessário:

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 de nível de modelo funciona com Cosmos3OmniPipeline, Cosmos3OmniModularPipeline e Cosmos3DistilledModularPipeline. O SeaCache é uma otimização aproximada e pode alterar as saídas geradas. Chame pipe.transformer.disable_cache() quando for necessário que cada etapa de desruidamento execute o transformador completo.

FirstBlockCache

FirstBlock Cache verifica o quanto as camadas iniciais do denoiser mudam de um timestep para o próximo. Se a mudança for pequena, o modelo ignora as camadas posteriores, mais custosas, e reutiliza a saída 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))

Cache TaylorSeer

TaylorSeer Cache acelera a inferência de difusão usando expansões em série de Taylor para aproximar e armazenar em cache as ativações intermediárias entre as etapas de desruido. O método prevê saídas futuras com base em cálculos passados, reutilizando-as em intervalos especificados para reduzir cálculos redundantes.

Esse mecanismo de cache entrega resultados robustos com um sobrecarga de memória adicional mínima. Para uma análise detalhada de desempenho, consulte nossas descobertas aqui.

Para habilitar o TaylorSeer Cache, crie um TaylorSeerCacheConfig e passe-o para o transformer do seu pipeline:

• cache_interval: Número de etapas para reutilizar saídas em cache antes de executar uma passagem completa
• disable_cache_before_step: Etapas iniciais que usam computações completas para coletar dados para aproximações
• max_order: Precisão da aproximação (em teoria, valores maiores melhoram a qualidade, mas aumentam o uso de memória; recomendamos definir como 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 a inferência ao pular blocos de transformador com base na magnitude da atualização residual. Ele observa que a magnitude das atualizações (Output - Input) decai de forma previsível ao longo do processo de difusão. Ao acumular um “orçamento de erro” com base em razões de magnitude pré-computadas, ele decide dinamicamente quando pular o cálculo e reutilizar a residual anterior.

MagCache depende de Magnitude Ratios (mag_ratios), que descrevem essa curva de decaimento. Essas proporções são específicas para o checkpoint do modelo e o scheduler. O FLUX_MAG_RATIOS incluído foi medido no FLUX.1; outros modelos, incluindo Flux2 e Flux2 Klein, precisam de sua própria execução de calibração.

Para usar o MagCache, você geralmente segue um processo de duas etapas: Calibração e Inferência.

1. Calibração: execute a inferência uma vez com calibrate=True. O hook medirá as magnitudes dos resíduos e imprimirá as proporções calculadas no console.
2. Inferência: Passe essas proporções para MagCacheConfig para habilitar a aceleração.

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 representam a curva de decaimento de magnitude intrínseca do modelo. Proporções calibradas para um número alto de etapas (por exemplo, 50) podem ser reutilizadas para contagens de etapas menores (por exemplo, 20). A implementação usa interpolação para mapear a curva para o número atual de etapas de inferência.

Para pipelines que executam Classifier-Free Guidance de forma sequencial (como o Kandinsky 5.0), o log de calibração pode imprimir dois arrays: um para a passagem Conditional e outro para a passagem Unconditional. Na maioria dos casos, você deve usar o primeiro array (Conditional).

Para pipelines que executam Classifier-Free Guidance de forma em lote (como SDXL ou Flux), os hidden_states processados pelo modelo contêm tanto os ramos condicionais quanto os incondicionais concatenados. O processo de calibração leva isso automaticamente em consideração, produzindo um único array de proporções que representa o comportamento conjunto. Você pode usar esse array resultante diretamente, sem modificações.