Work around black image bug on Mac 14.5 by forcing attention upcasting.

comfyanonymous · comfyanonymous · commit 8508df25691b · 2024-05-21T16:56:33.000-04:00
diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py
@@ -19,12 +19,13 @@
 import comfy.ops
 ops = comfy.ops.disable_weight_init
 
+FORCE_UPCAST_ATTENTION_DTYPE = model_management.force_upcast_attention_dtype()
 
 def get_attn_precision(attn_precision):
     if args.dont_upcast_attention:
         return None
-    if attn_precision is None and args.force_upcast_attention:
-        return torch.float32
+    if FORCE_UPCAST_ATTENTION_DTYPE is not None:
+        return FORCE_UPCAST_ATTENTION_DTYPE
     return attn_precision
 
 def exists(val):
diff --git a/comfy/model_management.py b/comfy/model_management.py
@@ -5,6 +5,7 @@
 import comfy.utils
 import torch
 import sys
+import platform
 
 class VRAMState(Enum):
     DISABLED = 0    #No vram present: no need to move models to vram
@@ -685,6 +686,18 @@ def pytorch_attention_flash_attention():
             return True
     return False
 
+def force_upcast_attention_dtype():
+    upcast = args.force_upcast_attention
+    try:
+        if platform.mac_ver()[0] in ['14.5']: #black image bug on OSX Sonoma 14.5
+            upcast = True
+    except:
+        pass
+    if upcast:
+        return torch.float32
+    else:
+        return None
+
 def get_free_memory(dev=None, torch_free_too=False):
     global directml_enabled
     if dev is None:
