452 lines
18 KiB
Python
452 lines
18 KiB
Python
# IMGNR-Utils/preview_image_base64
|
|
# Fixes: Transparant (RGBA) images (+ mask in/out)
|
|
# Fixes: Widget background now same as template node color
|
|
# Fixes: Zero size widget on spawn due to no image (using placeholder)
|
|
# New: Adhoc Save Node
|
|
# Updated: Renamed Sequence to Counter & Added Toggle
|
|
# Updated: Added Metadata Diff & Comparison Payload
|
|
# Fixed: JSON Serialization error (NaN)
|
|
# Fixed: Preview Node now correctly passes metadata for diffing
|
|
# Updated: Renamed Preview Ad-hoc Save - LastGen Compare (IMGNR)
|
|
# FINAL: Split into 3 distinct nodes (No Save, Ad-Hoc Save, Compare)
|
|
# FIXED: Workflow not embedded when triggering ad-hoc manual save outside execution loop
|
|
|
|
import os
|
|
import re
|
|
from datetime import datetime
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
import base64
|
|
import io
|
|
import json
|
|
import math
|
|
import sys
|
|
from server import PromptServer
|
|
from aiohttp import web
|
|
import folder_paths
|
|
from . import IMGNR_constants as C
|
|
|
|
# --- GLOBAL CACHE FOR MANUAL SAVES ---
|
|
# Maps node_id -> {"prompt": prompt, "extra_pnginfo": extra_pnginfo}
|
|
IMGNR_WORKFLOW_CACHE = {}
|
|
|
|
def format_comfy_string(text, prompt=None, extra_pnginfo=None):
|
|
if not isinstance(text, str):
|
|
return text
|
|
|
|
# 1. Resolve date variables
|
|
def format_date(match):
|
|
fmt = match.group(1)
|
|
fmt = fmt.replace("yyyy", "%Y").replace("yy", "%y")
|
|
fmt = fmt.replace("MM", "%m").replace("dd", "%d")
|
|
fmt = fmt.replace("HH", "%H").replace("mm", "%M").replace("ss", "%S")
|
|
try:
|
|
return datetime.now().strftime(fmt)
|
|
except Exception:
|
|
return match.group(0)
|
|
|
|
text = re.sub(r"%date:(.*?)%", format_date, text)
|
|
|
|
# 2. Resolve ComfyUI node variables (e.g. %Empty Latent Image.width%)
|
|
if prompt:
|
|
title_to_id = {}
|
|
class_to_id = {}
|
|
|
|
if extra_pnginfo and "workflow" in extra_pnginfo and "nodes" in extra_pnginfo["workflow"]:
|
|
for node in extra_pnginfo["workflow"]["nodes"]:
|
|
node_id = str(node.get("id"))
|
|
title = node.get("title", "")
|
|
class_type = node.get("type", "")
|
|
if title: title_to_id[title] = node_id
|
|
if class_type: class_to_id[class_type] = node_id
|
|
|
|
# Fallback for prompt-only situations
|
|
for node_id, node_data in prompt.items():
|
|
class_type = node_data.get("class_type", "")
|
|
if class_type and class_type not in class_to_id:
|
|
class_to_id[class_type] = str(node_id)
|
|
|
|
def resolve_var(match):
|
|
var_path = match.group(1)
|
|
parts = var_path.rsplit(".", 1)
|
|
if len(parts) != 2:
|
|
return match.group(0)
|
|
|
|
node_name, widget_name = parts
|
|
target_id = title_to_id.get(node_name) or class_to_id.get(node_name)
|
|
|
|
if target_id and target_id in prompt:
|
|
inputs = prompt[target_id].get("inputs", {})
|
|
if widget_name in inputs:
|
|
val = inputs[widget_name]
|
|
# Handle primitive links implicitly (grabs evaluated values from upstream)
|
|
if isinstance(val, list) and len(val) == 2:
|
|
link_id = str(val[0])
|
|
if link_id in prompt:
|
|
for k, v in prompt[link_id].get("inputs", {}).items():
|
|
if not isinstance(v, list):
|
|
return str(v)
|
|
return str(val)
|
|
return match.group(0)
|
|
|
|
text = re.sub(r"%([^%]+)%", resolve_var, text)
|
|
|
|
return text
|
|
|
|
# --- UTILITY: SAVE FUNCTION ---
|
|
def save_image_to_disk(image_data_base64, filename_prefix, counter, add_counter, filename_extras, overwrite, embed_workflow=False, prompt=None, extra_pnginfo=None, output_dir=""):
|
|
try:
|
|
# 1. Decode Image
|
|
if "," in image_data_base64:
|
|
image_data_base64 = image_data_base64.split(",")[1]
|
|
image_bytes = base64.b64decode(image_data_base64)
|
|
img = Image.open(io.BytesIO(image_bytes))
|
|
|
|
# 2. Handle Metadata (Workflow)
|
|
metadata = PngInfo()
|
|
if embed_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]))
|
|
|
|
# 3. Construct Filename
|
|
if not filename_prefix: filename_prefix = "ComfyUI"
|
|
|
|
# Catch manual ad-hoc dates and variables
|
|
filename_prefix = format_comfy_string(filename_prefix, prompt, extra_pnginfo)
|
|
|
|
extras_str = f"_{filename_extras}" if filename_extras and filename_extras.strip() else ""
|
|
|
|
full_output_dir = folder_paths.get_output_directory()
|
|
if output_dir:
|
|
full_output_dir = output_dir
|
|
|
|
# Resolve ComfyUI formatting natively (handles subfolders, %width%, etc.)
|
|
full_output_folder, filename_part, _, path_part, _ = folder_paths.get_save_image_path(filename_prefix, full_output_dir, img.width, img.height)
|
|
save_path = full_output_folder
|
|
|
|
if not os.path.exists(save_path):
|
|
os.makedirs(save_path)
|
|
|
|
# LOGIC: Add Counter or Not?
|
|
if add_counter:
|
|
counter_str = f"_{int(counter):05d}"
|
|
base_filename_no_ext = f"{filename_part}{counter_str}{extras_str}"
|
|
else:
|
|
base_filename_no_ext = f"{filename_part}{extras_str}"
|
|
|
|
file_extension = ".png"
|
|
full_file_path = os.path.join(save_path, base_filename_no_ext + file_extension)
|
|
|
|
# 4. Handle Overwrite & Determine Save Status
|
|
final_base_name = base_filename_no_ext
|
|
save_status = "new"
|
|
|
|
if os.path.exists(full_file_path):
|
|
if overwrite:
|
|
save_status = "overwritten"
|
|
else:
|
|
save_status = "saved_as"
|
|
safety_idx = 1
|
|
while os.path.exists(full_file_path):
|
|
new_base = f"{base_filename_no_ext}_{safety_idx:03d}"
|
|
full_file_path = os.path.join(save_path, new_base + file_extension)
|
|
final_base_name = new_base
|
|
safety_idx += 1
|
|
|
|
# 5. Save
|
|
img.save(full_file_path, pnginfo=metadata, compress_level=4)
|
|
|
|
# 6. Return Data
|
|
relative_path = os.path.relpath(full_file_path, folder_paths.get_output_directory())
|
|
|
|
# Increment counter ONLY if add_counter was active
|
|
next_counter = int(counter) + 1 if add_counter else int(counter)
|
|
|
|
# Construct the requested full filename string: path + final_base_name
|
|
if path_part:
|
|
returned_full_string = os.path.join(path_part, final_base_name).replace("\\", "/")
|
|
else:
|
|
returned_full_string = final_base_name
|
|
|
|
return True, full_file_path, relative_path, next_counter, returned_full_string, save_status
|
|
|
|
except Exception as e:
|
|
print(f"{C.ERR_PREFIX} Save Error: {e}")
|
|
return False, str(e), "", counter, "", "error"
|
|
|
|
# --- API: MANUAL SAVE ---
|
|
@PromptServer.instance.routes.post("/imgnr/save_manual")
|
|
async def imgnr_save_manual(request):
|
|
data = await request.json()
|
|
|
|
# Extract cached workflow metadata
|
|
node_id = str(data.get("node_id", ""))
|
|
cached_meta = IMGNR_WORKFLOW_CACHE.get(node_id, {})
|
|
|
|
success, full_path, rel_path, new_cnt, base_name, save_status = save_image_to_disk(
|
|
image_data_base64=data.get("image"),
|
|
filename_prefix=data.get("filename_prefix"),
|
|
counter=data.get("counter"),
|
|
add_counter=data.get("add_counter"),
|
|
filename_extras=data.get("filename_extras"),
|
|
overwrite=data.get("overwrite"),
|
|
embed_workflow=data.get("embed_workflow", True),
|
|
prompt=cached_meta.get("prompt", None),
|
|
extra_pnginfo=cached_meta.get("extra_pnginfo", None)
|
|
)
|
|
return web.json_response({
|
|
"success": success,
|
|
"message": full_path if success else rel_path,
|
|
"relative_path": rel_path,
|
|
"new_counter": new_cnt,
|
|
"save_status": save_status
|
|
})
|
|
|
|
|
|
# --- BASE CLASS LOGIC ---
|
|
class IMGNR_Preview_Base:
|
|
def tensor_to_pil(self, img_tensor):
|
|
if img_tensor.ndim == 3: img_tensor = img_tensor.unsqueeze(0)
|
|
arr = np.clip(255. * img_tensor[0].cpu().numpy(), 0, 255).astype(np.uint8)
|
|
|
|
# Explicitly enforce RGBA mode to prevent PIL from auto-stripping the alpha channel
|
|
if arr.shape[-1] == 4:
|
|
return Image.fromarray(arr, mode="RGBA")
|
|
return Image.fromarray(arr, mode="RGB")
|
|
|
|
def prepare_rgba(self, images, mask=None):
|
|
output_images = images
|
|
output_mask = None
|
|
|
|
if mask is not None:
|
|
if mask.ndim == 2: mask = mask.unsqueeze(0)
|
|
if mask.shape[0] == 1 and images.shape[0] > 1:
|
|
mask = mask.repeat(images.shape[0], 1, 1)
|
|
if mask.shape[1:] != images.shape[1:3]:
|
|
mask = torch.nn.functional.interpolate(
|
|
mask.unsqueeze(1), size=images.shape[1:3], mode='nearest'
|
|
).squeeze(1)
|
|
alpha = 1.0 - mask
|
|
output_images = torch.cat((images[..., :3], alpha.unsqueeze(-1)), dim=-1)
|
|
output_mask = mask
|
|
|
|
elif images.shape[-1] == 4:
|
|
output_images = images
|
|
output_mask = 1.0 - images[..., 3]
|
|
else:
|
|
output_mask = torch.zeros((images.shape[0], images.shape[1], images.shape[2]), dtype=torch.float32, device=images.device)
|
|
|
|
return output_images, output_mask
|
|
|
|
# Helper to recursively sanitize JSON (convert NaN/Inf to None)
|
|
def clean_json(self, data):
|
|
if isinstance(data, dict):
|
|
return {k: self.clean_json(v) for k, v in data.items()}
|
|
elif isinstance(data, list):
|
|
return [self.clean_json(v) for v in data]
|
|
elif isinstance(data, float):
|
|
if math.isnan(data) or math.isinf(data):
|
|
return None
|
|
return data
|
|
|
|
# Fixed signature: Now includes unique_id=None to handle ComfyUI's hidden inputs without crashing
|
|
def process_image(self, images, mask=None, filename_prefix="ComfyUI", counter=1, add_counter=True, filename_extras="", autosave=False, embed_workflow=True, overwrite=False, prompt=None, extra_pnginfo=None, unique_id=None):
|
|
ui_payload = []
|
|
current_cnt = counter
|
|
saved_rel_path = None
|
|
save_status = None
|
|
|
|
# Cache workflow for manual saves
|
|
if unique_id is not None:
|
|
node_id_str = str(unique_id[0]) if isinstance(unique_id, list) else str(unique_id)
|
|
IMGNR_WORKFLOW_CACHE[node_id_str] = {
|
|
"prompt": prompt,
|
|
"extra_pnginfo": extra_pnginfo
|
|
}
|
|
|
|
extras_str = f"_{filename_extras}" if filename_extras and filename_extras.strip() else ""
|
|
|
|
# Resolve ComfyUI formatting natively for the preview payload
|
|
width = images[0].shape[1] if (images is not None and len(images) > 0) else 0
|
|
height = images[0].shape[0] if (images is not None and len(images) > 0) else 0
|
|
|
|
resolved_prefix = format_comfy_string(filename_prefix, prompt, extra_pnginfo)
|
|
_, formatted_name, _, formatted_subfolder, _ = folder_paths.get_save_image_path(resolved_prefix, folder_paths.get_output_directory(), width, height)
|
|
resolved_filename_prefix = os.path.join(formatted_subfolder, formatted_name).replace("\\", "/") if formatted_subfolder else formatted_name
|
|
|
|
# Default string generation BEFORE attempting autosave (Keeps full path structure)
|
|
if add_counter:
|
|
cnt_str = f"_{int(current_cnt):05d}"
|
|
full_filename_str = f"{resolved_filename_prefix}{cnt_str}{extras_str}"
|
|
else:
|
|
full_filename_str = f"{resolved_filename_prefix}{extras_str}"
|
|
|
|
rgba_images, output_mask = self.prepare_rgba(images, mask)
|
|
|
|
if rgba_images is None:
|
|
empty = torch.zeros([1, 64, 64, 3])
|
|
return {"ui": {"imgnr_b64_previews": []}, "result": (empty, empty, filename_prefix, current_cnt, "")}
|
|
|
|
try:
|
|
pil_img = self.tensor_to_pil(rgba_images)
|
|
buffer = io.BytesIO()
|
|
pil_img.save(buffer, format="PNG", compress_level=4)
|
|
img_b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
data_uri = f"data:image/png;base64,{img_b64}"
|
|
|
|
if autosave:
|
|
_, _, saved_rel_path, next_cnt, saved_base_name, save_status = save_image_to_disk(
|
|
img_b64, filename_prefix, current_cnt, add_counter, filename_extras, overwrite,
|
|
embed_workflow, prompt, extra_pnginfo
|
|
)
|
|
current_cnt = next_cnt
|
|
# Overwrite the default string with the exact full saved final name (including _001 overrides)
|
|
full_filename_str = saved_base_name
|
|
|
|
# Flatten relevant metadata for easy JS Diffing
|
|
raw_meta = {}
|
|
if prompt is not None:
|
|
raw_meta["prompt"] = prompt
|
|
|
|
# SANITIZE: Ensure no NaNs get sent to frontend
|
|
meta_payload = self.clean_json(raw_meta)
|
|
|
|
ui_payload.append({
|
|
"image": data_uri,
|
|
"width": pil_img.width,
|
|
"height": pil_img.height,
|
|
"current_counter": current_cnt,
|
|
"saved_filename": saved_rel_path,
|
|
"save_status": save_status,
|
|
"params": {
|
|
"filename_prefix": filename_prefix,
|
|
"filename_extras": filename_extras,
|
|
"add_counter": add_counter,
|
|
"overwrite": overwrite,
|
|
"embed_workflow": embed_workflow
|
|
},
|
|
"meta": meta_payload # Passed to JS for A/B Diff
|
|
})
|
|
|
|
except Exception as e:
|
|
print(f"{C.ERR_PREFIX} [IMGNR_Preview] Error: {e}")
|
|
|
|
return {
|
|
"ui": {"imgnr_b64_previews": ui_payload},
|
|
"result": (rgba_images, output_mask, filename_prefix, current_cnt, full_filename_str)
|
|
}
|
|
|
|
|
|
# --- NODE 1: PREVIEW ONLY (NO SAVE) ---
|
|
class PreviewImageBase64Node(IMGNR_Preview_Base):
|
|
DESCRIPTION = "Displays input images without saving to disk."
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": { "images": ("IMAGE",) },
|
|
"optional": { "mask": ("MASK",) },
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("image", "mask")
|
|
FUNCTION = "run"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "IMGNR"
|
|
|
|
def run(self, images, mask=None, prompt=None, extra_pnginfo=None):
|
|
res = self.process_image(images, mask, prompt=prompt, extra_pnginfo=extra_pnginfo)
|
|
return {"ui": res["ui"], "result": (res["result"][0], res["result"][1])}
|
|
|
|
|
|
# --- NODE 2: AD-HOC SAVE (NO COMPARE) ---
|
|
class PreviewImageAdHocSaveNode(IMGNR_Preview_Base):
|
|
DESCRIPTION = """
|
|
Displays input images and allows Auto saving or Manual saving _after_ image generation.
|
|
Optimized memory (drops history payload).
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE", ),
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."}),
|
|
"counter": ("INT", {"default": 1, "min": 0, "max": 999999, "step": 1}),
|
|
"add_counter": ("BOOLEAN", {"default": True}),
|
|
"filename_extras": ("STRING", {"default": "", "tooltip": "Text to add after the counter.(no formatting)"}),
|
|
"autosave": ("BOOLEAN", {"default": False}),
|
|
"embed_workflow": ("BOOLEAN", {"default": True}),
|
|
"overwrite": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK", ),
|
|
},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", "unique_id": "UNIQUE_ID"},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "STRING", "INT", "STRING")
|
|
RETURN_NAMES = ("image", "mask", "filename_prefix", "counter_out", "full_filename")
|
|
FUNCTION = "run"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "IMGNR"
|
|
|
|
def run(self, **kwargs):
|
|
res = self.process_image(**kwargs)
|
|
return {"ui": res["ui"], "result": res["result"]}
|
|
|
|
|
|
# --- NODE 3: COMPARE NODE (SAVE + COMPARE) ---
|
|
class PreviewImageCompareNode(IMGNR_Preview_Base):
|
|
DESCRIPTION = """
|
|
Displays input images and allows Auto saving or Manual saving _after_ image generation.
|
|
Includes LastGen Comparison tools in the UI (Always On).
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE", ),
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."}),
|
|
"counter": ("INT", {"default": 1, "min": 0, "max": 999999, "step": 1}),
|
|
"add_counter": ("BOOLEAN", {"default": True}),
|
|
"filename_extras": ("STRING", {"default": "", "tooltip": "Text to add after the counter. (no formatting)"}),
|
|
"autosave": ("BOOLEAN", {"default": False}),
|
|
"embed_workflow": ("BOOLEAN", {"default": True}),
|
|
"overwrite": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK", ),
|
|
},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", "unique_id": "UNIQUE_ID"},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "STRING", "INT", "STRING")
|
|
RETURN_NAMES = ("image", "mask", "filename_prefix", "counter_out", "full_filename")
|
|
FUNCTION = "run"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "IMGNR"
|
|
|
|
def run(self, **kwargs):
|
|
res = self.process_image(**kwargs)
|
|
return {"ui": res["ui"], "result": res["result"]}
|
|
|
|
|
|
# --- REGISTER NODES ---
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PreviewImageBase64Node": PreviewImageBase64Node,
|
|
"PreviewImageAdHocSaveNode": PreviewImageAdHocSaveNode,
|
|
"PreviewImageCompareNode": PreviewImageCompareNode
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PreviewImageBase64Node": "Preview Image - No Save (IMGNR)",
|
|
"PreviewImageAdHocSaveNode": "Preview Image - Ad-hoc Save (IMGNR)",
|
|
"PreviewImageCompareNode": "Preview Compare LastGen (IMGNR)"
|
|
} |