feat: implement cache_context for pab and fastercache.

docs/source/en/optimization/cache.md

@@ -15,6 +15,11 @@ Caching accelerates inference by storing and reusing intermediate outputs of dif
 
 This guide shows you how to use the caching methods supported in Diffusers.
 
+Pyramid Attention Broadcast and FasterCache read the current timestep from the denoiser's `cache_context`.
+When writing a custom denoising loop, wrap each denoiser call with `model.cache_context("cond", timestep=t)`.
+Use a separate context name for each guidance branch, or `"cond_uncond"` for a combined batch, and call
+`model._reset_stateful_cache()` before starting a new generation. The pipeline examples below handle this for you.
+
 ## Pyramid Attention Broadcast
 
 [Pyramid Attention Broadcast (PAB)](https://huggingface.co/papers/2408.12588) is based on the observation that attention outputs aren't that different between successive timesteps of the generation process. The attention differences are smallest in the cross attention layers and are generally cached over a longer timestep range. This is followed by temporal attention and spatial attention layers.
@@ -36,7 +41,6 @@ pipeline.to("cuda")  # or "mps", "xpu", "cpu"
 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)
 ```
@@ -53,13 +57,12 @@ Set up and pass a [`FasterCacheConfig`] to a pipeline's transformer to enable it
 import torch
 from diffusers import CogVideoXPipeline, FasterCacheConfig
 
-pipe line= CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
+pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
 pipeline.to("cuda")  # or "mps", "xpu", "cpu"
 
 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),

docs/source/zh/optimization/cache.md

@@ -33,7 +33,6 @@ 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)
 ```
@@ -57,7 +56,6 @@ 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),

src/diffusers/hooks/pyramid_attention_broadcast.py

@@ -20,14 +20,14 @@
 
 from ..models.attention import AttentionModuleMixin
 from ..models.attention_processor import Attention, MochiAttention
-from ..utils import logging
+from ..utils import deprecate, logging
 from ._common import (
     _ATTENTION_CLASSES,
     _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS,
     _SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS,
     _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS,
 )
-from .hooks import HookRegistry, ModelHook
+from .hooks import CacheContext, HookRegistry, ModelHook, StateManager
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
@@ -83,12 +83,24 @@ class PyramidAttentionBroadcastConfig:
     temporal_attention_block_identifiers: tuple[str, ...] = _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS
     cross_attention_block_identifiers: tuple[str, ...] = _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS
 
-    current_timestep_callback: Callable[[], int] = None
+    current_timestep_callback: Callable[[], int] | None = None
 
     # TODO(aryan): add PAB for MLP layers (very limited speedup from testing with original codebase
     # so not added for now)
 
+    def __post_init__(self):
+        if self.current_timestep_callback is not None:
+            depr_message = (
+                "Passing `current_timestep_callback` to `PyramidAttentionBroadcastConfig` is deprecated. "
+                "Please use `cache_context(name, timestep=...)` to pass the current timestep instead. "
+                "See the caching documentation for guidance: https://huggingface.co/docs/diffusers/main/en/optimization/cache."
+            )
+            deprecate("current_timestep_callback", "0.45.0", depr_message)
+
     def __repr__(self) -> str:
+        callback_repr = ""
+        if self.current_timestep_callback is not None:
+            callback_repr = f"  current_timestep_callback={self.current_timestep_callback},\n"
         return (
             f"PyramidAttentionBroadcastConfig(\n"
             f"  spatial_attention_block_skip_range={self.spatial_attention_block_skip_range},\n"
@@ -100,7 +112,7 @@ def __repr__(self) -> str:
             f"  spatial_attention_block_identifiers={self.spatial_attention_block_identifiers},\n"
             f"  temporal_attention_block_identifiers={self.temporal_attention_block_identifiers},\n"
             f"  cross_attention_block_identifiers={self.cross_attention_block_identifiers},\n"
-            f"  current_timestep_callback={self.current_timestep_callback}\n"
+            f"{callback_repr}"
             ")"
         )
 
