Files
Hullabalo-ComfyUI-Loop/nodes.py
T
Hullabaloo 29269ffe64 Add files via upload
version 0.1.5 - so many changes
2025-04-30 22:38:42 +02:00

515 lines
20 KiB
Python

import time
import os
import shutil
import json
import node_helpers
import folder_paths
import numpy as np
import torch
import random
from PIL import Image, ImageOps
from PIL.PngImagePlugin import PngInfo
from comfy.cli_args import args
# a simple image loop for your workflow. MIT License. version 0.1.5
# https://github.com/Hullabalo/ComfyUI-Loop/
# last changes 4/30/2025 :
# Better integration with classic ComfyUI workflow, compatible inputs and outputs
# - added the node LoopImageSimple in a aim to replace the LoadImageSimple. A different flavor, that's something.
# - option given in loopimagesimple to choose the image name and folder, and/or use a classic image input and mask
# - some minors modifications into Save Image. Removed compatibility with jpg/jpeg ...
# - added the Crop Image node, woohoo javascript overdose, but it works ok with ComfyUI v0.3.10, ComfyUI_Frontend v1.6.15
# - now, 'Paste Image' hosts a optional cutting_mask input
# - small upgrade to 'Cut Image' node, added mask input and cutting_mask output
# TL;DR : I've reviewed everything.
# TODO : removing the shit out of the code, keeping the good.
class LoadImageSimple:
# DEPRECATED. Use vanilla Load Image with 'Loop Image' node instead
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"image": ("STRING",{"default": "image.png", "tooltip": "Name of an image file located in output folder."}),
}
}
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "STRING")
RETURN_NAMES = ("image", "mask", "height", "width", "path")
CATEGORY = "LOOP"
DESCRIPTION = "Load an image from output folder."
FUNCTION = "load_image"
def load_image(self, image):
image = os.path.join(self.output_dir, image)
img = node_helpers.pillow(Image.open, image)
output_images = []
output_masks = []
w, h = None, None
i = node_helpers.pillow(ImageOps.exif_transpose, img)
if i.mode == 'I':
i = i.point(lambda i: i * (1 / 255))
img = i.convert("RGB")
if len(output_images) == 0:
w = img.size[0]
h = img.size[1]
img = np.array(img).astype(np.float32) / 255.0
img = torch.from_numpy(img)[None,]
if 'A' in i.getbands():
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
mask = 1. - torch.from_numpy(mask)
else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
output_images.append(img)
output_masks.append(mask.unsqueeze(0))
output_image = output_images[0]
output_mask = output_masks[0]
return (output_image, output_mask, h, w, image)
@classmethod
def IS_CHANGED(s, image):
return float("NaN")
class LoopImageSimple:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",{}),
"loop_image": ("BOOLEAN", {"default": False, "tooltip": "Enable image loop mode. Disable (false) to load from image input"}),
"loop_mask": ("BOOLEAN", {"default": False, "tooltip": "Enable mask loop mode. Disable (false) to load from optional mask input"}),
"filename": ("STRING",{"default": "loop_image", "tooltip": "(facultative) Filename of loop image without extension (it's .png only). Use the name of an already existing image to load it from /output or one of its subfolders. Yeah, I know that it's a too long long long description for a simple field. :)"}),
"folder": ("STRING",{"default": "", "tooltip": "(facultative) Name a subfolder to load from or copy the input image into. If it does not exists, it will be created into /output"}),
},
"optional": {
"mask": ("MASK", {"tooltip": "Optional mask. Fire your workflow one time before connecting"}),
},
}
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "STRING")
RETURN_NAMES = ("image", "mask", "height", "width", "path")
CATEGORY = "LOOP"
DESCRIPTION = "Loop an image and mask. " \
"Image is copied from image input to the folder defined in corresponding field, " \
"and named as in filename field." \
"You can also use those two fields to load directly an image from /output/'folder'." \
"" \
"Use the loop_image and loop_mask to load a new mask and/or image, or to loop the same image and/or mask."
FUNCTION = "load_loop_image"
def load_loop_image(self, image, loop_image, loop_mask, filename, folder, mask=None):
# folder management
if folder != "":
folder_path = os.path.join(self.output_dir, folder)
if not os.path.exists(folder_path):
os.makedirs(folder_path)
else:
folder_path = self.output_dir
filename = f"loop_image.png" if filename == "" else filename+".png"
image_path = os.path.join(folder_path, filename)
# Check if the file already exists in loop folder, and that node is in a loop configuration
if os.path.exists(image_path) and loop_image:
# create a tensor from this file instead of using the input tensor
img = node_helpers.pillow(Image.open, image_path)
i = node_helpers.pillow(ImageOps.exif_transpose, img)
if i.mode == 'I':
i = i.point(lambda i: i * (1 / 255))
img = i.convert("RGB")
w = img.size[0]
h = img.size[1]
img = np.array(img).astype(np.float32) / 255.0
img = torch.from_numpy(img)[None,]
# by order, returns the optional input mask, or alpha from image, or an empty mask
if mask is not None and loop_mask is False:
mask_out = mask
elif 'A' in i.getbands():
mask_out = np.array(i.getchannel('A')).astype(np.float32) / 255.0
mask_out = 1. - torch.from_numpy(mask_out)
mask_out = mask_out.unsqueeze(0)
else:
mask_out = torch.zeros((h, w), dtype=torch.float32, device="cpu")
mask_out = mask_out.unsqueeze(0)
output_image = img
output_mask = mask_out
else:
# If the loop image has not been already created in /output/folder
pil_image = Image.fromarray((image[0].cpu().numpy() * 255).astype(np.uint8))
pil_image.save(image_path)
# using input image
output_image = image
h, w = image.shape[1], image.shape[2]
# by order, returns the optional input mask, or alpha from image, or an empty mask
if mask is not None and loop_mask is False:
output_mask = mask
elif image.shape[3] == 4: # if image is (RGBA), extract alpha channel
output_mask = image[0, :, :, 3:4].permute(2, 0, 1) # permute to (1, h, w)
else:
# create an empty mask
output_mask = torch.zeros((1, h, w), dtype=torch.float32, device="cpu")
return (output_image, output_mask, h, w, image_path)
@classmethod
def IS_CHANGED(cls, image, loop_image, loop_mask, filename, folder, mask):
return float("NaN")
class SaveImageSimple:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE", {"tooltip": "Image datas."}),
"image_path": ("STRING", {"default": "/path/to/image.ext", "defaultInput": True, "tooltip": "Full path of the saved image."}),
"save_steps": ("BOOLEAN", {"default": False, "tooltip": "Save a copy next to the saved image with a timestamp as suffix."}),
},
"optional": {
"mask": ("MASK", {"tooltip": "Optional mask to use as alpha channel."}),
},
"hidden": {
"prompt": "PROMPT",
"extra_pnginfo": "EXTRA_PNGINFO"
},
}
RETURN_TYPES = ()
FUNCTION = "save_images"
OUTPUT_NODE = True
CATEGORY = "LOOP"
DESCRIPTION = "Saves the input image, optionally add a mask as alpha channel."
def save_images(self, image, image_path, save_steps, prompt=None, extra_pnginfo=None, mask=None):
result = list()
filename = os.path.basename(image_path)
current_saving_path = os.path.dirname(image_path)
# Convert to numpy array as float32
i = (255. * image.cpu().numpy()).astype(np.float32)
# Handle image shape (1, H, W, C)
if len(i.shape) == 4 and i.shape[0] == 1:
i = i[0] # delete batch dimension
img = Image.fromarray(np.rint(i).clip(0, 255).astype(np.uint8)) # Convert to uint8 using np.rint for lossless output
# 4/29/2025 : removed the option for jpg/jpeg, compatibility is keeped with LoadImageSimple node only for png.
# save to specified path, overwriting source image
# if image_path.lower().endswith('.jpg') or image_path.lower().endswith('.jpeg'):
# img.save(image_path, quality=100)
# else:
# Add mask as alpha channel if provided
if mask is not None:
m = (255. * mask.cpu().numpy()).astype(np.uint8)
if len(m.shape) == 4:
m = m[0]
m = np.squeeze(m)
m = 255 - m
m = np.rint(m).clip(0, 255).astype(np.uint8) # m = np.clip(m, 0, 255).astype(np.uint8)
# Ensure mask size matches image size
if m.shape[:2] != (img.height, img.width):
print(f"Error : Mask size {m.shape[:2]} doesn't match image size {(img.height, img.width)}.")
# force mask resize
# if m.shape[:2] != (img.height, img.width):
# m = Image.fromarray(m)
# m = m.resize((img.width, img.height), resample=Image.BILINEAR)
# m = np.array(m)
else:
img = img.convert("RGBA")
alpha = Image.fromarray(m)
r, g, b = img.split()[:3]
img = Image.merge("RGBA", (r, g, b, alpha))
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]))
img.save(image_path, pnginfo=metadata, compress_level=0)
# optionally, save a incremented copy of the image,
# in the folder defined in Loop Image node (if undefined, default as /output)
if save_steps == True:
timestamp = f"{time.time():.6f}".replace(".", "")
base, ext = os.path.splitext(filename)
img_copy_name = f"{base}_{timestamp}{ext}"
img_copy_path = os.path.join(current_saving_path, img_copy_name)
shutil.copy(image_path, img_copy_path)
result.append({
"filename": filename,
"subfolder": os.path.basename(current_saving_path),
"type": "output"
})
return {"ui": {"images": result}}
class ImageCutLoop:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"x": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 8}),
"y": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 8}),
"width": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}),
"height": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}),
},
"optional": {
"mask": ("MASK", {"tooltip": "Optional mask."}),
},
}
RETURN_TYPES = ("IMAGE", "IMAGE", "MASK", "INT", "INT",)
RETURN_NAMES = ("source", "cutting", "cutting_mask", "x","y",)
FUNCTION = "cut"
CATEGORY = "LOOP"
DESCRIPTION = "Cut a part of an input image."
def cut(self, image, x, y, width, height, mask=None):
try:
_, img_height, img_width, _ = image.shape
# validate crop parameters
x = min(max(0, x), img_width - 1)
y = min(max(0, y), img_height - 1)
width = min(width, img_width - x)
height = min(height, img_height - y)
# crop image
cut = image[:, y:y+height, x:x+width, :]
# define cutting mask
if mask is not None:
cutting_mask = mask[:, y:y+width, x:x+height]
else:
cutting_mask = torch.zeros((1, height, width), dtype=torch.float32, device="cpu")
return (image, cut, cutting_mask, x, y)
except Exception as e:
print(f"Cut error: {str(e)}")
raise e
@classmethod
def IS_CHANGED(cls, image, mask, width, height, size):
return True
class ImageCropLoop:
"""
A ComfyUI node to generate an RGB image using pyvips.
Parameters:
- x (int): top left x origin of the crop.
- y (int): top left y origin of the crop.
- size (int): size of the crop (height = width). Multiple of 8
hint: if you need a size wider than 2048*2048 for your crop, change size widget "max" to anything:
your output image crop will never exceed the max size of your input image, whatever you choose or do.
Returns:
- source: source image
- cutting: image crop.
- cutting_mask: empty mask
- x: x origin
- y: y origin
- size: size of the crop
"""
def __init__(self):
self.temp_dir = folder_paths.get_temp_directory()
self.filename_prefix = "loop_crop_temp_" + ''.join(random.choice("abcdefghijklmnopqrstuvxyz") for _ in range(5))
self.type = "temp"
self.preview_size = 1024 # change this to something smaller/larger if you need a fast ui or a detailed preview. 1024 is just a compromise to keep the UI responsive
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE", {"tooltip": "Source image."}),
"x": ("INT", {"default": 0, "min": 0, "max": 32768, "tooltip": "x origin of the crop."}), # max 32768 is enought for 1073 Megapixels images. Much more than what you can effectively process :)
"y": ("INT", {"default": 0, "min": 0, "max": 32768, "tooltip": "y origin of the crop."}),
"size": ("INT", {"default":512, "min": 256, "max": 2048, "step": 8, "tooltip": "size of the crop. By 8 pixels increments."}), # step is 8 because inpaintModelConditionning output multiple of 8 latents (e.g. [[1, 16, 128, 128]] for a 1024*1024px batch of 1 image input)
},
"optional": {
"mask": ("MASK", {"tooltip": "Optional mask."}),
},
"hidden": {
}
}
RETURN_TYPES = ("IMAGE","IMAGE","MASK","INT","INT","INT")
RETURN_NAMES = ("source","cutting","cutting_mask","x","y","size")
FUNCTION = "click_and_crop"
CATEGORY = "LOOP"
DESCRIPTION = "Crop a clicked zone to the size dimensions. Output a cut image."
def click_and_crop(self, image, x, y, size, mask=None):
counter = 0
results = []
# get original dimensions (for preview, adjustement of crop size and passing to js ui)
_, oh, ow, _ = image.shape
# create and save preview image
pw_filename = f"{self.filename_prefix}_{ow}_{oh}_{counter:05}.png"
os.makedirs(self.temp_dir, exist_ok=True)
file_path = os.path.join(self.temp_dir, pw_filename)
pil_image = Image.fromarray((image[0].cpu().numpy() * 255).astype(np.uint8))
# resize preview
if self.preview_size < ow:
ratio = self.preview_size/ow
pil_image = pil_image.resize((int(ow*ratio), int(oh*ratio)), Image.LANCZOS)
pil_image.save(file_path)
results.append({
"filename": pw_filename,
"subfolder": "",
"type": self.type,
})
# adjust crop size if bigger than the image
if size > ow:
size = ow
print(f"crop size is larger than the width of the image, reducing to '{ow}'")
if size > oh:
size = oh
print(f"crop size is larger than the height of the image, reducing to '{oh}'")
# adjust coords x et y if out of the border
if x + size > ow:
x = max(0, ow - size)
print(f"x+crop size goes out of the borders, changing it to '{x}'")
if y + size > oh:
y = max(0, oh - size)
print(f"y+crop size goes out of the borders, changing it to '{y}'")
# extract tile
tile = image[:, y:y+size, x:x+size, :]
# define cutting mask
if mask is not None:
cutting_mask = mask[:, y:y+size, x:x+size]
else:
cutting_mask = torch.zeros((1, size, size), dtype=torch.float32, device="cpu")
return {"ui": {"images": results}, "result": (image, tile, cutting_mask, x, y, size)}
@classmethod
def IS_CHANGED(cls, image, mask, x, y, size):
return True
class ImagePasteLoop:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"source": ("IMAGE",),
"cutting": ("IMAGE",),
"x": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 8, "defaultInput": True}),
"y": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 8, "defaultInput": True}),
},
"optional": {
"cutting_mask": ("MASK", {"tooltip": "Optional cutting mask."}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "paste"
CATEGORY = "LOOP"
DESCRIPTION = "Replace a cutting in the original image."
def paste(self, source, cutting, x, y, cutting_mask=None):
try:
_, orig_height, orig_width, _ = source.shape
_, crop_height, crop_width, _ = cutting.shape
# create a copy of original image
paste_image = source.clone()
# check x, y parameters to assure they are bounded into the limits
x = min(max(0, x), orig_width - 1)
y = min(max(0, y), orig_height - 1)
# adjust dimensions if crop goes out of the borders
paste_width = min(crop_width, orig_width - x)
paste_height = min(crop_height, orig_height - y)
# if mask is None or empty
if cutting_mask is None or cutting_mask.max() == 0:
paste_image[:, y:y+paste_height, x:x+paste_width, :] = cutting[:, :paste_height, :paste_width, :]
else:
_, mask_height, mask_width = cutting_mask.shape # (batch, H, W)
if crop_height != mask_height or crop_width != mask_width:
raise ValueError("Mask and cropped image must have the same dimensions.")
# adjust mask and cropped image if needed
cutting = cutting[:, :paste_height, :paste_width, :]
cutting_mask = cutting_mask[:, :paste_height, :paste_width]
cutting_mask = cutting_mask.unsqueeze(-1)
# normalize mask (between 0 and 1 values)
cutting_mask = cutting_mask.float() / cutting_mask.max()
# apply mask
paste_image[:, y:y+paste_height, x:x+paste_width, :] = (
cutting * cutting_mask + paste_image[:, y:y+paste_height, x:x+paste_width, :] * (1 - cutting_mask)
)
return (paste_image,)
except Exception as e:
print(f"Error while pasting : {str(e)}")
raise e
NODE_CLASS_MAPPINGS = {
"LoadImageSimple": LoadImageSimple,
"LoopImageSimple": LoopImageSimple,
"SaveImageSimple": SaveImageSimple,
"ImageCutLoop": ImageCutLoop,
"ImageCropLoop": ImageCropLoop,
"ImagePasteLoop": ImagePasteLoop
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LoadImageSimple": "Load Image (LOOP) DEPRECATED : use Loop Image instead",
"LoopImageSimple": "Loop Image (LOOP)",
"SaveImageSimple": "Save Image (LOOP)",
"ImageCutLoop": "Cut Image (LOOP)",
"ImageCropLoop": "Crop Image (LOOP)",
"ImagePasteLoop": "Paste Image (LOOP)"
}