Diffusers documentation
Bộ nhớ đệm
Bộ nhớ đệm
Caching tăng tốc độ suy luận bằng cách lưu trữ và tái sử dụng các đầu ra trung gian của các lớp khác nhau, chẳng hạn như lớp attention và lớp feedforward, thay vì thực hiện toàn bộ phép tính ở mỗi bước suy luận. Nó cải thiện đáng kể tốc độ sinh ra nội dung với cái giá là sử dụng nhiều bộ nhớ hơn và không yêu cầu huấn luyện thêm.
Hướng dẫn này cho bạn cách sử dụng các phương pháp cache được hỗ trợ trong Diffusers.
Phát sóng Pyramid Attention
Pyramid Attention Broadcast (PAB) dựa trên quan sát rằng đầu ra của attention không khác biệt nhiều giữa các bước thời gian liên tiếp trong quá trình sinh. Sự khác biệt của attention là nhỏ nhất ở các lớp cross attention và thường được cache trên một khoảng bước thời gian dài hơn. Tiếp theo là các lớp temporal attention và spatial attention.
Không phải tất cả các mô hình video đều có ba loại chú ý (cross, temporal và spatial)!
PAB có thể được kết hợp với các kỹ thuật khác như song song hóa chuỗi và song song hóa dẫn hướng không bộ phân loại (song song hóa dữ liệu) để tạo video gần như thời gian thực.
Cài đặt và truyền một PyramidAttentionBroadcastConfig cho transformer của pipeline để kích hoạt tính năng này. Tham số spatial_attention_block_skip_range kiểm soát tần suất bỏ qua các phép tính attention trong các khối spatial attention, còn spatial_attention_timestep_skip_range là khoảng các timestep cần bỏ qua. Hãy chọn khoảng phù hợp vì khoảng nhỏ hơn có thể làm giảm tốc độ suy luận, trong khi khoảng lớn hơn có thể làm giảm chất lượng sinh thành.
import torch
from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig
pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
pipeline.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,
)
pipeline.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 vì sự khác biệt đầu ra giữa các bước thời gian liên tiếp là rất nhỏ.
Phương pháp này cũng có thể chọn bỏ qua dự đoán nhánh không điều kiện, khi sử dụng classifier-free guidance để lấy mẫu (phổ biến ở hầu hết các mô hình cơ bản), và ước lượng nó từ dự đoán nhánh có điều kiện nếu có sự dư thừa đáng kể trong các đầu ra latent được dự đoán giữa các bước thời gian liên tiếp.
Cài đặt và truyền một FasterCacheConfig cho transformer của pipeline để bật tính năng này.
import torch
from diffusers import CogVideoXPipeline, FasterCacheConfig
pipe line= CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
pipeline.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",
)
pipeline.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 của chỉ báo vẫn nằm dưới ngưỡng, nó sẽ bỏ qua chồng khối transformer đắt đỏ và dự đoán đầu ra của chúng từ các phần dư được lưu cache. Xây dựng chỉ báo từ các visual latents tạo thành đầu ra được sinh ra. Bao gồm các khung điều kiện sạch khi chúng là một phần của quỹ đạo đầu ra đó, như trong sinh ảnh-đến-video và video-đến-video. Loại trừ các visual hints riêng biệt điều kiện hóa quá trình sinh nhưng không phải là một phần của đầu ra. Điều kiện văn bản bị loại trừ vì nó không phải là visual latent.
Các gói Transfer của Cosmos 3 đóng gói các gợi ý điều khiển (control hints) dưới dạng các chuỗi thị giác riêng biệt, do đó adapter của nó loại trừ chúng khỏi bộ chỉ báo (indicator). Các nhánh Control-CFG so sánh cùng một quỹ đạo đầu ra trong khi vẫn giữ lại các phần dư (residuals) được đệm (cached) của riêng chúng. Các gợi ý điều khiển vẫn điều kiện hóa (condition) transformer.
Cài đặt cung cấp các bộ thích ứng tích hợp sẵn cho các mô hình sau:
- Cosmos 3 là tích hợp được tối ưu và đo lường hiệu năng chính. Nó lưu bộ nhớ đệm toàn bộ chồng giải mã thông qua ranh giới chuẩn hóa hậu (post-normalization).
- Wan T2V sử dụng đường dẫn khối lặp chung ở chế độ eager. Việc tích hợp này minh họa cách một video transformer single-stream khác có thể cung cấp raw vision latents cho SeaCache; đây không phải là khẳng định rằng các tham số cache giống nhau là tối ưu cho Wan hay rằng các biến thể Wan khác được hỗ trợ.
Các video transformer khác có thể tích hợp với đường dẫn chung khi chúng sử dụng CacheMixin, công khai một danh sách khối lặp lại được công nhận, và đăng ký bố cục đầu vào/đầu ra của khối trong TransformerBlockRegistry. Pipeline phải đi vào một cache_context cho mỗi lần gọi transformer, gắn step_index, sigma và num_inference_steps, và sử dụng các tên ngữ cảnh riêng biệt cho các quỹ đạo độc lập như hướng dẫn có điều kiện và không có điều kiện. Truyền một raw_vision_callback trả về các visual latents tạo thành đầu ra được sinh ra khi không có adapter tích hợp sẵn. Kiểm tra chất lượng đầu ra và tinh chỉnh các tham số cache cho từng model và scheduler; kết quả hỗ trợ và benchmark không tự động chuyển từ Cosmos 3.
Cosmos 3
SeaCache bị tắt theo mặc định. Hãy bật nó trên transformer; vòng lặp khử nhiễu của Cosmos 3 gắn bước scheduler đang hoạt động, sigma và số bước vào mỗi cuộc gọi cache_context, vì vậy không cần thêm cấu hình nào:
from diffusers import Cosmos3OmniPipeline, SeaCacheConfig
pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano")
pipe.transformer.enable_cache(SeaCacheConfig(threshold=0.2, max_consecutive_cached=2))API ở cấp độ mô hình này hoạt động với Cosmos3OmniPipeline, Cosmos3OmniModularPipeline và Cosmos3DistilledModularPipeline. SeaCache là một tối ưu hóa xấp xỉ và có thể làm thay đổi kết quả đầu ra được tạo ra. Hãy gọi pipe.transformer.disable_cache() khi bạn cần mỗi bước khử nhiễu thực thi đầy đủ transformer.
FirstBlockCache
FirstBlock Cache kiểm tra mức độ thay đổi của các lớp đầu tiên của denoiser từ một bước thời gian sang bước tiếp theo. Nếu sự thay đổi là nhỏ, mô hình sẽ bỏ qua các lớp đắt đỏ hơn ở phía sau và tái sử dụng kết quả từ lần trước.
import torch
from diffusers import DiffusionPipeline
from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig
pipeline = DiffusionPipeline.from_pretrained(
"Qwen/Qwen-Image", dtype=torch.bfloat16
)
apply_first_block_cache(pipeline.transformer, 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 để xấp xỉ và lưu cache các kích hoạt trung gian giữa các bước khử nhiễu. Phương pháp này dự đoán các đầu ra tương lai dựa trên các phép tính trước đó, tái sử dụng chúng ở các khoảng thời gian được chỉ định để giảm thiểu các phép tính trùng lặp.
Cơ chế lưu trữ bộ nhớ đệm này mang lại kết quả ấn tượng với chi phí bộ nhớ bổ sung tối thiểu. Để xem phân tích hiệu suất chi tiết, hãy xem kết quả của chúng tôi tại đây.
Để bật TaylorSeer Cache, hãy tạo một TaylorSeerCacheConfig và truyền nó vào transformer của pipeline:
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: Độ chính xác của phép xấp xỉ (về mặt lý thuyết, giá trị cao hơn sẽ cải thiện chất lượng nhưng làm tăng mức sử dụng bộ nhớ, tuy nhiên chúng tôi khuyến nghị nên đặt thành1)
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 tăng tốc suy luận bằng cách bỏ qua các khối transformer dựa trên độ lớn của bản cập nhật phần dư. Nó quan sát thấy rằng độ lớn của các bản cập nhật (Output - Input) suy giảm một cách có thể dự đoán được trong quá trình khuếch tán. Bằng cách tích lũy một “ngân sách lỗi” dựa trên các tỷ lệ độ lớn được tính trước, nó quyết định động khi nào nên bỏ qua việc tính toán và tái sử dụng phần dư trước đó.
MagCache dựa vào Magnitude Ratios (mag_ratios), mô tả đường cong suy giảm này. Các tỷ lệ này đặc thù cho từng checkpoint mô hình và scheduler. Bộ FLUX_MAG_RATIOS đi kèm đượ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.
Để sử dụng MagCache, bạn thường tuân theo quy trình hai bước: Hiệu chuẩn và Suy luận.
- 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 các tỷ lệ đã tính ra console. - Suy luận: Truyền các tỷ lệ này vào
MagCacheConfigđể bật tăng tốc.
import torch
from diffusers import FluxPipeline, MagCacheConfig
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)
# Logs will print something like: "MagCache Calibration Results: [1.0, 1.37, 0.97, 0.87]"
# 2. Inference Step
# Apply the specific ratios obtained from calibration for optimized speed.
# Note: For Flux models, you can also import defaults:
# from diffusers.hooks.mag_cache import 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]
mag_ratiosđại diện cho đường cong suy giảm biên độ nội tại của mô hình. Các tỷ lệ được hiệu chuẩn cho số bước cao (ví dụ: 50) có thể được tái sử dụng cho số bước thấp hơn (ví dụ: 20). Việc triển khai sử dụng nội suy để ánh xạ đường cong sang số bước suy luận hiện tại.
Đối với các pipeline chạy Classifier-Free Guidance tuần tự (như Kandinsky 5.0), log hiệu chuẩn có thể in ra hai mảng: một cho lượt chạy Conditional và một cho lượt chạy Unconditional. Trong hầu hết các trường hợp, bạn nên sử dụng mảng đầu tiên (Conditional).
Update on GitHubĐối với các pipeline chạy Classifier-Free Guidance theo cách batched (như SDXL hoặc Flux),
hidden_statesđược xử lý bởi mô hình chứa cả nhánh có điều kiện và không có điều kiện được nối lại với nhau. Quá trình hiệu chuẩn tự động tính đến điều này, tạo ra một mảng duy nhất các tỷ số đại diện cho hành vi chung. Bạn có thể sử dụng trực tiếp mảng kết quả này mà không cần chỉnh sửa.