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)