Bộ nhớ đệm

Caching tái sử dụng đầu ra của các lớp trung gian giữa các bước khử nhiễu để tăng tốc suy luận. Phương pháp này sử dụng nhiều bộ nhớ hơn và không cần huấn luyện. Hãy bật một phương pháp bằng cách đặt cấu hình trên một transformer hỗ trợ caching.

Chọn một phương pháp cache

Chọn phương pháp dựa trên mức độ cấu hình bạn sẽ thiết lập và mức độ phù hợp cần thiết. Mỗi phương pháp cân bằng giữa tốc độ, bộ nhớ và mức độ khớp của đầu ra so với lần chạy không có cache theo những cách khác nhau. Hãy so sánh đầu ra có và không có cache trên các prompt của riêng bạn trước khi tin dùng một phương pháp.

Phát sóng Pyramid Attention

Pyramid Attention Broadcast (PAB) xấp xỉ attention qua các bước khử nhiễu bằng cách tái sử dụng kết quả attention cho một số khối và bước thời gian thay vì tính lại ở mỗi bước. Config tách biệt các loại attention (spatial, temporal, cross) khi mô hình có chúng. Không phải mọi mô hình video đều hiển thị cả ba loại, và chỉ đặt các khoảng phù hợp với các khối bạn có.

Mỗi loại sử dụng *_attention_block_skip_range (tần suất tính lại so với tái sử dụng trong cửa sổ) và *_attention_timestep_skip_range (các bước thời gian khử nhiễu nào có thể bỏ qua). Bạn phải truyền current_timestep_callback để hook có thể đọc bước thời gian hiện tại của pipeline. Việc bỏ qua rộng hơn hoặc quyết liệt hơn thường đồng nghĩa với tốc độ cao hơn và rủi ro về chất lượng lớn hơn.

Truyền một PyramidAttentionBroadcastConfig để bật tính năng này.

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 lưu trữ và tái sử dụng các đặc trưng chú ý tương tự như PAB. Nó cũng có thể bỏ qua nhánh không điều kiện dưới classifier-free guidance và ước lượng nó từ nhánh có điều kiện khi các latent liên tiếp đủ dư thừa.

Truyền một FasterCacheConfig để bật tính năng này. Giống như PAB, hãy đặt *_attention_block_skip_range và *_attention_timestep_skip_range cho các loại attention mà bạn có, cùng với các tùy chọn skip cho nhánh CFG khi bạn muốn.

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 so sánh các chỉ báo Spectral Evolution Aware (SEA) giữa các bước khử nhiễu liên tiếp. Khi sự thay đổi tích lũy nằm dưới ngưỡng, phương pháp này bỏ qua chồng khối transformer và dự đoán đầu ra từ các phần dư được lưu cache. Phương pháp này mang tính xấp xỉ và được thiết kế cho việc tạo video.

Các bộ thích ứng tích hợp sẵn cho SeaCache bao gồm:

• Cosmos 3 là tích hợp chính được tối ưu và đánh chuẩn.
• Wan T2V sử dụng đường dẫn khối lặp chung như một bản demo cho cách cung cấp các latents thị giác thô cho SeaCache. Các tham số cache tương tự có thể không áp dụng được cho Wan hoặc các biến thể khác của Wan.

Bật SeaCache trên transformer. Vòng lặp khử nhiễu của Cosmos 3 gắn bước scheduler, sigma và số bước vào mỗi cache_context, nên không cần thêm tham số nào khác.

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 có thể làm thay đổi kết quả đầu ra. Gọi pipe.transformer.disable_cache() khi bạn cần mỗi bước thực hiện đầy đủ transformer. Lệnh enable tương tự hoạt động với Cosmos3OmniPipeline, Cosmos3OmniModularPipeline và Cosmos3DistilledModularPipeline.

Để sử dụng SeaCache với một video transformer khác, hãy làm theo Thêm caching vào một model mới. SeaCache cũng cần step_index, sigma và num_inference_steps trong mỗi cache_context, và một raw_vision_callback khi không có adapter tích hợp sẵn. Điều chỉnh các tham số theo từng model và scheduler.

FirstBlockCache

FirstBlockCacheConfig kiểm tra mức độ thay đổi của các lớp đầu tiên của denoiser từ một timestep sang timestep tiếp theo. Nếu mức thay đổi nhỏ, mô hình sẽ bỏ qua các lớp phía sau tốn kém và tái sử dụng kết quả từ lần trước.

Bật tính năng này thông qua enable_cache để disable_cache và is_cache_enabled luôn đồng bộ. Giá trị mặc định của threshold là 0.05. Một giá trị cao hơn, chẳng hạn 0.2, sẽ bỏ qua thường xuyên hơn để tăng tốc độ, nhưng chất lượng sinh có thể giảm.

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 Cache

