update

Author: a-r-r-o-w

src/diffusers/models/hooks.py

@@ -117,45 +117,72 @@ def reset_state(self, module):
 class PyramidAttentionBroadcastHook(ModelHook):
     def __init__(
         self,
-        skip_range: int,
-        timestep_range: Tuple[int, int],
-        timestep_callback: Callable[[], Union[torch.LongTensor, int]],
+        skip_callback: Callable[[torch.nn.Module], bool],
+        # skip_range: int,
+        # timestep_range: Tuple[int, int],
+        # timestep_callback: Callable[[], Union[torch.LongTensor, int]],
     ) -> None:
         super().__init__()
 
-        self.skip_range = skip_range
-        self.timestep_range = timestep_range
-        self.timestep_callback = timestep_callback
+        # self.skip_range = skip_range
+        # self.timestep_range = timestep_range
+        # self.timestep_callback = timestep_callback
+        self.skip_callback = skip_callback
 
-        self.attention_cache = None
+        self.cache = None
         self._iteration = 0
 
     def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any:
         args, kwargs = module._diffusers_hook.pre_forward(module, *args, **kwargs)
 
-        current_timestep = self.timestep_callback()
-        is_within_timestep_range = self.timestep_range[0] < current_timestep < self.timestep_range[1]
-        should_compute_attention = self._iteration % self.skip_range == 0
+        # current_timestep = self.timestep_callback()
+        # is_within_timestep_range = self.timestep_range[0] < current_timestep < self.timestep_range[1]
+        # should_compute_attention = self._iteration % self.skip_range == 0
 
-        if not is_within_timestep_range or should_compute_attention:
-            output = module._old_forward(*args, **kwargs)
-        else:
-            output = self.attention_cache
+        # if not is_within_timestep_range or should_compute_attention:
+        #     output = module._old_forward(*args, **kwargs)
+        # else:
+        #     output = self.attention_cache
 
-        self._iteration = self._iteration + 1
+        if self.cache is not None and self.skip_callback(module):
+            output = self.cache
+        else:
+            output = module._old_forward(*args, **kwargs)
 
         return module._diffusers_hook.post_forward(module, output)
 
     def post_forward(self, module: torch.nn.Module, output: Any) -> Any:
-        self.attention_cache = output
+        self.cache = output
         return output
 
     def reset_state(self, module: torch.nn.Module) -> torch.nn.Module:
-        self.attention_cache = None
+        self.cache = None
         self._iteration = 0
         return module
 
 
+class LayerSkipHook(ModelHook):
+    def __init__(self, skip_: Callable[[torch.nn.Module], bool]) -> None:
+        super().__init__()
+
+        self.skip_callback = skip_
+    
+    def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any:
+        args, kwargs = module._diffusers_hook.pre_forward(module, *args, **kwargs)
+
+        if self.skip_callback(module):
+            # We want to skip this layer, so we have to return the input of the current layer
+            # as output of the next layer. But at this point, we don't have information about
+            # the arguments required by next layer. Even if we did, order matters unless we
+            # always pass kwargs. But that is not the case usually with hidden_states, encoder_hidden_states,
+            # temb, etc. TODO(aryan): implement correctly later
+            output = None
+        else:
+            output = module._old_forward(*args, **kwargs)
+        
+        return module._diffusers_hook.post_forward(module, output)
+
+
 def add_hook_to_module(module: torch.nn.Module, hook: ModelHook, append: bool = False):
     r"""
     Adds a hook to a given module. This will rewrite the `forward` method of the module to include the hook, to remove

src/diffusers/pipelines/allegro/pipeline_allegro.py

@@ -38,7 +38,6 @@
 )
 from ...utils.torch_utils import randn_tensor
 from ...video_processor import VideoProcessor
-from ..pyramid_broadcast_utils import PyramidAttentionBroadcastMixin
 from .pipeline_output import AllegroPipelineOutput
 
 
@@ -132,7 +131,7 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-class AllegroPipeline(DiffusionPipeline, PyramidAttentionBroadcastMixin):
+class AllegroPipeline(DiffusionPipeline):
     r"""
     Pipeline for text-to-video generation using Allegro.
 

src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py

@@ -29,7 +29,6 @@
 from ...utils import logging, replace_example_docstring
 from ...utils.torch_utils import randn_tensor
 from ...video_processor import VideoProcessor
