From 6c23854f54edb51643026b7b10d1ae9e34cfd3d6 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Wed, 22 May 2024 13:56:28 -0400
Subject: [PATCH] Fix OSX latent2rgb previews.

---
 comfy/model_management.py | 10 ++++++++--
 comfy/ops.py              |  2 +-
 latent_preview.py         |  3 ++-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index fbfbb7e4..ef36a2c4 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -630,8 +630,14 @@ def supports_dtype(device, dtype): #TODO
 def device_supports_non_blocking(device):
     if is_device_mps(device):
         return False #pytorch bug? mps doesn't support non blocking
+    return True
+
+def device_should_use_non_blocking(device):
+    if not device_supports_non_blocking(device):
+        return False
     return False
-    # return True #TODO: figure out why this causes issues
+    # return True #TODO: figure out why this causes memory issues on Nvidia and possibly others
+
 
 def cast_to_device(tensor, device, dtype, copy=False):
     device_supports_cast = False
@@ -643,7 +649,7 @@ def cast_to_device(tensor, device, dtype, copy=False):
         elif is_intel_xpu():
             device_supports_cast = True
 
-    non_blocking = device_supports_non_blocking(device)
+    non_blocking = device_should_use_non_blocking(device)
 
     if device_supports_cast:
         if copy:
diff --git a/comfy/ops.py b/comfy/ops.py
index eb650768..7ebb3dd2 100644
--- a/comfy/ops.py
+++ b/comfy/ops.py
@@ -21,7 +21,7 @@ import comfy.model_management
 
 def cast_bias_weight(s, input):
     bias = None
-    non_blocking = comfy.model_management.device_supports_non_blocking(input.device)
+    non_blocking = comfy.model_management.device_should_use_non_blocking(input.device)
     if s.bias is not None:
         bias = s.bias.to(device=input.device, dtype=input.dtype, non_blocking=non_blocking)
         if s.bias_function is not None:
diff --git a/latent_preview.py b/latent_preview.py
index 05d750e2..dae9beb6 100644
--- a/latent_preview.py
+++ b/latent_preview.py
@@ -4,6 +4,7 @@ import struct
 import numpy as np
 from comfy.cli_args import args, LatentPreviewMethod
 from comfy.taesd.taesd import TAESD
+import comfy.model_management
 import folder_paths
 import comfy.utils
 import logging
@@ -43,7 +44,7 @@ class Latent2RGBPreviewer(LatentPreviewer):
         latents_ubyte = (((latent_image + 1) / 2)
                             .clamp(0, 1)  # change scale from -1..1 to 0..1
                             .mul(0xFF)  # to 0..255
-                            ).to(device="cpu", dtype=torch.uint8, non_blocking=True)
+                            ).to(device="cpu", dtype=torch.uint8, non_blocking=comfy.model_management.device_supports_non_blocking(latent_image.device))
 
         return Image.fromarray(latents_ubyte.numpy())