캐싱

Caching은 denoising 단계 간에 중간 레이어의 출력을 재사용하여 추론 속도를 높입니다. 메모리 사용량이 증가하지만 학습이 필요하지 않습니다. Caching을 지원하는 transformer에서 config를 통해 해당 방법을 활성화합니다.

cache 방법 선택

설정할 config의 양과 필요한 적합도에 따라 방법을 선택하세요. 각 방법은 속도, 메모리, 캐시 없는 실행과 출력의 일치 정도를 서로 다른 방식으로 균형 있게 조정합니다. 특정 방법에 의존하기 전에 자체 프롬프트로 캐시 사용 시와 미사용 시의 출력을 비교해 보세요.

Pyramid Attention Broadcast

Pyramid Attention Broadcast (PAB)은 모든 단계에서 다시 계산하는 대신 일부 블록과 타임스텝에 대해 attention 출력을 재사용하여 denoising 단계 간 attention을 근사합니다. 모델에 해당 항목이 있는 경우 Config는 attention 종류(spatial, temporal, cross)를 구분합니다. 모든 video 모델이 세 가지 모두를 노출하지 않으므로, 보유 중인 블록에 해당하는 범위만 설정하세요.

각 유형은 *_attention_block_skip_range(윈도우 내에서 재계산과 재사용의 빈도)과 *_attention_timestep_skip_range(어떤 디노이징 타임스텝을 건너뛸 수 있는지)를 사용합니다. 훅이 파이프라인의 현재 타임스텝을 읽을 수 있도록 current_timestep_callback을 전달해야 합니다. 더 넓거나 공격적인 스킵은 일반적으로 더 높은 속도와 더 큰 품질 리스크를 의미합니다.

활성화하려면 PyramidAttentionBroadcastConfig를 전달하세요.

import torch
from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig

pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

config = PyramidAttentionBroadcastConfig(
    spatial_attention_block_skip_range=2,
    spatial_attention_timestep_skip_range=(100, 800),
    current_timestep_callback=lambda: pipe.current_timestep,
)
pipe.transformer.enable_cache(config)

FasterCache

FasterCache는 PAB와 유사하게 attention feature를 캐싱하고 재사용합니다. 또한 classifier-free guidance 하의 unconditional branch를 건너뛰고, 연속된 latents가 충분히 중복될 경우 conditional branch로부터 이를 추정할 수도 있습니다.

활성화하려면 FasterCacheConfig를 전달하세요. PAB와 마찬가지로, 사용 중인 attention 종류에 대해 *_attention_block_skip_range와 *_attention_timestep_skip_range를 설정하고, 필요할 경우 CFG-branch skip 옵션도 함께 설정하세요.

import torch
from diffusers import CogVideoXPipeline, FasterCacheConfig

pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

config = FasterCacheConfig(
    spatial_attention_block_skip_range=2,
    spatial_attention_timestep_skip_range=(-1, 681),
    current_timestep_callback=lambda: pipe.current_timestep,
    attention_weight_callback=lambda _: 0.3,
    unconditional_batch_skip_range=5,
    unconditional_batch_timestep_skip_range=(-1, 781),
    tensor_format="BFCHW",
)
pipe.transformer.enable_cache(config)

SeaCache

SeaCache는 연속된 노이즈 제거 단계 간에 Spectral Evolution Aware(SEA) 지표를 비교합니다. 누적된 변화가 임계값 미만으로 유지되면 transformer 블록 스택을 건너뛰고 캐시된 잔차로부터 출력을 예측합니다. 이 방법은 근사적이며 비디오 생성을 위해 설계되었습니다.

SeaCache의 내장 어댑터에는 다음이 포함됩니다:

• Cosmos 3는 주로 최적화되고 벤치마크된 통합입니다.
• Wan T2V는 SeaCache에 원시 비전 latents를 제공하는 방법을 시연하기 위해 범용 반복 블록 경로를 사용합니다. 동일한 캐시 매개변수가 Wan 또는 다른 Wan 변형에 그대로 적용되지 않을 수 있습니다.

transformer에서 SeaCache를 활성화합니다. Cosmos 3의 denoising 루프는 각 cache_context에 scheduler step, sigma, step count를 연결하므로 추가 매개변수가 필요하지 않습니다.

