Buckets:

mishig's picture
|
download
raw
29.6 kB
# 缓存方法
缓存方法通过存储并复用特定层(如注意力层和前馈层)的中间输出,而非在每个推理步骤中重新计算,从而加速扩散 Transformer。
## CacheMixin[[diffusers.CacheMixin]]
#### diffusers.CacheMixin[[diffusers.CacheMixin]]
```python
diffusers.CacheMixin()
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/cache_utils.py#L23)
用于在扩散模型上启用/禁用缓存技术的类。
支持的缓存技术:
- [金字塔注意力广播](https://huggingface.co/papers/2408.12588)
- [FasterCache](https://huggingface.co/papers/2410.19355)
- [FirstBlockCache](https://github.com/chengzeyi/ParaAttention/blob/7a266123671b55e7e5a2fe9af3121f07a36afc78/README.md#first-block-cache-our-dynamic-caching)
- [SeaCache](https://huggingface.co/papers/2602.18993)
#### cache_context[[diffusers.CacheMixin.cache_context]]
```python
cache_context(name: str, **kwargs)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/cache_utils.py#L181)
提供缓存管理信息的上下文管理器。
`name` 是去噪调用的名称,通常为 `"cond"` 或 `"uncond"`。关键字参数用于描述去噪循环所处的位置——有关接受的字段,请参阅 `CacheContext`,例如 `cache_context("cond", step_index=i, sigma=sigma)`。
#### enable_cache[[diffusers.CacheMixin.enable_cache]]
```python
enable_cache(config)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/cache_utils.py#L40)
**参数:**
config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | SeaCacheConfig | TextKVCacheConfig`):用于应用缓存技术的配置。当前支持的缓存技术包括:- [PyramidAttentionBroadcastConfig](/docs/diffusers/v0.43.0/zh/api/cache#diffusers.PyramidAttentionBroadcastConfig) - [FasterCacheConfig](/docs/diffusers/v0.43.0/zh/api/cache#diffusers.FasterCacheConfig) - [FirstBlockCacheConfig](/docs/diffusers/v0.43.0/zh/api/cache#diffusers.FirstBlockCacheConfig) - [SeaCacheConfig](/docs/diffusers/v0.43.0/zh/api/cache#diffusers.SeaCacheConfig) - `TextKVCacheConfig`
在模型上启用缓存技术。
示例:
```python
>>> 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]]
#### diffusers.PyramidAttentionBroadcastConfig[[diffusers.PyramidAttentionBroadcastConfig]]
```python
diffusers.PyramidAttentionBroadcastConfig(spatial_attention_block_skip_range: int | None = None, temporal_attention_block_skip_range: int | None = None, cross_attention_block_skip_range: int | None = None, 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', '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)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/pyramid_attention_broadcast.py#L40)
**参数:**
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)`):在 cross-attention 层中需要跳过的时间步范围。如果当前时间步位于指定范围内,则条件性地跳过注意力计算。
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[[diffusers.apply_pyramid_attention_broadcast]]
```python
diffusers.apply_pyramid_attention_broadcast(module: Module, config: PyramidAttentionBroadcastConfig)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/pyramid_attention_broadcast.py#L181)
**参数:**
module (`torch.nn.Module`):要应用 Pyramid Attention Broadcast 的模块。
config (`PyramidAttentionBroadcastConfig | None`, `optional`, 默认为 `None`) : 用于 Pyramid Attention Broadcast 的配置。
将 [Pyramid Attention Broadcast](https://huggingface.co/papers/2408.12588) 应用于指定的 pipeline。
PAB 是一种注意力近似方法,它利用不同时间步之间注意力状态的相似性来降低注意力计算的开销。论文的核心结论是:跨时间步的交叉注意力层中注意力相似性最高,其次是时间层和空间层。这使得交叉注意力层可以比时间层和空间层更频繁地跳过注意力计算。因此,应用 PAB 可以加速推理过程。
示例:
```python
>>> 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]]
#### diffusers.FasterCacheConfig[[diffusers.FasterCacheConfig]]
```python
diffusers.FasterCacheConfig(spatial_attention_block_skip_range: int = 2, temporal_attention_block_skip_range: int | None = None, spatial_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.1, alpha_high_frequency: float = 1.1, unconditional_batch_skip_range: int = 5, unconditional_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] = None, low_frequency_weight_callback: typing.Callable[[torch.nn.Module], float] = None, high_frequency_weight_callback: typing.Callable[[torch.nn.Module], float] = None, tensor_format: str = 'BCFHW', is_guidance_distilled: bool = False, current_timestep_callback: typing.Callable[[], int] = None, _unconditional_conditional_input_kwargs_identifiers: list[str] = ('hidden_states', 'encoder_hidden_states', 'timestep', 'attention_mask', 'encoder_attention_mask'))
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/faster_cache.py#L50)
**参数:**
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 蒸馏。如果模型经过 guidance 蒸馏,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](https://huggingface.co/papers/2410.19355) 的配置。
#### diffusers.apply_faster_cache[[diffusers.apply_faster_cache]]
```python
diffusers.apply_faster_cache(module: Module, config: FasterCacheConfig)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/faster_cache.py#L486)
**参数:**
module (`torch.nn.Module`):要应用 FasterCache 的 PyTorch 模块。通常,这应该是 Diffusers 支持的 Transformer 架构,例如 `CogVideoXTransformer3DModel`,但外部实现也可能有效。
config (`FasterCacheConfig`):用于 FasterCache 的配置。
将 [FasterCache](https://huggingface.co/papers/2410.19355) 应用于指定的 pipeline。
示例:
```python
>>> 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]]
#### diffusers.FirstBlockCacheConfig[[diffusers.FirstBlockCacheConfig]]
```python
diffusers.FirstBlockCacheConfig(threshold: float = 0.05)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/first_block_cache.py#L33)
**参数:**
threshold (`float`,默认值为 `0.05`):用于判断是否需要执行模型所有层的前向传播的阈值。较高的阈值通常意味着前向传播经过的层数更少,推理速度更快,但可能导致生成质量下降。较低的阈值可能无法带来显著的生成速度提升。该阈值与第一个 Transformer 块中当前输出与缓存输出之间残差的绝对均值差进行比较。如果该差值低于阈值,则跳过前向传播。
[First Block Cache](https://github.com/chengzeyi/ParaAttention/blob/7a266123671b55e7e5a2fe9af3121f07a36afc78/README.md#first-block-cache-our-dynamic-caching) 的配置。
#### diffusers.apply_first_block_cache[[diffusers.apply_first_block_cache]]
```python
diffusers.apply_first_block_cache(module: Module, config: FirstBlockCacheConfig)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/first_block_cache.py#L193)
**参数:**
module (`torch.nn.Module`):要应用 FBCache 的 PyTorch 模块。通常,这应该是 Diffusers 支持的 Transformer 架构,例如 `CogVideoXTransformer3DModel`,但外部实现也可能有效。
config (`FirstBlockCacheConfig`):用于应用 FBCache 方法的配置。
将 [First Block Cache](https://github.com/chengzeyi/ParaAttention/blob/4de137c5b96416489f06e43e19f2c14a772e28fd/README.md#first-block-cache-our-dynamic-caching) 应用于指定模块。
First Block Cache 基于 [TeaCache](https://huggingface.co/papers/2411.19108) 的思想构建。它在更广泛的模型范围内实现起来更为简单,并首先作为实验目的被集成。
示例:
```python
>>> 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]]
#### diffusers.TaylorSeerCacheConfig[[diffusers.TaylorSeerCacheConfig]]
```python
diffusers.TaylorSeerCacheConfig(cache_interval: int = 5, disable_cache_before_step: int = 3, disable_cache_after_step: int | None = None, max_order: int = 1, taylor_factors_dtype: typing.Optional[torch.dtype] = torch.bfloat16, skip_predict_identifiers: list[str] | None = None, cache_identifiers: list[str] | None = None, use_lite_mode: bool = False)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/taylorseer_cache.py#L26)
**参数:**
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`,则只会对匹配的模块进行挂钩。
- 如果两者均未提供,则默认会挂钩所有注意力类模块。
未激活和激活使用的示例:
```py
def forward(x):
x = self.module1(x) # inactive module: returns zeros tensor based on shape recorded during full compute
x = self.module2(x) # active module: caches output here, avoiding recomputation of prior steps
return x
```
#### bfloat16[[diffusers.TaylorSeerCacheConfig.taylor_factors_dtype]]
```python
bfloat16(*args, **kwargs)
```
点分路径(例如 `torch.float32`)的模拟值:属性访问链和调用行为如同透传装饰器,`repr` 返回点分路径,将其用作基类时会替换为普通的 `type` 基类(PEP 560 `__mro_entries__`),因此真实的子类保留正常的元类,且 `inspect.signature` 读取的是它们真实的 `__init__`,而非模拟对象的。
#### diffusers.apply_taylorseer_cache[[diffusers.apply_taylorseer_cache]]
```python
diffusers.apply_taylorseer_cache(module: Module, config: TaylorSeerCacheConfig)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/taylorseer_cache.py#L261)
**参数:**
module (torch.nn.Module):要应用钩子的模型子树。
config TaylorSeerCacheConfig:缓存的配置。
将 TaylorSeer 缓存应用于给定的流水线(通常是 transformer / UNet)。
此函数会挂钩模型中指定的模块,以便根据提供的配置启用缓存或跳过操作,从而减少扩散去噪循环中的冗余计算。
示例:
```python
>>> 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]]
#### diffusers.MagCacheConfig[[diffusers.MagCacheConfig]]
```python
diffusers.MagCacheConfig(threshold: float = 0.06, max_skip_steps: int = 3, retention_ratio: float = 0.2, num_inference_steps: int = 28, mag_ratios: typing.Union[torch.Tensor, typing.List[float], NoneType] = None, calibrate: bool = False)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/mag_cache.py#L84)
**参数:**
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](https://github.com/Zehong-Ma/MagCache) 的配置。
#### diffusers.apply_mag_cache[[diffusers.apply_mag_cache]]
```python
diffusers.apply_mag_cache(module: Module, config: MagCacheConfig)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/mag_cache.py#L398)
**参数:**
module (`torch.nn.Module`):要应用 MagCache 的模块。
config (`MagCacheConfig`):MagCache 的配置。
将 MagCache 应用于给定模块(通常是 Transformer)。
## SeaCacheConfig[[diffusers.SeaCacheConfig]]
#### diffusers.SeaCacheConfig[[diffusers.SeaCacheConfig]]
```python
diffusers.SeaCacheConfig(threshold: float = 0.25, residual_order: int = 1, retention_steps: int = 1, cache_end_steps: int = 1, max_consecutive_cached: int = 2, power_exp: float = 3.0, raw_vision_callback: typing.Callable[[torch.nn.Module, tuple[typing.Any, ...], dict[str, typing.Any]], list[torch.Tensor] | None] = None)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/sea_cache.py#L38)
**参数:**
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](https://huggingface.co/papers/2602.18993) 的配置。
SeaCache 默认处于禁用状态,只有在将此配置传递给 `transformer.enable_cache(config)` 后才会激活。
SeaCache 会对比调度器步骤之间的 Spectral-Evolution-Aware (SEA) 指标。如果它们的累积相对变化保持在 `threshold` 以下,重复的 transformer 块堆栈将被替换为缓存的残差。特定于模型的适配器可以将残差边界扩展到这些块之外。其他视频 transformer 可以通过注册其块元数据并提供 `raw_vision_callback` 来使用通用路径;其 pipeline 必须为每个 transformer 轨迹进入 `cache_context`。
示例:
```python
>>> from diffusers import Cosmos3OmniPipeline, SeaCacheConfig
>>> pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano")
>>> pipe.transformer.enable_cache(SeaCacheConfig())
```
#### diffusers.apply_sea_cache[[diffusers.apply_sea_cache]]
```python
diffusers.apply_sea_cache(module: Module, config: SeaCacheConfig)
```
[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/hooks/sea_cache.py#L980)
**参数:**
module (`torch.nn.Module`):要缓存的 Transformer 模块。
config (`SeaCacheConfig`):SeaCache 配置。
将 SeaCache 应用于受支持的 transformer。
SeaCache 会复用 Transformer 重复块堆栈的缓存残差。模型适配器可以自定义指示器提取和残差边界;通用路径使用已注册的块元数据以及 `raw_vision_callback`。针对每个模型,都必须验证兼容性和缓存参数。

Xet Storage Details

Size:
29.6 kB
·
Xet hash:
f9acf65e83704081835f355a62996bcfe0798e6364d4a5a3c087d90825ad8df9

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.