503 lines
15 KiB
Python
503 lines
15 KiB
Python
import os
|
|
import json
|
|
import numpy as np
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
import folder_paths
|
|
from datetime import datetime
|
|
|
|
|
|
SAMPLERS = {
|
|
"euler": "Euler",
|
|
"euler_cfg_pp": "Euler CFG++",
|
|
"euler_ancestral": "Euler a",
|
|
"euler_ancestral_cfg_pp": "Euler a CFG++",
|
|
"heun": "Heun",
|
|
"heunpp2": "HeunPP2",
|
|
"dpm_2": "DPM2",
|
|
"dpm_2_ancestral": "DPM2 a",
|
|
"lms": "LMS",
|
|
"dpm_fast": "DPM fast",
|
|
"dpm_adaptive": "DPM adaptive",
|
|
"dpmpp_2s_ancestral": "DPM++ 2S a",
|
|
"dpmpp_sde": "DPM++ SDE",
|
|
"dpmpp_sde_gpu": "DPM++ SDE",
|
|
"dpmpp_2m": "DPM++ 2M",
|
|
"dpmpp_2m_sde": "DPM++ 2M SDE",
|
|
"dpmpp_2m_sde_gpu": "DPM++ 2M SDE",
|
|
"dpmpp_2m_sde_heun": "DPM++ 2M SDE Heun",
|
|
"dpmpp_2m_sde_heun_gpu": "DPM++ 2M SDE Heun",
|
|
"dpmpp_3m_sde": "DPM++ 3M SDE",
|
|
"dpmpp_3m_sde_gpu": "DPM++ 3M SDE",
|
|
"ipndm": "IPNDM",
|
|
"ipndm_v": "IPNDM_V",
|
|
"deis": "DEIS",
|
|
}
|
|
|
|
SCHEDULERS = {
|
|
"simple": "Simple",
|
|
"sgm_uniform": "SGM Uniform",
|
|
"karras": "Karras",
|
|
"exponential": "Exponential",
|
|
"normal": "Normal",
|
|
"kl_optimal": "KL Optimal",
|
|
"align_your_steps": "Align Your Steps"
|
|
}
|
|
|
|
|
|
def get_node_by_type(prompt, node_type):
|
|
return [v for k, v in prompt.items() if v.get("class_type") == node_type]
|
|
|
|
|
|
def build_a1111_meta(pipe, prompt, w, h, model):
|
|
pipe_meta = pipe.get("meta", {})
|
|
p_text = pipe_meta.get("positive_text", "")
|
|
n_text = pipe_meta.get("negative_text", "")
|
|
seed = pipe.get("seed", 0)
|
|
|
|
if not prompt:
|
|
meta = f"{p_text}\nNegative prompt: {n_text}\n"
|
|
# width and height here should be divided by the upscale value
|
|
meta += f"Seed: {seed}, Size: {w}x{h}, Model: {model}"
|
|
return meta
|
|
|
|
lora_meta = ""
|
|
stacker = get_node_by_type(prompt, "Lora Stacker (LoraManager)")
|
|
if stacker:
|
|
v_stack = stacker[0].get("inputs", {})
|
|
loras = v_stack.get("lora_stack", [])
|
|
for lora in loras:
|
|
if isinstance(lora, list) and len(lora) >= 3:
|
|
name, strength = lora[0], lora[1]
|
|
name = name.split("/")[-1].rsplit(".", 1)[0]
|
|
lora_meta += f"<lora:{name}:{strength}>, "
|
|
|
|
steps, sampler, sched, cfg = "20", "euler_ancestral", "karras", "7.0"
|
|
clip_skip = 2
|
|
model_hash = "unknown"
|
|
|
|
loader = get_node_by_type(prompt, "LoaderFullPipe")
|
|
if loader:
|
|
v_loader = loader[0].get("inputs", {})
|
|
clip_skip = abs(v_loader.get("clip_skip", clip_skip))
|
|
ckpt_name = v_loader.get("ckpt_name", "")
|
|
if ckpt_name:
|
|
import hashlib
|
|
model_hash = hashlib.sha256(ckpt_name.encode()).hexdigest()[:10]
|
|
|
|
base_nodes = get_node_by_type(prompt, "BaseNode")
|
|
if base_nodes:
|
|
v = base_nodes[0].get("inputs", {})
|
|
sampler = v.get("sampler_name", sampler)
|
|
sampler = SAMPLERS.get(sampler, sampler)
|
|
sched = v.get("scheduler", sched)
|
|
sched = SCHEDULERS.get(sched, sched)
|
|
steps = v.get("steps", steps)
|
|
cfg = v.get("cfg", cfg)
|
|
|
|
meta = f"{lora_meta}{p_text}\nNegative prompt: {n_text}\n"
|
|
meta += (f"Steps: {steps}, Sampler: {sampler}, Schedule type: {sched}, "
|
|
f"CFG scale: {cfg}, Seed: {seed}, Size: {w}x{h}, "
|
|
f"Model hash: {model_hash}, Model: {model}, "
|
|
f"Clip skip: {clip_skip}")
|
|
|
|
upscale = get_node_by_type(prompt, "UpscaleNode")
|
|
if upscale:
|
|
v = upscale[0].get("inputs", {})
|
|
upscale_model = v.get('upscale_model').split('/')[-1]
|
|
meta += (f", Hires upscale: {v.get('scale_by')}, "
|
|
f"Hires steps: {v.get('steps')}, "
|
|
f"Denoising strength: {v.get('denoise')}, "
|
|
f"Hires upscaler: {upscale_model}")
|
|
|
|
detailers = get_node_by_type(prompt, "DetailerPipeNode")
|
|
for i, det in enumerate(detailers):
|
|
v = det.get("inputs", {})
|
|
bbox_model = v.get('bbox_model').split('/')[-1]
|
|
sfx = " 2nd" if i == 1 else ""
|
|
meta += (f", ADetailer model{sfx}: {bbox_model}, "
|
|
f"ADetailer confidence{sfx}: {v.get('threshold')}, "
|
|
f"ADetailer denoising strength{sfx}: {v.get('denoise')}")
|
|
|
|
nestedDetailer_nodes = get_node_by_type(prompt, "NestedDetailerPipeNode")
|
|
if nestedDetailer_nodes:
|
|
v = nestedDetailer_nodes[0].get("inputs", {})
|
|
bbox_model = []
|
|
bbox_model.append(v.get('face_model').split('/')[-1])
|
|
bbox_model.append(v.get('eyes_pair_model').split('/')[-1])
|
|
denoise = []
|
|
denoise.append(v.get('face_denoise'))
|
|
denoise.append(v.get('eye_denoise'))
|
|
for i in range(0, 2):
|
|
sfx = " 2nd" if i == 1 else ""
|
|
meta += (f", ADetailer model{sfx}: {bbox_model[i]}, "
|
|
f"ADetailer confidence{sfx}: {v.get('threshold')}, "
|
|
f"ADetailer denoising strength{sfx}: {denoise[i]}")
|
|
|
|
return meta
|
|
|
|
|
|
def replace_variables(text, full_pipe, width, height, timestamp):
|
|
"""
|
|
Replace variables in text with actual values.
|
|
|
|
Args:
|
|
text: String containing variables like %date, %time, etc.
|
|
full_pipe: Dictionary containing pipe data
|
|
width: Image width
|
|
height: Image height
|
|
timestamp: datetime object for consistent timestamps
|
|
|
|
Returns:
|
|
String with variables replaced
|
|
"""
|
|
pipe_meta = full_pipe.get('meta', {})
|
|
ckpt_name = pipe_meta.get('ckpt_name', 'unknown')
|
|
if '/' in ckpt_name:
|
|
ckpt_name = ckpt_name.split('/')[-1]
|
|
|
|
replacements = {
|
|
'%date': timestamp.strftime('%Y-%m-%d'),
|
|
'%time': timestamp.strftime('%Y-%m-%d-%H%M%S'),
|
|
'%model': ckpt_name,
|
|
'%width': str(width),
|
|
'%height': str(height),
|
|
'%seed': str(full_pipe.get('seed', 0)),
|
|
}
|
|
|
|
result = text
|
|
for var, value in replacements.items():
|
|
result = result.replace(var, value)
|
|
|
|
return result
|
|
|
|
|
|
class SaveFullPipe:
|
|
"""
|
|
Save image from FULL_PIPE with metadata and display preview.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
tooltip = (
|
|
"Available variables: %date, %time, %model, "
|
|
"%width, %height, %seed"
|
|
)
|
|
|
|
return {
|
|
"required": {
|
|
"full_pipe": ("FULL_PIPE",),
|
|
"filename_prefix": (
|
|
"STRING",
|
|
{
|
|
"default": "%time_%seed",
|
|
"tooltip": tooltip
|
|
}
|
|
),
|
|
"path": (
|
|
"STRING",
|
|
{
|
|
"default": "%date",
|
|
"tooltip": tooltip
|
|
}
|
|
),
|
|
"extension": (
|
|
["png", "jpg", "jpeg", "webp"], {"default": "png"}),
|
|
"a1111_metadata": ("BOOLEAN", {"default": True}),
|
|
"comfyui_workflow": ("BOOLEAN", {"default": True}),
|
|
"show_preview": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",)
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO"
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "save_images"
|
|
CATEGORY = "mudknight/save"
|
|
DESCRIPTION = (
|
|
"Save images from full pipe or optional image input."
|
|
"Optionally saves ComfyUI and A1111 metadata."
|
|
"ComfyUI samplers are mapped to their A1111 equivalents."
|
|
)
|
|
|
|
def save_images(
|
|
self,
|
|
full_pipe,
|
|
filename_prefix,
|
|
path,
|
|
extension,
|
|
a1111_metadata,
|
|
comfyui_workflow,
|
|
show_preview,
|
|
image=None,
|
|
prompt=None,
|
|
extra_pnginfo=None
|
|
):
|
|
# Extract data from pipe
|
|
if image is None:
|
|
image = full_pipe.get("image")
|
|
if image is None:
|
|
return {"ui": {"text": ["No image in pipe"]}}
|
|
|
|
ckpt_name = (
|
|
full_pipe.get("meta", {})
|
|
.get("ckpt_name", "unknown")
|
|
.split('/')[-1]
|
|
)
|
|
|
|
# Get image dimensions
|
|
batch_size = image.shape[0]
|
|
height = image.shape[1]
|
|
width = image.shape[2]
|
|
|
|
# Create consistent timestamp for this batch
|
|
timestamp = datetime.now()
|
|
|
|
# Replace variables in path
|
|
path = replace_variables(
|
|
path, full_pipe, width, height, timestamp
|
|
)
|
|
|
|
# Create output directory
|
|
output_dir = folder_paths.get_output_directory()
|
|
if path:
|
|
output_dir = os.path.join(output_dir, path)
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
results = []
|
|
for i, img_tensor in enumerate(image):
|
|
# Convert tensor to PIL Image
|
|
img_array = (
|
|
255.0 * img_tensor.cpu().numpy()
|
|
).astype(np.uint8)
|
|
pil_image = Image.fromarray(img_array)
|
|
|
|
# Replace variables in filename
|
|
filename = replace_variables(
|
|
filename_prefix, full_pipe, width, height, timestamp
|
|
)
|
|
if batch_size > 1:
|
|
filename = f"{filename}_{i:04d}"
|
|
filename = f"{filename}.{extension}"
|
|
|
|
# Prepare metadata
|
|
metadata = PngInfo()
|
|
|
|
# Create A1111-style parameters text
|
|
if a1111_metadata:
|
|
meta_text = build_a1111_meta(full_pipe, prompt, width,
|
|
height, ckpt_name)
|
|
metadata.add_text("parameters", meta_text)
|
|
|
|
# Embed workflow if requested (PNG only)
|
|
# IMPORTANT: workflow must come before prompt
|
|
if comfyui_workflow and extension == "png":
|
|
if extra_pnginfo is not None:
|
|
for key, value in extra_pnginfo.items():
|
|
metadata.add_text(
|
|
key,
|
|
json.dumps(value)
|
|
)
|
|
if prompt is not None:
|
|
metadata.add_text(
|
|
"prompt",
|
|
json.dumps(prompt)
|
|
)
|
|
|
|
# Save image
|
|
filepath = os.path.join(output_dir, filename)
|
|
save_kwargs = {}
|
|
|
|
if extension == "png":
|
|
save_kwargs["pnginfo"] = metadata
|
|
save_kwargs["compress_level"] = 4
|
|
|
|
elif extension in {"jpg", "jpeg"}:
|
|
save_kwargs["quality"] = 95
|
|
|
|
elif extension == "webp":
|
|
save_kwargs["quality"] = 95
|
|
save_kwargs["method"] = 6
|
|
|
|
pil_image.save(filepath, **save_kwargs)
|
|
|
|
results.append({
|
|
"filename": filename,
|
|
"subfolder": path,
|
|
"type": "output"
|
|
})
|
|
|
|
# Return preview images
|
|
return {
|
|
"ui": {"images": results if show_preview else []},
|
|
"result": (image,)
|
|
}
|
|
|
|
|
|
class SaveGif:
|
|
"""
|
|
Save a batch of images from FULL_PIPE or image input as a
|
|
looping GIF.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
tooltip = (
|
|
"Available variables: %date, %time, %model, "
|
|
"%width, %height, %seed"
|
|
)
|
|
|
|
return {
|
|
"required": {
|
|
"full_pipe": ("FULL_PIPE",),
|
|
"filename_prefix": (
|
|
"STRING",
|
|
{
|
|
"default": "%time_%seed",
|
|
"tooltip": tooltip
|
|
}
|
|
),
|
|
"path": (
|
|
"STRING",
|
|
{
|
|
"default": "%date",
|
|
"tooltip": tooltip
|
|
}
|
|
),
|
|
"frame_duration": ("FLOAT", {
|
|
"default": 0.15,
|
|
"min": 0.01,
|
|
"max": 10.0,
|
|
"step": 0.01,
|
|
}),
|
|
"scale": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.1,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
}),
|
|
"gif_quality": ("INT", {
|
|
"default": 128,
|
|
"min": 1,
|
|
"max": 256,
|
|
}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "save_gif"
|
|
CATEGORY = "mudknight/save"
|
|
DESCRIPTION = "Save a batch of images as a looping GIF."
|
|
|
|
def save_gif(
|
|
self,
|
|
full_pipe,
|
|
filename_prefix,
|
|
path,
|
|
frame_duration,
|
|
scale,
|
|
gif_quality,
|
|
image=None,
|
|
):
|
|
# Fall back to pipe image if no input provided
|
|
if image is None:
|
|
image = full_pipe.get("image")
|
|
if image is None:
|
|
return {"ui": {"text": ["No image in pipe"]}}
|
|
|
|
# Get dimensions from tensor (B, H, W, C)
|
|
height = image.shape[1]
|
|
width = image.shape[2]
|
|
|
|
timestamp = datetime.now()
|
|
|
|
# Resolve output directory
|
|
resolved_path = replace_variables(
|
|
path, full_pipe, width, height, timestamp
|
|
)
|
|
output_dir = folder_paths.get_output_directory()
|
|
if resolved_path:
|
|
output_dir = os.path.join(output_dir, resolved_path)
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# Resolve filename
|
|
filename = replace_variables(
|
|
filename_prefix, full_pipe, width, height, timestamp
|
|
)
|
|
filepath = os.path.join(output_dir, f"{filename}.gif")
|
|
|
|
# Convert each frame tensor to a PIL image, applying scale
|
|
target_w = max(1, int(width * scale))
|
|
target_h = max(1, int(height * scale))
|
|
rgb_frames = []
|
|
for img_tensor in image:
|
|
arr = (255.0 * img_tensor.cpu().numpy()).astype(np.uint8)
|
|
pil_frame = Image.fromarray(arr)
|
|
if scale != 1.0:
|
|
pil_frame = pil_frame.resize(
|
|
(target_w, target_h), Image.LANCZOS
|
|
)
|
|
rgb_frames.append(pil_frame)
|
|
|
|
if not rgb_frames:
|
|
return {"ui": {"text": ["No frames to save"]}}
|
|
|
|
# Build a shared palette from all frames combined so that
|
|
# colours stay consistent across the GIF animation
|
|
combined_w = target_w * len(rgb_frames)
|
|
combined = Image.new("RGB", (combined_w, target_h))
|
|
for i, f in enumerate(rgb_frames):
|
|
combined.paste(f, (i * target_w, 0))
|
|
palette_image = combined.quantize(colors=gif_quality)
|
|
|
|
# Apply the shared palette to each frame individually
|
|
frames = [
|
|
f.quantize(palette=palette_image) for f in rgb_frames
|
|
]
|
|
|
|
# frame_duration is in seconds; Pillow uses milliseconds
|
|
duration_ms = int(frame_duration * 1000)
|
|
|
|
# Save as looping GIF (loop=0 means infinite)
|
|
frames[0].save(
|
|
filepath,
|
|
format="GIF",
|
|
save_all=True,
|
|
append_images=frames[1:],
|
|
loop=0,
|
|
duration=duration_ms,
|
|
)
|
|
|
|
return {
|
|
"ui": {
|
|
"images": [{
|
|
"filename": f"{filename}.gif",
|
|
"subfolder": resolved_path,
|
|
"type": "output",
|
|
}]
|
|
}
|
|
}
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SaveFullPipe": SaveFullPipe,
|
|
"SaveGif": SaveGif,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SaveFullPipe": "Save (Full Pipe)",
|
|
"SaveGif": "Save GIF",
|
|
}
|