CacheDiT

CacheDiT tăng tốc các pipeline DiT bằng cách tái sử dụng đầu ra của các khối transformer giữa các bước khử nhiễu. Nó không cần huấn luyện và hỗ trợ hầu hết các pipeline DiT của Diffusers, bao gồm Flux, Qwen-Image, Wan và HunyuanVideo. Diffusers cũng có cache tích hợp sẵn mà không yêu cầu phụ thuộc bổ sung.

Cài đặt CacheDiT từ PyPI.

pip install -U cache-dit

Gọi cache_dit.supported_pipelines() để liệt kê các nhóm pipeline mà CacheDiT hỗ trợ.

import cache_dit

cache_dit.supported_pipelines()

Bật bộ nhớ đệm

Gọi cache_dit.enable_cache trên một pipeline để lưu cache với các cài đặt mặc định, sau đó chạy pipeline như bình thường.

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 cũng hoạt động với torch.compile. Hãy biên dịch transformer sau khi bạn gọi enable_cache. Xem tài liệu compile để biết các cài đặt giúp tránh việc biên dịch lại với các hình dạng đầu vào động.

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

Gọi cache_dit.summary sau suy luận để ghi log số bước đã được cache và sự khác biệt phần dư giữa các bước.

stats = cache_dit.summary(pipeline)

Gọi cache_dit.disable_cache để khôi phục pipeline ban đầu.

cache_dit.disable_cache(pipeline)

Cấu hình bộ nhớ đệm

DBCache (Dual Block Cache) tính toán n khối đầu tiên (Fn) ở mỗi bước. Khi đầu ra của chúng hầu như không thay đổi so với bước trước, nó tái sử dụng đầu ra được lưu trong bộ nhớ đệm cho các khối còn lại, và có thể tính lại n khối cuối cùng (Bn) để hiệu chỉnh. Bộ hiệu chuẩn TaylorSeer dự đoán đầu ra được lưu trong bộ nhớ đệm từ các bước trước thay vì tái sử dụng nó nguyên vẹn.

enable_cache mặc định dùng DBCache với 8 khối đầu luôn được tính toán (F8B0) và 8 bước khởi động không có cache. Để đánh đổi tốc độ lấy chất lượng, hãy tăng Fn_compute_blocks hoặc giảm residual_diff_threshold (mặc định 0.08). Để đạt chất lượng tốt nhất ở tỷ lệ cache cao, hãy thêm bộ hiệu chuẩn 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),
)

Đối với các pipeline được hỗ trợ, CacheDiT đã biết liệu CFG có chạy như một bước forward riêng biệt hay không. Đối với các model khác, hãy đặt enable_separate_cfg=True trong DBCacheConfig nếu model chạy các bước điều kiện và không điều kiện riêng biệt, hoặc đặt False nếu model gộp chúng lại hoặc không sử dụng CFG.

Xem tài liệu thiết kế DBCache để biết cách các khối Fn và Bn hoạt động, và các bài kiểm tra hiệu năng bộ nhớ đệm để xem các con số về tốc độ và chất lượng.

Các bước tiếp theo

Đối với các pipeline mà CacheDiT chưa hỗ trợ, hãy xem tài liệu BlockAdapter.
CacheDiT cũng hỗ trợ song song ngữ cảnh và lượng tử hóa.

Update on GitHub