update

Author: a-r-r-o-w

src/diffusers/models/attention_processor.py

@@ -930,6 +930,8 @@ def __init__(
         self.out_dim = out_dim if out_dim is not None else query_dim
         self.out_context_dim = out_context_dim if out_context_dim else query_dim
         self.context_pre_only = context_pre_only
+        # TODO(aryan): Maybe try to improve the checks in PAB instead
+        self.is_cross_attention = False
 
         self.heads = out_dim // dim_head if out_dim is not None else heads
 

src/diffusers/pipelines/pyramid_attention_broadcast_utils.py

@@ -18,7 +18,7 @@
 
 import torch.nn as nn
 
-from ..models.attention_processor import Attention
+from ..models.attention_processor import Attention, MochiAttention
 from ..models.hooks import ModelHook, add_hook_to_module
 from ..utils import logging
 from .pipeline_utils import DiffusionPipeline
@@ -27,7 +27,7 @@
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
-_ATTENTION_CLASSES = (Attention,)
+_ATTENTION_CLASSES = (Attention, MochiAttention)
 
 _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "single_transformer_blocks")
 _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS = ("temporal_transformer_blocks",)
@@ -175,8 +175,10 @@ def apply_pyramid_attention_broadcast(
     for name, module in denoiser.named_modules():
         if not isinstance(module, _ATTENTION_CLASSES):
             continue
-        if isinstance(module, Attention):
+        if isinstance(module, (Attention)):
             _apply_pyramid_attention_broadcast_on_attention_class(pipeline, name, module, config)
+        if isinstance(module, MochiAttention):
+            _apply_pyramid_attention_broadcast_on_mochi_attention_class(pipeline, name, module, config)
 
 
 def apply_pyramid_attention_broadcast_on_module(
@@ -263,6 +265,13 @@ def skip_callback(module: nn.Module) -> bool:
     return True
 
 
+def _apply_pyramid_attention_broadcast_on_mochi_attention_class(
+    pipeline: DiffusionPipeline, name: str, module: MochiAttention, config: PyramidAttentionBroadcastConfig
+) -> bool:
+    # The same logic as Attention class works here, so just use that for now
+    return _apply_pyramid_attention_broadcast_on_attention_class(pipeline, name, module, config)
+
+
 class PyramidAttentionBroadcastHook(ModelHook):
     r"""A hook that applies Pyramid Attention Broadcast to a given module."""