Add Save Image Dumas node
This commit is contained in:
@@ -1,5 +1,8 @@
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
@@ -7,6 +10,124 @@ 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"):
|
||||
@@ -106,10 +227,154 @@ class DumasImageCompareNode:
|
||||
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",
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user