Diffusers 文档
缓存方法
并获得增强的文档体验
开始使用
缓存方法
缓存方法通过存储和重用特定层（如注意力层和前馈层）的中间输出，而不是在每个推理步骤重新计算它们，从而加速扩散 Transformer。
CacheMixin
用于启用/禁用扩散模型上的缓存技术的类。
支持的缓存技术
enable_cache
< source >( config )
在模型上启用缓存技术。
示例
>>> 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
class diffusers.PyramidAttentionBroadcastConfig
< source >( spatial_attention_block_skip_range: typing.Optional[int] = None temporal_attention_block_skip_range: typing.Optional[int] = None cross_attention_block_skip_range: typing.Optional[int] = None spatial_attention_timestep_skip_range: typing.Tuple[int, int] = (100, 800) temporal_attention_timestep_skip_range: typing.Tuple[int, int] = (100, 800) cross_attention_timestep_skip_range: typing.Tuple[int, int] = (100, 800) spatial_attention_block_identifiers: typing.Tuple[str, ...] = ('blocks', 'transformer_blocks', 'single_transformer_blocks') temporal_attention_block_identifiers: typing.Tuple[str, ...] = ('temporal_transformer_blocks',) cross_attention_block_identifiers: typing.Tuple[str, ...] = ('blocks', 'transformer_blocks') 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, ...], 默认为("blocks", "transformer_blocks")) — 用于匹配层名称以确定该层是否为空间注意力层的标识符。 - temporal_attention_block_identifiers (
Tuple[str, ...], 默认为("temporal_transformer_blocks",)) — 用于匹配层名称以确定该层是否为时间注意力层的标识符。 - cross_attention_block_identifiers (
Tuple[str, ...], 默认为("blocks", "transformer_blocks")) — 用于匹配层名称以确定该层是否为交叉注意力层的标识符。
金字塔注意力广播的配置。
diffusers.apply_pyramid_attention_broadcast
< source >( module: Module config: PyramidAttentionBroadcastConfig )
将金字塔注意力广播应用于给定的管道。
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
class diffusers.FasterCacheConfig
< source >( spatial_attention_block_skip_range: int = 2 temporal_attention_block_skip_range: typing.Optional[int] = None spatial_attention_timestep_skip_range: typing.Tuple[int, int] = (-1, 681) temporal_attention_timestep_skip_range: typing.Tuple[int, int] = (-1, 681) low_frequency_weight_update_timestep_range: typing.Tuple[int, int] = (99, 901) high_frequency_weight_update_timestep_range: typing.Tuple[int, int] = (-1, 301) alpha_low_frequency: float = 1.1 alpha_high_frequency: float = 1.1 unconditional_batch_skip_range: int = 5 unconditional_batch_timestep_skip_range: typing.Tuple[int, int] = (-1, 641) spatial_attention_block_identifiers: typing.Tuple[str, ...] = ('^blocks.*attn', '^transformer_blocks.*attn', '^single_transformer_blocks.*attn') temporal_attention_block_identifiers: typing.Tuple[str, ...] = ('^temporal_transformer_blocks.*attn',) attention_weight_callback: typing.Callable[[torch.nn.modules.module.Module], float] = None low_frequency_weight_callback: typing.Callable[[torch.nn.modules.module.Module], float] = None high_frequency_weight_callback: typing.Callable[[torch.nn.modules.module.Module], float] = None tensor_format: str = 'BCFHW' is_guidance_distilled: bool = False current_timestep_callback: typing.Callable[[], int] = None _unconditional_conditional_input_kwargs_identifiers: typing.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) — 模型是否经过指导蒸馏。如果模型经过指导蒸馏，FasterCache 不会在去噪器级别应用以跳过无条件分支计算（因为不存在）。 - _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
< 来源 >( module: Module config: FasterCacheConfig )
将 FasterCache 应用于给定的管道。
示例
>>> 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)