Buckets:

mishig's picture
|
download
raw
3.67 kB

CacheDiT

CacheDiT 通过在不同去噪步骤之间复用 transformer block 的输出,从而加速 DiT 流水线。它无需训练,并支持大多数 Diffusers DiT 流水线,包括 Flux、Qwen-Image、Wan 和 HunyuanVideo。Diffusers 还提供了 内置缓存,无需额外依赖。

从 PyPI 安装 CacheDiT。

pip install -U cache-dit

调用 cache_dit.supported_pipelines() 以列出 CacheDiT 支持的 pipeline 系列。

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 之后再编译 transformer。有关避免在动态输入形状下重新编译的设置,请参阅 compile 文档。

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

在推理后调用 cache_dit.summary,以记录缓存的步数以及各步之间的残差差异。

stats = cache_dit.summary(pipeline)

调用 cache_dit.disable_cache 以恢复原始 pipeline。

cache_dit.disable_cache(pipeline)

配置缓存

DBCache (双块缓存)在每一步计算前 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 是否作为独立的前向传播运行。对于其他模型,如果模型分别运行条件路径和无条件路径,请在 DBCacheConfig 中设置 enable_separate_cfg=True;如果模型将它们融合或不使用 CFG,则设置为 False。

请参阅 DBCache 设计 文档,了解 Fn 和 Bn 块的工作原理,并查看 缓存基准测试 以获取速度和质量的数值。

后续步骤

Xet Storage Details

Size:
3.67 kB
·
Xet hash:
7f933dc30e0460b61324a332a7e06c5113c5efcf32d07dbfd35f5c6f30e688a0

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