From 86755c9f8020939b6c5e328b3b84689170a7ae7f Mon Sep 17 00:00:00 2001
From: "Lt.Dr.Data" <lt.dr.data@gmail.com>
Date: Thu, 8 Jun 2023 17:34:49 +0900
Subject: [PATCH] remove image_opt use latent_preview

---
 comfy/utils.py    | 18 ------------------
 latent_preview.py |  4 ++--
 nodes.py          | 38 +++++++++++++++++---------------------
 3 files changed, 19 insertions(+), 41 deletions(-)

diff --git a/comfy/utils.py b/comfy/utils.py
index cbcb47506..585ebda51 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -1,6 +1,5 @@
 import torch
 import math
-from PIL import Image
 import struct
 
 def load_torch_file(ckpt, safe_load=False):
@@ -221,20 +220,3 @@ class ProgressBar:
 
     def update(self, value):
         self.update_absolute(self.current + value)
-
-
-def latent_to_rgb(latent_tensor):
-    latent_rgb_factors = torch.tensor([
-        #   R     G      B
-        [0.298, 0.207, 0.208],  # L1
-        [0.187, 0.286, 0.173],  # L2
-        [-0.158, 0.189, 0.264],  # L3
-        [-0.184, -0.271, -0.473],  # L4
-    ], device="cpu")
-    rgb = torch.einsum('...lhw,lr -> ...rhw', latent_tensor.cpu().float(), latent_rgb_factors)
-    tensor = (((rgb + 1) / 2)
-              .clamp(0, 1)  # change scale from -1..1 to 0..1
-              .mul(0xFF)    # to 0..255
-              .byte())
-
-    return Image.fromarray(tensor.movedim(1, -1)[0].cpu().numpy())
diff --git a/latent_preview.py b/latent_preview.py
index ef6c201b6..323e3f375 100644
--- a/latent_preview.py
+++ b/latent_preview.py
@@ -69,10 +69,10 @@ class Latent2RGBPreviewer(LatentPreviewer):
         return Image.fromarray(latents_ubyte.numpy())
 
 
-def get_previewer(device):
+def get_previewer(device, force=False):
     previewer = None
     method = args.preview_method
-    if method != LatentPreviewMethod.NoPreviews:
+    if method != LatentPreviewMethod.NoPreviews or force:
         # TODO previewer methods
         taesd_decoder_path = folder_paths.get_full_path("vae_approx", "taesd_decoder.pth")
 
diff --git a/nodes.py b/nodes.py
index ee7320e0f..8b75155b6 100644
--- a/nodes.py
+++ b/nodes.py
@@ -299,7 +299,6 @@ class SavePreviewLatent(SaveLatent):
     def INPUT_TYPES(s):
         return {"required": {"samples": ("LATENT", ),
                              "filename_prefix": ("STRING", {"default": "latents/ComfyUI"}), },
-                "optional": {"image_opt": ("IMAGE", ), },
                 "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                 }
     RETURN_TYPES = ()
@@ -328,33 +327,30 @@ class SavePreviewLatent(SaveLatent):
         image.save(image_path, format='png', exif=exif_bytes, pnginfo=metadata, optimize=True)
 
     @staticmethod
-    def prepare_preview(image, latent_tensor):
+    def prepare_preview(latent_tensor):
         lower_bound = 128
         upper_bound = 512
 
-        if image is None:
-            image = comfy.utils.latent_to_rgb(latent_tensor).convert("P", palette=Image.ADAPTIVE, colors=256)
+        device = comfy.model_management.get_torch_device()
 
-            min_size = min(image.size[0], image.size[1])
-            max_size = max(image.size[0], image.size[1])
+        previewer = latent_preview.get_previewer(device, True)
 
-            scale_factor = 1
-            if max_size > upper_bound:
-                scale_factor = upper_bound/max_size
+        image = previewer.decode_latent_to_preview(latent_tensor)
+        min_size = min(image.size[0], image.size[1])
+        max_size = max(image.size[0], image.size[1])
 
-            # prevent too small preview
-            if min_size*scale_factor < lower_bound:
-                scale_factor = lower_bound/min_size
+        scale_factor = 1
+        if max_size > upper_bound:
+            scale_factor = upper_bound/max_size
 
-            w = int(image.size[0] * scale_factor)
-            h = int(image.size[1] * scale_factor)
+        # prevent too small preview
+        if min_size*scale_factor < lower_bound:
+            scale_factor = lower_bound/min_size
 
-            image = image.resize((w, h), resample=Image.NEAREST)
+        w = int(image.size[0] * scale_factor)
+        h = int(image.size[1] * scale_factor)
 
-        else:
-            # don't resize if provide preview image intentionally
-            i = 255. * image[0].cpu().numpy()
-            image = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+        image = image.resize((w, h), resample=Image.NEAREST)
 
         return SavePreviewLatent.attach_format_text(image)
 
@@ -378,11 +374,11 @@ class SavePreviewLatent(SaveLatent):
 
         return new_image
 
-    def save_preview_latent(self, samples, filename_prefix="ComfyUI", image_opt=None, prompt=None, extra_pnginfo=None):
+    def save_preview_latent(self, samples, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
 
         # load preview
-        preview = SavePreviewLatent.prepare_preview(image_opt, samples['samples'])
+        preview = SavePreviewLatent.prepare_preview(samples['samples'])
 
         # support save metadata for latent sharing
         file = f"{filename}_{counter:05}_.latent.png"