from diffusers import Cosmos3OmniPipeline, SeaCacheConfig

pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano")
pipe.transformer.enable_cache(SeaCacheConfig(threshold=0.2, max_consecutive_cached=2))

SeaCache는 출력을 변경할 수 있습니다. 모든 단계에서 전체 transformer를 실행해야 하는 경우 pipe.transformer.disable_cache()를 호출하세요. 동일한 enable 호출은 Cosmos3OmniPipeline, Cosmos3OmniModularPipeline, Cosmos3DistilledModularPipeline에서도 작동합니다.

다른 비디오 transformer에서 SeaCache를 사용하려면 새 모델에 캐싱 추가를 따르십시오. SeaCache는 모든 cache_context에 step_index, sigma, num_inference_steps가 필요하며, 내장 어댑터가 없는 경우 raw_vision_callback도 필요합니다. 모델과 스케줄러별로 매개변수를 조정하십시오.

FirstBlockCache

FirstBlockCacheConfig는 denoiser의 초기 레이어가 timestep마다 얼마나 변하는지 확인합니다. 변화가 작으면 모델은 비용이 높은 후속 레이어를 건너뛰고 이전 출력을 재사용합니다.

enable_cache를 통해 활성화하여 disable_cache와 is_cache_enabled가 동기화되도록 하십시오. 기본 threshold 값은 0.05입니다. 0.2와 같은 더 높은 값을 사용하면 더 자주 건너뛰어 속도를 높일 수 있지만, 생성 품질이 저하될 수 있습니다.

import torch
from diffusers import DiffusionPipeline, FirstBlockCacheConfig

pipe = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16
)
pipe.transformer.enable_cache(FirstBlockCacheConfig(threshold=0.2))

TaylorSeer 캐시

TaylorSeer Cache는 디노이징 단계에 걸친 테일러 급수 확장을 통해 diffusion 추론을 가속화합니다. 이전 단계의 활성화 값에서 후속 단계의 활성화 값을 예측하고, 해당 예측 값을 여러 단계에 걸쳐 재사용하여 transformer의 전체 연산량을 줄입니다.

• cache_interval: 전체 forward pass를 수행하기 전에 캐시된 출력을 재사용할 스텝 수
• disable_cache_before_step: 근사치를 위한 데이터를 수집하기 위해 전체 계산을 사용하는 초기 단계
• max_order: 더 높은 Taylor order는 더 정확할 수 있지만 더 많은 메모리를 사용합니다. 변경할 특별한 이유가 없는 한 1으로 유지하세요.

import torch
from diffusers import FluxPipeline, TaylorSeerCacheConfig

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

config = TaylorSeerCacheConfig(
    cache_interval=5,
    max_order=1,
    disable_cache_before_step=10,
    taylor_factors_dtype=torch.bfloat16,
)
pipe.transformer.enable_cache(config)

MagCache

MagCache는 잔차 업데이트의 크기를 기준으로 transformer 블록을 건너뜁니다. 업데이트 크기는 노이즈 제거 과정에서 예측 가능하게 감소하며, MagCache는 사전 계산된 크기 비율(mag_ratios)로부터 에러 버짓을 추적하여 재사용이 안전한 시점을 결정합니다. 이러한 비율은 체크포인트와 스케줄러에 따라 달라집니다. 높은 스텝 수에서 계산된 비율은 더 적은 스텝 수로 보간할 수 있습니다.

번들된 FLUX_MAG_RATIOS는 FLUX.1에서 측정되었습니다. Flux2 및 Flux2 Klein을 포함한 다른 모델은 자체 캘리브레이션 실행이 필요합니다.

MagCache는 두 단계로 진행됩니다:

1. Calibration: calibrate=True로 추론을 한 번 실행합니다. 훅이 잔차의 크기를 측정하고 계산된 비율을 출력합니다.
2. Inference: 교정 캐시를 비활성화한 후, 해당 비율을 MagCacheConfig에 전달하여 가속합니다.

