From 5265d27b936e7745d88ca4fcb071fa53c11278a5 Mon Sep 17 00:00:00 2001 From: "Salvador E. Tropea" Date: Wed, 5 Nov 2025 16:47:55 -0300 Subject: [PATCH] [Added] Node to save images to arbitrary paths --- src/nodes/helpers.py | 30 +++++++++++++++++++++++++++--- src/nodes/nodes_img.py | 38 ++++++++++++++++++++++++++++++++------ 2 files changed, 59 insertions(+), 9 deletions(-) diff --git a/src/nodes/helpers.py b/src/nodes/helpers.py index 907e74b..3395e79 100644 --- a/src/nodes/helpers.py +++ b/src/nodes/helpers.py @@ -1,7 +1,8 @@ +import json import numpy as np import os -from PIL import Image, ImageOps, ImageSequence -from PIL import ImageFile, UnidentifiedImageError +from PIL import Image, ImageOps, ImageSequence, ImageFile, UnidentifiedImageError +from PIL.PngImagePlugin import PngInfo import torch from . import main_logger @@ -80,7 +81,7 @@ class CustomLoadImage(object): return (output_image, output_mask) -def load_image_wrapper(file_name, embed_transparency, disp_name=None): +def load_image_wrapper(file_name, embed_transparency, disp_name=None, show_preview=True): disp_name = disp_name or file_name # --- REUSE ComfyUI's LoadImage LOGIC --- @@ -109,6 +110,9 @@ def load_image_wrapper(file_name, embed_transparency, disp_name=None): # Concatenate image and mask into (b, h, w, 4) image_with_alpha = torch.cat([image, 1.0 - mask], dim=-1) result = (image_with_alpha, mask) + # No preview + if not show_preview: + return result # This information is for the preview, as we are an output node and we return images # they will be displayed in our node. Quite simple. if os.path.isabs(file_name): @@ -132,3 +136,23 @@ def load_image_wrapper(file_name, embed_transparency, disp_name=None): # Re-raise to make the error visible in ComfyUI raise IOError(f"Could not load the image file '{disp_name}' using the standard loader. " "It may be corrupt or in an unsupported format.") from e + + +def save_image(images, filenames, prompt=None, extra_pnginfo=None, compress_level=4): + if isinstance(filenames, str): + filenames = [filenames] + B = images.shape[0] + if len(filenames) != B: + raise ValueError(f"{B} images provided but only {len(filenames)} file names") + for batch_number, (image, filename) in enumerate(zip(images, filenames)): + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + + logger.debug(f"Saving {image.shape[1]}x{image.shape[0]} image to {filename}") + img.save(filename, pnginfo=metadata, compress_level=compress_level) diff --git a/src/nodes/nodes_img.py b/src/nodes/nodes_img.py index 89e9518..fd3d0dc 100644 --- a/src/nodes/nodes_img.py +++ b/src/nodes/nodes_img.py @@ -25,7 +25,7 @@ from typing import Optional # We are the main source, so we use the main_logger from . import main_logger -from .helpers import load_image_wrapper +from .helpers import load_image_wrapper, save_image try: from folder_paths import get_input_directory, get_output_directory from comfy import model_management @@ -286,16 +286,40 @@ class ImageLoad: DESCRIPTION = ("Loads an image from any path") UNIQUE_NAME = "SET_ImageLoad" DISPLAY_NAME = "Load Image from Path" - # This node stores a result to disk. So this IS an output node. - # It can be used without connecting any other node. - # Declaring it as output helps with the preview mechanism. - OUTPUT_NODE = True def execute(self, file_name: str, embed_transparency: bool = False): if not os.path.exists(file_name): raise ValueError(f"File '{file_name}' not found") - return load_image_wrapper(file_name, embed_transparency) + return load_image_wrapper(file_name, embed_transparency, show_preview=False) + + +class ImageSave: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE", {"tooltip": "The images to save."}), + "filename": ("STRING", {"default": "ComfyUI", "tooltip": "The file name for the image"}) + }, + "hidden": { + "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" + }, + } + + RETURN_TYPES = () + FUNCTION = "execute" + + OUTPUT_NODE = True + + CATEGORY = BASE_CATEGORY + "/" + IO_CATEGORY + DESCRIPTION = ("Saves an image to an arbitrary path") + UNIQUE_NAME = "SET_ImageSave" + DISPLAY_NAME = "Save Image to Path" + + def execute(self, image, filename, prompt=None, extra_pnginfo=None): + save_image(image, filename, prompt, extra_pnginfo) + return () class ImageDataset: @@ -407,6 +431,8 @@ class ImageDataset: results.append(str(dest_path)) references.append(ref_filename if ref_dir else "") + if not len(images): + raise ValueError("No images to process") logger.info(f"Found {len(images)} images") logger.debug(images)