CacheDiT

CacheDiT ускоряет пайплайны DiT за счёт повторного использования выходов блоков трансформера на шагах денойзинга. Обучение не требуется. Поддерживаются большинство пайплайнов DiT в Diffusers, включая Flux, Qwen-Image, Wan и HunyuanVideo. В Diffusers также есть встроенный кэширование, которое не требует дополнительных зависимостей.

Установите CacheDiT из PyPI.

pip install -U cache-dit

Вызовите cache_dit.supported_pipelines(), чтобы получить список семейств пайплайнов, поддерживаемых CacheDiT.

import cache_dit

cache_dit.supported_pipelines()

Включить кэширование

Вызовите cache_dit.enable_cache для конвейера, чтобы кэшировать его с настройками по умолчанию, а затем запустите конвейер в обычном режиме.

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 также работает с torch.compile. Компилируйте трансформер после вызова enable_cache. См. документацию по compile для настроек, позволяющих избежать повторной компиляции при динамических формах входных данных.

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

Вызовите cache_dit.summary после инференса, чтобы зафиксировать количество закэшированных шагов и остаточные различия между шагами.

stats = cache_dit.summary(pipeline)

Вызовите cache_dit.disable_cache, чтобы восстановить исходный конвейер.

cache_dit.disable_cache(pipeline)

Настройка кэша

DBCache (Dual Block Cache) вычисляет первые n блоков (Fn) на каждом шаге. Если их выход почти не меняется по сравнению с предыдущим шагом, кэшированный выход переиспользуется для оставшихся блоков, при этом последние n блоков (Bn) могут быть пересчитаны для коррекции. Калибратор TaylorSeer предсказывает кэшированный выход из более ранних шагов, вместо того чтобы переиспользовать его без изменений.

enable_cache по умолчанию использует DBCache, при этом первые 8 блоков всегда вычисляются (F8B0), а также выполняется 8 шагов прогрева без кэширования. Чтобы пожертвовать скоростью ради качества, увеличьте Fn_compute_blocks или уменьшите residual_diff_threshold (по умолчанию 0.08). Для наилучшего качества при высоких скоростях кэширования добавьте калибровщик 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),
)

Для поддерживаемых пайплайнов CacheDiT уже знает, выполняется ли CFG как отдельный прямой проход. Для других моделей установите enable_separate_cfg=True в DBCacheConfig, если модель выполняет условный и безусловный проходы отдельно, или False, если она объединяет их или не использует CFG.

См. документацию проектирования DBCache о том, как работают блоки Fn и Bn, а также бенчмарки кэша для показателей скорости и качества.

Дальнейшие шаги

Для пайплайнов, которые CacheDiT пока не поддерживает, см. документацию BlockAdapter.
CacheDiT также поддерживает контекстную параллельность и квантование.

Update on GitHub