1929 lines
62 KiB
Python
1929 lines
62 KiB
Python
import base64
|
|
import json
|
|
import math
|
|
import random
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import torch
|
|
|
|
try:
|
|
import piexif.helper
|
|
import piexif
|
|
from .exif.exif import read_info_from_image_stealth
|
|
|
|
piexif_loaded = True
|
|
except ImportError:
|
|
piexif_loaded = False
|
|
|
|
from .imgio.converter import PILHandlingHodes
|
|
from .autonode import node_wrapper, get_node_names_mappings, validate, anytype, PILImage
|
|
import time
|
|
import os
|
|
import shutil
|
|
from PIL import Image
|
|
from PIL import ImageOps
|
|
from PIL import ImageEnhance
|
|
from PIL.PngImagePlugin import PngInfo
|
|
try:
|
|
import folder_paths
|
|
except ModuleNotFoundError:
|
|
folder_paths = None
|
|
try:
|
|
from comfy.cli_args import args
|
|
except ModuleNotFoundError:
|
|
# Allow importing this module outside a full ComfyUI install (e.g. unit tests).
|
|
class _Args:
|
|
disable_metadata = True
|
|
|
|
args = _Args()
|
|
import filelock
|
|
import tempfile
|
|
|
|
fundamental_classes = []
|
|
fundamental_node = node_wrapper(fundamental_classes)
|
|
|
|
|
|
@fundamental_node
|
|
class SleepNodeAny:
|
|
FUNCTION = "sleep"
|
|
RETURN_TYPES = (anytype,)
|
|
CATEGORY = "Misc"
|
|
custom_name = "SleepNode"
|
|
|
|
@staticmethod
|
|
def sleep(interval, inputs):
|
|
time.sleep(interval)
|
|
return (inputs,)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"interval": ("FLOAT", {"default": 0.0}),
|
|
},
|
|
"optional": {
|
|
"inputs": (anytype, {"default": 0.0}),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class SleepNodeImage:
|
|
FUNCTION = "sleep"
|
|
RETURN_TYPES = (anytype,)
|
|
CATEGORY = "Misc"
|
|
custom_name = "Sleep (Image tunnel)"
|
|
|
|
@staticmethod
|
|
def sleep(interval, image):
|
|
time.sleep(interval)
|
|
return (image,)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"interval": ("FLOAT", {"default": 0.0}),
|
|
"image": (anytype,),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ErrorNode:
|
|
FUNCTION = "raise_error"
|
|
RETURN_TYPES = ("STRING",)
|
|
CATEGORY = "Misc"
|
|
custom_name = "ErrorNode"
|
|
|
|
@staticmethod
|
|
def raise_error(error_msg="Error"):
|
|
raise Exception("Error: {}".format(error_msg))
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"error_msg": ("STRING", {"default": "Error"}),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class CurrentTimestamp:
|
|
"""
|
|
Returns the current Unix timestamp or a formatted time string.
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
def generate(self, format_string):
|
|
if format_string.strip() == "":
|
|
# return Unix timestamp
|
|
return (int(time.time()),)
|
|
else:
|
|
# return formatted date/time
|
|
return (time.strftime(format_string, time.localtime()),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"format_string": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"display": "text",
|
|
"comment": "Leave blank for raw timestamp, or use format directives like '%Y-%m-%d %H:%M:%S'",
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",) # or ("INT",) if returning raw int timestamp
|
|
FUNCTION = "generate"
|
|
CATEGORY = "Logic Gates"
|
|
custom_name = "Current Timestamp"
|
|
|
|
|
|
@fundamental_node
|
|
class DebugComboInputNode:
|
|
FUNCTION = "debug_combo_input"
|
|
RETURN_TYPES = ("STRING",)
|
|
CATEGORY = "Misc"
|
|
custom_name = "Debug Combo Input"
|
|
|
|
@staticmethod
|
|
def debug_combo_input(input1):
|
|
print(input1)
|
|
return (input1,)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"input1": (["0", "1", "2"], {"default": "0"}),
|
|
}
|
|
}
|
|
|
|
|
|
# https://github.com/comfyanonymous/ComfyUI/blob/340177e6e85d076ab9e222e4f3c6a22f1fb4031f/custom_nodes/example_node.py.example#L18
|
|
@fundamental_node
|
|
class TextPreviewNode:
|
|
"""
|
|
Can't display text but it makes always changed state
|
|
"""
|
|
|
|
FUNCTION = "text_preview"
|
|
RETURN_TYPES = ()
|
|
CATEGORY = "Misc"
|
|
custom_name = "Text Preview"
|
|
RESULT_NODE = True
|
|
OUTPUT_NODE = True
|
|
|
|
def text_preview(self, text):
|
|
print(text)
|
|
# below does not work, why?
|
|
return {"ui": {"text": str(text)}}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text": (anytype, {"default": "text", "type": "output"}),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, *args, **kwargs):
|
|
return float("nan")
|
|
|
|
|
|
@fundamental_node
|
|
class ParseExifNode:
|
|
"""
|
|
Parses exif data from image
|
|
"""
|
|
|
|
FUNCTION = "parse_exif"
|
|
RETURN_TYPES = ("STRING",)
|
|
CATEGORY = "Misc"
|
|
custom_name = "Parse Exif"
|
|
|
|
@staticmethod
|
|
def parse_exif(image):
|
|
return (read_info_from_image_stealth(image),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
|
|
def throw_if_parent_or_root_access(path):
|
|
if ".." in path or path.startswith("/") or path.startswith("\\"):
|
|
raise RuntimeError("Tried to access parent or root directory")
|
|
if path.startswith("~"):
|
|
raise RuntimeError("Tried to access home directory")
|
|
if os.path.isabs(path):
|
|
raise RuntimeError("Path cannot be absolute")
|
|
|
|
|
|
@fundamental_node
|
|
class SaveImageCustomNode:
|
|
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",),
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
|
|
"subfolder_dir": ("STRING", {"default": ""}),
|
|
},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",) # Filename
|
|
FUNCTION = "save_images"
|
|
|
|
OUTPUT_NODE = True
|
|
RESULT_NODE = True
|
|
CATEGORY = "image"
|
|
custom_name = "Save Image Custom Node"
|
|
|
|
def save_images(
|
|
self,
|
|
images,
|
|
filename_prefix="ComfyUI",
|
|
subfolder_dir="",
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
# `images` can be None or empty in some edge cases.
|
|
if images is None:
|
|
return {"ui": {"images": []}, "outputs": {"images": ""}}
|
|
if isinstance(images, torch.Tensor) and images.shape[0] == 0:
|
|
return {"ui": {"images": []}, "outputs": {"images": ""}}
|
|
if isinstance(images, (list, tuple)) and len(images) == 0:
|
|
return {"ui": {"images": []}, "outputs": {"images": ""}}
|
|
|
|
filename_prefix += self.prefix_append
|
|
throw_if_parent_or_root_access(filename_prefix)
|
|
throw_if_parent_or_root_access(subfolder_dir)
|
|
output_dir = os.path.join(self.output_dir, subfolder_dir)
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = (
|
|
folder_paths.get_save_image_path(
|
|
filename_prefix, output_dir, images[0].shape[1], images[0].shape[0]
|
|
)
|
|
)
|
|
results = list()
|
|
for image in images:
|
|
i = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
metadata = None
|
|
if not args.disable_metadata:
|
|
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]))
|
|
|
|
file = f"{filename}_{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}
|
|
)
|
|
counter += 1
|
|
|
|
return {"ui": {"images": results}, "outputs": {"images": file.rstrip(".png")}}
|
|
|
|
|
|
@fundamental_node
|
|
class SaveTextCustomNode:
|
|
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": {
|
|
"text": (anytype,),
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
|
|
"subfolder_dir": ("STRING", {"default": ""}),
|
|
"filename": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",) # Filename
|
|
FUNCTION = "save_text"
|
|
custom_name = "Save Text Custom Node"
|
|
CATEGORY = "text"
|
|
RESULT_NODE = True
|
|
OUTPUT_NODE = True
|
|
|
|
def save_text(self, text, filename_prefix="ComfyUI", subfolder_dir="", filename=""):
|
|
text = str(text)
|
|
throw_if_parent_or_root_access(filename_prefix)
|
|
throw_if_parent_or_root_access(subfolder_dir)
|
|
assert (
|
|
len(text) > 0 and len(filename) > 0
|
|
), "Text and filename must be non-empty"
|
|
filename_prefix += self.prefix_append
|
|
output_dir = os.path.join(self.output_dir, subfolder_dir)
|
|
filename_merged = filename_prefix + filename + ".txt"
|
|
full_output_folder, subfolder, actual_filename = output_dir, "", filename_merged
|
|
results = list()
|
|
file = actual_filename
|
|
with open(os.path.join(full_output_folder, file), "w") as f:
|
|
f.write(text)
|
|
results.append({"filename": file, "subfolder": subfolder, "type": self.type})
|
|
|
|
return {"ui": {"texts": results}, "outputs": {"images": file.rstrip(".txt")}}
|
|
|
|
|
|
@fundamental_node
|
|
class CommaRejoinNode:
|
|
FUNCTION = "comma_rejoin"
|
|
RETURN_TYPES = ("STRING",)
|
|
CATEGORY = "text"
|
|
custom_name = "Comma Rejoin"
|
|
|
|
@staticmethod
|
|
def comma_rejoin(text, split_separator=",", join_separator=", "):
|
|
# Split (by comma or given separator), strip items, then join using the
|
|
# join separator verbatim (e.g. ", ").
|
|
parts = str(text).split(split_separator) if split_separator else [str(text)]
|
|
stripped = [part.strip() for part in parts]
|
|
stripped = [part for part in stripped if part != ""]
|
|
return (join_separator.join(stripped),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"default": ""}),
|
|
},
|
|
"optional": {
|
|
"split_separator": ("STRING", {"default": ","}),
|
|
"join_separator": ("STRING", {"default": ", "}),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class DumpTextJsonlNode:
|
|
"""
|
|
Appends text to a JSONL file (one JSON object per line).
|
|
Each line will have the structure: { "<keyname>": "<text_item>" }
|
|
|
|
For concurrency safety, this node uses filelock to block
|
|
concurrent writes to the same file.
|
|
"""
|
|
|
|
FUNCTION = "dump_text_jsonl"
|
|
RETURN_TYPES = ("STRING",) # We return the filename for convenience
|
|
CATEGORY = "text"
|
|
custom_name = "Dump Text JSONL Node"
|
|
RESULT_NODE = True
|
|
OUTPUT_NODE = True
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output" # for consistent UI listing
|
|
self.prefix_append = ""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""
|
|
text can be a single string or a list of strings.
|
|
If it's a list, each item is appended as a separate line.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"text": (anytype,), # Single string or list of strings
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
|
|
"subfolder_dir": ("STRING", {"default": ""}),
|
|
"filename": ("STRING", {"default": "dump.jsonl"}),
|
|
"keyname": ("STRING", {"default": "text"}),
|
|
},
|
|
}
|
|
|
|
def dump_text_jsonl(
|
|
self,
|
|
text,
|
|
filename_prefix="ComfyUI",
|
|
subfolder_dir="",
|
|
filename="dump.jsonl",
|
|
keyname="text",
|
|
):
|
|
# Security checks to avoid writing outside of the ComfyUI output folder
|
|
throw_if_parent_or_root_access(filename_prefix)
|
|
throw_if_parent_or_root_access(subfolder_dir)
|
|
|
|
# Build the actual output path
|
|
filename_prefix += self.prefix_append # If you want to append something
|
|
output_dir = os.path.join(self.output_dir, subfolder_dir)
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
final_filename = filename_prefix + "_" + filename
|
|
full_path = os.path.join(output_dir, final_filename)
|
|
lock_path = full_path + ".lock"
|
|
|
|
# Ensure we can safely write concurrently
|
|
with filelock.FileLock(lock_path, timeout=10):
|
|
with open(full_path, "a", encoding="utf-8") as f:
|
|
# If `text` is a list, write each element as its own JSON line
|
|
if isinstance(text, list):
|
|
for item in text:
|
|
# Convert each item to string, just to be safe
|
|
line = {keyname: str(item)}
|
|
f.write(json.dumps(line, ensure_ascii=False) + "\n")
|
|
else:
|
|
# Single string input
|
|
line = {keyname: str(text)}
|
|
f.write(json.dumps(line, ensure_ascii=False) + "\n")
|
|
|
|
# Return data for UI usage
|
|
results = [
|
|
{"filename": final_filename, "subfolder": subfolder_dir, "type": self.type}
|
|
]
|
|
return {
|
|
"ui": {"texts": results},
|
|
"outputs": {"filename": final_filename},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ConcatGridNode:
|
|
"""
|
|
Concatenate multiple images in a row, a column, or a square-like grid
|
|
using either resizing or padding to match dimensions.
|
|
|
|
direction:
|
|
- "horizontal": line up side by side
|
|
- "vertical": stack top to bottom
|
|
- "square-like": arrange images in an NxN grid (where N = ceil(sqrt(#images)))
|
|
|
|
match_method:
|
|
- "resize": scale images so their matching dimension is the same
|
|
(height for horizontal, width for vertical, or cell-size for square-like)
|
|
- "pad": keep original size but add transparent padding so the matching dimension is the same
|
|
"""
|
|
|
|
FUNCTION = "concat_grid"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Concat Grid (Batch to single grid)"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"direction": (
|
|
["horizontal", "vertical", "square-like"],
|
|
{"default": "horizontal"},
|
|
),
|
|
"match_method": (["resize", "pad"], {"default": "resize"}),
|
|
}
|
|
}
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def concat_grid(images, direction="horizontal", match_method="resize"):
|
|
# 1) Convert images input to a list of PIL RGBA images
|
|
# - If it's a torch.Tensor with shape (B, C, H, W) or a single image, unify into list.
|
|
if not (
|
|
isinstance(images, torch.Tensor) and len(images.shape) == 4
|
|
) and not isinstance(images, (list, tuple)):
|
|
images = [images]
|
|
|
|
converted = PILHandlingHodes.handle_input(images) # returns PIL or list of PIL
|
|
if isinstance(converted, list):
|
|
pil_images = [img.convert("RGBA") for img in converted]
|
|
else:
|
|
pil_images = [converted.convert("RGBA")]
|
|
|
|
if len(pil_images) == 0:
|
|
raise RuntimeError("No images provided to Concat Grid")
|
|
|
|
# 2) Handle the three layout directions
|
|
if direction == "horizontal":
|
|
# --- Horizontal layout ---
|
|
max_height = max(img.height for img in pil_images)
|
|
processed = []
|
|
for img in pil_images:
|
|
if match_method == "resize":
|
|
# Scale the image so that height == max_height
|
|
if img.height == 0:
|
|
raise RuntimeError("Encountered an image of zero height.")
|
|
ratio = max_height / float(img.height)
|
|
new_w = int(img.width * ratio)
|
|
new_h = max_height
|
|
new_img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
|
else: # "pad"
|
|
# Create a new image with the same width but max_height
|
|
new_img = Image.new("RGBA", (img.width, max_height), (0, 0, 0, 0))
|
|
new_img.paste(img, (0, 0))
|
|
processed.append(new_img)
|
|
|
|
total_width = sum(im.width for im in processed)
|
|
out = Image.new("RGBA", (total_width, max_height), (0, 0, 0, 0))
|
|
x_offset = 0
|
|
for im in processed:
|
|
out.paste(im, (x_offset, 0))
|
|
x_offset += im.width
|
|
|
|
elif direction == "vertical":
|
|
# --- Vertical layout ---
|
|
max_width = max(img.width for img in pil_images)
|
|
processed = []
|
|
for img in pil_images:
|
|
if match_method == "resize":
|
|
if img.width == 0:
|
|
raise RuntimeError("Encountered an image of zero width.")
|
|
ratio = max_width / float(img.width)
|
|
new_w = max_width
|
|
new_h = int(img.height * ratio)
|
|
new_img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
|
else: # "pad"
|
|
new_img = Image.new("RGBA", (max_width, img.height), (0, 0, 0, 0))
|
|
new_img.paste(img, (0, 0))
|
|
processed.append(new_img)
|
|
|
|
total_height = sum(im.height for im in processed)
|
|
out = Image.new("RGBA", (max_width, total_height), (0, 0, 0, 0))
|
|
y_offset = 0
|
|
for im in processed:
|
|
out.paste(im, (0, y_offset))
|
|
y_offset += im.height
|
|
|
|
else: # direction == "square-like"
|
|
# --- Square-like NxN grid ---
|
|
count = len(pil_images)
|
|
# Determine grid size
|
|
num_cols = int(math.ceil(math.sqrt(count)))
|
|
num_rows = int(math.ceil(count / num_cols))
|
|
|
|
# Find maximum width/height among images
|
|
max_width = max(img.width for img in pil_images)
|
|
max_height = max(img.height for img in pil_images)
|
|
|
|
processed = []
|
|
for img in pil_images:
|
|
if match_method == "resize":
|
|
# Here we forcibly resize each image to (max_width, max_height)
|
|
# (which may distort if aspect ratios differ).
|
|
new_img = img.resize(
|
|
(max_width, max_height), Image.Resampling.LANCZOS
|
|
)
|
|
else: # "pad"
|
|
# Keep original size but create a new RGBA canvas so each cell is (max_w, max_h)
|
|
new_img = Image.new("RGBA", (max_width, max_height), (0, 0, 0, 0))
|
|
new_img.paste(img, (0, 0))
|
|
processed.append(new_img)
|
|
|
|
# Create the final output canvas
|
|
grid_width = num_cols * max_width
|
|
grid_height = num_rows * max_height
|
|
out = Image.new("RGBA", (grid_width, grid_height), (0, 0, 0, 0))
|
|
|
|
# Paste images in row-major order
|
|
idx = 0
|
|
for row in range(num_rows):
|
|
for col in range(num_cols):
|
|
if idx >= count:
|
|
break # no more images
|
|
x_offset = col * max_width
|
|
y_offset = row * max_height
|
|
out.paste(processed[idx], (x_offset, y_offset))
|
|
idx += 1
|
|
|
|
return (out,)
|
|
|
|
|
|
@fundamental_node
|
|
class ConcatTwoImagesNode:
|
|
"""
|
|
Concatenate exactly two images (imageA, imageB).
|
|
|
|
direction:
|
|
- "horizontal": line them up side by side
|
|
- "vertical": place them top to bottom
|
|
|
|
match_method:
|
|
- "resize": scale images so their matching dimension is the same
|
|
(height for horizontal, width for vertical)
|
|
- "pad": keep original size but pad them so the matching dimension is the same
|
|
"""
|
|
|
|
FUNCTION = "concat_two_images"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Concat 2 Images to Grid"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"imageA": ("IMAGE",),
|
|
"imageB": ("IMAGE",),
|
|
"direction": (["horizontal", "vertical"], {"default": "horizontal"}),
|
|
"match_method": (["resize", "pad"], {"default": "resize"}),
|
|
}
|
|
}
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def concat_two_images(
|
|
imageA, imageB, direction="horizontal", match_method="resize"
|
|
):
|
|
# Convert input to PIL images (RGBA to preserve alpha if needed)
|
|
pilA = PILHandlingHodes.handle_input(imageA)
|
|
if isinstance(pilA, list):
|
|
raise RuntimeError(
|
|
"Expected a single image for imageA, grid only supports two images"
|
|
)
|
|
pilB = PILHandlingHodes.handle_input(imageB)
|
|
if isinstance(pilB, list):
|
|
raise RuntimeError(
|
|
"Expected a single image for imageB, grid only supports two images"
|
|
)
|
|
if direction == "horizontal":
|
|
# We want to unify heights
|
|
max_h = max(pilA.height, pilB.height)
|
|
|
|
if match_method == "resize":
|
|
# Scale each image so their heights match
|
|
def scale_height(img, target_h):
|
|
if img.height == 0:
|
|
raise RuntimeError("Encountered an image with zero height.")
|
|
ratio = target_h / float(img.height)
|
|
new_w = int(img.width * ratio)
|
|
new_h = target_h
|
|
return img.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
|
|
|
pilA = scale_height(pilA, max_h)
|
|
pilB = scale_height(pilB, max_h)
|
|
|
|
else: # match_method == "pad"
|
|
# Pad images with transparent background so they share the same height
|
|
def pad_height(img, target_h):
|
|
new_img = Image.new("RGBA", (img.width, target_h), (0, 0, 0, 0))
|
|
new_img.paste(img, (0, 0))
|
|
return new_img
|
|
|
|
pilA = pad_height(pilA, max_h)
|
|
pilB = pad_height(pilB, max_h)
|
|
|
|
total_width = pilA.width + pilB.width
|
|
out = Image.new("RGBA", (total_width, max_h), (0, 0, 0, 0))
|
|
# Paste images side by side
|
|
out.paste(pilA, (0, 0))
|
|
out.paste(pilB, (pilA.width, 0))
|
|
|
|
else:
|
|
# direction == "vertical"
|
|
# We want to unify widths
|
|
max_w = max(pilA.width, pilB.width)
|
|
|
|
if match_method == "resize":
|
|
# Scale each image so their widths match
|
|
def scale_width(img, target_w):
|
|
if img.width == 0:
|
|
raise RuntimeError("Encountered an image with zero width.")
|
|
ratio = target_w / float(img.width)
|
|
new_w = target_w
|
|
new_h = int(img.height * ratio)
|
|
return img.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
|
|
|
pilA = scale_width(pilA, max_w)
|
|
pilB = scale_width(pilB, max_w)
|
|
|
|
else: # match_method == "pad"
|
|
# Pad images with transparent background so they share the same width
|
|
def pad_width(img, target_w):
|
|
new_img = Image.new("RGBA", (target_w, img.height), (0, 0, 0, 0))
|
|
new_img.paste(img, (0, 0))
|
|
return new_img
|
|
|
|
pilA = pad_width(pilA, max_w)
|
|
pilB = pad_width(pilB, max_w)
|
|
|
|
total_height = pilA.height + pilB.height
|
|
out = Image.new("RGBA", (max_w, total_height), (0, 0, 0, 0))
|
|
# Paste images top to bottom
|
|
out.paste(pilA, (0, 0))
|
|
out.paste(pilB, (0, pilA.height))
|
|
|
|
return (out,)
|
|
|
|
|
|
@fundamental_node
|
|
class SaveCustomJPGNode:
|
|
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",),
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
|
|
"subfolder_dir": ("STRING", {"default": ""}),
|
|
},
|
|
"optional": {
|
|
"quality": ("INT", {"default": 95}),
|
|
"optimize": ("BOOLEAN", {"default": True}),
|
|
"metadata_string": ("STRING", {"default": ""}),
|
|
},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",) # Filename
|
|
FUNCTION = "save_images"
|
|
|
|
OUTPUT_NODE = True
|
|
RESULT_NODE = True
|
|
|
|
CATEGORY = "image"
|
|
custom_name = "Save Custom JPG Node"
|
|
|
|
def save_images(
|
|
self,
|
|
images,
|
|
filename_prefix="ComfyUI",
|
|
subfolder_dir="",
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
quality=95,
|
|
optimize=True,
|
|
metadata_string="",
|
|
):
|
|
if images is None:
|
|
return {"ui": {"images": []}, "outputs": {"images": ""}}
|
|
if isinstance(images, torch.Tensor) and images.shape[0] == 0:
|
|
return {"ui": {"images": []}, "outputs": {"images": ""}}
|
|
if isinstance(images, (list, tuple)) and len(images) == 0:
|
|
return {"ui": {"images": []}, "outputs": {"images": ""}}
|
|
if not isinstance(images, (list, tuple, torch.Tensor)):
|
|
images = [images]
|
|
|
|
throw_if_parent_or_root_access(filename_prefix)
|
|
throw_if_parent_or_root_access(subfolder_dir)
|
|
|
|
filename_prefix += self.prefix_append
|
|
output_dir = os.path.join(self.output_dir, subfolder_dir)
|
|
filelock_path = os.path.join(output_dir, filename_prefix + ".lock")
|
|
|
|
results = []
|
|
for image in images:
|
|
if isinstance(image, torch.Tensor):
|
|
if image.device.type != "cpu":
|
|
image = image.cpu()
|
|
image = 255.0 * image.numpy()
|
|
clipped = np.clip(image, 0, 255).astype(np.uint8)
|
|
if clipped.shape[0] <= 3:
|
|
clipped = np.transpose(clipped, (1, 2, 0))
|
|
img = Image.fromarray(clipped)
|
|
else:
|
|
img = PILHandlingHodes.handle_input(image)
|
|
|
|
metadata = {}
|
|
if not args.disable_metadata:
|
|
if prompt is not None:
|
|
metadata["prompt"] = json.dumps(prompt)
|
|
if extra_pnginfo is not None:
|
|
for x in extra_pnginfo:
|
|
metadata[x] = json.dumps(extra_pnginfo[x])
|
|
|
|
if metadata_string:
|
|
metadata = {"metadata": metadata_string}
|
|
|
|
exif_bytes = None
|
|
if piexif_loaded:
|
|
exif_bytes = piexif.dump(
|
|
{
|
|
"Exif": {
|
|
piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(
|
|
json.dumps(metadata), encoding="unicode"
|
|
)
|
|
},
|
|
}
|
|
)
|
|
|
|
with filelock.FileLock(filelock_path, timeout=10):
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = (
|
|
folder_paths.get_save_image_path(
|
|
filename_prefix, output_dir, img.size[1], img.size[0]
|
|
)
|
|
)
|
|
counter_len = len(str(len(images)))
|
|
file = f"{filename}_{str(counter).zfill(max(5, counter_len))}_.jpg"
|
|
|
|
with tempfile.NamedTemporaryFile(
|
|
suffix=".jpg", delete=False
|
|
) as tmpfile:
|
|
tmp_path = tmpfile.name
|
|
img.save(tmp_path, "JPEG", quality=quality, optimize=optimize)
|
|
|
|
if piexif_loaded and exif_bytes:
|
|
piexif.insert(exif_bytes, tmp_path)
|
|
|
|
final_path = os.path.join(full_output_folder, file)
|
|
shutil.copy2(tmp_path, final_path)
|
|
os.remove(tmp_path)
|
|
|
|
results.append(
|
|
{
|
|
"filename": os.path.join(full_output_folder, file),
|
|
"subfolder": subfolder_dir,
|
|
"type": self.type,
|
|
}
|
|
)
|
|
|
|
return {
|
|
"ui": {"images": results},
|
|
"outputs": {
|
|
"images": os.path.join(full_output_folder, file).rstrip(".jpg")
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class SaveImageWebpCustomNode:
|
|
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",),
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
|
|
"subfolder_dir": ("STRING", {"default": ""}),
|
|
},
|
|
"optional": {
|
|
"quality": ("INT", {"default": 100}),
|
|
"lossless": ("BOOLEAN", {"default": False}),
|
|
"compression": ("INT", {"default": 4}),
|
|
"optimize": ("BOOLEAN", {"default": False}),
|
|
"metadata_string": ("STRING", {"default": ""}),
|
|
"optional_additional_metadata": ("STRING", {"default": ""}),
|
|
},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",) # Filename
|
|
FUNCTION = "save_images"
|
|
|
|
OUTPUT_NODE = True
|
|
RESULT_NODE = True
|
|
|
|
CATEGORY = "image"
|
|
custom_name = "Save Image Webp Node"
|
|
|
|
def save_images(
|
|
self,
|
|
images,
|
|
filename_prefix="ComfyUI",
|
|
subfolder_dir="",
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
quality=100,
|
|
lossless=False,
|
|
compression=4,
|
|
optimize=False,
|
|
metadata_string="",
|
|
optional_additional_metadata="",
|
|
):
|
|
if images is None: # sometimes images is empty
|
|
return {"ui": {"images": []}, "outputs": {"images": ""}}
|
|
if isinstance(images, torch.Tensor) and images.shape[0] == 0:
|
|
return {"ui": {"images": []}, "outputs": {"images": ""}}
|
|
if isinstance(images, (list, tuple)) and len(images) == 0:
|
|
return {"ui": {"images": []}, "outputs": {"images": ""}}
|
|
if not isinstance(images, (list, tuple, torch.Tensor)):
|
|
images = [images]
|
|
throw_if_parent_or_root_access(filename_prefix)
|
|
throw_if_parent_or_root_access(subfolder_dir)
|
|
filename_prefix += self.prefix_append
|
|
output_dir = os.path.join(self.output_dir, subfolder_dir)
|
|
filelock_path = os.path.join(output_dir, filename_prefix + ".lock")
|
|
|
|
results = list()
|
|
for image in images:
|
|
if isinstance(image, torch.Tensor):
|
|
if image.device.type != "cpu":
|
|
image = image.cpu()
|
|
image = 255.0 * image.numpy()
|
|
clipped = np.clip(image, 0, 255).astype(np.uint8)
|
|
if clipped.shape[0] == 3:
|
|
clipped = np.transpose(clipped, (1, 2, 0)) # [1216, 832, 3]
|
|
# if len(shape) is 4 and first dimension is 1, remove it (batch size)
|
|
if clipped.shape[0] == 1 and len(clipped.shape) == 4:
|
|
clipped = clipped[0]
|
|
# print(clipped.shape)
|
|
img = Image.fromarray(clipped)
|
|
else:
|
|
img = PILHandlingHodes.handle_input(image)
|
|
metadata = None
|
|
if not args.disable_metadata:
|
|
metadata = {}
|
|
if prompt is not None:
|
|
metadata["prompt"] = json.dumps(prompt)
|
|
if extra_pnginfo is not None:
|
|
for x in extra_pnginfo:
|
|
metadata[x] = json.dumps(extra_pnginfo[x])
|
|
if metadata_string: # override metadata
|
|
metadata = {}
|
|
metadata["metadata"] = metadata_string
|
|
if optional_additional_metadata:
|
|
metadata["optional_additional_metadata"] = optional_additional_metadata
|
|
if piexif_loaded:
|
|
exif_bytes = piexif.dump(
|
|
{
|
|
"Exif": {
|
|
piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(
|
|
json.dumps(metadata) or "", encoding="unicode"
|
|
)
|
|
},
|
|
}
|
|
)
|
|
|
|
with filelock.FileLock(
|
|
filelock_path, timeout=10
|
|
): # timeout 10 seconds should be enough for most cases
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = (
|
|
folder_paths.get_save_image_path(
|
|
filename_prefix, output_dir, img.size[1], img.size[0]
|
|
)
|
|
)
|
|
counter_len = len(str(len(images))) # for padding
|
|
# file = f"{filename}_{counter:05}_.webp"
|
|
file = f"{filename}_{str(counter).zfill(max(5, counter_len))}_.webp"
|
|
with tempfile.NamedTemporaryFile(
|
|
suffix=".webp", delete=False
|
|
) as tmpfile:
|
|
tmp_path = tmpfile.name
|
|
img.save(
|
|
tmp_path,
|
|
"WEBP",
|
|
pnginfo=metadata,
|
|
compress_level=compression,
|
|
quality=quality,
|
|
lossless=lossless,
|
|
optimize=optimize,
|
|
)
|
|
if piexif_loaded:
|
|
piexif.insert(exif_bytes, tmp_path)
|
|
final_path = os.path.join(full_output_folder, file)
|
|
shutil.copy2(tmp_path, final_path)
|
|
os.remove(tmp_path)
|
|
|
|
results.append(
|
|
{
|
|
"filename": os.path.join(full_output_folder, file),
|
|
"subfolder": subfolder_dir,
|
|
"type": self.type,
|
|
}
|
|
)
|
|
|
|
return {
|
|
"ui": {"images": results},
|
|
"outputs": {
|
|
"images": os.path.join(full_output_folder, file).rstrip(".webp")
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ComposeRGBAImageFromMask:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
"invert": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "compose"
|
|
CATEGORY = "image"
|
|
custom_name = "Compose RGBA Image From Mask"
|
|
|
|
@staticmethod
|
|
def compose(image, mask, invert):
|
|
if invert:
|
|
mask = 1.0 - mask
|
|
|
|
# Ensure mask has shape (batch_size, height, width, 1)
|
|
mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1], 1))
|
|
# check devices, move to cpu
|
|
if hasattr(image, "device"):
|
|
image = image.cpu()
|
|
if hasattr(mask, "device"):
|
|
mask = mask.cpu()
|
|
# Resize mask to match image dimensions if necessary
|
|
if (
|
|
image.shape[0] != mask.shape[0]
|
|
or image.shape[1] != mask.shape[1]
|
|
or image.shape[2] != mask.shape[2]
|
|
):
|
|
# Resize mask to match image dimensions
|
|
mask = torch.nn.functional.interpolate(
|
|
mask.permute(0, 3, 1, 2),
|
|
size=(image.shape[1], image.shape[2]),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
).permute(0, 2, 3, 1)
|
|
|
|
num_channels = image.shape[-1]
|
|
if num_channels == 3:
|
|
rgba_image = torch.cat((image, mask), dim=-1)
|
|
elif num_channels == 4:
|
|
rgba_image = image.clone()
|
|
rgba_image[:, :, :, 3:] = mask
|
|
else:
|
|
raise ValueError("Image must have 3 (RGB) or 4 (RGBA) channels")
|
|
|
|
return (rgba_image,)
|
|
|
|
|
|
@fundamental_node
|
|
class ResizeImageNode:
|
|
FUNCTION = "resize_image"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Resize Image"
|
|
|
|
constants = {
|
|
"NEAREST": Image.Resampling.NEAREST,
|
|
"LANCZOS": Image.Resampling.LANCZOS,
|
|
"BICUBIC": Image.Resampling.BICUBIC,
|
|
}
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def resize_image(image, width, height, method):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
return (image.resize((width, height), ResizeImageNode.constants[method]),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"width": ("INT", {"default": 512}),
|
|
"height": ("INT", {"default": 512}),
|
|
"method": (["NEAREST", "LANCZOS", "BICUBIC"],),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ResizeImageResolution:
|
|
FUNCTION = "resize_image_resolution"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Resize Image With Resolution"
|
|
|
|
constants = {
|
|
"NEAREST": Image.Resampling.NEAREST,
|
|
"LANCZOS": Image.Resampling.LANCZOS,
|
|
"BICUBIC": Image.Resampling.BICUBIC,
|
|
}
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def resize_image_resolution(image, resolution, method):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
image_width, image_height = image.size
|
|
total_pixels = image_width * image_height
|
|
if total_pixels == 0:
|
|
raise RuntimeError("Image has no pixels")
|
|
if resolution < 256:
|
|
raise RuntimeError("Resolution must be positive and at least 256")
|
|
# get ratio
|
|
target_pixels = resolution**2
|
|
ratio = (target_pixels / total_pixels) ** 0.5
|
|
target_width = int(image_width * ratio)
|
|
target_height = int(image_height * ratio)
|
|
return (
|
|
image.resize(
|
|
(target_width, target_height), ResizeImageResolution.constants[method]
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"resolution": ("INT", {"default": 512}),
|
|
"method": (["NEAREST", "LANCZOS", "BICUBIC"],),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ResizeImageEnsuringMultiple:
|
|
FUNCTION = "resize_image_ensuring_multiple"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Resize Image Ensuring W/H Multiple"
|
|
|
|
constants = {
|
|
"NEAREST": Image.Resampling.NEAREST,
|
|
"LANCZOS": Image.Resampling.LANCZOS,
|
|
"BICUBIC": Image.Resampling.BICUBIC,
|
|
}
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def resize_image_ensuring_multiple(image, multiple, method):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
image_width, image_height = image.size
|
|
total_pixels = image_width * image_height
|
|
if total_pixels == 0:
|
|
raise RuntimeError("Image has no pixels")
|
|
target_width = (image_width // multiple) * multiple
|
|
target_height = (image_height // multiple) * multiple
|
|
return (
|
|
image.resize(
|
|
(target_width, target_height), ResizeImageResolution.constants[method]
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"multiple": ("INT", {"default": 32}),
|
|
"method": (["NEAREST", "LANCZOS", "BICUBIC"],),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ResizeImageResolutionIfBigger:
|
|
FUNCTION = "resize_image_resolution_if_bigger"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Resize Image With Resolution If Bigger"
|
|
|
|
constants = {
|
|
"NEAREST": Image.Resampling.NEAREST,
|
|
"LANCZOS": Image.Resampling.LANCZOS,
|
|
"BICUBIC": Image.Resampling.BICUBIC,
|
|
}
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def resize_image_resolution_if_bigger(image, resolution, method):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
image_width, image_height = image.size
|
|
total_pixels = image_width * image_height
|
|
if total_pixels == 0:
|
|
raise RuntimeError("Image has no pixels")
|
|
if total_pixels <= resolution**2:
|
|
return (image,)
|
|
# get ratio
|
|
target_pixels = resolution**2
|
|
ratio = target_pixels / total_pixels
|
|
target_width = int(image_width * ratio)
|
|
target_height = int(image_height * ratio)
|
|
return (
|
|
image.resize(
|
|
(target_width, target_height),
|
|
ResizeImageResolutionIfBigger.constants[method],
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"resolution": ("INT", {"default": 512}),
|
|
"method": (["NEAREST", "LANCZOS", "BICUBIC"],),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ResizeImageResolutionIfSmaller:
|
|
FUNCTION = "resize_image_resolution_if_smaller"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Resize Image With Resolution If Smaller"
|
|
|
|
constants = {
|
|
"NEAREST": Image.Resampling.NEAREST,
|
|
"LANCZOS": Image.Resampling.LANCZOS,
|
|
"BICUBIC": Image.Resampling.BICUBIC,
|
|
}
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def resize_image_resolution_if_smaller(image, resolution, method):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
image_width, image_height = image.size
|
|
total_pixels = image_width * image_height
|
|
if total_pixels == 0:
|
|
raise RuntimeError("Image has no pixels")
|
|
if total_pixels >= resolution**2:
|
|
return (image,)
|
|
# get ratio
|
|
target_pixels = resolution**2
|
|
ratio = target_pixels / total_pixels
|
|
target_width = int(image_width * ratio)
|
|
target_height = int(image_height * ratio)
|
|
return (
|
|
image.resize(
|
|
(target_width, target_height),
|
|
ResizeImageResolutionIfSmaller.constants[method],
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"resolution": ("INT", {"default": 512}),
|
|
"method": (["NEAREST", "LANCZOS", "BICUBIC"],),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class Base64DecodeNode:
|
|
FUNCTION = "base64_decode"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Base64 Decode to Image"
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def base64_decode(base64_string):
|
|
image = PILHandlingHodes.handle_input(
|
|
base64_string
|
|
) # automatically converts to PIL image
|
|
return (image,)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"base64_string": ("STRING",),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ImageFromURLNode:
|
|
FUNCTION = "url_download"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Download Image from URL"
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def url_download(url):
|
|
if not url.startswith("http"): # for security reasons
|
|
raise RuntimeError(
|
|
"Strict URL check is required, however the URL does not start with http"
|
|
)
|
|
image = PILHandlingHodes.handle_input(url) # automatically downloads image
|
|
return (image,)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"url": ("STRING",),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class RandomImageFromFolderNode:
|
|
FUNCTION = "random_image_from_folder"
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "path")
|
|
CATEGORY = "image"
|
|
custom_name = "Random Image From Folder"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return float("nan")
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def random_image_from_folder(folder, extension, recursive, seed=0):
|
|
folder = str(folder).strip()
|
|
if not folder:
|
|
raise ValueError("folder must be a non-empty path")
|
|
|
|
folder_path = Path(folder).expanduser()
|
|
if not folder_path.is_dir():
|
|
raise ValueError(f"folder is not a directory: {folder_path}")
|
|
|
|
ext = str(extension).strip().lower().lstrip(".")
|
|
if ext == "all":
|
|
extensions = {".jpg", ".jpeg", ".png", ".webp"}
|
|
elif ext == "jpg":
|
|
extensions = {".jpg", ".jpeg"}
|
|
else:
|
|
extensions = {f".{ext}"}
|
|
|
|
candidates = []
|
|
if recursive:
|
|
for suffix in extensions:
|
|
candidates.extend(folder_path.rglob(f"*{suffix}"))
|
|
else:
|
|
for suffix in extensions:
|
|
candidates.extend(folder_path.glob(f"*{suffix}"))
|
|
|
|
candidates = [p for p in candidates if p.is_file()]
|
|
if not candidates:
|
|
raise FileNotFoundError(
|
|
f"No images found in '{folder_path}' with extension '{extension}'"
|
|
)
|
|
|
|
candidates = sorted(set(candidates), key=lambda p: str(p).lower())
|
|
rng = random.Random(seed)
|
|
chosen = candidates[rng.randrange(len(candidates))]
|
|
|
|
with Image.open(chosen) as img:
|
|
img = img.convert("RGB")
|
|
|
|
return (img, str(chosen))
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"folder": ("STRING", {"default": ""}),
|
|
"extension": (["all", "jpg", "png", "webp"], {"default": "all"}),
|
|
"recursive": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 2**63 - 1}),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class Base64EncodeNode:
|
|
FUNCTION = "base64_encode"
|
|
RETURN_TYPES = ("STRING",)
|
|
CATEGORY = "image"
|
|
custom_name = "Image to Base64 Encode"
|
|
|
|
@staticmethod
|
|
def base64_encode(image, quality, format, gzip_compress):
|
|
image = PILHandlingHodes.to_base64(image, quality, format, gzip_compress)
|
|
return (image,)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
"quality": ("INT", {"default": 100}),
|
|
"format": (["PNG", "WEBP", "JPG"], {"default": "PNG"}),
|
|
"gzip_compress": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class StringToBase64Node:
|
|
FUNCTION = "string_to_base64"
|
|
RETURN_TYPES = ("STRING",)
|
|
CATEGORY = "image"
|
|
custom_name = "String to Base64 Encode"
|
|
|
|
@staticmethod
|
|
def string_to_base64(string, gzip_compress):
|
|
return (PILHandlingHodes.string_to_base64(string, gzip_compress),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"string": ("STRING",),
|
|
},
|
|
"optional": {
|
|
"gzip_compress": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class Base64ToStringNode:
|
|
FUNCTION = "base64_to_string"
|
|
RETURN_TYPES = ("STRING",)
|
|
CATEGORY = "image"
|
|
custom_name = "Base64 to String Decode"
|
|
|
|
@staticmethod
|
|
def base64_to_string(base64_string):
|
|
return (PILHandlingHodes.maybe_gzip_base64_to_string(base64_string),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"base64_string": ("STRING",),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class InvertImageNode:
|
|
FUNCTION = "invert_image"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Invert Image"
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def invert_image(image):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
return (ImageOps.invert(image),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ResizeScaleImageNode:
|
|
FUNCTION = "resize_scale_image"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Resize Scale Image"
|
|
|
|
constants = {
|
|
"NEAREST": Image.Resampling.NEAREST,
|
|
"LANCZOS": Image.Resampling.LANCZOS,
|
|
"BICUBIC": Image.Resampling.BICUBIC,
|
|
}
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def resize_scale_image(image, scale, method):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
if scale < 0:
|
|
raise RuntimeError("Scale must be positive")
|
|
return (
|
|
image.resize(
|
|
(int(image.width * scale), int(image.height * scale)),
|
|
ResizeScaleImageNode.constants[method],
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"scale": ("INT", {"default": 2}),
|
|
"method": (["NEAREST", "LANCZOS", "BICUBIC"],),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ResizeShortestToNode:
|
|
FUNCTION = "resize_shortest_to"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Resize Shortest To"
|
|
|
|
constants = {
|
|
"NEAREST": Image.Resampling.NEAREST,
|
|
"LANCZOS": Image.Resampling.LANCZOS,
|
|
"BICUBIC": Image.Resampling.BICUBIC,
|
|
}
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def resize_shortest_to(image, size, method):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
if size < 0:
|
|
raise RuntimeError("Size must be positive")
|
|
if image.width < image.height:
|
|
return (
|
|
image.resize(
|
|
(size, int(image.height * size / image.width)),
|
|
ResizeShortestToNode.constants[method],
|
|
),
|
|
)
|
|
else:
|
|
return (
|
|
image.resize(
|
|
(int(image.width * size / image.height), size),
|
|
ResizeShortestToNode.constants[method],
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"size": ("INT", {"default": 512}),
|
|
"method": (["NEAREST", "LANCZOS", "BICUBIC"],),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ResizeLongestToNode:
|
|
FUNCTION = "resize_longest_to"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Resize Longest To"
|
|
|
|
constants = {
|
|
"NEAREST": Image.Resampling.NEAREST,
|
|
"LANCZOS": Image.Resampling.LANCZOS,
|
|
"BICUBIC": Image.Resampling.BICUBIC,
|
|
}
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def resize_longest_to(image, size, method):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
if size < 0:
|
|
raise RuntimeError("Size must be positive")
|
|
if image.width > image.height:
|
|
return (
|
|
image.resize(
|
|
(size, int(image.height * size / image.width)),
|
|
ResizeLongestToNode.constants[method],
|
|
),
|
|
)
|
|
else:
|
|
return (
|
|
image.resize(
|
|
(int(image.width * size / image.height), size),
|
|
ResizeLongestToNode.constants[method],
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"size": ("INT", {"default": 512}),
|
|
"method": (["NEAREST", "LANCZOS", "BICUBIC"],),
|
|
},
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ConvertGreyscaleNode:
|
|
FUNCTION = "convert_greyscale"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Convert Greyscale"
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def convert_greyscale(image):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
greyscale_image = image.convert("L")
|
|
# 3 channel greyscale image
|
|
return (greyscale_image.convert("RGB"),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class RotateImageNode:
|
|
FUNCTION = "rotate_image"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Rotate Image"
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def rotate_image(image, angle):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
return (image.rotate(angle),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"angle": ("INT", {"default": 0}),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class BrightnessNode:
|
|
FUNCTION = "brightness"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Brightness"
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def brightness(image, factor):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
enhancer = ImageEnhance.Brightness(image)
|
|
return (enhancer.enhance(factor),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"factor": ("FLOAT", {"default": 1.0}),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ContrastNode:
|
|
FUNCTION = "contrast"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Contrast"
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def contrast(image, factor):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
enhancer = ImageEnhance.Contrast(image)
|
|
return (enhancer.enhance(factor),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"factor": ("FLOAT", {"default": 1.0}),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class SharpnessNode:
|
|
FUNCTION = "sharpness"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Sharpness"
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def sharpness(image, factor):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
enhancer = ImageEnhance.Sharpness(image)
|
|
return (enhancer.enhance(factor),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"factor": ("FLOAT", {"default": 1.0}),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ColorNode:
|
|
FUNCTION = "color"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Color"
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def color(image, factor):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
enhancer = ImageEnhance.Color(image)
|
|
return (enhancer.enhance(factor),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"factor": ("FLOAT", {"default": 1.0}),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ConvertRGBNode:
|
|
FUNCTION = "convert_rgb"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Convert RGB"
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def convert_rgb(image):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
return (image.convert("RGB"),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class FFTNode:
|
|
FUNCTION = "fft_image"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "FFT Image"
|
|
|
|
@staticmethod
|
|
def fft_image(image: Image.Image, mask_radius: int) -> Image.Image:
|
|
"""
|
|
Applies an FFT-based low-pass filter to an input PIL image.
|
|
|
|
Args:
|
|
image (Image.Image): Input PIL image to filter.
|
|
mask_radius (int): Radius of the low-pass circular mask.
|
|
|
|
Returns:
|
|
Image.Image: The filtered image as a PIL Image.
|
|
"""
|
|
# Convert image to numpy array
|
|
images = PILHandlingHodes.handle_input(image)
|
|
results = []
|
|
if isinstance(images, list):
|
|
for image in images:
|
|
image_np = np.array(image.convert("RGB"))
|
|
|
|
# Compute FFT for each channel and shift to center
|
|
fft_channels = [
|
|
np.fft.fftshift(np.fft.fft2(image_np[:, :, channel]))
|
|
for channel in range(3)
|
|
]
|
|
|
|
# Create low-pass filter mask
|
|
rows, cols = image_np.shape[:2]
|
|
crow, ccol = rows // 2, cols // 2
|
|
mask = np.zeros((rows, cols), dtype=np.uint8)
|
|
y, x = np.ogrid[-crow : rows - crow, -ccol : cols - ccol]
|
|
mask_area = x**2 + y**2 <= mask_radius**2
|
|
mask[mask_area] = 1
|
|
|
|
# Apply mask and perform inverse FFT
|
|
filtered_channels = [
|
|
np.abs(np.fft.ifft2(np.fft.ifftshift(channel * mask)))
|
|
for channel in fft_channels
|
|
]
|
|
|
|
# Combine channels and convert back to image format
|
|
filtered_image_np = np.stack(filtered_channels, axis=-1)
|
|
filtered_image_np = np.clip(filtered_image_np, 0, 255).astype(np.uint8)
|
|
results.append(
|
|
PILHandlingHodes.handle_output_as_tensor(
|
|
Image.fromarray(filtered_image_np)
|
|
)
|
|
)
|
|
return (results,)
|
|
else:
|
|
image_np = np.array(image.convert("RGB"))
|
|
|
|
# Compute FFT for each channel and shift to center
|
|
fft_channels = [
|
|
np.fft.fftshift(np.fft.fft2(image_np[:, :, channel]))
|
|
for channel in range(3)
|
|
]
|
|
|
|
# Create low-pass filter mask
|
|
rows, cols = image_np.shape[:2]
|
|
crow, ccol = rows // 2, cols // 2
|
|
mask = np.zeros((rows, cols), dtype=np.uint8)
|
|
y, x = np.ogrid[-crow : rows - crow, -ccol : cols - ccol]
|
|
mask_area = x**2 + y**2 <= mask_radius**2
|
|
mask[mask_area] = 1
|
|
|
|
# Apply mask and perform inverse FFT
|
|
filtered_channels = [
|
|
np.abs(np.fft.ifft2(np.fft.ifftshift(channel * mask)))
|
|
for channel in fft_channels
|
|
]
|
|
|
|
# Combine channels and convert back to image format
|
|
filtered_image_np = np.stack(filtered_channels, axis=-1)
|
|
filtered_image_np = np.clip(filtered_image_np, 0, 255).astype(np.uint8)
|
|
return (
|
|
PILHandlingHodes.handle_output_as_tensor(
|
|
Image.fromarray(filtered_image_np)
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"mask_radius": ("INT", {"default": 50}),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class GetImageInfoNode:
|
|
FUNCTION = "get_image_info"
|
|
RETURN_TYPES = ("WIDTH", "HEIGHT", "TOTAL_PIXELS")
|
|
CATEGORY = "image"
|
|
custom_name = "Get Image Info"
|
|
|
|
@staticmethod
|
|
def get_image_info(image):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
width, height = image.size
|
|
return (width, height, width * height)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
|
|
@fundamental_node
|
|
class ThresholdNode:
|
|
FUNCTION = "threshold"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "image"
|
|
custom_name = "Threshold image with value"
|
|
|
|
@staticmethod
|
|
@PILHandlingHodes.output_wrapper
|
|
def threshold(image, threshold):
|
|
image = PILHandlingHodes.handle_input(image)
|
|
return (image.point(lambda p: p > threshold and 255),)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"threshold": ("INT", {"default": 128}),
|
|
}
|
|
}
|
|
|
|
|
|
CLASS_MAPPINGS, CLASS_NAMES = get_node_names_mappings(fundamental_classes)
|
|
validate(fundamental_classes)
|