import json import os import re import numpy as np from pathlib import Path from datetime import datetime from PIL import Image from PIL.PngImagePlugin import PngInfo import folder_paths import comfy.samplers class AodhMetadataImageSaver: def __init__(self): self.output_dir = folder_paths.get_output_directory() self.type = "output" self.prefix_append = "" @classmethod def INPUT_TYPES(s): return { "required": { "images": ("IMAGE", ), "filename_prefix": ("STRING", {"default": "ComfyUI"}), }, "optional": { "save_directory": ("STRING", {"default": "%Y-%m-%d"}), # Generation parameters "positive_prompt": ("STRING", {"default": "", "multiline": True}), "negative_prompt": ("STRING", {"default": "", "multiline": True}), "checkpoint_name": ("STRING", {"default": ""}), "vae_name": ("STRING", {"default": ""}), "clip_skip": ("INT", {"default": -2}), "lora_name": ("STRING", {"default": ""}), "lora_strength": ("FLOAT", {"default": 1.0}), # Sampling "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "steps": ("INT", {"default": 20}), "cfg": ("FLOAT", {"default": 8.0}), "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ), "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), # Resolution "width": ("INT", {"default": 512}), "height": ("INT", {"default": 512}), "upscale_factor": ("FLOAT", {"default": 1.0}), "upscaler_name": ("STRING", {"default": ""}), "hires_steps": ("INT", {"default": 10}), "denoise": ("FLOAT", {"default": 0.0}), # Character / App Metadata "character_name": ("STRING", {"default": ""}), "tags": ("STRING", {"default": ""}), "rating": (["safe", "questionable", "explicit"], {"default": "safe"}), # Extension "extension": (["png", "jpg", "webp"], {"default": "png"}), "include_workflow": ("BOOLEAN", {"default": True}), }, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("images",) FUNCTION = "save_images" OUTPUT_NODE = True CATEGORY = "AODH Pack/Image" def save_images(self, images, filename_prefix="ComfyUI", save_directory="%Y-%m-%d", positive_prompt="", negative_prompt="", checkpoint_name="", vae_name="", clip_skip=-2, lora_name="", lora_strength=1.0, seed=0, steps=20, cfg=8.0, sampler_name="euler", scheduler="normal", width=512, height=512, upscale_factor=1.0, upscaler_name="", hires_steps=10, denoise=0.0, character_name="", tags="", rating="safe", extension="png", include_workflow=True, prompt=None, extra_pnginfo=None): now = datetime.now() if save_directory: # Handle date:yyyy-MM-dd format if "date:" in save_directory: def replace_date(match): fmt = match.group(1) # Convert common JS-like date formats to python strftime fmt = fmt.replace("yyyy", "%Y").replace("MM", "%m").replace("dd", "%d") fmt = fmt.replace("HH", "%H").replace("mm", "%M").replace("ss", "%S") return now.strftime(fmt) save_directory = re.sub(r"date:([\w\-\:]+)", replace_date, save_directory) # Also handle direct strftime patterns if any % is present if "%" in save_directory: save_directory = now.strftime(save_directory) filename_prefix = os.path.join(save_directory, filename_prefix) full_output_folder, filename, counter, subfolder, filename_prefix = \ folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) results = list() for image in images: i = 255. * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) metadata = PngInfo() # 1. Construct A1111 Parameters String # Format: Prompt # Negative prompt: ... # Steps: ..., Sampler: ..., CFG scale: ..., Seed: ..., Size: ...x..., Model: ..., ... parameters_text = f"{positive_prompt}\nNegative prompt: {negative_prompt}\n" parameters_text += f"Steps: {steps}, Sampler: {sampler_name}, Schedule: {scheduler}, CFG scale: {cfg}, Seed: {seed}, " parameters_text += f"Size: {width}x{height}, Model: {checkpoint_name}" if clip_skip != -2: parameters_text += f", Clip skip: {clip_skip}" if upscale_factor > 1.0: parameters_text += f", Hires upscale: {upscale_factor}, Hires steps: {hires_steps}, Hires upscaler: {upscaler_name}, Denoising strength: {denoise}" if lora_name: display_lora_name = Path(lora_name).name parameters_text += f", Lora: {display_lora_name}" metadata.add_text("parameters", parameters_text) # 2. Construct ComfyUI Metadata JSON (following spec) comfy_meta = { "_description": "Extended metadata specific to ComfyUI workflow", "workflow_name": "AODH Generator", # Could be dynamic if we parsed extra_pnginfo "workflow_version": "1.0.0", "generation": { "checkpoint": checkpoint_name, "vae": vae_name, "clip_skip": clip_skip, "lora": [{"name": Path(lora_name).name, "strength_model": lora_strength}] if lora_name else [] }, "sampling": { "sampler": sampler_name, "scheduler": scheduler, "steps": steps, "cfg": cfg, "seed": seed, "batch_size": images.shape[0] }, "resolution": { "width": width, "height": height, "upscale_factor": upscale_factor, "upscaler": upscaler_name, "hires_steps": hires_steps, "denoise_strength": denoise }, "prompt_structure": { "positive": {"full": positive_prompt}, "negative": {"full": negative_prompt} }, "workflow": {} } # Embed raw workflow if available if include_workflow: if prompt is not None: comfy_meta["workflow"]["execution"] = prompt if extra_pnginfo is not None and "workflow" in extra_pnginfo: comfy_meta["workflow"]["nodes"] = extra_pnginfo["workflow"] # 3. App Metadata app_meta = { "tags": [t.strip() for t in tags.split(",") if t.strip()], "character": character_name, "rating": rating } # Add to main JSON full_metadata = { "parameters": parameters_text, "source": "comfyui", "comfyui": True, "comfyui_metadata": comfy_meta, "app_metadata": app_meta } # Embed as "comment" or a specific key. # The spec says "A ComfyUI node should embed this metadata in the 'parameters' PNG text chunk." # But usually JSON goes into a separate chunk or we rely on the parser to extract it from 'parameters' if it was formatted there. # However, standard ComfyUI puts workflow in "workflow" and "prompt" chunks. # The spec implies a specific JSON structure. We'll add it as a separate text chunk "comfyui_metadata_json" # OR we can try to append it to parameters, but A1111 parsers might break. # Let's put the full JSON in a "aodh_metadata" chunk for safety, # and ALSO try to adhere to standard Comfy behavior by keeping workflow/prompt chunks. if include_workflow: 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])) # Add our custom full spec JSON metadata.add_text("aodh_metadata", json.dumps(full_metadata, indent=2)) file = f"{filename}_{counter:05}_.{extension}" img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=4) results.append({ "filename": file, "subfolder": subfolder, "type": self.type }) counter += 1 return { "ui": { "images": results }, "result": (images,) }