Skip to content

Commit

Permalink
Work around black image bug on Mac 14.5 by forcing attention upcasting.
Browse files Browse the repository at this point in the history
  • Loading branch information
comfyanonymous committed May 21, 2024
1 parent 83d969e commit 8508df2
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 2 deletions.
5 changes: 3 additions & 2 deletions comfy/ldm/modules/attention.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
13 changes: 13 additions & 0 deletions comfy/model_management.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 8508df2

Please sign in to comment.