From 248f1039e96c7e049a0f07072b4dcc4a1142278b Mon Sep 17 00:00:00 2001 From: Aryan Date: Wed, 15 Jan 2025 11:31:03 +0100 Subject: [PATCH] apply suggestions from review --- src/diffusers/hooks/pyramid_attention_broadcast.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/diffusers/hooks/pyramid_attention_broadcast.py b/src/diffusers/hooks/pyramid_attention_broadcast.py index d5ee4568d8bf..6d7964395929 100644 --- 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(