Cache

O cache reutiliza as saídas das camadas intermediárias entre as etapas de remoção de ruído para acelerar a inferência. Ele usa mais memória e não requer treinamento. Ative um método com uma configuração em um transformer que suporte cache.

Escolha um método de cache

Escolha um método com base na quantidade de configuração que você definirá e no ajuste necessário. Cada método equilibra velocidade, memória e o quão de perto as saídas correspondem à execução sem cache de maneira diferente. Compare as saídas com e sem o cache nos seus próprios prompts antes de confiar em um método.

Pyramid Attention Broadcast

Pyramid Attention Broadcast (PAB) aproxima a atenção entre as etapas de desruidamento reutilizando as saídas de atenção para alguns blocos e timesteps, em vez de recalcular a cada etapa. O Config separa os tipos de atenção (espacial, temporal, cruzada) quando o modelo os possui. Nem todo modelo de vídeo expõe os três tipos, e defina apenas os intervalos que correspondem aos blocos que você possui.

Cada tipo usa um *_attention_block_skip_range (com que frequência recalcular versus reutilizar dentro da janela) e um *_attention_timestep_skip_range (quais timesteps de desruidamento podem ser pulados). Você deve passar current_timestep_callback para que o hook possa ler o timestep atual do pipeline. Pulos mais amplos ou agressivos geralmente significam mais velocidade e mais risco de qualidade.

Passe um PyramidAttentionBroadcastConfig para ativá-lo.

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 faz o cache e reutiliza recursos de atenção semelhantes aos de PAB. Ele também pode pular o ramo incondicional sob a orientação livre de classificador e estimá-lo a partir do ramo condicional quando os latentes sucessivos forem redundantes o suficiente.

Passe um FasterCacheConfig para ativá-lo. Assim como no PAB, defina *_attention_block_skip_range e *_attention_timestep_skip_range para os tipos de atenção que você possui, além das opções de skip do ramo CFG, quando desejá-las.

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 os indicadores Spectral Evolution Aware (SEA) entre as etapas sucessivas de remoção de ruído. Quando a mudança acumulada permanece abaixo de um limite, ele ignora a pilha de blocos do transformer e prevê a saída a partir dos resíduos em cache. O método é aproximado e projetado para geração de vídeo.

Os adaptadores integrados para SeaCache incluem:

• Cosmos 3 é a integração principal otimizada e benchmarked.
• Wan T2V usa o caminho de bloco repetido genérico como demonstração de como fornecer os latentes de visão brutos ao SeaCache. Os mesmos parâmetros de cache podem não se transferir para o Wan ou outras variantes do Wan.

Ative o SeaCache no transformer. O loop de desruidamento do Cosmos 3 anexa o passo do agendador, o sigma e a contagem de passos a cada cache_context, portanto, nenhum parâmetro 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))

SeaCache pode alterar as saídas. Chame pipe.transformer.disable_cache() quando precisar que cada etapa execute o transformador completo. A mesma chamada de habilitação funciona com Cosmos3OmniPipeline, Cosmos3OmniModularPipeline e Cosmos3DistilledModularPipeline.

Para usar o SeaCache com outro video transformer, siga Adicionar cache a um novo modelo. O SeaCache também precisa de step_index, sigma e num_inference_steps em cada cache_context, e de um raw_vision_callback quando não existe um adaptador embutido. Ajuste os parâmetros conforme o modelo e o agendador.

FirstBlockCache

FirstBlockCacheConfig 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, que são mais custosas, e reutiliza a saída anterior.

Ative-o por meio de enable_cache para que disable_cache e is_cache_enabled permaneçam sincronizados. O valor padrão de threshold é 0.05. Um valor maior, como 0.2, ignora com mais frequência para ganhar velocidade, mas a qualidade da geração pode diminuir.

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

Cache TaylorSeer

TaylorSeer Cache acelera a inferência de difusão com expansões em série de Taylor ao longo das etapas de desruido. Ele prevê as ativações de etapas posteriores a partir das anteriores e reutiliza essas previsões por várias etapas, de modo que o transformer realize menos trabalho completo.

• 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: Ordens de Taylor mais altas podem ser mais precisas, mas usam mais memória. Mantenha esse valor em 1, a menos que você tenha um motivo para alterá-lo.

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 ignora blocos de transformador com base na magnitude da atualização residual. As magnitudes de atualização decaem de forma previsível ao longo do desruidamento, e o MagCache rastreia um orçamento de erro a partir de razões de magnitude pré-computadas (mag_ratios) para decidir quando a reutilização é segura. Essas razões são específicas do checkpoint e do agendador. Razões de um número alto de passos podem ser interpoladas para um número menor de passos.

Os FLUX_MAG_RATIOS incluídos foram medidos no FLUX.1. Outros modelos, incluindo Flux2 e Flux2 Klein, precisam de sua própria execução de calibração.

O MagCache segue duas etapas:

1. Calibração: execute a inferência uma vez com calibrate=True. O hook mede as magnitudes residuais e imprime as proporções calculadas.
2. Inferência: desative o cache de calibração e, em seguida, passe essas proporções para MagCacheConfig para aceleração.

A orientação sem classificador (classifier-free guidance) pode afetar a calibração. Pipelines que usam CFG verdadeiro com contextos sequenciais, como o Flux quando true_cfg_scale > 1, entram em cache_context("cond") e cache_context("uncond") separadamente. A calibração pode imprimir um array por contexto, mas você deve usar o array condicional na maioria dos casos. Pipelines que agrupam o CFG concatenando entradas condicionais e incondicionais (por exemplo, CogVideoX) produzem um array conjunto único que você pode usar diretamente.

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 de KV de Texto

TextKVCacheConfig calcula as projeções de chave e valor do texto uma única vez e as reutiliza em cada etapa de desruido, em vez de recalculá-las. apply_text_kv_cache() atualmente faz o hook apenas em NucleusMoEImageTransformerBlock. Ative-o com 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]

Adicionar cache a um novo modelo

Os métodos de cache são anexados a um modelo por meio de hooks, e esses hooks dependem da classe do modelo, dos seus blocos de transformer e do loop de desruidamento do pipeline. Configure os três antes de chamar enable_cache em um novo modelo.

1. Herde de CacheMixin na classe do transformer. Isso adiciona enable_cache, disable_cache e cache_context.

2. Registre o layout de saída do bloco do transformer em _register_transformer_blocks_metadata em hooks/_helpers.py. Métodos que pulam ou reutilizam blocos inteiros, como FirstBlockCache, MagCache e SeaCache, usam esse registro para localizar os estados ocultos nas saídas do bloco.

   Copiado

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

3. Envolva cada chamada do transformer no loop de desruidamento do pipeline com cache_context. Use nomes separados, como "cond" e "uncond", quando a classifier-free guidance executa os ramos condicional e incondicional como chamadas separadas, para que cada ramo mantenha seu próprio estado de cache.

   Copiado

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

PAB e FasterCache também precisam de current_timestep_callback para que os hooks possam ler o timestep atual. Verifique a qualidade da saída no novo modelo, pois as configurações de cache ajustadas para um modelo geralmente não se transferem para outro.