381 lines
13 KiB
Python
381 lines
13 KiB
Python
import json
|
|
import os
|
|
import random
|
|
import re
|
|
import time
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
import folder_paths
|
|
|
|
|
|
_MEDIA_EXT_RE = re.compile(
|
|
r"\.(png|jpe?g|webp|gif|bmp|tiff?|avif|mp4|mov|webm|mkv|m4v)$", re.IGNORECASE
|
|
)
|
|
_DATE_TOKEN_RE = re.compile(r"%date:([^%]+)%")
|
|
|
|
|
|
def _clean_input_token_value(value):
|
|
cleaned = ""
|
|
if value is not None:
|
|
cleaned = value if isinstance(value, str) else str(value)
|
|
cleaned = _MEDIA_EXT_RE.sub("", cleaned.strip())
|
|
cleaned = cleaned.replace("\\", "_").replace("/", "_")
|
|
return cleaned
|
|
|
|
|
|
def _expand_date_tokens(value):
|
|
if not isinstance(value, str) or "%date:" not in value:
|
|
return value
|
|
|
|
now = time.localtime()
|
|
|
|
def pad(number, width):
|
|
return str(number).zfill(width)
|
|
|
|
def repl(match):
|
|
fmt = match.group(1)
|
|
|
|
def swap(token_match):
|
|
token = token_match.group(0)
|
|
if token == "yyyy":
|
|
return pad(now.tm_year, 4)
|
|
if token == "yy":
|
|
return str(now.tm_year)[-2:]
|
|
if token == "MM":
|
|
return pad(now.tm_mon, 2)
|
|
if token == "M":
|
|
return str(now.tm_mon)
|
|
if token == "dd":
|
|
return pad(now.tm_mday, 2)
|
|
if token == "d":
|
|
return str(now.tm_mday)
|
|
if token in ("hh", "HH"):
|
|
return pad(now.tm_hour, 2)
|
|
if token in ("h", "H"):
|
|
return str(now.tm_hour)
|
|
if token == "mm":
|
|
return pad(now.tm_min, 2)
|
|
if token == "m":
|
|
return str(now.tm_min)
|
|
if token == "ss":
|
|
return pad(now.tm_sec, 2)
|
|
if token == "s":
|
|
return str(now.tm_sec)
|
|
return token
|
|
|
|
return re.sub(r"yyyy|yy|MM|M|dd|d|hh|h|HH|H|mm|m|ss|s", swap, fmt)
|
|
|
|
return _DATE_TOKEN_RE.sub(repl, value)
|
|
|
|
|
|
def _expand_native_tokens(value):
|
|
if not isinstance(value, str) or "%" not in value:
|
|
return value
|
|
|
|
now = time.localtime()
|
|
replacements = (
|
|
("%year%", f"{now.tm_year:04}"),
|
|
("%month%", f"{now.tm_mon:02}"),
|
|
("%day%", f"{now.tm_mday:02}"),
|
|
("%hour%", f"{now.tm_hour:02}"),
|
|
("%minute%", f"{now.tm_min:02}"),
|
|
("%second%", f"{now.tm_sec:02}"),
|
|
)
|
|
for token, replacement in replacements:
|
|
value = value.replace(token, replacement)
|
|
return value
|
|
|
|
|
|
def _safe_pattern(value):
|
|
value = str(value or "").replace("\\", "/")
|
|
value = re.sub(r'[<>:"|?*]', "_", value)
|
|
value = re.sub(r"/{2,}", "/", value).strip(" /.") or "image_%counter%"
|
|
return value
|
|
|
|
|
|
def _next_counter(directory, filename_template):
|
|
os.makedirs(directory, exist_ok=True)
|
|
if "%counter%" not in filename_template:
|
|
return 1
|
|
|
|
parts = filename_template.split("%counter%")
|
|
highest = 0
|
|
for entry in os.listdir(directory):
|
|
if not entry.startswith(parts[0]) or not entry.endswith(parts[-1]):
|
|
continue
|
|
middle = entry[len(parts[0]):]
|
|
if parts[-1]:
|
|
middle = middle[: -len(parts[-1])]
|
|
if middle.isdigit():
|
|
highest = max(highest, int(middle))
|
|
return highest + 1
|
|
|
|
|
|
def _build_pnginfo(prompt=None, extra_pnginfo=None):
|
|
try:
|
|
pnginfo = Image.PngImagePlugin.PngInfo()
|
|
except AttributeError:
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
pnginfo = PngInfo()
|
|
if prompt is not None:
|
|
pnginfo.add_text("prompt", json.dumps(prompt))
|
|
if isinstance(extra_pnginfo, dict):
|
|
for key, value in extra_pnginfo.items():
|
|
pnginfo.add_text(str(key), json.dumps(value))
|
|
return pnginfo
|
|
|
|
|
|
def _tensor_image_to_pil_image(tensor):
|
|
image_tensor = tensor[0]
|
|
if hasattr(image_tensor, "mul") and hasattr(image_tensor, "clamp"):
|
|
image_array = image_tensor.mul(255).clamp(0, 255)
|
|
if hasattr(image_array, "byte"):
|
|
image_array = image_array.byte()
|
|
image_array = image_array.cpu().numpy()
|
|
return Image.fromarray(image_array)
|
|
|
|
image_array = 255.0 * image_tensor.cpu().numpy()
|
|
return Image.fromarray(np.clip(image_array, 0, 255).astype(np.uint8))
|
|
|
|
|
|
class DumasImageCompareNode:
|
|
DESCRIPTION = (
|
|
"Dumas Image Compare shows the difference between two images directly on "
|
|
"the node. Connect one or two IMAGE inputs to compare before/after "
|
|
"results, model variants, or processing stages without breaking a "
|
|
"workflow when one branch is bypassed."
|
|
)
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("new image",)
|
|
FUNCTION = "compare_images"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "Dumas/Image"
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_temp_directory()
|
|
self.type = "temp"
|
|
self.prefix_append = "_dumascmp_" + "".join(
|
|
random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(5)
|
|
)
|
|
self.compress_level = 4
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"optional": {
|
|
"image1": (
|
|
"IMAGE",
|
|
{
|
|
"tooltip": (
|
|
"First image to compare. Optional so muted or bypassed "
|
|
"branches do not trigger a missing-input error."
|
|
)
|
|
},
|
|
),
|
|
"image2": (
|
|
"IMAGE",
|
|
{
|
|
"tooltip": (
|
|
"Second image to compare. Optional so the node can still "
|
|
"display a single available image."
|
|
)
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
def compare_images(self, image1=None, image2=None):
|
|
present = []
|
|
if image1 is not None:
|
|
present.append((1, image1))
|
|
if image2 is not None:
|
|
present.append((2, image2))
|
|
results = []
|
|
|
|
if present:
|
|
first_tensor = present[0][1]
|
|
prefix = "dumas_compare" + self.prefix_append
|
|
first_image = first_tensor[0]
|
|
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
|
|
prefix,
|
|
self.output_dir,
|
|
first_image.shape[1],
|
|
first_image.shape[0],
|
|
)
|
|
join_path = os.path.join
|
|
for slot, tensor in present:
|
|
image = _tensor_image_to_pil_image(tensor)
|
|
file_name = f"{filename}_{counter:05}_.png"
|
|
image.save(
|
|
join_path(full_output_folder, file_name),
|
|
compress_level=self.compress_level,
|
|
)
|
|
results.append(
|
|
{
|
|
"filename": file_name,
|
|
"subfolder": subfolder,
|
|
"type": self.type,
|
|
"slot": slot,
|
|
}
|
|
)
|
|
counter += 1
|
|
|
|
new_image = image2 if image2 is not None else image1
|
|
return {"ui": {"images": results}, "result": (new_image,)}
|
|
|
|
|
|
class DumasSaveImageNode:
|
|
DESCRIPTION = (
|
|
"Dumas Save Image writes images to any folder, with filename tokens "
|
|
"such as %input%, %input2%, %date:yyyy-MM-dd%, %counter%, %width%, "
|
|
"%height%, and %batch_num%."
|
|
)
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save_images"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "Dumas/Image"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE", {"tooltip": "Image batch to save."}),
|
|
"folder": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": False,
|
|
"tooltip": (
|
|
"Destination folder. Leave empty to save in ComfyUI's "
|
|
"output directory."
|
|
),
|
|
},
|
|
),
|
|
"pattern": (
|
|
"STRING",
|
|
{
|
|
"default": "image_%date:yyyy-MM-dd%_%counter%",
|
|
"multiline": False,
|
|
"tooltip": (
|
|
"Filename pattern with optional subfolders. Tokens: "
|
|
"%input%, %input2%, %date:yyyy-MM-dd%, %counter%, "
|
|
"%width%, %height%, %batch_num%."
|
|
),
|
|
},
|
|
),
|
|
"format": (["png", "jpg"], {"default": "png"}),
|
|
"quality": (
|
|
"INT",
|
|
{"default": 100, "min": 1, "max": 100, "step": 1},
|
|
),
|
|
"embed_workflow": ("BOOLEAN", {"default": True}),
|
|
"save_on_run": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"name": (
|
|
"STRING",
|
|
{
|
|
"forceInput": True,
|
|
"tooltip": "Optional text inserted by the %input% token.",
|
|
},
|
|
),
|
|
"name_2": (
|
|
"STRING",
|
|
{
|
|
"forceInput": True,
|
|
"tooltip": "Optional text inserted by the %input2% token.",
|
|
},
|
|
),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **_kwargs):
|
|
return float("nan")
|
|
|
|
def save_images(
|
|
self,
|
|
images,
|
|
folder,
|
|
pattern,
|
|
format,
|
|
quality,
|
|
embed_workflow,
|
|
save_on_run,
|
|
name=None,
|
|
name_2=None,
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
if not save_on_run:
|
|
return {"ui": {"images": []}}
|
|
|
|
width = int(images.shape[2])
|
|
height = int(images.shape[1])
|
|
output_dir = folder_paths.get_output_directory()
|
|
target_dir = os.path.abspath(folder.strip()) if str(folder or "").strip() else output_dir
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
|
|
resolved_pattern = str(pattern or "image_%date:yyyy-MM-dd%_%counter%")
|
|
resolved_pattern = resolved_pattern.replace("%input%", _clean_input_token_value(name))
|
|
resolved_pattern = resolved_pattern.replace("%input2%", _clean_input_token_value(name_2))
|
|
resolved_pattern = _expand_date_tokens(resolved_pattern)
|
|
resolved_pattern = _expand_native_tokens(resolved_pattern)
|
|
resolved_pattern = resolved_pattern.replace("%width%", str(width))
|
|
resolved_pattern = resolved_pattern.replace("%height%", str(height))
|
|
resolved_pattern = _safe_pattern(resolved_pattern)
|
|
|
|
extension = ".jpg" if format == "jpg" else ".png"
|
|
quality = max(1, min(100, int(quality)))
|
|
ui_images = []
|
|
|
|
for batch_index in range(images.shape[0]):
|
|
frame_pattern = resolved_pattern.replace("%batch_num%", str(batch_index))
|
|
frame_parts = [part for part in frame_pattern.split("/") if part]
|
|
sub_dirs = frame_parts[:-1]
|
|
filename_template = (frame_parts[-1] if frame_parts else "image_%counter%") + extension
|
|
frame_dir = os.path.join(target_dir, *sub_dirs)
|
|
counter = _next_counter(frame_dir, filename_template)
|
|
filename = filename_template.replace("%counter%", str(counter).zfill(3))
|
|
image = _tensor_image_to_pil_image(images[batch_index : batch_index + 1])
|
|
full_path = os.path.join(frame_dir, filename)
|
|
|
|
if format == "jpg":
|
|
image = image.convert("RGB")
|
|
image.save(full_path, "JPEG", quality=quality)
|
|
else:
|
|
pnginfo = None
|
|
if embed_workflow:
|
|
pnginfo = _build_pnginfo(prompt=prompt, extra_pnginfo=extra_pnginfo)
|
|
image.save(full_path, "PNG", pnginfo=pnginfo)
|
|
|
|
if os.path.commonpath([output_dir, full_path]) == output_dir:
|
|
subfolder = os.path.relpath(frame_dir, output_dir)
|
|
ui_images.append(
|
|
{
|
|
"filename": filename,
|
|
"subfolder": "" if subfolder == "." else subfolder.replace("\\", "/"),
|
|
"type": "output",
|
|
}
|
|
)
|
|
|
|
return {"ui": {"images": ui_images}}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DumasImageCompare": DumasImageCompareNode,
|
|
"DumasSaveImage": DumasSaveImageNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DumasImageCompare": "Dumas Image Compare",
|
|
"DumasSaveImage": "Save Image Dumas",
|
|
}
|