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

This commit is contained in:
comfyanonymous 2024-05-21 16:56:33 -04:00
parent 83d969e397
commit 8508df2569
2 changed files with 16 additions and 2 deletions

View File

@ -19,12 +19,13 @@ from comfy.cli_args import args
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):

View File

@ -5,6 +5,7 @@ from comfy.cli_args import args
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: