211 lines
9.3 KiB
Python
211 lines
9.3 KiB
Python
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,) }
|