From 27e35f13faa79c90dddae45eea324baffcc13ace Mon Sep 17 00:00:00 2001
From: drhead <1313496+drhead@users.noreply.github.com>
Date: Sun, 19 May 2024 18:34:09 -0400
Subject: [PATCH] handle non blocking better and case of single image

---
 modules/sd_samplers_common.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index dd9e10a8b..cd9b912d7 100644
--- a/modules/sd_samplers_common.py
+++ b/modules/sd_samplers_common.py
@@ -72,10 +72,6 @@ def decode_first_stage(model, x):
     return samples_to_images_tensor(x, approx_index, model)
 
 
-def sample_to_image(samples, index=0, approximation=None):
-    return single_sample_to_image(samples[index], approximation)
-
-
 if torch.cuda.is_available():
     lp_stream = torch.cuda.Stream()
     live_preview_stream_context = torch.cuda.stream(lp_stream)
@@ -83,9 +79,17 @@ else:
     lp_stream = None
     live_preview_stream_context = nullcontext()
 
+def sample_to_image(samples, index=0, approximation=None):
+    with live_preview_stream_context:
+        sample = single_sample_to_image(samples[index], approximation, non_blocking=lp_stream is not None)
+    if lp_stream is not None:
+        lp_stream.synchronize()
+    return Image.fromarray(sample.numpy())
+
+
 def samples_to_image_grid(samples, approximation=None):
     with live_preview_stream_context:
-        sample_tensors = [single_sample_to_image(sample, approximation, non_blocking=True) for sample in samples]
+        sample_tensors = [single_sample_to_image(sample, approximation, non_blocking=lp_stream is not None) for sample in samples]
     if lp_stream is not None:
         lp_stream.synchronize()
     return images.image_grid([Image.fromarray(sample.numpy()) for sample in sample_tensors])