diff --git a/README.md b/README.md index ea26e04..38b9243 100644 --- a/README.md +++ b/README.md @@ -10,11 +10,17 @@ This custom node extends ComfyUI's image loading functionality with filename out - **Compatible**: Maintains all original functionality (IMAGE and MASK outputs) ### LoadImageFolder -- **Folder Loading**: Load all images from a selected folder +- **Folder Loading**: Load all images from a specified folder path - **Batch Processing**: Returns all images concatenated as a single tensor -- **Filename Tracking**: Returns a comma-separated list of all loaded filenames +- **Filename Tracking**: Returns a list of all loaded filenames - **Error Handling**: Gracefully handles corrupted or unsupported image files +### SaveImageWithFilename +- **Enhanced Save Image Node**: Based on the original ComfyUI SaveImage node +- **Custom Filenames**: Accepts single filenames or comma-separated lists of filenames +- **Flexible Naming**: Can use provided filenames or fall back to default naming +- **Batch Support**: Handles multiple images with corresponding filenames + ## Installation 1. Place `nodes.py` in your ComfyUI `custom_nodes` directory @@ -31,11 +37,22 @@ This custom node extends ComfyUI's image loading functionality with filename out - `filename`: The filename of the loaded image ### LoadImageFolder -- **Input**: Select a folder from the dropdown +- **Input**: Enter a folder path as a string - **Outputs**: - `image`: All images from the folder concatenated as a single tensor - `mask`: All masks concatenated as a single tensor - - `filenames`: Comma-separated list of all loaded filenames + - `filenames`: List of all loaded filenames + +### SaveImageWithFilename +- **Inputs**: + - `images`: The images to save (IMAGE tensor) + - `filenames`: Single filename or comma-separated list of filenames (optional) + - `filename_prefix`: Prefix for default naming (optional) +- **Behavior**: + - If filenames are provided, uses them for the corresponding images + - If no filenames or fewer filenames than images, uses default naming for remaining images + - Automatically adds .png extension to filenames + - Saves to ComfyUI output directory ## Requirements @@ -50,8 +67,11 @@ This custom node extends ComfyUI's image loading functionality with filename out - If no valid images are found in a folder, empty tensors will be returned - All images in a folder must have the same dimensions for proper concatenation - Error messages are printed to console for any files that fail to load +- The SaveImageWithFilename node preserves original filenames when possible +- If filenames contain extensions, they will be replaced with .png ## Based On - [ComfyUI LoadImage Node](https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py) +- [ComfyUI SaveImage Node](https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py) - [Issue #8699](https://github.com/comfyanonymous/ComfyUI/issues/8699) - Request for filename output functionality \ No newline at end of file diff --git a/__init__.py b/__init__.py index a33997b..48b3d62 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,13 @@ -from .nodes import LoadImageWithFilename, LoadImageFolder +from .nodes import LoadImageWithFilename, LoadImageFolder, SaveImageWithFilename NODE_CLASS_MAPPINGS = { "LoadImageWithFilename": LoadImageWithFilename, "LoadImageFolder": LoadImageFolder, + "SaveImageWithFilename": SaveImageWithFilename, } NODE_DISPLAY_NAME_MAPPINGS = { "LoadImageWithFilename": "Load Image With Filename", "LoadImageFolder": "Load Image Folder", + "SaveImageWithFilename": "Save Image With Filename", } \ No newline at end of file diff --git a/nodes.py b/nodes.py index 43525f0..392b95a 100644 --- a/nodes.py +++ b/nodes.py @@ -2,7 +2,9 @@ import os import hashlib import torch import numpy as np +import json from PIL import Image, ImageOps, ImageSequence +from PIL.PngImagePlugin import PngInfo import folder_paths import node_helpers @@ -231,4 +233,117 @@ class LoadImageFolder: return "Path does not exist: {}".format(folder_path) if not os.path.isdir(folder_path): return "Not a directory: {}".format(folder_path) + return True + + +class SaveImageWithFilename: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "" + self.compress_level = 4 + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", {"tooltip": "The images to save."}), + "filenames": ("STRING", {"default": "", "tooltip": "Single filename or comma-separated list of filenames. If empty, will use default naming."}), + "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."}) + }, + "hidden": { + "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" + }, + } + + RETURN_TYPES = () + FUNCTION = "save_images" + + OUTPUT_NODE = True + + CATEGORY = "image" + DESCRIPTION = "Saves the input images with specified filenames to your ComfyUI output directory." + + def save_images(self, images, filenames="", filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): + filename_prefix += self.prefix_append + + # Parse filenames - handle both single filename and comma-separated list + filename_list = [] + if filenames: + # Split by comma and strip whitespace + filename_list = [f.strip() for f in filenames.split(",") if f.strip()] + + # If no filenames provided or not enough filenames, use default naming + if not filename_list or len(filename_list) < len(images): + # Use default naming for remaining images + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) + + results = list() + for (batch_number, image) in enumerate(images): + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + metadata = None + 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])) + + # Use provided filename if available, otherwise use default naming + if batch_number < len(filename_list): + # Use the provided filename + provided_filename = filename_list[batch_number] + # Remove extension if present and add .png + base_name = os.path.splitext(provided_filename)[0] + file = f"{base_name}.png" + else: + # Use default naming + filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) + file = f"{filename_with_batch_num}_{counter:05}_.png" + + img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + else: + # Use provided filenames for all images + results = list() + for (batch_number, image) in enumerate(images): + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + metadata = None + 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])) + + # Use the provided filename + provided_filename = filename_list[batch_number] + # Remove extension if present and add .png + base_name = os.path.splitext(provided_filename)[0] + file = f"{base_name}.png" + + # Save to output directory + img.save(os.path.join(self.output_dir, file), pnginfo=metadata, compress_level=self.compress_level) + results.append({ + "filename": file, + "subfolder": "", + "type": self.type + }) + + return {"ui": {"images": results}} + + @classmethod + def IS_CHANGED(s, images, filenames, filename_prefix): + return hashlib.sha256(str(images).encode()).hexdigest() + + @classmethod + def VALIDATE_INPUTS(s, images, filenames, filename_prefix): + if not images: + return "No images provided" return True \ No newline at end of file