diff --git a/__init__.py b/__init__.py index 0d18bc8..7f526e7 100644 --- a/__init__.py +++ b/__init__.py @@ -31,6 +31,7 @@ from .loaders.caption_image_loader import CaptionImageLoader from .loaders.load_image_metadata import LoadImageAndExtractPrompt from .loaders.folder_image_metadata import FolderImageAndExtractPrompt from .loaders.folder_image_metadata_by_name import FolderImageMetadataByName +from .loaders.save_image_with_tag import AnotherSaveImageWithTag from .loaders.load_gif_frames import LoadGifFrames, RemapGifFrames from .loaders.batch_image_list import BatchToImageList from .loaders.folder_image_loader import FolderImageLoader @@ -167,6 +168,8 @@ NODE_CLASS_MAPPINGS = { "LoadImageAndExtractPrompt": LoadImageAndExtractPrompt, "FolderImageAndExtractPrompt": FolderImageAndExtractPrompt, "FolderImageMetadataByName": FolderImageMetadataByName, + "AnotherSaveImageWithTag": AnotherSaveImageWithTag, + "FolderImageLoader": FolderImageLoader, "ImageListToBatch": ImageListToBatch, "IndicesListTo50": IndicesListTo50, } @@ -239,6 +242,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "LoadImageAndExtractPrompt": "Load Image and Extract Prompt", "FolderImageAndExtractPrompt": "Folder Image and Extract Prompt", "FolderImageMetadataByName": "Folder Metadata by Node Name", + "AnotherSaveImageWithTag": "Another Save Image with Tag", + "FolderImageLoader": "Folder Image Loader", "ImageListToBatch": "Image List To Multi Batch", "IndicesListTo 50": "Indices List To 50 Inputs", } diff --git a/loaders/folder_image_metadata_by_name.py b/loaders/folder_image_metadata_by_name.py index b5dd9db..88e06e0 100644 --- a/loaders/folder_image_metadata_by_name.py +++ b/loaders/folder_image_metadata_by_name.py @@ -29,6 +29,10 @@ class FolderImageMetadataByName: if not metadata: return "" + # 0. Check for custom embedded tag FIRST (best for batches) + if 'another_tag' in metadata: + return str(metadata['another_tag']) + node_name_clean = node_name.lower().strip() # 1. Try 'prompt' (Execution JSON) diff --git a/loaders/save_image_with_tag.py b/loaders/save_image_with_tag.py new file mode 100644 index 0000000..c5b11a9 --- /dev/null +++ b/loaders/save_image_with_tag.py @@ -0,0 +1,70 @@ +import os +import json +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo +import folder_paths +import comfy.utils + +class AnotherSaveImageWithTag: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "custom_tag": ("STRING", {"forceInput": True}), + "filename_prefix": ("STRING", {"default": "AnotherTag"}), + }, + } + + INPUT_IS_LIST = True + RETURN_TYPES = () + FUNCTION = "save_images" + OUTPUT_NODE = True + CATEGORY = "AnotherUtils/loaders" + + def save_images(self, images, custom_tag, filename_prefix): + filename_prefix = filename_prefix[0] if isinstance(filename_prefix, list) else filename_prefix + + # Determine number of items to process + num_images = len(images) + num_tags = len(custom_tag) + + results = list() + for i in range(num_images): + img_tensor = images[i] + # Handle list mapping for tags + tag_value = custom_tag[i] if i < num_tags else custom_tag[-1] + + # Convert tensor to PIL + i_numpy = 255. * img_tensor.cpu().numpy() + img = Image.fromarray(np.clip(i_numpy, 0, 255).astype(np.uint8)) + + # ComfyUI Filename generation + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, img.width, img.height) + file = f"{filename}_{counter:05}_.png" + + # Prepare Metadata + metadata = PngInfo() + # Standard ComfyUI stuff (not strictly required if we only want our tag, + # but usually good for compatibility) + # However, we ONLY care about the custom tag for recovery. + metadata.add_text("another_tag", str(tag_value)) + + # Save + save_path = os.path.join(full_output_folder, file) + img.save(save_path, pnginfo=metadata, compress_level=4) + + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + counter += 1 + + return { "ui": { "images": results } }