Initial import
This commit is contained in:
Yolan
@@ -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.
|
||||
+74
-76
@@ -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"
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user