Use basic attention implementation for small inputs on old pytorch.

comfyanonymous · Jan 9, 2024 · 6a7bc35 · 6a7bc35
1 parent b3b5ddb
commit 6a7bc35
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py
@@ -351,8 +351,11 @@ def attention_pytorch(q, k, v, heads, mask=None):
 optimized_attention_masked = optimized_attention
 
 def optimized_attention_for_device(device, mask=False, small_input=False):
-    if small_input and model_management.pytorch_attention_enabled():
-        return attention_pytorch #TODO: need to confirm but this is probably slightly faster for small inputs in all cases
+    if small_input:
+        if model_management.pytorch_attention_enabled():
+            return attention_pytorch #TODO: need to confirm but this is probably slightly faster for small inputs in all cases
+        else:
+            return attention_basic
 
     if device == torch.device("cpu"):
         return attention_sub_quad