CacheDiT

CacheDiT acelera los pipelines de DiT reutilizando las salidas de los bloques de transformer entre los pasos de desruido. No requiere entrenamiento y es compatible con la mayoría de los pipelines DiT de Diffusers, incluidos Flux, Qwen-Image, Wan y HunyuanVideo. Diffusers también cuenta con cache integrada que no requiere una dependencia adicional.

Instale CacheDiT desde PyPI.

pip install -U cache-dit

Llama a cache_dit.supported_pipelines() para listar las familias de pipelines que admite CacheDiT.

import cache_dit

cache_dit.supported_pipelines()

Activar la caché

Llama a cache_dit.enable_cache en un pipeline para cachearlo con la configuración predeterminada y luego ejecuta el pipeline como de costumbre.

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]

CacheDiT también funciona con torch.compile. Compila el transformer después de llamar a enable_cache. Consulta la documentación de compile para conocer las configuraciones que evitan la recompilación con formas de entrada dinámicas.

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

Llama a cache_dit.summary después de la inferencia para registrar cuántos pasos se almacenaron en caché y las diferencias residuales entre pasos.

stats = cache_dit.summary(pipeline)

Llama a cache_dit.disable_cache para restaurar el pipeline original.

cache_dit.disable_cache(pipeline)

Configura la caché

DBCache (Dual Block Cache) calcula los primeros n bloques (Fn) en cada paso. Cuando su salida apenas cambia respecto al paso anterior, reutiliza la salida en caché para los bloques restantes y puede recalcular los últimos n bloques (Bn) para corregirla. El calibrador TaylorSeer predice la salida en caché a partir de pasos anteriores en lugar de reutilizarla tal cual.

enable_cache usa DBCache por defecto, con los primeros 8 bloques siempre calculados (F8B0) y 8 pasos de calentamiento sin caché. Para priorizar la calidad a costa de la velocidad, aumente Fn_compute_blocks o reduzca residual_diff_threshold (valor predeterminado 0.08). Para obtener la mejor calidad con altas tasas de caché, agregue el 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 los pipelines compatibles, CacheDiT ya sabe si CFG se ejecuta como una pasada de inferencia independiente. Para otros modelos, establece enable_separate_cfg=True en DBCacheConfig si el modelo ejecuta las pasadas condicionales y no condicionales por separado, o False si las fusiona o no utiliza CFG.

Consulte la documentación del diseño de DBCache para ver cómo funcionan los bloques Fn y Bn, y los benchmarks de caché para conocer las cifras de velocidad y calidad.

Próximos pasos

• Para los pipelines que CacheDiT aún no admite, consulte la documentación de BlockAdapter.
• CacheDiT también admite paralelismo de contexto y cuantización.