缓存

缓存通过存储和重用不同层的中间输出（例如注意力层和前馈层）来加速推理，而不是在每个推理步骤中执行完整的计算。它以消耗更多内存为代价显著提高了生成速度，并且不需要额外的训练。

本指南向您展示如何使用 Diffusers 中支持的缓存方法。

Pyramid Attention Broadcast

金字塔注意力广播（Pyramid Attention Broadcast, PAB）基于这样一个观察：生成过程中连续时间步之间的注意力输出差异不大。交叉注意力层中的注意力差异最小，通常在更长的时间步范围内进行缓存。其次是时间注意力和空间注意力层。

并非所有视频模型都具有三种类型的注意力（交叉、时间和空间）！

PAB 可以与其他技术（如序列并行和无分类器指导并行（数据并行））相结合，实现近乎实时的视频生成。

设置并传递一个 PyramidAttentionBroadcastConfig 到 pipeline 的 transformer 中以启用它。spatial_attention_block_skip_range 控制在空间注意力块中跳过注意力计算的频率，而 spatial_attention_timestep_skip_range 是要跳过的时间步范围。请注意选择一个合适的范围，因为较小的间隔可能导致较慢的推理速度，而较大的间隔可能导致较低的生成质量。

import torch
from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig

pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
pipeline.to("cuda")

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 缓存和重用注意力特征，类似于 PAB，因为每个连续时间步的输出差异很小。

当使用无分类器指导进行采样时（在大多数基础模型中很常见），如果连续时间步之间预测的潜在输出存在显著冗余，此方法也可能选择跳过无条件分支预测，并从有条件分支预测中估计它。

设置并传递一个 FasterCacheConfig 到 pipeline 的 transformer 中以启用它。

import torch
from diffusers import CogVideoXPipeline, FasterCacheConfig

pipe line= CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
pipeline.to("cuda")

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)

< > 在 GitHub 上更新