Buckets:

mishig's picture
|
download
raw
5.66 kB

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, а также бенчмарки кэша для показателей скорости и качества.

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

Xet Storage Details

Size:
5.66 kB
·
Xet hash:
50720e8329c130fbe40b0e4d89b789bf7c66aad572f11eaca5401518c2ad8097

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.