1lch2
66dbf20ef0
Remove the seed input parameter and automatically resolve the seed value from the first KSampler node in the workflow prompt. Supports following link chains for connected seed inputs (Seed (rgthree), ImpactInt, etc).
247 lines
8.9 KiB
Python
247 lines
8.9 KiB
Python
import os
|
|
import json
|
|
import re
|
|
import numpy as np
|
|
from datetime import datetime
|
|
|
|
import folder_paths
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
from comfy.cli_args import args
|
|
|
|
|
|
def _is_link(obj):
|
|
"""Check if a value is a ComfyUI link: [node_id_str, output_slot_index]."""
|
|
return (isinstance(obj, list) and len(obj) == 2
|
|
and isinstance(obj[0], str)
|
|
and isinstance(obj[1], (int, float)))
|
|
|
|
|
|
def _resolve_seed_from_prompt(prompt):
|
|
"""Find the seed value from the first KSampler node in the workflow prompt.
|
|
|
|
Follows link chains to resolve connected seed inputs from common seed nodes
|
|
(Seed (rgthree), ImpactInt, etc).
|
|
"""
|
|
if not prompt:
|
|
return 0
|
|
|
|
# Find first KSampler node
|
|
ksampler = None
|
|
for node_info in prompt.values():
|
|
if node_info.get("class_type") in ("KSampler", "KSamplerAdvanced"):
|
|
ksampler = node_info
|
|
break
|
|
if not ksampler:
|
|
return 0
|
|
|
|
seed_input = ksampler.get("inputs", {}).get("seed")
|
|
return _resolve_value(prompt, seed_input)
|
|
|
|
|
|
def _resolve_value(prompt, value, visited=None):
|
|
"""Resolve a prompt input value, following links recursively.
|
|
|
|
Handles:
|
|
- Direct int values
|
|
- Links to Seed (rgthree) -> inputs.seed
|
|
- Links to ImpactInt -> inputs.value
|
|
- Links to nodes with 'seed' or 'value' int inputs
|
|
"""
|
|
if visited is None:
|
|
visited = set()
|
|
|
|
if isinstance(value, int):
|
|
return value
|
|
|
|
if not _is_link(value):
|
|
return 0
|
|
|
|
node_id = value[0]
|
|
if node_id in visited:
|
|
return 0
|
|
visited.add(node_id)
|
|
|
|
node_info = prompt.get(node_id)
|
|
if not node_info:
|
|
return 0
|
|
|
|
inputs = node_info.get("inputs", {})
|
|
class_type = node_info.get("class_type", "")
|
|
|
|
# Known seed source nodes
|
|
if class_type == "Seed (rgthree)":
|
|
return _resolve_value(prompt, inputs.get("seed"), visited)
|
|
if class_type == "ImpactInt":
|
|
return _resolve_value(prompt, inputs.get("value"), visited)
|
|
|
|
# Generic: try 'seed' then 'value' inputs
|
|
for key in ("seed", "value"):
|
|
if key in inputs:
|
|
val = inputs[key]
|
|
if isinstance(val, int):
|
|
return val
|
|
if _is_link(val):
|
|
return _resolve_value(prompt, val, visited)
|
|
|
|
return 0
|
|
|
|
|
|
def _create_webp_metadata(img, prompt, extra_pnginfo):
|
|
"""Creates EXIF metadata for WebP images. Reference: WAS_Node_Suite.py:7452-7464"""
|
|
exif_data = img.getexif()
|
|
if prompt is not None:
|
|
exif_data[0x010f] = "Prompt:" + json.dumps(prompt) # EXIF 0x010f = Make
|
|
if extra_pnginfo is not None:
|
|
workflow_metadata = ''
|
|
for key, value in extra_pnginfo.items():
|
|
workflow_metadata += json.dumps(value)
|
|
exif_data[0x010e] = "Workflow:" + workflow_metadata # EXIF 0x010e = ImageDescription
|
|
return exif_data.tobytes()
|
|
|
|
|
|
class SaveImageAsSDWebUI:
|
|
"""Save images in SD WebUI style with date-based subfolders and workflow metadata.
|
|
|
|
Seed is automatically resolved from the first KSampler node in the workflow.
|
|
Fallback to 0 if no KSampler is found.
|
|
|
|
Prefix rules (output_dir = ComfyUI/output):
|
|
"txt2img/" -> output/txt2img/{date}/00001-seed.ext
|
|
"abc" -> output/{date}/abc00001-seed.ext
|
|
"imgEdit/qwen" -> output/imgEdit/{date}/qwen00001-seed.ext
|
|
"" -> output/{date}/00001-seed.ext
|
|
"abc/def/" -> output/abc/def/{date}/00001-seed.ext
|
|
"abc/def/114514" -> output/abc/def/{date}/11451400001-seed.ext
|
|
|
|
{date} = yyyy-mm-dd, always the direct parent folder of the file.
|
|
Slash at end = directory only, no filename prefix.
|
|
No slash = filename prefix only, no subdirectory.
|
|
Middle slash(s) = nested directories above date folder.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output"
|
|
self.compress_level = 4
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE", {"tooltip": "The images to save."}),
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": (
|
|
"Filename prefix and/or output subdirectory.\n"
|
|
"A date folder (yyyy-mm-dd) is always created as the direct parent.\n"
|
|
"Seed is auto-resolved from the first KSampler node (fallback: 0).\n\n"
|
|
"Examples (output dir = ComfyUI/output):\n"
|
|
" 'txt2img/' -> output/txt2img/{date}/00001-seed.ext\n"
|
|
" 'abc' -> output/{date}/abc00001-seed.ext\n"
|
|
" 'imgEdit/qwen' -> output/imgEdit/{date}/qwen00001-seed.ext\n"
|
|
" '' -> output/{date}/00001-seed.ext\n"
|
|
" 'abc/def/' -> output/abc/def/{date}/00001-seed.ext\n"
|
|
" 'abc/def/114514' -> output/abc/def/{date}/11451400001-seed.ext"
|
|
)}),
|
|
"image_format": (["png", "jpg", "webp"], {"default": "png", "tooltip": "Image save format."}),
|
|
"quality": ("INT", {"default": 100, "min": 0, "max": 100, "tooltip": "Save quality for JPG/WebP formats (0-100). Ignored for PNG."}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save_images"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "image"
|
|
|
|
def save_images(self, images, filename_prefix="ComfyUI", image_format="png", quality=100, prompt=None, extra_pnginfo=None):
|
|
seed = _resolve_seed_from_prompt(prompt)
|
|
# Split prefix: "dir/name" -> dir="dir", name="name"
|
|
# "dir/" -> dir="dir", name=""
|
|
# "name" -> dir="", name="name"
|
|
prefix = filename_prefix
|
|
if '/' in prefix or '\\' in prefix:
|
|
prefix = prefix.replace('\\', '/')
|
|
last_sep = prefix.rfind('/')
|
|
folder_parts = prefix[:last_sep].split('/')
|
|
name_part = prefix[last_sep + 1:]
|
|
else:
|
|
folder_parts = []
|
|
name_part = prefix
|
|
|
|
date_str = datetime.now().strftime("%Y-%m-%d")
|
|
output_folder = os.path.join(self.output_dir, *folder_parts, date_str)
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
|
|
# Safety check: prevent path traversal
|
|
if os.path.commonpath((self.output_dir, os.path.abspath(output_folder))) != self.output_dir:
|
|
raise Exception("Saving outside the output folder is not allowed.")
|
|
|
|
# Scan for max counter
|
|
if name_part:
|
|
pattern = re.compile(r'^' + re.escape(name_part) + r'(\d{5})-\d+\.\w+$')
|
|
else:
|
|
pattern = re.compile(r'^(\d{5})-\d+\.\w+$')
|
|
max_counter = 0
|
|
for f in os.listdir(output_folder):
|
|
m = pattern.match(f)
|
|
if m:
|
|
max_counter = max(max_counter, int(m.group(1)))
|
|
counter = max_counter + 1
|
|
|
|
subfolder = os.path.join(*folder_parts, date_str) if folder_parts else date_str
|
|
|
|
results = []
|
|
for (batch_number, image) in enumerate(images):
|
|
i = 255. * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
|
|
metadata = None
|
|
if not args.disable_metadata:
|
|
if image_format == "png":
|
|
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]))
|
|
elif image_format == "webp":
|
|
metadata = _create_webp_metadata(img, prompt, extra_pnginfo)
|
|
|
|
file = f"{name_part}{counter:05d}-{seed}.{image_format}"
|
|
filepath = os.path.join(output_folder, file)
|
|
|
|
save_kwargs = {}
|
|
if image_format == "png":
|
|
save_kwargs["pnginfo"] = metadata
|
|
save_kwargs["compress_level"] = self.compress_level
|
|
elif image_format == "webp":
|
|
save_kwargs["exif"] = metadata
|
|
save_kwargs["quality"] = quality
|
|
if img.mode == "RGBA":
|
|
img = img.convert("RGB")
|
|
elif image_format == "jpg":
|
|
save_kwargs["quality"] = quality
|
|
if img.mode == "RGBA":
|
|
img = img.convert("RGB")
|
|
|
|
img.save(filepath, **save_kwargs)
|
|
results.append({
|
|
"filename": file,
|
|
"subfolder": subfolder,
|
|
"type": self.type,
|
|
})
|
|
counter += 1
|
|
|
|
return {"ui": {"images": results}}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SaveImageAsSDWebUI": SaveImageAsSDWebUI,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SaveImageAsSDWebUI": "Save Image (SD WebUI Style)",
|
|
}
|