ComfyUI/comfy_extras/v3/nodes_images.py

728 lines
25 KiB
Python

import hashlib
import os
import numpy as np
import torch
from PIL import Image, ImageOps, ImageSequence
import comfy.utils
import folder_paths
import node_helpers
import nodes
from comfy_api.v3 import io, ui
from server import PromptServer
class GetImageSize(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="GetImageSize_V3",
display_name="Get Image Size _V3",
description="Returns width and height of the image, and passes it through unchanged.",
category="image",
inputs=[
io.Image.Input("image"),
],
outputs=[
io.Int.Output(display_name="width"),
io.Int.Output(display_name="height"),
io.Int.Output(display_name="batch_size"),
],
hidden=[io.Hidden.unique_id],
)
@classmethod
def execute(cls, image) -> io.NodeOutput:
height = image.shape[1]
width = image.shape[2]
batch_size = image.shape[0]
if cls.hidden.unique_id:
PromptServer.instance.send_progress_text(
f"width: {width}, height: {height}\n batch size: {batch_size}", cls.hidden.unique_id
)
return io.NodeOutput(width, height, batch_size)
class ImageAddNoise(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="ImageAddNoise_V3",
display_name="Image Add Noise _V3",
category="image",
inputs=[
io.Image.Input("image"),
io.Int.Input(
"seed",
default=0,
min=0,
max=0xFFFFFFFFFFFFFFFF,
control_after_generate=True,
tooltip="The random seed used for creating the noise.",
),
io.Float.Input("strength", default=0.5, min=0.0, max=1.0, step=0.01),
],
outputs=[io.Image.Output()],
)
@classmethod
def execute(cls, image, seed, strength) -> io.NodeOutput:
generator = torch.manual_seed(seed)
s = torch.clip(
(image + strength * torch.randn(image.size(), generator=generator, device="cpu").to(image)),
min=0.0,
max=1.0,
)
return io.NodeOutput(s)
class ImageCrop(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="ImageCrop_V3",
display_name="Image Crop _V3",
category="image/transform",
inputs=[
io.Image.Input("image"),
io.Int.Input("width", default=512, min=1, max=nodes.MAX_RESOLUTION, step=1),
io.Int.Input("height", default=512, min=1, max=nodes.MAX_RESOLUTION, step=1),
io.Int.Input("x", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
io.Int.Input("y", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
],
outputs=[io.Image.Output()],
)
@classmethod
def execute(cls, image, width, height, x, y) -> io.NodeOutput:
x = min(x, image.shape[2] - 1)
y = min(y, image.shape[1] - 1)
to_x = width + x
to_y = height + y
return io.NodeOutput(image[:, y:to_y, x:to_x, :])
class ImageFlip(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="ImageFlip_V3",
display_name="Image Flip _V3",
category="image/transform",
inputs=[
io.Image.Input("image"),
io.Combo.Input("flip_method", options=["x-axis: vertically", "y-axis: horizontally"]),
],
outputs=[io.Image.Output()],
)
@classmethod
def execute(cls, image, flip_method) -> io.NodeOutput:
if flip_method.startswith("x"):
image = torch.flip(image, dims=[1])
elif flip_method.startswith("y"):
image = torch.flip(image, dims=[2])
return io.NodeOutput(image)
class ImageFromBatch(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="ImageFromBatch_V3",
display_name="Image From Batch _V3",
category="image/batch",
inputs=[
io.Image.Input("image"),
io.Int.Input("batch_index", default=0, min=0, max=4095),
io.Int.Input("length", default=1, min=1, max=4096),
],
outputs=[io.Image.Output()],
)
@classmethod
def execute(cls, image, batch_index, length) -> io.NodeOutput:
s_in = image
batch_index = min(s_in.shape[0] - 1, batch_index)
length = min(s_in.shape[0] - batch_index, length)
s = s_in[batch_index : batch_index + length].clone()
return io.NodeOutput(s)
class ImageRotate(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="ImageRotate_V3",
display_name="Image Rotate _V3",
category="image/transform",
inputs=[
io.Image.Input("image"),
io.Combo.Input("rotation", options=["none", "90 degrees", "180 degrees", "270 degrees"]),
],
outputs=[io.Image.Output()],
)
@classmethod
def execute(cls, image, rotation) -> io.NodeOutput:
rotate_by = 0
if rotation.startswith("90"):
rotate_by = 1
elif rotation.startswith("180"):
rotate_by = 2
elif rotation.startswith("270"):
rotate_by = 3
return io.NodeOutput(torch.rot90(image, k=rotate_by, dims=[2, 1]))
class ImageStitch(io.ComfyNode):
"""Upstreamed from https://github.com/kijai/ComfyUI-KJNodes"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="ImageStitch_V3",
display_name="Image Stitch _V3",
description="Stitches image2 to image1 in the specified direction. "
"If image2 is not provided, returns image1 unchanged. "
"Optional spacing can be added between images.",
category="image/transform",
inputs=[
io.Image.Input("image1"),
io.Combo.Input("direction", options=["right", "down", "left", "up"], default="right"),
io.Boolean.Input("match_image_size", default=True),
io.Int.Input("spacing_width", default=0, min=0, max=1024, step=2),
io.Combo.Input("spacing_color", options=["white", "black", "red", "green", "blue"], default="white"),
io.Image.Input("image2", optional=True),
],
outputs=[io.Image.Output()],
)
@classmethod
def execute(cls, image1, direction, match_image_size, spacing_width, spacing_color, image2=None) -> io.NodeOutput:
if image2 is None:
return io.NodeOutput(image1)
# Handle batch size differences
if image1.shape[0] != image2.shape[0]:
max_batch = max(image1.shape[0], image2.shape[0])
if image1.shape[0] < max_batch:
image1 = torch.cat([image1, image1[-1:].repeat(max_batch - image1.shape[0], 1, 1, 1)])
if image2.shape[0] < max_batch:
image2 = torch.cat([image2, image2[-1:].repeat(max_batch - image2.shape[0], 1, 1, 1)])
# Match image sizes if requested
if match_image_size:
h1, w1 = image1.shape[1:3]
h2, w2 = image2.shape[1:3]
aspect_ratio = w2 / h2
if direction in ["left", "right"]:
target_h, target_w = h1, int(h1 * aspect_ratio)
else: # up, down
target_w, target_h = w1, int(w1 / aspect_ratio)
image2 = comfy.utils.common_upscale(
image2.movedim(-1, 1), target_w, target_h, "lanczos", "disabled"
).movedim(1, -1)
color_map = {
"white": 1.0,
"black": 0.0,
"red": (1.0, 0.0, 0.0),
"green": (0.0, 1.0, 0.0),
"blue": (0.0, 0.0, 1.0),
}
color_val = color_map[spacing_color]
# When not matching sizes, pad to align non-concat dimensions
if not match_image_size:
h1, w1 = image1.shape[1:3]
h2, w2 = image2.shape[1:3]
pad_value = 0.0
if not isinstance(color_val, tuple):
pad_value = color_val
if direction in ["left", "right"]:
# For horizontal concat, pad heights to match
if h1 != h2:
target_h = max(h1, h2)
if h1 < target_h:
pad_h = target_h - h1
pad_top, pad_bottom = pad_h // 2, pad_h - pad_h // 2
image1 = torch.nn.functional.pad(
image1, (0, 0, 0, 0, pad_top, pad_bottom), mode="constant", value=pad_value
)
if h2 < target_h:
pad_h = target_h - h2
pad_top, pad_bottom = pad_h // 2, pad_h - pad_h // 2
image2 = torch.nn.functional.pad(
image2, (0, 0, 0, 0, pad_top, pad_bottom), mode="constant", value=pad_value
)
else: # up, down
# For vertical concat, pad widths to match
if w1 != w2:
target_w = max(w1, w2)
if w1 < target_w:
pad_w = target_w - w1
pad_left, pad_right = pad_w // 2, pad_w - pad_w // 2
image1 = torch.nn.functional.pad(
image1, (0, 0, pad_left, pad_right), mode="constant", value=pad_value
)
if w2 < target_w:
pad_w = target_w - w2
pad_left, pad_right = pad_w // 2, pad_w - pad_w // 2
image2 = torch.nn.functional.pad(
image2, (0, 0, pad_left, pad_right), mode="constant", value=pad_value
)
# Ensure same number of channels
if image1.shape[-1] != image2.shape[-1]:
max_channels = max(image1.shape[-1], image2.shape[-1])
if image1.shape[-1] < max_channels:
image1 = torch.cat(
[
image1,
torch.ones(
*image1.shape[:-1],
max_channels - image1.shape[-1],
device=image1.device,
),
],
dim=-1,
)
if image2.shape[-1] < max_channels:
image2 = torch.cat(
[
image2,
torch.ones(
*image2.shape[:-1],
max_channels - image2.shape[-1],
device=image2.device,
),
],
dim=-1,
)
# Add spacing if specified
if spacing_width > 0:
spacing_width = spacing_width + (spacing_width % 2) # Ensure even
if direction in ["left", "right"]:
spacing_shape = (
image1.shape[0],
max(image1.shape[1], image2.shape[1]),
spacing_width,
image1.shape[-1],
)
else:
spacing_shape = (
image1.shape[0],
spacing_width,
max(image1.shape[2], image2.shape[2]),
image1.shape[-1],
)
spacing = torch.full(spacing_shape, 0.0, device=image1.device)
if isinstance(color_val, tuple):
for i, c in enumerate(color_val):
if i < spacing.shape[-1]:
spacing[..., i] = c
if spacing.shape[-1] == 4: # Add alpha
spacing[..., 3] = 1.0
else:
spacing[..., : min(3, spacing.shape[-1])] = color_val
if spacing.shape[-1] == 4:
spacing[..., 3] = 1.0
# Concatenate images
images = [image2, image1] if direction in ["left", "up"] else [image1, image2]
if spacing_width > 0:
images.insert(1, spacing)
concat_dim = 2 if direction in ["left", "right"] else 1
return io.NodeOutput(torch.cat(images, dim=concat_dim))
class LoadImage(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoadImage_V3",
display_name="Load Image _V3",
category="image",
inputs=[
io.Combo.Input(
"image",
upload=io.UploadType.image,
image_folder=io.FolderType.input,
options=cls.get_files_options(),
),
],
outputs=[
io.Image.Output(),
io.Mask.Output(),
],
)
@classmethod
def get_files_options(cls) -> list[str]:
target_dir = folder_paths.get_input_directory()
files = [f for f in os.listdir(target_dir) if os.path.isfile(os.path.join(target_dir, f))]
return sorted(folder_paths.filter_files_content_types(files, ["image"]))
@classmethod
def execute(cls, image) -> io.NodeOutput:
img = node_helpers.pillow(Image.open, folder_paths.get_annotated_filepath(image))
output_images = []
output_masks = []
w, h = None, None
excluded_formats = ["MPO"]
for i in ImageSequence.Iterator(img):
i = node_helpers.pillow(ImageOps.exif_transpose, i)
if i.mode == "I":
i = i.point(lambda i: i * (1 / 255))
image = i.convert("RGB")
if len(output_images) == 0:
w = image.size[0]
h = image.size[1]
if image.size[0] != w or image.size[1] != h:
continue
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
if "A" in i.getbands():
mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0
mask = 1.0 - torch.from_numpy(mask)
elif i.mode == "P" and "transparency" in i.info:
mask = np.array(i.convert("RGBA").getchannel("A")).astype(np.float32) / 255.0
mask = 1.0 - torch.from_numpy(mask)
else:
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
output_images.append(image)
output_masks.append(mask.unsqueeze(0))
if len(output_images) > 1 and img.format not in excluded_formats:
output_image = torch.cat(output_images, dim=0)
output_mask = torch.cat(output_masks, dim=0)
else:
output_image = output_images[0]
output_mask = output_masks[0]
return io.NodeOutput(output_image, output_mask)
@classmethod
def fingerprint_inputs(s, image):
image_path = folder_paths.get_annotated_filepath(image)
m = hashlib.sha256()
with open(image_path, "rb") as f:
m.update(f.read())
return m.digest().hex()
@classmethod
def validate_inputs(s, image):
if not folder_paths.exists_annotated_filepath(image):
return "Invalid image file: {}".format(image)
return True
class LoadImageOutput(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoadImageOutput_V3",
display_name="Load Image (from Outputs) _V3",
description="Load an image from the output folder. "
"When the refresh button is clicked, the node will update the image list "
"and automatically select the first image, allowing for easy iteration.",
category="image",
inputs=[
io.Combo.Input(
"image",
upload=io.UploadType.image,
image_folder=io.FolderType.output,
remote=io.RemoteOptions(
route="/internal/files/output",
refresh_button=True,
control_after_refresh="first",
),
),
],
outputs=[
io.Image.Output(),
io.Mask.Output(),
],
)
@classmethod
def execute(cls, image) -> io.NodeOutput:
img = node_helpers.pillow(Image.open, folder_paths.get_annotated_filepath(image))
output_images = []
output_masks = []
w, h = None, None
excluded_formats = ["MPO"]
for i in ImageSequence.Iterator(img):
i = node_helpers.pillow(ImageOps.exif_transpose, i)
if i.mode == "I":
i = i.point(lambda i: i * (1 / 255))
image = i.convert("RGB")
if len(output_images) == 0:
w = image.size[0]
h = image.size[1]
if image.size[0] != w or image.size[1] != h:
continue
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
if "A" in i.getbands():
mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0
mask = 1.0 - torch.from_numpy(mask)
elif i.mode == "P" and "transparency" in i.info:
mask = np.array(i.convert("RGBA").getchannel("A")).astype(np.float32) / 255.0
mask = 1.0 - torch.from_numpy(mask)
else:
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
output_images.append(image)
output_masks.append(mask.unsqueeze(0))
if len(output_images) > 1 and img.format not in excluded_formats:
output_image = torch.cat(output_images, dim=0)
output_mask = torch.cat(output_masks, dim=0)
else:
output_image = output_images[0]
output_mask = output_masks[0]
return io.NodeOutput(output_image, output_mask)
@classmethod
def fingerprint_inputs(s, image):
image_path = folder_paths.get_annotated_filepath(image)
m = hashlib.sha256()
with open(image_path, "rb") as f:
m.update(f.read())
return m.digest().hex()
@classmethod
def validate_inputs(s, image):
if not folder_paths.exists_annotated_filepath(image):
return "Invalid image file: {}".format(image)
return True
class PreviewImage(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="PreviewImage_V3",
display_name="Preview Image _V3",
description="Preview the input images.",
category="image",
inputs=[
io.Image.Input("images", tooltip="The images to preview."),
],
hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
is_output_node=True,
)
@classmethod
def execute(cls, images) -> io.NodeOutput:
return io.NodeOutput(ui=ui.PreviewImage(images, cls=cls))
class RepeatImageBatch(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="RepeatImageBatch_V3",
display_name="Repeat Image Batch _V3",
category="image/batch",
inputs=[
io.Image.Input("image"),
io.Int.Input("amount", default=1, min=1, max=4096),
],
outputs=[io.Image.Output()],
)
@classmethod
def execute(cls, image, amount) -> io.NodeOutput:
return io.NodeOutput(image.repeat((amount, 1, 1, 1)))
class ResizeAndPadImage(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="ResizeAndPadImage_V3",
display_name="Resize and Pad Image _V3",
category="image/transform",
inputs=[
io.Image.Input("image"),
io.Int.Input("target_width", default=512, min=1, max=nodes.MAX_RESOLUTION, step=1),
io.Int.Input("target_height", default=512, min=1, max=nodes.MAX_RESOLUTION, step=1),
io.Combo.Input("padding_color", options=["white", "black"]),
io.Combo.Input("interpolation", options=["area", "bicubic", "nearest-exact", "bilinear", "lanczos"]),
],
outputs=[io.Image.Output()],
)
@classmethod
def execute(cls, image, target_width, target_height, padding_color, interpolation) -> io.NodeOutput:
batch_size, orig_height, orig_width, channels = image.shape
scale_w = target_width / orig_width
scale_h = target_height / orig_height
scale = min(scale_w, scale_h)
new_width = int(orig_width * scale)
new_height = int(orig_height * scale)
image_permuted = image.permute(0, 3, 1, 2)
resized = comfy.utils.common_upscale(image_permuted, new_width, new_height, interpolation, "disabled")
pad_value = 0.0 if padding_color == "black" else 1.0
padded = torch.full(
(batch_size, channels, target_height, target_width), pad_value, dtype=image.dtype, device=image.device
)
y_offset = (target_height - new_height) // 2
x_offset = (target_width - new_width) // 2
padded[:, :, y_offset : y_offset + new_height, x_offset : x_offset + new_width] = resized
return io.NodeOutput(padded.permute(0, 2, 3, 1))
class SaveAnimatedPNG(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SaveAnimatedPNG_V3",
display_name="Save Animated PNG _V3",
category="image/animation",
inputs=[
io.Image.Input("images"),
io.String.Input("filename_prefix", default="ComfyUI"),
io.Float.Input("fps", default=6.0, min=0.01, max=1000.0, step=0.01),
io.Int.Input("compress_level", default=4, min=0, max=9),
],
hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
is_output_node=True,
)
@classmethod
def execute(cls, images, fps, compress_level, filename_prefix="ComfyUI") -> io.NodeOutput:
return io.NodeOutput(
ui=ui.ImageSaveHelper.get_save_animated_png_ui(
images=images,
filename_prefix=filename_prefix,
cls=cls,
fps=fps,
compress_level=compress_level,
)
)
class SaveAnimatedWEBP(io.ComfyNode):
COMPRESS_METHODS = {"default": 4, "fastest": 0, "slowest": 6}
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SaveAnimatedWEBP_V3",
display_name="Save Animated WEBP _V3",
category="image/animation",
inputs=[
io.Image.Input("images"),
io.String.Input("filename_prefix", default="ComfyUI"),
io.Float.Input("fps", default=6.0, min=0.01, max=1000.0, step=0.01),
io.Boolean.Input("lossless", default=True),
io.Int.Input("quality", default=80, min=0, max=100),
io.Combo.Input("method", options=list(cls.COMPRESS_METHODS.keys())),
],
hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
is_output_node=True,
)
@classmethod
def execute(cls, images, fps, filename_prefix, lossless, quality, method) -> io.NodeOutput:
return io.NodeOutput(
ui=ui.ImageSaveHelper.get_save_animated_webp_ui(
images=images,
filename_prefix=filename_prefix,
cls=cls,
fps=fps,
lossless=lossless,
quality=quality,
method=cls.COMPRESS_METHODS.get(method)
)
)
class SaveImage(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SaveImage_V3",
display_name="Save Image _V3",
description="Saves the input images to your ComfyUI output directory.",
category="image",
inputs=[
io.Image.Input(
"images",
tooltip="The images to save.",
),
io.String.Input(
"filename_prefix",
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=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
is_output_node=True,
)
@classmethod
def execute(cls, images, filename_prefix="ComfyUI") -> io.NodeOutput:
return io.NodeOutput(
ui=ui.ImageSaveHelper.get_save_images_ui(images, filename_prefix=filename_prefix, cls=cls, compress_level=4)
)
NODES_LIST: list[type[io.ComfyNode]] = [
GetImageSize,
ImageAddNoise,
ImageCrop,
ImageFlip,
ImageFromBatch,
ImageRotate,
ImageStitch,
LoadImage,
LoadImageOutput,
PreviewImage,
RepeatImageBatch,
ResizeAndPadImage,
SaveAnimatedPNG,
SaveAnimatedWEBP,
SaveImage,
]