import json import os import re from typing import Any, Dict, Optional import numpy as np import folder_paths # type: ignore from ..services.service_registry import ServiceRegistry from ..metadata_collector.metadata_processor import MetadataProcessor from ..metadata_collector import get_metadata from PIL import Image, PngImagePlugin import piexif import logging logger = logging.getLogger(__name__) class SaveImageLM: NAME = "Save Image (LoraManager)" CATEGORY = "Lora Manager/utils" DESCRIPTION = "Save images with embedded generation metadata in compatible format" def __init__(self): self.output_dir = folder_paths.get_output_directory() self.type = "output" self.prefix_append = "" self.compress_level = 4 self.counter = 0 # Add pattern format regex for filename substitution pattern_format = re.compile(r"(%[^%]+%)") @classmethod def INPUT_TYPES(cls): return { "required": { "images": ("IMAGE",), "filename_prefix": ( "STRING", { "default": "ComfyUI", "tooltip": "Base filename for saved images. Supports format patterns like %seed%, %width%, %height%, %model%, etc.", }, ), "file_format": ( ["png", "jpeg", "webp"], { "tooltip": "Image format to save as. PNG preserves quality, JPEG is smaller, WebP balances size and quality." }, ), }, "optional": { "lossless_webp": ( "BOOLEAN", { "default": False, "tooltip": "When enabled, saves WebP images with lossless compression. Results in larger files but no quality loss.", }, ), "quality": ( "INT", { "default": 100, "min": 1, "max": 100, "tooltip": "Compression quality for JPEG and lossy WebP formats (1-100). Higher values mean better quality but larger files.", }, ), "embed_workflow": ( "BOOLEAN", { "default": False, "tooltip": "Embeds the complete workflow data into the image metadata. Only works with PNG and WebP formats.", }, ), "add_counter_to_filename": ( "BOOLEAN", { "default": True, "tooltip": "Adds an incremental counter to filenames to prevent overwriting previous images.", }, ), }, "hidden": { "id": "UNIQUE_ID", "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("images",) FUNCTION = "process_image" OUTPUT_NODE = True def get_lora_hash(self, lora_name): """Get the lora hash from cache""" scanner = ServiceRegistry.get_service_sync("lora_scanner") # Use the new direct filename lookup method if scanner is not None: hash_value = scanner.get_hash_by_filename(lora_name) if hash_value: return hash_value return None def get_checkpoint_hash(self, checkpoint_path): """Get the checkpoint hash from cache""" scanner = ServiceRegistry.get_service_sync("checkpoint_scanner") if not checkpoint_path: return None # Extract basename without extension checkpoint_name = os.path.basename(checkpoint_path) checkpoint_name = os.path.splitext(checkpoint_name)[0] # Try direct filename lookup first if scanner is not None: hash_value = scanner.get_hash_by_filename(checkpoint_name) if hash_value: return hash_value return None def format_metadata(self, metadata_dict): """Format metadata in the requested format similar to userComment example""" if not metadata_dict: return "" # Helper function to only add parameter if value is not None def add_param_if_not_none(param_list, label, value): if value is not None: param_list.append(f"{label}: {value}") # Extract the prompt and negative prompt prompt = metadata_dict.get("prompt", "") negative_prompt = metadata_dict.get("negative_prompt", "") # Extract loras from the prompt if present loras_text = metadata_dict.get("loras", "") lora_hashes = {} # If loras are found, add them on a new line after the prompt if loras_text: prompt_with_loras = f"{prompt}\n{loras_text}" # Extract lora names from the format lora_matches = re.findall(r"]+)>", loras_text) # Get hash for each lora for lora_name, strength in lora_matches: hash_value = self.get_lora_hash(lora_name) if hash_value: lora_hashes[lora_name] = hash_value else: prompt_with_loras = prompt # Format the first part (prompt and loras) metadata_parts = [prompt_with_loras] # Add negative prompt if negative_prompt: metadata_parts.append(f"Negative prompt: {negative_prompt}") # Format the second part (generation parameters) params = [] # Add standard parameters in the correct order if "steps" in metadata_dict: add_param_if_not_none(params, "Steps", metadata_dict.get("steps")) # Combine sampler and scheduler information sampler_name = None scheduler_name = None if "sampler" in metadata_dict: sampler = metadata_dict.get("sampler") # Convert ComfyUI sampler names to user-friendly names sampler_mapping = { "euler": "Euler", "euler_ancestral": "Euler a", "dpm_2": "DPM2", "dpm_2_ancestral": "DPM2 a", "heun": "Heun", "dpm_fast": "DPM fast", "dpm_adaptive": "DPM adaptive", "lms": "LMS", "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", "ddim": "DDIM", } sampler_name = sampler_mapping.get(sampler, sampler) if "scheduler" in metadata_dict: scheduler = metadata_dict.get("scheduler") scheduler_mapping = { "normal": "Simple", "karras": "Karras", "exponential": "Exponential", "sgm_uniform": "SGM Uniform", "sgm_quadratic": "SGM Quadratic", } scheduler_name = scheduler_mapping.get(scheduler, scheduler) # Add combined sampler and scheduler information if sampler_name: if scheduler_name: params.append(f"Sampler: {sampler_name} {scheduler_name}") else: params.append(f"Sampler: {sampler_name}") # CFG scale (Use guidance if available, otherwise fall back to cfg_scale or cfg) if "guidance" in metadata_dict: add_param_if_not_none(params, "CFG scale", metadata_dict.get("guidance")) elif "cfg_scale" in metadata_dict: add_param_if_not_none(params, "CFG scale", metadata_dict.get("cfg_scale")) elif "cfg" in metadata_dict: add_param_if_not_none(params, "CFG scale", metadata_dict.get("cfg")) # Seed if "seed" in metadata_dict: add_param_if_not_none(params, "Seed", metadata_dict.get("seed")) # Size if "size" in metadata_dict: add_param_if_not_none(params, "Size", metadata_dict.get("size")) # Model info if "checkpoint" in metadata_dict: # Ensure checkpoint is a string before processing checkpoint = metadata_dict.get("checkpoint") if checkpoint is not None: # Get model hash model_hash = self.get_checkpoint_hash(checkpoint) # Extract basename without path checkpoint_name = os.path.basename(checkpoint) # Remove extension if present checkpoint_name = os.path.splitext(checkpoint_name)[0] # Add model hash if available if model_hash: params.append( f"Model hash: {model_hash[:10]}, Model: {checkpoint_name}" ) else: params.append(f"Model: {checkpoint_name}") # Add LoRA hashes if available if lora_hashes: lora_hash_parts = [] for lora_name, hash_value in lora_hashes.items(): lora_hash_parts.append(f"{lora_name}: {hash_value[:10]}") if lora_hash_parts: params.append(f'Lora hashes: "{", ".join(lora_hash_parts)}"') # Combine all parameters with commas metadata_parts.append(", ".join(params)) # Join all parts with a new line return "\n".join(metadata_parts) # credit to nkchocoai # Add format_filename method to handle pattern substitution def format_filename(self, filename, metadata_dict): """Format filename with metadata values""" if not metadata_dict: return filename result = re.findall(self.pattern_format, filename) for segment in result: parts = segment.replace("%", "").split(":") key = parts[0] if key == "seed" and "seed" in metadata_dict: filename = filename.replace(segment, str(metadata_dict.get("seed", ""))) elif key == "width" and "size" in metadata_dict: size = metadata_dict.get("size", "x") w = size.split("x")[0] if isinstance(size, str) else size[0] filename = filename.replace(segment, str(w)) elif key == "height" and "size" in metadata_dict: size = metadata_dict.get("size", "x") h = size.split("x")[1] if isinstance(size, str) else size[1] filename = filename.replace(segment, str(h)) elif key == "pprompt" and "prompt" in metadata_dict: prompt = metadata_dict.get("prompt", "").replace("\n", " ") if len(parts) >= 2: length = int(parts[1]) prompt = prompt[:length] filename = filename.replace(segment, prompt.strip()) elif key == "nprompt" and "negative_prompt" in metadata_dict: prompt = metadata_dict.get("negative_prompt", "").replace("\n", " ") if len(parts) >= 2: length = int(parts[1]) prompt = prompt[:length] filename = filename.replace(segment, prompt.strip()) elif key == "model": model_value = metadata_dict.get("checkpoint") if isinstance(model_value, (bytes, os.PathLike)): model_value = str(model_value) if not isinstance(model_value, str) or not model_value: model = "model_unavailable" else: model = os.path.splitext(os.path.basename(model_value))[0] if len(parts) >= 2: length = int(parts[1]) model = model[:length] filename = filename.replace(segment, model) elif key == "date": from datetime import datetime now = datetime.now() date_table = { "yyyy": f"{now.year:04d}", "yy": f"{now.year % 100:02d}", "MM": f"{now.month:02d}", "dd": f"{now.day:02d}", "hh": f"{now.hour:02d}", "mm": f"{now.minute:02d}", "ss": f"{now.second:02d}", } if len(parts) >= 2: date_format = parts[1] for k, v in date_table.items(): date_format = date_format.replace(k, v) filename = filename.replace(segment, date_format) else: date_format = "yyyyMMddhhmmss" for k, v in date_table.items(): date_format = date_format.replace(k, v) filename = filename.replace(segment, date_format) return filename def save_images( self, images, filename_prefix, file_format, id, prompt=None, extra_pnginfo=None, lossless_webp=True, quality=100, embed_workflow=False, add_counter_to_filename=True, ): """Save images with metadata""" results = [] # Get metadata using the metadata collector raw_metadata = get_metadata() metadata_dict = MetadataProcessor.to_dict(raw_metadata, id) metadata = self.format_metadata(metadata_dict) # Process filename_prefix with pattern substitution filename_prefix = self.format_filename(filename_prefix, metadata_dict) # Get initial save path info once for the batch full_output_folder, filename, counter, subfolder, processed_prefix = ( folder_paths.get_save_image_path( filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0] ) ) # Create directory if it doesn't exist if not os.path.exists(full_output_folder): os.makedirs(full_output_folder, exist_ok=True) # Process each image with incrementing counter for i, image in enumerate(images): # Convert the tensor image to numpy array img = 255.0 * image.cpu().numpy() img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) # Generate filename with counter if needed base_filename = filename if add_counter_to_filename: # Use counter + i to ensure unique filenames for all images in batch current_counter = counter + i base_filename += f"_{current_counter:05}_" # Set file extension and prepare saving parameters file: str save_kwargs: Dict[str, Any] pnginfo: Optional[PngImagePlugin.PngInfo] = None if file_format == "png": file = base_filename + ".png" file_extension = ".png" # Remove "optimize": True to match built-in node behavior save_kwargs = {"compress_level": self.compress_level} pnginfo = PngImagePlugin.PngInfo() elif file_format == "jpeg": file = base_filename + ".jpg" file_extension = ".jpg" save_kwargs = {"quality": quality, "optimize": True} elif file_format == "webp": file = base_filename + ".webp" file_extension = ".webp" # Add optimization param to control performance save_kwargs = { "quality": quality, "lossless": lossless_webp, "method": 0, } else: raise ValueError(f"Unsupported file format: {file_format}") # Full save path file_path = os.path.join(full_output_folder, file) # Save the image with metadata try: if file_format == "png": assert pnginfo is not None if metadata: pnginfo.add_text("parameters", metadata) if embed_workflow and extra_pnginfo is not None: workflow_json = json.dumps(extra_pnginfo["workflow"]) pnginfo.add_text("workflow", workflow_json) save_kwargs["pnginfo"] = pnginfo img.save(file_path, format="PNG", **save_kwargs) elif file_format == "jpeg": # For JPEG, use piexif if metadata: try: exif_dict = { "Exif": { piexif.ExifIFD.UserComment: b"UNICODE\0" + metadata.encode("utf-16be") } } exif_bytes = piexif.dump(exif_dict) save_kwargs["exif"] = exif_bytes except Exception as e: logger.error(f"Error adding EXIF data: {e}") img.save(file_path, format="JPEG", **save_kwargs) elif file_format == "webp": try: # For WebP, use piexif for metadata exif_dict = {} if metadata: exif_dict["Exif"] = { piexif.ExifIFD.UserComment: b"UNICODE\0" + metadata.encode("utf-16be") } # Add workflow if needed if embed_workflow and extra_pnginfo is not None: workflow_json = json.dumps(extra_pnginfo["workflow"]) exif_dict["0th"] = { piexif.ImageIFD.ImageDescription: "Workflow:" + workflow_json } exif_bytes = piexif.dump(exif_dict) save_kwargs["exif"] = exif_bytes except Exception as e: logger.error(f"Error adding EXIF data: {e}") img.save(file_path, format="WEBP", **save_kwargs) results.append( {"filename": file, "subfolder": subfolder, "type": self.type} ) except Exception as e: logger.error(f"Error saving image: {e}") return results def process_image( self, images, id, filename_prefix="ComfyUI", file_format="png", prompt=None, extra_pnginfo=None, lossless_webp=True, quality=100, embed_workflow=False, add_counter_to_filename=True, ): """Process and save image with metadata""" # Make sure the output directory exists os.makedirs(self.output_dir, exist_ok=True) # If images is already a list or array of images, do nothing; otherwise, convert to list if isinstance(images, (list, np.ndarray)): pass else: # Ensure images is always a list of images if len(images.shape) == 3: # Single image (height, width, channels) images = [images] else: # Multiple images (batch, height, width, channels) images = [img for img in images] # Save all images results = self.save_images( images, filename_prefix, file_format, id, prompt, extra_pnginfo, lossless_webp, quality, embed_workflow, add_counter_to_filename, ) return (images,)