CacheDiT

CacheDiT acelera os pipelines DiT reutilizando as saídas dos blocos de transformer entre as etapas de desruido. Ele não requer treinamento e suporta a maioria dos pipelines DiT do Diffusers, incluindo Flux, Qwen-Image, Wan e HunyuanVideo. O Diffusers também possui cache integrado, que não exige uma dependência extra.

Instale o CacheDiT do PyPI.

pip install -U cache-dit

Chame cache_dit.supported_pipelines() para listar as famílias de pipelines que o CacheDiT suporta.

import cache_dit

cache_dit.supported_pipelines()

Ativar o cache

Chame cache_dit.enable_cache em um pipeline para ativá-lo com as configurações padrão e, em seguida, execute o pipeline normalmente.

import torch
import cache_dit
from diffusers import FluxPipeline

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", dtype=torch.bfloat16
).to("cuda")
cache_dit.enable_cache(pipeline)

image = pipeline(
    "A cat holding a sign that says hello world", num_inference_steps=28
).images[0]

O CacheDiT também funciona com torch.compile. Compile o transformer depois de chamar enable_cache. Consulte a documentação de compile para configurações que evitam recompilação com formas de entrada dinâmicas.

pipeline.transformer = torch.compile(pipeline.transformer)

Chame cache_dit.summary após a inferência para registrar quantas etapas foram armazenadas em cache e as diferenças residuais entre as etapas.

stats = cache_dit.summary(pipeline)

Chame cache_dit.disable_cache para restaurar o pipeline original.

cache_dit.disable_cache(pipeline)

Configure o cache

O DBCache (Dual Block Cache) calcula os primeiros n blocos (Fn) a cada etapa. Quando a sua saída muda pouco em relação à etapa anterior, ele reutiliza a saída em cache para os blocos restantes e pode recalcular os últimos n blocos (Bn) para corrigi-la. O calibrador TaylorSeer prevê a saída em cache a partir de etapas anteriores, em vez de reutilizá-la como está.

enable_cache usa DBCache por padrão, com os primeiros 8 blocos sempre calculados (F8B0) e 8 passos de aquecimento sem cache. Para trocar velocidade por qualidade, aumente Fn_compute_blocks ou diminua residual_diff_threshold (padrão 0.08). Para a melhor qualidade em altas taxas de cache, adicione o calibrador TaylorSeer.

from cache_dit import DBCacheConfig, TaylorSeerCalibratorConfig

cache_dit.enable_cache(
    pipeline,
    cache_config=DBCacheConfig(
        max_warmup_steps=8,
        Fn_compute_blocks=8,
        Bn_compute_blocks=0,
        residual_diff_threshold=0.12,
    ),
    calibrator_config=TaylorSeerCalibratorConfig(taylorseer_order=1),
)

Para os pipelines suportados, o CacheDiT já sabe se o CFG é executado como uma passagem de forward separada. Para outros modelos, defina enable_separate_cfg=True em DBCacheConfig se o modelo executar as passagens condicional e incondicional separadamente, ou False se ele as fundir ou não usar CFG.

Consulte a documentação do design do DBCache para entender como os blocos Fn e Bn funcionam, e os benchmarks de cache para números de velocidade e qualidade.

Próximas etapas

• Para pipelines que o CacheDiT ainda não suporta, consulte a documentação do BlockAdapter.
• O CacheDiT também oferece suporte a paralelismo de contexto e quantização.