From 72ef54fd492527d93bc4cd6569fc1aadca23ab72 Mon Sep 17 00:00:00 2001 From: Yolan Date: Sat, 21 Dec 2024 10:52:32 -0800 Subject: [PATCH] Initial import --- README.md | 20 +++++-- __init__.py | 150 ++++++++++++++++++++++++++-------------------------- 2 files changed, 90 insertions(+), 80 deletions(-) diff --git a/README.md b/README.md index 011c239..77671f9 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,26 @@ # Comfy UI Node Template -This is a template for creating custom nodes for the Comfy UI stable diffusion client. +This is a custom node for the Comfy UI stable diffusion client. ## Description -This Python script is an optional add-on to the Comfy UI stable diffusion client. It introduces quality of life improvements by providing variable nodes and shared global variables. +The `SaveImageARGB16PNG` node provides functionality for saving images as uncompressed PNG files with ARGB16 precision. +This node is particularly useful for workflows that require high-quality image saving with metadata such as prompts and additional PNG info. + +## Inputs +- **images**: A tensor of images to be saved. Each image is expected to be in a format compatible with ARGB16. +- **filename_prefix**: A string used as the prefix for the saved filenames. This can include dynamic formatting options. +- **prompt** (optional): Metadata to embed in the PNG file. +- **extra_pnginfo** (optional): Additional metadata to include in the PNG file as key-value pairs. + +## Outputs +- Saves images to the Comfy UI output directory with metadata embedded in the PNG files. +- Returns a UI-compatible dictionary containing details about the saved images. ## Getting Started -Import into the custom nodes directory of your Comfy UI client. +Import this script into the custom nodes directory of your Comfy UI client. ## Dependencies -ComfyUI +- ComfyUI +- Pillow (for PNG processing) ## License This project is licensed under the MIT License. \ No newline at end of file diff --git a/__init__.py b/__init__.py index 175202a..63b01a0 100644 --- a/__init__.py +++ b/__init__.py @@ -1,95 +1,93 @@ -class Example: - """ - A example node +import torch - Class methods - ------------- - INPUT_TYPES (dict): - Tell the main program input parameters of nodes. +import os +import sys +import json - Attributes - ---------- - RETURN_TYPES (`tuple`): - The type of each element in the output tulple. - RETURN_NAMES (`tuple`): - Optional: The name of each output in the output tulple. - FUNCTION (`str`): - The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() - OUTPUT_NODE ([`bool`]): - If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. - The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. - Assumed to be False if not present. - CATEGORY (`str`): - The category the node should appear in the UI. - execute(s) -> tuple || None: - The entry point method. The name of this method must be the same as the value of property `FUNCTION`. - For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. - """ +sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy")) + +import folder_paths + +class SaveImageARGB16PNG: def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - """ - Return a dictionary which contains config for all input fields. - Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". - Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. - The type can be a list for selection. + import os + import json + try: + from PIL import Image + self.Image = Image + except ImportError: + raise ImportError("Pillow module not found. Please install it to save PNG images.") - Returns: `dict`: - - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` - - Value input_fields (`dict`): Contains input fields config: - * Key field_name (`string`): Name of a entry-point method's argument - * Value field_config (`tuple`): - + First value is a string indicate the type of field or a list for selection. - + Secound value is a config for type "INT", "STRING" or "FLOAT". - """ + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "" + self.compress_level = 4 + + @classmethod + def INPUT_TYPES(cls): return { "required": { - "image": ("IMAGE",), - "int_field": ("INT", { - "default": 0, - "min": 0, #Minimum value - "max": 4096, #Maximum value - "step": 64 #Slider's step - }), - "float_field": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), - "print_to_screen": (["enable", "disable"],), - "string_field": ("STRING", { - "multiline": False, #True if you want the field to look like the one on the ClipTextEncode node - "default": "Hello World!" - }), + "images": ("IMAGE",), + "filename_prefix": ("STRING", {"default": "ComfyUI"}) + }, + "hidden": { + "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" }, } - RETURN_TYPES = ("IMAGE",) - #RETURN_NAMES = ("image_output_name",) + RETURN_TYPES = () + FUNCTION = "savepng" + OUTPUT_NODE = True + CATEGORY = "Marigold" - FUNCTION = "test" + def savepng(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): + import numpy as np + import os + import re + from PIL.PngImagePlugin import PngInfo - #OUTPUT_NODE = False + filename_prefix += self.prefix_append + 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 = [] - CATEGORY = "Example" + for batch_number, image in enumerate(images): + image_np = image.cpu().numpy() + image_np = (image_np * 65535).astype(np.uint16) # Scale to 16-bit range - def test(self, image, string_field, int_field, float_field, print_to_screen): - if print_to_screen == "enable": - print(f"""Your input contains: - string_field aka input text: {string_field} - int_field: {int_field} - float_field: {float_field} - """) - #do some processing on the image, in this example I just invert it - image = 1.0 - image - return (image,) + if image_np.shape[-1] == 4: + mode = "RGBA" + else: + mode = "RGB" + + image_pil = self.Image.fromarray(image_np, mode=mode) + + 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])) + + filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) + file = f"{filename_with_batch_num}_{counter:05}_.png" + + image_pil.save(os.path.join(full_output_folder, file), format="PNG", compress_level=self.compress_level, pnginfo=metadata) + + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + counter += 1 + + return { "ui": { "images": results } } -# A dictionary that contains all nodes you want to export with their names -# NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { - "Example": Example + "SaveImageARGB16PNG": SaveImageARGB16PNG } - -# A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { - "Example": "Example Node" + "SaveImageARGB16PNG": "SaveImageARGB16PNG" }