-from ..pyramid_broadcast_utils import PyramidAttentionBroadcastMixin
 from .pipeline_output import CogVideoXPipelineOutput
 
 
@@ -138,7 +137,7 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin, PyramidAttentionBroadcastMixin):
+class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
     r"""
     Pipeline for text-to-video generation using CogVideoX.
 

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py

@@ -30,7 +30,6 @@
 from ...utils import logging, replace_example_docstring
 from ...utils.torch_utils import randn_tensor
 from ...video_processor import VideoProcessor
-from ..pyramid_broadcast_utils import PyramidAttentionBroadcastMixin
 from .pipeline_output import CogVideoXPipelineOutput
 
 
@@ -145,7 +144,7 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-class CogVideoXFunControlPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin, PyramidAttentionBroadcastMixin):
+class CogVideoXFunControlPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
     r"""
     Pipeline for controlled text-to-video generation using CogVideoX Fun.
 

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py

@@ -34,7 +34,6 @@
 )
 from ...utils.torch_utils import randn_tensor
 from ...video_processor import VideoProcessor
-from ..pyramid_broadcast_utils import PyramidAttentionBroadcastMixin
 from .pipeline_output import CogVideoXPipelineOutput
 
 
@@ -154,7 +153,7 @@ def retrieve_latents(
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-class CogVideoXImageToVideoPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin, PyramidAttentionBroadcastMixin):
+class CogVideoXImageToVideoPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
     r"""
     Pipeline for image-to-video generation using CogVideoX.
 

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py

@@ -30,7 +30,6 @@
 from ...utils import logging, replace_example_docstring
 from ...utils.torch_utils import randn_tensor
 from ...video_processor import VideoProcessor
-from ..pyramid_broadcast_utils import PyramidAttentionBroadcastMixin
 from .pipeline_output import CogVideoXPipelineOutput
 
 
@@ -160,7 +159,7 @@ def retrieve_latents(
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-class CogVideoXVideoToVideoPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin, PyramidAttentionBroadcastMixin):
+class CogVideoXVideoToVideoPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
     r"""
     Pipeline for video-to-video generation using CogVideoX.
 

src/diffusers/pipelines/latte/pipeline_latte.py

@@ -37,7 +37,6 @@
 )
 from ...utils.torch_utils import is_compiled_module, randn_tensor
 from ...video_processor import VideoProcessor
-from ..pyramid_broadcast_utils import PyramidAttentionBroadcastMixin
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
@@ -133,7 +132,7 @@ class LattePipelineOutput(BaseOutput):
     frames: torch.Tensor
 
 
