缓存方法
缓存方法通过存储并复用特定层(如注意力层和前馈层)的中间输出,而非在每个推理步骤中重新计算,从而加速扩散 Transformer。
CacheMixin
用于在扩散模型上启用/禁用缓存技术的类。
支持的缓存技术:
提供缓存管理信息的上下文管理器。
name 是去噪调用的名称,通常为 "cond" 或 "uncond"。关键字参数描述了去噪循环的位置——有关接受的字段,请参阅 CacheContext,例如 cache_context("cond", step_index=i, sigma=sigma)。
enable_cache
< source >( config )
参数
- config (
PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | SeaCacheConfig | TextKVCacheConfig) — 用于应用缓存技术的配置。当前支持的缓存技术包括:- PyramidAttentionBroadcastConfig
- FasterCacheConfig
- FirstBlockCacheConfig
- SeaCacheConfig
TextKVCacheConfig
在模型上启用缓存技术。
示例:
>>> import torch
>>> from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig
>>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
>>> pipe.to("cuda")
>>> 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)PyramidAttentionBroadcastConfig
类 diffusers.PyramidAttentionBroadcastConfig
< source >( spatial_attention_block_skip_range: int | None = Nonetemporal_attention_block_skip_range: int | None = Nonecross_attention_block_skip_range: int | None = Nonespatial_attention_timestep_skip_range: tuple[int, int] = (100, 800)temporal_attention_timestep_skip_range: tuple[int, int] = (100, 800)cross_attention_timestep_skip_range: tuple[int, int] = (100, 800)spatial_attention_block_identifiers: tuple[str, ...] = ('blocks', 'transformer_blocks', 'single_transformer_blocks', 'layers', 'visual_transformer_blocks')temporal_attention_block_identifiers: tuple[str, ...] = ('temporal_transformer_blocks',)cross_attention_block_identifiers: tuple[str, ...] = ('blocks', 'transformer_blocks', 'layers')current_timestep_callback: typing.Callable[[], int] = None )
参数
- spatial_attention_block_skip_range (
int,可选,默认为None) — 在计算用于复用的注意力状态之前,跳过特定空间注意力广播的次数。如果将此值设置为N,则注意力计算将被跳过N - 1次(即复用旧的注意力状态),然后再次计算新的注意力状态。 - temporal_attention_block_skip_range (
int,可选,默认为None) — 在重新计算可复用的注意力状态之前,跳过特定时间注意力广播的次数。如果将此值设置为N,则注意力计算将被跳过N - 1次(即复用旧的注意力状态),然后再次计算新的注意力状态。 - cross_attention_block_skip_range (
int,可选,默认为None) — 在计算用于复用的注意力状态之前,跳过特定交叉注意力广播的次数。如果将此值设置为N,则注意力计算将被跳过N - 1次(即复用旧的注意力状态),然后再次计算新的注意力状态。 - spatial_attention_timestep_skip_range (
tuple[int, int],默认为(100, 800)) — 空间注意力层中需要跳过的时间步范围。如果当前时间步位于指定范围内,则条件性地跳过注意力计算。 - temporal_attention_timestep_skip_range (
tuple[int, int],默认为(100, 800)) — 时间注意力层中需要跳过的时间步范围。如果当前时间步位于指定范围内,则条件性地跳过注意力计算。 - cross_attention_timestep_skip_range (
tuple[int, int],默认为(100, 800)) — 交叉注意力层中需要跳过的时间步范围。如果当前时间步位于指定范围内,则条件性地跳过注意力计算。 - spatial_attention_block_identifiers (
tuple[str, ...]) — 用于与层名称进行匹配,以判断该层是否为空间注意力层的标识符。 - temporal_attention_block_identifiers (
tuple[str, ...]) — 用于与层名称进行匹配,以判断该层是否为时间注意力层的标识符。 - cross_attention_block_identifiers (
tuple[str, ...]) — 用于与层名称进行匹配,以判断该层是否为交叉注意力层的标识符。
Pyramid Attention Broadcast 的配置。
diffusers.apply_pyramid_attention_broadcast
< source >( module: Moduleconfig: PyramidAttentionBroadcastConfig )
将 Pyramid Attention Broadcast 应用于指定的 pipeline。
PAB 是一种注意力近似方法,它利用不同时间步之间注意力状态的相似性来降低注意力计算的开销。该论文的核心结论是:不同时间步之间,交叉注意力层的注意力相似性最高,其次是时间维度和空间维度层。这使得交叉注意力层可以比时间维度和空间维度层更频繁地跳过注意力计算。因此,应用 PAB 可以加速推理过程。
示例:
>>> import torch
>>> from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
>>> from diffusers.utils import export_to_video
>>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
>>> pipe.to("cuda")
>>> config = PyramidAttentionBroadcastConfig(
... spatial_attention_block_skip_range=2,
... spatial_attention_timestep_skip_range=(100, 800),
... current_timestep_callback=lambda: pipe.current_timestep,
... )
>>> apply_pyramid_attention_broadcast(pipe.transformer, config)FasterCacheConfig
类 diffusers.FasterCacheConfig
< source >( spatial_attention_block_skip_range: int = 2temporal_attention_block_skip_range: int | None = Nonespatial_attention_timestep_skip_range: tuple[int, int] = (-1, 681)temporal_attention_timestep_skip_range: tuple[int, int] = (-1, 681)low_frequency_weight_update_timestep_range: tuple[int, int] = (99, 901)high_frequency_weight_update_timestep_range: tuple[int, int] = (-1, 301)alpha_low_frequency: float = 1.1alpha_high_frequency: float = 1.1unconditional_batch_skip_range: int = 5unconditional_batch_timestep_skip_range: tuple[int, int] = (-1, 641)spatial_attention_block_identifiers: tuple[str, ...] = ('^blocks.*attn', '^transformer_blocks.*attn', '^single_transformer_blocks.*attn')temporal_attention_block_identifiers: tuple[str, ...] = ('^temporal_transformer_blocks.*attn',)attention_weight_callback: typing.Callable[[torch.nn.Module], float] = Nonelow_frequency_weight_callback: typing.Callable[[torch.nn.Module], float] = Nonehigh_frequency_weight_callback: typing.Callable[[torch.nn.Module], float] = Nonetensor_format: str = 'BCFHW'is_guidance_distilled: bool = Falsecurrent_timestep_callback: typing.Callable[[], int] = None_unconditional_conditional_input_kwargs_identifiers: list[str] = ('hidden_states', 'encoder_hidden_states', 'timestep', 'attention_mask', 'encoder_attention_mask') )
参数
- spatial_attention_block_skip_range (
int,默认为2) — 每隔N次迭代计算一次注意力状态。如果将其设置为N,则在重新计算新的注意力状态之前,注意力计算将被跳过N - 1次(即复用缓存的注意力状态)。 - temporal_attention_block_skip_range (
int, 可选,默认为None) — 每隔N次迭代计算一次注意力状态。如果将此值设置为N,则在重新计算新的注意力状态之前,注意力计算将被跳过N - 1次(即复用缓存的注意力状态)。 - spatial_attention_timestep_skip_range (
tuple[float, float],默认为(-1, 681)) — 可以跳过空间注意力计算而不会显著降低质量的步长范围。此范围需由用户根据底层模型确定。元组中的第一个值为下界,第二个值为上界。通常,用于去噪的扩散步长位于 0 到 1000 的反向范围内(即去噪从步长 1000 开始,到步长 0 结束)。对于默认值,这意味着只有在去噪步长达到 681 之后,才会应用空间注意力计算的跳过,并持续到去噪过程结束。 - temporal_attention_timestep_skip_range (
tuple[float, float],可选,默认为None) — 可以跳过时间注意力计算而不会显著降低质量的步长范围。此范围需由用户根据底层模型确定。元组中的第一个值为下界,第二个值为上界。通常,用于去噪的扩散步长位于 0 到 1000 的反向范围内(即去噪从步长 1000 开始,到步长 0 结束)。 - low_frequency_weight_update_timestep_range (
tuple[int, int],默认为(99, 901)) — 应用低频权重缩放更新的时间步范围。元组中的第一个值是时间步范围的下界,第二个值是上界。更新回调函数仅在此范围内被调用。 - high_frequency_weight_update_timestep_range (
tuple[int, int],默认为(-1, 301)) — 应用高频权重缩放更新的时间步范围。元组中的第一个值是时间步范围的下界,第二个值是上界。更新回调函数仅在此范围内被调用。 - alpha_low_frequency (
float, 默认值为1.1) — 用于缩放低频更新的权重。此参数用于近似从条件分支输出到无条件分支的转换。 - alpha_high_frequency (
float, 默认值为1.1) — 用于缩放高频更新的权重。此参数用于从条件分支输出近似无条件分支。 - unconditional_batch_skip_range (
int,默认为5) — 每N次迭代处理一次无条件分支。如果将此值设置为N,则在重新计算新的无条件分支状态之前,无条件分支计算将被跳过N - 1次(即复用缓存的无条件分支状态)。 - unconditional_batch_timestep_skip_range (
tuple[float, float],默认为(-1, 641)) — 可以跳过无条件分支计算而不会显著降低质量的步长范围。此范围需由用户根据底层模型确定。元组中的第一个值为下界,第二个值为上界。 - spatial_attention_block_identifiers (
tuple[str, ...],默认为("blocks.*attn1", "transformer_blocks.*attn1", "single_transformer_blocks.*attn1")) — 用于匹配模型中空间注意力块的标识符。如果块的名称包含这些标识符中的任意一个,FasterCache 将应用于该块。这些标识符可以是完整的层名称、部分层名称或正则表达式模式。匹配操作始终使用正则表达式匹配。 - temporal_attention_block_identifiers (
tuple[str, ...],默认为("temporal_transformer_blocks.*attn1",)) — 用于匹配模型中时间注意力块的标识符。如果块的名称包含这些标识符中的任意一个,FasterCache 将应用于该块。这些标识符可以是完整的层名称、部分层名称或正则表达式模式。匹配操作始终使用正则表达式匹配。 - attention_weight_callback (
Callable[[torch.nn.Module], float],默认为None) — 用于确定缩放注意力输出权重的回调函数。该函数应以注意力模块作为输入,并返回一个浮点值。此参数用于从条件分支的输出中近似无条件分支。如果未提供,则所有时间步的默认权重为 0.5。通常,如论文所述,随着推理的进行,该权重应从 0 逐渐增加到 1。建议用户进行实验,并提供自定义的权重调度方案,该方案应考虑推理步数以及去噪过程中底层模型的行为。 - low_frequency_weight_callback (
Callable[[torch.nn.Module], float],默认为None) — 用于确定缩放低频更新权重的回调函数。如果未提供,则默认权重为 1.1,适用于指定范围内的时间步(如论文中所述)。 - high_frequency_weight_callback (
Callable[[torch.nn.Module], float],默认为None) — 用于确定缩放高频更新权重的回调函数。如果未提供,则默认权重为 1.1,适用于指定范围内的时间步(如论文中所述)。 - tensor_format (
str,默认为"BCFHW") — 输入张量的格式。应为"BCFHW"、"BFCHW"或"BCHW"之一。该格式用于拆分单独的潜在帧,以便计算低频和高频分量。 - is_guidance_distilled (
bool,默认为False) — 指示模型是否经过 guidance distillation。如果模型经过 guidance distillation,FasterCache 将不会在 denoiser 层面应用,以跳过无条件分支的计算(因为不存在该分支)。 - _unconditional_conditional_input_kwargs_identifiers (
list[str],默认为("hidden_states", "encoder_hidden_states", "timestep", "attention_mask", "encoder_attention_mask")) — 用于匹配包含按批次拼接的无条件输入和有条件输入的输入 kwargs 的标识符。如果输入 kwargs 的名称包含这些标识符中的任意一个,FasterCache 会将输入拆分为无条件分支和有条件分支。这必须是一个列表,其中包含确切包含按批次拼接的无条件输入和有条件输入的输入 kwargs 名称。
FasterCache 的配置。
diffusers.apply_faster_cache
< source >( module: Moduleconfig: FasterCacheConfig )
将 FasterCache 应用于指定的 pipeline。
示例:
>>> import torch
>>> from diffusers import CogVideoXPipeline, FasterCacheConfig, apply_faster_cache
>>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
>>> pipe.to("cuda")
>>> config = FasterCacheConfig(
... spatial_attention_block_skip_range=2,
... spatial_attention_timestep_skip_range=(-1, 681),
... low_frequency_weight_update_timestep_range=(99, 641),
... high_frequency_weight_update_timestep_range=(-1, 301),
... spatial_attention_block_identifiers=["transformer_blocks"],
... attention_weight_callback=lambda _: 0.3,
... tensor_format="BFCHW",
... )
>>> apply_faster_cache(pipe.transformer, config)FirstBlockCacheConfig
类 diffusers.FirstBlockCacheConfig
< source >( threshold: float = 0.05 )
First Block Cache 的配置。
diffusers.apply_first_block_cache
< source >( module: Moduleconfig: FirstBlockCacheConfig )
将 First Block Cache 应用于指定模块。
First Block Cache 基于 TeaCache 的思想构建。它在更广泛的模型范围内实现起来要简单得多,并且最初是出于实验目的而集成的。
示例:
>>> import torch
>>> from diffusers import CogView4Pipeline
>>> from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig
>>> pipe = CogView4Pipeline.from_pretrained("THUDM/CogView4-6B", torch_dtype=torch.bfloat16)
>>> pipe.to("cuda")
>>> apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=0.2))
>>> prompt = "A photo of an astronaut riding a horse on mars"
>>> image = pipe(prompt, generator=torch.Generator().manual_seed(42)).images[0]
>>> image.save("output.png")TaylorSeerCacheConfig
类 diffusers.TaylorSeerCacheConfig
< source >( cache_interval: int = 5disable_cache_before_step: int = 3disable_cache_after_step: int | None = Nonemax_order: int = 1taylor_factors_dtype: typing.Optional[torch.dtype] = torch.bfloat16skip_predict_identifiers: list[str] | None = Nonecache_identifiers: list[str] | None = Noneuse_lite_mode: bool = False )
参数
- cache_interval (
int,默认为5) — 完整计算步骤之间的间隔。在一次完整计算后,缓存的(预测的)输出将在接下来的这么多去噪步骤中被复用,直到通过新的完整前向传播进行刷新。 - disable_cache_before_step (
int,默认为3) — 在此去噪步索引之前禁用缓存,即对初始步骤(0 到 disable_cache_before_step - 1)执行完整计算,以收集用于泰勒级数近似的数据。在这些步骤中,会更新泰勒因子,但不会应用缓存/预测。缓存从该步骤开始。 - disable_cache_after_step (
int, 可选,默认为None) — 禁用缓存的去噪步骤索引。如果设置,对于大于或等于此值的步骤,所有模块都将执行完整计算,而不进行预测或状态更新,从而在需要时确保后期阶段的准确性。 - max_order (
int, 默认为1) — 用于近似模块输出的泰勒级数展开的最高阶数。更高的阶数可提供更好的近似效果,但会增加计算和内存开销。 - taylor_factors_dtype (
torch.dtype,默认为torch.bfloat16) — 用于存储和计算泰勒级数因子的数据类型。较低精度可减少内存占用,但可能影响稳定性;较高精度可提高准确性,但会消耗更多内存。 - skip_predict_identifiers (
list[str],可选,默认为None) — 用于在“cache”模式下将模块名称标记为“skip”的正则表达式模式(使用re.fullmatch)。在此模式下, 模块在初始或刷新步骤中会完整计算,但在预测步骤中返回一个零张量(与记录的形状匹配),从而以较低成本跳过计算。 - cache_identifiers (
list[str],可选,默认为None) — 用于指定哪些模块名称应置于泰勒级数缓存模式的正则表达式模式(使用re.fullmatch),在该模式下,输出会被近似计算并缓存以供复用。 - use_lite_mode (
bool,可选,默认为False) — 启用轻量级 TaylorSeer 变体,通过应用预定义模式来跳过和缓存(例如跳过块并缓存投影),从而最小化内存使用。这将覆盖任何自定义的inactive_identifiers或active_identifiers。
TaylorSeer 缓存的配置。参见:https://huggingface.co/papers/2503.06923
注意事项:
- 模式使用
re.fullmatch对模块名称进行匹配。 - 如果提供了
skip_predict_identifiers或cache_identifiers,则只会对匹配的模块进行挂钩。 - 如果两者均未提供,则默认会挂钩所有注意力类模块。
diffusers.apply_taylorseer_cache
< source >( module: Moduleconfig: TaylorSeerCacheConfig )
将 TaylorSeer 缓存应用于给定的流水线(通常是 transformer / UNet)。
此函数会挂钩模型中指定的模块,以便根据提供的配置启用缓存或跳过,从而减少扩散去噪循环中的冗余计算。
示例:
>>> import torch
>>> from diffusers import FluxPipeline, TaylorSeerCacheConfig
>>> pipe = FluxPipeline.from_pretrained(
... "black-forest-labs/FLUX.1-dev",
... torch_dtype=torch.bfloat16,
... )
>>> pipe.to("cuda")
>>> config = TaylorSeerCacheConfig(
... cache_interval=5,
... max_order=1,
... disable_cache_before_step=3,
... taylor_factors_dtype=torch.float32,
... )
>>> pipe.transformer.enable_cache(config)MagCacheConfig
类 diffusers.MagCacheConfig
< source >( threshold: float = 0.06max_skip_steps: int = 3retention_ratio: float = 0.2num_inference_steps: int = 28mag_ratios: typing.Union[torch.Tensor, typing.List[float], NoneType] = Nonecalibrate: bool = False )
参数
- threshold (
float, 默认为0.06) — 累积误差的阈值。如果累积误差低于此阈值,则跳过该块的计算。较高的阈值允许更激进的跳过(速度更快),但可能会降低质量。 - max_skip_steps (
int,默认为3) — 可跳过的最大连续步数(论文中的 K)。 - retention_ratio (
float,默认为0.2) — 在初始步骤中禁用跳步的比例,以确保稳定性。例如,如果num_inference_steps为 28 且retention_ratio为 0.2,则前 6 个步骤永远不会被跳过。 - num_inference_steps (
int,默认为28) — 流水线中使用的推理步数。此参数用于正确插值mag_ratios。 - mag_ratios (
torch.Tensor, 可选) — 模型的预计算幅度比率。这些比率依赖于检查点。如果未提供,您必须设置calibrate=True以针对您的特定模型进行计算。对于 Flux 模型,您可以使用diffusers.hooks.mag_cache.FLUX_MAG_RATIOS。 - calibrate (
bool,默认为False) — 如果为 True,则启用校准模式。在此模式下,不会跳过任何块。相反,钩子会计算当前运行的幅度比率,并在结束时记录它们。使用此选项可为新模型或调度器获取mag_ratios。
MagCache 的配置。
diffusers.apply_mag_cache
< source >( module: Moduleconfig: MagCacheConfig )
将 MagCache 应用于给定模块(通常是 Transformer)。
SeaCacheConfig
类 diffusers.SeaCacheConfig
< source >( threshold: float = 0.25residual_order: int = 1retention_steps: int = 1cache_end_steps: int = 1max_consecutive_cached: int = 2power_exp: float = 3.0raw_vision_callback: typing.Callable[[torch.nn.Module, tuple[typing.Any, ...], dict[str, typing.Any]], list[torch.Tensor] | None] = None )
参数
- threshold (
float,默认为0.25) — 累积的相对 L1 预算。值越大,越频繁地复用缓存。 - residual_order (
int,默认为1) — 用于预测 transformer 残差的顺序。0直接复用最近的残差,1则从最近两次完整执行进行线性外推。 - retention_steps (
int,默认为1) — 始终完整执行的初始调度器步数。 - cache_end_steps (
int,默认为1) — 始终完整执行的最后几个调度器步骤的数量。 - max_consecutive_cached (
int,默认为2) — 在强制完整执行之前,每个缓存上下文允许的最大连续残差复用次数。0表示禁用此限制。 - power_exp (
float,默认为3.0) — SEA 干净信号功率先验的指数。SeaCache 对视频特征使用3.0。 - raw_vision_callback (
Callable, 可选) — 高级模型适配器,返回构成生成输出的视觉潜在变量,每个的形状为(C, T, H, W)。在输出轨迹中包含干净的 conditioning 帧,但排除不属于输出的独立视觉提示。如果省略,则使用内置适配器(如果可用)。
SeaCache 的配置。
SeaCache 默认处于禁用状态,只有在将此配置传递给 transformer.enable_cache(config) 后才会激活。
SeaCache 比较调度器步骤之间的频谱演化感知(SEA)指标。如果它们的累积相对变化保持在 threshold 以下,则重复的 Transformer 块堆栈将被替换为缓存的残差。特定于模型的适配器可以将残差边界扩展到这些块之外。其他视频 Transformer 可以通过注册其块元数据并提供 raw_vision_callback 来使用通用路径;其流水线必须为每个 Transformer 轨迹进入 cache_context。
diffusers.apply_sea_cache
< source >( module: Moduleconfig: SeaCacheConfig )
将 SeaCache 应用于受支持的 transformer。
SeaCache 会复用 Transformer 重复块堆栈的缓存残差。模型适配器可以自定义指示符提取和残差边界;通用路径使用已注册的块元数据和 raw_vision_callback。必须针对每个模型验证兼容性和缓存参数。