feat: add AODH Image Saver (Metadata), Lora Selector, Checkpoint Selector, and various node improvements
This commit is contained in:
Aodhan Collins
9
nodes/metadata_saver/__init__.py
Normal file
9
nodes/metadata_saver/__init__.py
Normal file
@@ -0,0 +1,9 @@
|
||||
from .metadata_saver import AodhMetadataImageSaver
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AodhMetadataImageSaver": AodhMetadataImageSaver
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AodhMetadataImageSaver": "AODH Image Saver (Metadata)"
|
||||
}
|
||||
210
nodes/metadata_saver/metadata_saver.py
Normal file
210
nodes/metadata_saver/metadata_saver.py
Normal file
@@ -0,0 +1,210 @@
|
||||
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,) }
|
||||
Reference in New Issue
Block a user