-class LattePipeline(DiffusionPipeline, PyramidAttentionBroadcastMixin):
+class LattePipeline(DiffusionPipeline):
     r"""
     Pipeline for text-to-video generation using Latte.
 

src/diffusers/pipelines/pyramid_broadcast_utils.py renamed to src/diffusers/pipelines/pyramid_attention_broadcast_utils.py

@@ -12,18 +12,105 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional, Tuple
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Tuple, Type, TypeVar
 
 import torch.nn as nn
 
 from ..models.attention_processor import Attention
 from ..models.hooks import PyramidAttentionBroadcastHook, add_hook_to_module, remove_hook_from_module
 from ..utils import logging
+from .pipeline_utils import DiffusionPipeline
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
+_ATTENTION_CLASSES = (Attention,)
+
+_SPATIAL_ATTENTION_BLOCK_IDENTIFIERS = ["blocks", "transformer_blocks"]
+_TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS = ["temporal_transformer_blocks"]
+_CROSS_ATTENTION_BLOCK_IDENTIFIERS = ["blocks", "transformer_blocks"]
+
+
+@dataclass
+class PyramidAttentionBroadcastConfig:
+    spatial_attention_block_skip = None
+    temporal_attention_block_skip = None
+    cross_attention_block_skip = None
+    
+    spatial_attention_timestep_skip_range = (100, 800)
+    temporal_attention_timestep_skip_range = (100, 800)
+    cross_attention_timestep_skip_range = (100, 800)
+
+    spatial_attention_block_identifiers = _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS
+    temporal_attention_block_identifiers = _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS
+    cross_attention_block_identifiers = _CROSS_ATTENTION_BLOCK_IDENTIFIERS
+
+
+class PyramidAttentionBroadcastState:
+    iteration = 0
+
+
+def apply_pyramid_attention_broadcast(
+    pipeline: DiffusionPipeline,
+    config: Optional[PyramidAttentionBroadcastConfig] = None,
+    denoiser: Optional[nn.Module] = None,
+):
+    if config is None:
+        config = PyramidAttentionBroadcastConfig()
+    
+    if config.spatial_attention_block_skip is None and config.temporal_attention_block_skip is None and config.cross_attention_block_skip is None:
+        logger.warning(
+            "Pyramid Attention Broadcast requires one or more of `spatial_attention_block_skip`, `temporal_attention_block_skip` "
+            "or `cross_attention_block_skip` parameters to be set to an integer, not `None`. Defaulting to using `spatial_attention_block_skip=2`. "
+            "To avoid this warning, please set one of the above parameters."
+        )
+        config.spatial_attention_block_skip = 2
+    
+    if denoiser is None:
+        denoiser = pipeline.transformer if hasattr(pipeline, "transformer") else pipeline.unet
+        
+    for name, module in denoiser.named_modules():
+        if not isinstance(module, _ATTENTION_CLASSES):
+            continue
+        if isinstance(module, Attention):
+            _apply_pyramid_attention_broadcast_on_attention_class(pipeline, name, module, config)
+
+
+# def apply_pyramid_attention_broadcast_spatial(module: TypeVar[_ATTENTION_CLASSES], config: PyramidAttentionBroadcastConfig):
+#     hook = PyramidAttentionBroadcastHook(skip_callback=)
+#     add_hook_to_module(module)
+
+
+def _apply_pyramid_attention_broadcast_on_attention_class(pipeline: DiffusionPipeline, name: str, module: Attention, config: PyramidAttentionBroadcastConfig):
+    # Similar check as PEFT to determine if a string layer name matches a module name
+    is_spatial_self_attention = (
+        any(f"{identifier}." in name or identifier == name for identifier in config.spatial_attention_block_identifiers)
+        and config.spatial_attention_timestep_skip_range is not None
+        and not module.is_cross_attention
+    )
+    is_temporal_self_attention = (
+        any(f"{identifier}." in name or identifier == name for identifier in config.temporal_attention_block_identifiers)
+        and config.temporal_attention_timestep_skip_range is not None
+        and not module.is_cross_attention
+    )
+    is_cross_attention = (
+        any(f"{identifier}." in name or identifier == name for identifier in config.cross_attention_block_identifiers)
+        and config.cross_attention_timestep_skip_range is not None
+        and not module.is_cross_attention
+    )
+
+    if is_spatial_self_attention:
+        apply_pyramid_attention_broadcast_spatial(module, config)
+    elif is_temporal_self_attention:
+        apply_pyramid_attention_broadcast_temporal(module, config)
+    elif is_cross_attention:
+        apply_pyramid_attention_broadcast_cross(module, config)
+    else:
+        logger.warning(f"Unable to apply Pyramid Attention Broadcast to the selected layer: {name}.")
+
+
 class PyramidAttentionBroadcastMixin:
     r"""Mixin class for [Pyramid Attention Broadcast](https://www.arxiv.org/abs/2408.12588)."""
 

tests/pipelines/cogvideo/test_cogvideox.py

@@ -21,7 +21,7 @@
 from transformers import AutoTokenizer, T5EncoderModel
 
 from diffusers import AutoencoderKLCogVideoX, CogVideoXPipeline, CogVideoXTransformer3DModel, DDIMScheduler
-from diffusers.pipelines.pyramid_broadcast_utils import PyramidAttentionBroadcastAttentionProcessorWrapper
+from diffusers.pipelines.pyramid_attention_broadcast_utils import PyramidAttentionBroadcastAttentionProcessorWrapper
 from diffusers.utils.testing_utils import (
     enable_full_determinism,
     numpy_cosine_similarity_distance,

tests/pipelines/cogvideo/test_cogvideox_image2video.py

@@ -22,7 +22,7 @@
 from transformers import AutoTokenizer, T5EncoderModel
 
 from diffusers import AutoencoderKLCogVideoX, CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel, DDIMScheduler
-from diffusers.pipelines.pyramid_broadcast_utils import PyramidAttentionBroadcastAttentionProcessorWrapper
+from diffusers.pipelines.pyramid_attention_broadcast_utils import PyramidAttentionBroadcastAttentionProcessorWrapper
 from diffusers.utils import load_image
 from diffusers.utils.testing_utils import (
     enable_full_determinism,