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)", }