From a6decf1e620907347c9c5d8c815172f349b19c21 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 1 Aug 2024 16:18:14 -0400
Subject: [PATCH] Fix bfloat16 potentially not being enabled on mps.

---
 comfy/model_management.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index de4bd44..b4f32d6 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -897,7 +897,10 @@ def should_use_bf16(device=None, model_params=0, prioritize_performance=True, ma
     if directml_enabled:
         return False
 
-    if cpu_mode() or mps_mode():
+    if mps_mode():
+        return True
+
+    if cpu_mode():
         return False
 
     if is_intel_xpu():