apply suggestions from review

huggingface · Jan 15, 2025 · 248f103 · 248f103
1 parent 40fc7a5
commit 248f103
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/diffusers/hooks/pyramid_attention_broadcast.py b/src/diffusers/hooks/pyramid_attention_broadcast.py
@@ -221,8 +221,7 @@ def apply_pyramid_attention_broadcast(
             # cannot be applied to this layer. For custom layers, users can extend this functionality and implement
             # their own PAB logic similar to `_apply_pyramid_attention_broadcast_on_attention_class`.
             continue
-        if isinstance(submodule, (Attention, MochiAttention)):
-            _apply_pyramid_attention_broadcast_on_attention_class(name, submodule, config)
+        _apply_pyramid_attention_broadcast_on_attention_class(name, submodule, config)
 
 
 def _apply_pyramid_attention_broadcast_on_attention_class(