TaylorSeer Cache tăng tốc suy luận diffusion bằng cách sử dụng khai triển chuỗi Taylor qua các bước khử nhiễu. Nó dự đoán các kích hoạt ở các bước sau từ các bước trước và tái sử dụng những dự đoán đó trong một số bước, nhờ đó transformer thực hiện ít công việc toàn phần hơn.

• cache_interval: Số bước để tái sử dụng các đầu ra được lưu trong bộ nhớ đệm trước khi thực hiện một lần truyền tiến đầy đủ
• disable_cache_before_step: Các bước ban đầu sử dụng tính toán đầy đủ để thu thập dữ liệu cho các phép xấp xỉ
• max_order: Bậc Taylor cao hơn có thể chính xác hơn nhưng sử dụng nhiều bộ nhớ hơn. Giữ giá trị này ở 1 trừ khi bạn có lý do để thay đổi.

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 bỏ qua các khối transformer dựa trên độ lớn của cập nhật phần dư. Độ lớn cập nhật suy giảm một cách có thể dự đoán được trong quá trình khử nhiễu, và MagCache theo dõi một ngân sách lỗi từ các tỷ lệ độ lớn được tính trước (mag_ratios) để quyết định khi nào việc tái sử dụng là an toàn. Các tỷ lệ này phụ thuộc vào checkpoint và bộ lập lịch cụ thể. Các tỷ lệ từ số bước cao có thể được nội suy xuống số bước ít hơn.

Các FLUX_MAG_RATIOS được đóng gói đã được đo trên FLUX.1. Các mô hình khác, bao gồm Flux2 và Flux2 Klein, cần chạy hiệu chuẩn riêng.

MagCache tuân theo hai bước:

1. Hiệu chuẩn: Chạy suy luận một lần với calibrate=True. Hook sẽ đo độ lớn của phần dư và in ra các tỷ lệ đã tính toán.
2. Suy luận: Tắt bộ đệm hiệu chuẩn, sau đó truyền các tỷ lệ đó vào MagCacheConfig để tăng tốc.

Classifier-free guidance có thể ảnh hưởng đến hiệu chuẩn. Các pipeline sử dụng CFG thực sự với ngữ cảnh tuần tự, chẳng hạn như Flux khi true_cfg_scale > 1, sẽ lần lượt đi vào cache_context("cond") và cache_context("uncond"). Hiệu chuẩn có thể in ra một mảng cho mỗi ngữ cảnh, nhưng trong hầu hết các trường hợp, bạn nên sử dụng mảng điều kiện. Các pipeline thực hiện CFG theo lô bằng cách nối các đầu vào có điều kiện và không điều kiện (ví dụ: CogVideoX) sẽ tạo ra một mảng liên hợp duy nhất mà bạn có thể sử dụng trực tiếp.

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]

Bộ đệm KV văn bản

TextKVCacheConfig tính toán các phép chiếu key và value của văn bản một lần và tái sử dụng chúng ở mỗi bước khử nhiễu thay vì tính lại. apply_text_kv_cache() hiện chỉ gắn vào NucleusMoEImageTransformerBlock. Bật tính năng này bằng 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]

Thêm bộ nhớ đệm vào một mô hình mới

Các phương thức cache được gắn vào mô hình thông qua các hook, và các hook này phụ thuộc vào lớp mô hình, các khối transformer của nó và vòng lặp khử nhiễu của pipeline. Hãy thiết lập cả ba thành phần này trước khi gọi enable_cache trên một mô hình mới.

1. Kế thừa từ CacheMixin trong lớp transformer. Nó thêm enable_cache, disable_cache và cache_context.

2. Đăng ký bố cục đầu ra của khối transformer trong _register_transformer_blocks_metadata trong hooks/_helpers.py. Các phương thức bỏ qua hoặc tái sử dụng toàn bộ khối, chẳng hạn như FirstBlockCache, MagCache và SeaCache, sử dụng nó để tìm các trạng thái ẩn trong đầu ra của khối.

   Đã sao chép

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

3. Bọc mỗi lần gọi transformer trong vòng lặp khử nhiễu của pipeline bằng cache_context. Sử dụng các tên riêng biệt, chẳng hạn như "cond" và "uncond", khi classifier-free guidance chạy các nhánh có điều kiện và không điều kiện như các lần gọi riêng biệt, để mỗi nhánh giữ trạng thái cache của riêng mình.

   Đã sao chép

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

PAB và FasterCache cũng cần current_timestep_callback để các hook có thể đọc timestep hiện tại. Kiểm tra chất lượng đầu ra trên mô hình mới, vì các cài đặt cache được tinh chỉnh cho một mô hình thường không áp dụng được cho mô hình khác.