@@ -141,7 +153,10 @@ class PyramidAttentionBroadcastHook(ModelHook):
     _is_stateful = True
 
     def __init__(
-        self, timestep_skip_range: tuple[int, int], block_skip_range: int, current_timestep_callback: Callable[[], int]
+        self,
+        timestep_skip_range: tuple[int, int],
+        block_skip_range: int,
+        current_timestep_callback: Callable[[], int] | None = None,
     ) -> None:
         super().__init__()
 
@@ -150,31 +165,37 @@ def __init__(
         self.current_timestep_callback = current_timestep_callback
 
     def initialize_hook(self, module):
-        self.state = PyramidAttentionBroadcastState()
+        self.state_manager = StateManager(PyramidAttentionBroadcastState)
         return module
 
     def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any:
-        is_within_timestep_range = (
-            self.timestep_skip_range[0] < self.current_timestep_callback() < self.timestep_skip_range[1]
-        )
+        if self.state_manager._context is None and self.current_timestep_callback is not None:
+            self.state_manager.set_context(CacheContext(name="inference"))
+        timestep = self.state_manager.context.timestep
+        if timestep is None and self.current_timestep_callback is not None:
+            timestep = self.current_timestep_callback()
+        if timestep is None:
+            raise ValueError("Pyramid Attention Broadcast requires `cache_context(name, timestep=...)`.")
+        state = self.state_manager.get_state()
+        is_within_timestep_range = self.timestep_skip_range[0] < timestep < self.timestep_skip_range[1]
         should_compute_attention = (
-            self.state.cache is None
-            or self.state.iteration == 0
+            state.cache is None
+            or state.iteration == 0
             or not is_within_timestep_range
-            or self.state.iteration % self.block_skip_range == 0
+            or state.iteration % self.block_skip_range == 0
         )
 
         if should_compute_attention:
             output = self.fn_ref.original_forward(*args, **kwargs)
         else:
-            output = self.state.cache
+            output = state.cache
 
-        self.state.cache = output
-        self.state.iteration += 1
+        state.cache = output
+        state.iteration += 1
         return output
 
     def reset_state(self, module: torch.nn.Module) -> None:
-        self.state.reset()
+        self.state_manager.reset()
         return module
 
 
@@ -207,16 +228,10 @@ def apply_pyramid_attention_broadcast(module: torch.nn.Module, config: PyramidAt
     >>> 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)
     ```
     """
-    if config.current_timestep_callback is None:
-        raise ValueError(
-            "The `current_timestep_callback` function must be provided in the configuration to apply Pyramid Attention Broadcast."
-        )
-
     if (
         config.spatial_attention_block_skip_range is None
         and config.temporal_attention_block_skip_range is None
@@ -291,7 +306,7 @@ def _apply_pyramid_attention_broadcast_hook(
     module: Attention | MochiAttention,
     timestep_skip_range: tuple[int, int],
     block_skip_range: int,
-    current_timestep_callback: Callable[[], int],
+    current_timestep_callback: Callable[[], int] | None = None,
 ):
     r"""
     Apply [Pyramid Attention Broadcast](https://huggingface.co/papers/2408.12588) to a given torch.nn.Module.
@@ -306,8 +321,6 @@ def _apply_pyramid_attention_broadcast_hook(
             The number of times a specific attention broadcast is skipped before computing the attention states to
             re-use. If this is set to the value `N`, the attention computation will be skipped `N - 1` times (i.e., old
             attention states will be reused) before computing the new attention states again.
-        current_timestep_callback (`Callable[[], int]`):
-            A callback function that returns the current inference timestep.
     """
     registry = HookRegistry.check_if_exists_or_initialize(module)
     hook = PyramidAttentionBroadcastHook(timestep_skip_range, block_skip_range, current_timestep_callback)

src/diffusers/models/cache_utils.py

@@ -62,7 +62,6 @@ def enable_cache(self, config) -> None:
         >>> 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)
         ```

src/diffusers/modular_pipelines/cosmos/denoise.py

@@ -233,6 +233,7 @@ def __call__(
             }
             with components.transformer.cache_context(
                 pass_name,
+                timestep=t,
                 step_index=i,
                 sigma=float(components.scheduler.sigmas[i]),
                 num_inference_steps=components.scheduler.num_inference_steps,
@@ -797,9 +798,11 @@ def intermediate_outputs(self) -> list[OutputParam]:
         return [OutputParam("velocity", type_hint=torch.Tensor, description="Predicted (masked) transfer velocity.")]
 
     @staticmethod
-    def _forward(components, static, vision_tokens, vision_timesteps, context_name, step, sigma, num_inference_steps):
+    def _forward(
+        components, static, vision_tokens, vision_timesteps, context_name, step, timestep, sigma, num_inference_steps
+    ):
         with components.transformer.cache_context(
-            context_name, step_index=step, sigma=sigma, num_inference_steps=num_inference_steps
+            context_name, step_index=step, timestep=timestep, sigma=sigma, num_inference_steps=num_inference_steps
         ):
             preds_vision, _, _ = components.transformer(
                 input_ids=static["input_ids"],
@@ -847,6 +850,7 @@ def __call__(
             block_state.vision_timesteps,
             "cond",
             step=i,
+            timestep=t,
             sigma=float(components.scheduler.sigmas[i]),
             num_inference_steps=components.scheduler.num_inference_steps,
         )
@@ -860,6 +864,7 @@ def __call__(
                 block_state.vision_timesteps,
                 "cond_no_control",
                 step=i,
+                timestep=t,
                 sigma=float(components.scheduler.sigmas[i]),
                 num_inference_steps=components.scheduler.num_inference_steps,
             )
@@ -873,6 +878,7 @@ def __call__(
                 block_state.vision_timesteps,
                 "uncond",
                 step=i,
+                timestep=t,
                 sigma=float(components.scheduler.sigmas[i]),
                 num_inference_steps=components.scheduler.num_inference_steps,
             )

src/diffusers/modular_pipelines/echo/denoise.py

@@ -233,34 +233,37 @@ def __call__(
         audio_context = block_state.connector_audio_prompt_embeds.to(transformer_dtype)
         context_mask = block_state.connector_attention_mask
 
-        velocity_video, velocity_audio = components.transformer(
-            hidden_states=block_state.latent_model_input.to(transformer_dtype),
-            audio_hidden_states=block_state.audio_latent_model_input.to(transformer_dtype),
-            encoder_hidden_states=video_context,
-            audio_encoder_hidden_states=audio_context,
-            timestep=block_state.model_video_timestep,
-            audio_timestep=block_state.model_audio_timestep,
-            # Echo deliberately uses the first target token as its global video sigma. When a clean first frame is
-            # present this is zero, while the audio branch keeps the current DMD sigma. `use_cross_timestep=True`
-            # exchanges these values in the cross-modal modulation blocks, matching the released wrapper.
-            sigma=block_state.video_timestep[:, 0],
-            audio_sigma=block_state.audio_timestep[:, 0],
-            encoder_attention_mask=context_mask,
-            audio_encoder_attention_mask=context_mask,
-            num_frames=block_state.latent_num_frames,
-            height=block_state.latent_height,
-            width=block_state.latent_width,
-            fps=block_state.frame_rate,
-            audio_num_frames=block_state.audio_num_frames,
-            video_coords=block_state.model_video_coords,
-            audio_coords=block_state.model_audio_coords,
-            isolate_modalities=False,
-            spatio_temporal_guidance_blocks=None,
-            perturbation_mask=None,
-            use_cross_timestep=True,
-            attention_kwargs=block_state.attention_kwargs,
-            return_dict=False,
-        )
+        with components.transformer.cache_context(
+            "inference", timestep=sigma * components.transformer.config.timestep_scale_multiplier
+        ):
+            velocity_video, velocity_audio = components.transformer(
+                hidden_states=block_state.latent_model_input.to(transformer_dtype),
+                audio_hidden_states=block_state.audio_latent_model_input.to(transformer_dtype),
+                encoder_hidden_states=video_context,
+                audio_encoder_hidden_states=audio_context,
+                timestep=block_state.model_video_timestep,
+                audio_timestep=block_state.model_audio_timestep,
+                # Echo deliberately uses the first target token as its global video sigma. When a clean first frame is
+                # present this is zero, while the audio branch keeps the current DMD sigma. `use_cross_timestep=True`
+                # exchanges these values in the cross-modal modulation blocks, matching the released wrapper.
+                sigma=block_state.video_timestep[:, 0],
+                audio_sigma=block_state.audio_timestep[:, 0],
+                encoder_attention_mask=context_mask,
+                audio_encoder_attention_mask=context_mask,
+                num_frames=block_state.latent_num_frames,
+                height=block_state.latent_height,
+                width=block_state.latent_width,
+                fps=block_state.frame_rate,
+                audio_num_frames=block_state.audio_num_frames,
+                video_coords=block_state.model_video_coords,
+                audio_coords=block_state.model_audio_coords,
+                isolate_modalities=False,
+                spatio_temporal_guidance_blocks=None,
+                perturbation_mask=None,
+                use_cross_timestep=True,
+                attention_kwargs=block_state.attention_kwargs,
+                return_dict=False,
+            )
         velocity_video = velocity_video[:, block_state.memory_video_token_count :]
         velocity_audio = velocity_audio[:, block_state.memory_audio_token_count :]
         timestep_scale = float(components.transformer.config.timestep_scale_multiplier)

src/diffusers/modular_pipelines/flux/denoise.py

@@ -93,17 +93,18 @@ def inputs(self) -> list[tuple[str, Any]]:
     def __call__(
         self, components: FluxModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
     ) -> tuple[FluxModularPipeline, BlockState]:
-        noise_pred = components.transformer(
-            hidden_states=block_state.latents,
-            timestep=t.flatten() / 1000,
-            guidance=block_state.guidance,
-            encoder_hidden_states=block_state.prompt_embeds,
-            pooled_projections=block_state.pooled_prompt_embeds,
-            joint_attention_kwargs=block_state.joint_attention_kwargs,
-            txt_ids=block_state.txt_ids,
-            img_ids=block_state.img_ids,
-            return_dict=False,
-        )[0]
+        with components.transformer.cache_context("inference", timestep=t):
+            noise_pred = components.transformer(
+                hidden_states=block_state.latents,
+                timestep=t.flatten() / 1000,
+                guidance=block_state.guidance,
+                encoder_hidden_states=block_state.prompt_embeds,
+                pooled_projections=block_state.pooled_prompt_embeds,
+                joint_attention_kwargs=block_state.joint_attention_kwargs,
+                txt_ids=block_state.txt_ids,
+                img_ids=block_state.img_ids,
+                return_dict=False,
+            )[0]
         block_state.noise_pred = noise_pred
 
         return components, block_state
@@ -182,17 +183,18 @@ def __call__(
             latent_model_input = torch.cat([latent_model_input, image_latents], dim=1)
 
         timestep = t.expand(latents.shape[0]).to(latents.dtype)
-        noise_pred = components.transformer(
-            hidden_states=latent_model_input,
-            timestep=timestep / 1000,
-            guidance=block_state.guidance,
-            encoder_hidden_states=block_state.prompt_embeds,
-            pooled_projections=block_state.pooled_prompt_embeds,
-            joint_attention_kwargs=block_state.joint_attention_kwargs,
-            txt_ids=block_state.txt_ids,
-            img_ids=block_state.img_ids,
-            return_dict=False,
-        )[0]
+        with components.transformer.cache_context("inference", timestep=t):
+            noise_pred = components.transformer(
+                hidden_states=latent_model_input,
+                timestep=timestep / 1000,
+                guidance=block_state.guidance,
+                encoder_hidden_states=block_state.prompt_embeds,
+                pooled_projections=block_state.pooled_prompt_embeds,
+                joint_attention_kwargs=block_state.joint_attention_kwargs,
+                txt_ids=block_state.txt_ids,
+                img_ids=block_state.img_ids,
+                return_dict=False,
+            )[0]
         noise_pred = noise_pred[:, : latents.size(1)]
         block_state.noise_pred = noise_pred