Buckets:

mishig's picture
|
download
raw
4.28 kB

CacheDiT

CacheDiT은 디노이징 단계 간에 트랜스포머 블록 출력을 재사용하여 DiT 파이프라인의 속도를 높입니다. 학습이 필요하지 않으며, Flux, Qwen-Image, Wan, HunyuanVideo를 포함한 대부분의 Diffusers DiT 파이프라인을 지원합니다. Diffusers에는 추가적인 의존성이 필요하지 않은 내장 캐싱 기능도 있습니다.

PyPI에서 CacheDiT를 설치합니다.

pip install -U cache-dit

CacheDiT가 지원하는 파이프라인 계열을 나열하려면 cache_dit.supported_pipelines()를 호출하세요.

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를 호출하여 원래의 파이프라인을 복원합니다.

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가 별도의 forward pass로 실행되는지 여부를 이미 알고 있습니다. 다른 모델의 경우, 모델이 조건부 패스와 무조건부 패스를 별도로 실행하는지, 아니면 이를 융합하거나 CFG를 사용하지 않는지에 따라 DBCacheConfig에서 enable_separate_cfg=True 또는 False로 설정하십시오.

Fn 및 Bn 블록의 작동 방식은 DBCache 설계 문서를, 속도 및 품질 수치는 캐시 벤치마크를 참고하세요.

다음 단계

Xet Storage Details

Size:
4.28 kB
·
Xet hash:
c7ece5118ab66f928908139d856cfececebb9ea1d52037e1cda19b3f20a353b0

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