Support attn2==None for xformers (#1759)

32a5d70c · Anton Lozhkov · GitHub · 429e5449 · 32a5d70c
Unverified Commit 32a5d70c authored Dec 19, 2022 by Anton Lozhkov Committed by GitHub Dec 19, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

src/diffusers/models/attention.py src/diffusers/models/attention.py +2 -1

No files found.
--- a/src/diffusers/models/attention.py
+++ b/src/diffusers/models/attention.py
@@ -473,7 +473,8 @@ class BasicTransformerBlock(nn.Module):
            except Exception as e:
                raise e
            self.attn1._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
-            self.attn2._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
+            if self.attn2 is not None:
+                self.attn2._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers

    def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, attention_mask=None):
        # 1. Self-Attention