From c7b8e7250d03aaa3b2fc3b3b41f1a852385362e8 Mon Sep 17 00:00:00 2001
From: drhead <1313496+drhead@users.noreply.github.com>
Date: Sun, 9 Mar 2025 19:42:19 -0400
Subject: [PATCH 1/3] make conds use non-blocking transfers

---
 comfy/conds.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/comfy/conds.py b/comfy/conds.py
index 211fb8d57..775785d37 100644
--- a/comfy/conds.py
+++ b/comfy/conds.py
@@ -1,6 +1,7 @@
 import torch
 import math
 import comfy.utils
+from comfy.model_management import device_should_use_non_blocking
 
 
 class CONDRegular:
@@ -10,8 +11,14 @@ class CONDRegular:
     def _copy_with(self, cond):
         return self.__class__(cond)
 
+    def _pin_memory(self, cond):
+        if cond.device == torch.device('cpu'):
+            return cond.pin_memory()
+        else:
+            return cond
+
     def process_cond(self, batch_size, device, **kwargs):
-        return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device))
+        return self._copy_with(comfy.utils.repeat_to_batch_size(self._pin_memory(self.cond), batch_size).to(device, non_blocking=device_should_use_non_blocking(device))
 
     def can_concat(self, other):
         if self.cond.shape != other.cond.shape:

From 8e5f33cc9c254959826497580819239857c97fef Mon Sep 17 00:00:00 2001
From: drhead <1313496+drhead@users.noreply.github.com>
Date: Sun, 9 Mar 2025 19:45:29 -0400
Subject: [PATCH 2/3] better way of gating non blocking use

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

diff --git a/comfy/conds.py b/comfy/conds.py
index 775785d37..e643d6830 100644
--- a/comfy/conds.py
+++ b/comfy/conds.py
@@ -18,7 +18,10 @@ class CONDRegular:
             return cond
 
     def process_cond(self, batch_size, device, **kwargs):
-        return self._copy_with(comfy.utils.repeat_to_batch_size(self._pin_memory(self.cond), batch_size).to(device, non_blocking=device_should_use_non_blocking(device))
+        if device_should_use_non_blocking(device):
+            return self._copy_with(comfy.utils.repeat_to_batch_size(self._pin_memory(self.cond), batch_size).to(device, non_blocking=True)
+        else:
+            return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device)
 
     def can_concat(self, other):
         if self.cond.shape != other.cond.shape:

From d175cbd3151e97dbfdfa701e59211e0a14d75f01 Mon Sep 17 00:00:00 2001
From: drhead <1313496+drhead@users.noreply.github.com>
Date: Sun, 9 Mar 2025 19:57:51 -0400
Subject: [PATCH 3/3] fix syntax

---
 comfy/conds.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/comfy/conds.py b/comfy/conds.py
index e643d6830..5fac44ef5 100644
--- a/comfy/conds.py
+++ b/comfy/conds.py
@@ -19,9 +19,9 @@ class CONDRegular:
 
     def process_cond(self, batch_size, device, **kwargs):
         if device_should_use_non_blocking(device):
-            return self._copy_with(comfy.utils.repeat_to_batch_size(self._pin_memory(self.cond), batch_size).to(device, non_blocking=True)
+            return self._copy_with(comfy.utils.repeat_to_batch_size(self._pin_memory(self.cond), batch_size).to(device, non_blocking=True))
         else:
-            return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device)
+            return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device))
 
     def can_concat(self, other):
         if self.cond.shape != other.cond.shape: