diff --git a/latent_preview.py b/latent_preview.py index 0e7bdc8c1..ba01169c5 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -3,6 +3,7 @@ from PIL import Image, ImageOps from io import BytesIO import struct import numpy as np +import comfy.latent_formats as latent_formats from comfy.cli_args import args, LatentPreviewMethod from comfy.taesd.taesd import TAESD @@ -63,7 +64,18 @@ class Latent2RGBPreviewer(LatentPreviewer): return Image.fromarray(latents_ubyte.numpy()) -def get_previewer(device, latent_format, force=False): +def get_previewer(device, latent_format=None, latent=None, force=False): + # If the latent_format parameter is not provided, fallback to assuming SD15 format. + # Ultimately, it seems that the format information should be included in the latent itself. + if latent_format is None: + if latent is not None and 'format' in latent: + if latent['format'] == 'SDXL': + latent_format = latent_formats.SDXL() + else: + latent_format = latent_formats.SD15() + else: + latent_format = latent_formats.SD15() + previewer = None method = args.preview_method if method != LatentPreviewMethod.NoPreviews or force: diff --git a/nodes.py b/nodes.py index aa0290e3b..e6317b187 100644 --- a/nodes.py +++ b/nodes.py @@ -329,12 +329,12 @@ class SavePreviewLatent(SaveLatent): image.save(image_path, format='png', exif=exif_bytes, pnginfo=metadata, optimize=True) @staticmethod - def prepare_preview(latent_tensor): + def prepare_preview(latent): lower_bound = 128 upper_bound = 256 - previewer = latent_preview.get_previewer("cpu", force=True) - image = previewer.decode_latent_to_preview(latent_tensor) + previewer = latent_preview.get_previewer("cpu", latent=latent, force=True) + image = previewer.decode_latent_to_preview(latent['samples']) min_size = min(image.size[0], image.size[1]) max_size = max(image.size[0], image.size[1]) @@ -377,7 +377,7 @@ class SavePreviewLatent(SaveLatent): 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(samples['samples']) + preview = SavePreviewLatent.prepare_preview(samples) # support save metadata for latent sharing file = f"{filename}_{counter:05}_.latent.png"