Classifier-free guidance는 캘리브레이션에 영향을 줄 수 있습니다. Flux와 같이 true_cfg_scale > 1일 때 순차적 컨텍스트를 사용하는 true CFG를 적용하는 파이프라인은 cache_context("cond")과 cache_context("uncond")를 각각 입력합니다. 캘리브레이션은 컨텍스트별로 하나의 배열을 출력할 수 있지만, 대부분의 경우 조건부 배열을 사용해야 합니다. 조건부 및 무조건부 입력을 연결하여 CFG를 배치하는 파이프라인(예: CogVideoX)은 바로 사용할 수 있는 단일 결합된 배열을 생성합니다.

import torch
from diffusers import FluxPipeline, MagCacheConfig
from diffusers.hooks.mag_cache import FLUX_MAG_RATIOS

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

# 1. Calibration Step
# Run full inference to measure model behavior.
calib_config = MagCacheConfig(calibrate=True, num_inference_steps=4)
pipe.transformer.enable_cache(calib_config)

# Run a prompt to trigger calibration
pipe("A cat playing chess", num_inference_steps=4)
# Prints: [MagCache] Calibration Complete. Copy these values to MagCacheConfig(mag_ratios=...):

# 2. Inference Step
# Disable calibration hooks before enabling MagCache for inference.
pipe.transformer.disable_cache()

# Apply ratios from calibration, or use the Flux defaults:
# mag_ratios=FLUX_MAG_RATIOS
mag_config = MagCacheConfig(
    mag_ratios=[1.0, 1.37, 0.97, 0.87],
    num_inference_steps=4
)

pipe.transformer.enable_cache(mag_config)

image = pipe("A cat playing chess", num_inference_steps=4).images[0]

텍스트 KV 캐시

TextKVCacheConfig은 텍스트 키 및 값 투영을 한 번만 계산한 후, 재계산하지 않고 모든 디노이징 단계에서 재사용합니다. apply_text_kv_cache()는 현재 NucleusMoEImageTransformerBlock에만 연결되어 있습니다. enable_cache를 통해 활성화할 수 있습니다.

import torch
from diffusers import NucleusMoEImagePipeline, TextKVCacheConfig

pipe = NucleusMoEImagePipeline.from_pretrained(
    "NucleusAI/NucleusMoE-Image", dtype=torch.bfloat16
)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

pipe.transformer.enable_cache(TextKVCacheConfig())

image = pipe("A cat holding a sign that says hello world", num_inference_steps=50).images[0]

새 모델에 캐싱 추가

Cache 메서드는 훅을 통해 모델에 연결되며, 이 훅은 모델 클래스, 해당 transformer 블록, 그리고 파이프라인의 denoising 루프에 의존합니다. 새 모델에서 enable_cache를 호출하기 전에 이 세 가지를 모두 설정하세요.

1. transformer 클래스에서 CacheMixin을 상속하세요. enable_cache, disable_cache, cache_context가 추가됩니다.

2. hooks/_helpers.py의 _register_transformer_blocks_metadata에 transformer 블록의 출력 레이아웃을 등록합니다. FirstBlockCache, MagCache, SeaCache와 같이 블록 전체를 건너뛰거나 재사용하는 메서드는 이를 사용하여 블록 출력에서 hidden states를 찾습니다.

   복사됨

   TransformerBlockRegistry.register(
       model_class=MyTransformerBlock,
       metadata=TransformerBlockMetadata(
           return_hidden_states_index=0,
           return_encoder_hidden_states_index=None,
       ),
   )

3. cache_context으로 파이프라인의 디노이징 루프 내 각 transformer 호출을 감싸세요. classifier-free guidance가 조건부 및 무조건부 분기를 별도의 호출로 실행하는 경우, 각 분기가 자체 캐시 상태를 유지하도록 "cond"와 "uncond"와 같은 서로 다른 이름을 사용하세요.

   복사됨

   with self.transformer.cache_context("cond"):
       noise_pred = self.transformer(hidden_states=latents, timestep=timestep, ...)[0]

PAB와 FasterCache도 훅이 현재 timestep을 읽을 수 있도록 current_timestep_callback이 필요합니다. 캐시 설정은 하나의 모델에 맞게 조정된 것이 다른 모델로 잘 이식되지 않는 경우가 많으므로, 새 모델에서 출력 품질을 확인하세요.