[Added] Node to save images to arbitrary paths

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)

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)