Support attn2==None for xformers (huggingface#1759)

0xdigiscore · Dec 19, 2022 · 32a5d70 · 32a5d70
1 parent 429e544
commit 32a5d70
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py
@@ -473,7 +473,8 @@ def set_use_memory_efficient_attention_xformers(self, use_memory_efficient_atten
             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