diff --git a/__init__.py b/__init__.py index c697e12f..008219c8 100644 --- a/__init__.py +++ b/__init__.py @@ -2,13 +2,13 @@ from .py.lora_manager import LoraManager from .py.nodes.lora_loader import LoraManagerLoader from .py.nodes.trigger_word_toggle import TriggerWordToggle from .py.nodes.lora_stacker import LoraStacker -# from .py.nodes.save_image import SaveImage +from .py.nodes.save_image import SaveImage NODE_CLASS_MAPPINGS = { LoraManagerLoader.NAME: LoraManagerLoader, TriggerWordToggle.NAME: TriggerWordToggle, LoraStacker.NAME: LoraStacker, - # SaveImage.NAME: SaveImage + SaveImage.NAME: SaveImage } WEB_DIRECTORY = "./web/comfyui" diff --git a/py/nodes/save_image.py b/py/nodes/save_image.py index 09fec5e4..63c7b925 100644 --- a/py/nodes/save_image.py +++ b/py/nodes/save_image.py @@ -1,16 +1,43 @@ import json +import os +import asyncio +import re +import numpy as np +import time from server import PromptServer # type: ignore +import folder_paths # type: ignore +from ..services.lora_scanner import LoraScanner +from ..config import config +from ..workflow.parser import WorkflowParser +from PIL import Image, PngImagePlugin +import piexif +from io import BytesIO class SaveImage: NAME = "Save Image (LoraManager)" CATEGORY = "Lora Manager/utils" - DESCRIPTION = "Experimental node to display image preview and print prompt and extra_pnginfo" + 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 @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), + "filename_prefix": ("STRING", {"default": "ComfyUI"}), + "file_format": (["png", "jpeg", "webp"],), + }, + "optional": { + "lossless_webp": ("BOOLEAN", {"default": True}), + "quality": ("INT", {"default": 100, "min": 1, "max": 100}), + "save_workflow_json": ("BOOLEAN", {"default": False}), + "add_counter_to_filename": ("BOOLEAN", {"default": True}), }, "hidden": { "prompt": "PROMPT", @@ -18,24 +45,252 @@ class SaveImage: }, } - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("image",) + RETURN_TYPES = ("IMAGE", "STRING") + RETURN_NAMES = ("image", "filename") FUNCTION = "process_image" + OUTPUT_NODE = True - def process_image(self, image, prompt=None, extra_pnginfo=None): - # Print the prompt information - print("SaveImage Node - Prompt:") + async def get_lora_hash(self, lora_name): + """Get the lora hash from cache""" + scanner = await LoraScanner.get_instance() + cache = await scanner.get_cached_data() + + for item in cache.raw_data: + if item.get('file_name') == lora_name: + return item.get('sha256') + return None + + async def format_metadata(self, parsed_workflow): + """Format metadata in the requested format similar to userComment example""" + if not parsed_workflow: + return "" + + # Extract the prompt and negative prompt + prompt = parsed_workflow.get('prompt', '') + negative_prompt = parsed_workflow.get('negative_prompt', '') + + # Extract loras from the prompt if present + loras_text = parsed_workflow.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 = await 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 parsed_workflow: + params.append(f"Steps: {parsed_workflow.get('steps')}") + + if 'sampler' in parsed_workflow: + sampler = parsed_workflow.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) + params.append(f"Sampler: {sampler_name}") + + if 'scheduler' in parsed_workflow: + scheduler = parsed_workflow.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) + params.append(f"Schedule type: {scheduler_name}") + + # CFG scale (cfg in parsed_workflow) + if 'cfg_scale' in parsed_workflow: + params.append(f"CFG scale: {parsed_workflow.get('cfg_scale')}") + elif 'cfg' in parsed_workflow: + params.append(f"CFG scale: {parsed_workflow.get('cfg')}") + + # Seed + if 'seed' in parsed_workflow: + params.append(f"Seed: {parsed_workflow.get('seed')}") + + # Size + if 'size' in parsed_workflow: + params.append(f"Size: {parsed_workflow.get('size')}") + + # Model info + if 'checkpoint' in parsed_workflow: + # Extract basename without path + checkpoint = os.path.basename(parsed_workflow.get('checkpoint', '')) + # Remove extension if present + checkpoint = os.path.splitext(checkpoint)[0] + params.append(f"Model: {checkpoint}") + + # 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}") + + 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) + + def save_images(self, images, filename_prefix, file_format, prompt=None, extra_pnginfo=None, + lossless_webp=True, quality=100, save_workflow_json=False, add_counter_to_filename=True): + """Save images with metadata""" + results = [] + + # Parse the workflow using the WorkflowParser + parser = WorkflowParser() if prompt: - print(json.dumps(prompt, indent=2)) + parsed_workflow = parser.parse_workflow(prompt) else: - print("No prompt information available") + parsed_workflow = {} + + # Get or create metadata asynchronously + metadata = asyncio.run(self.format_metadata(parsed_workflow)) - # Print the extra_pnginfo - print("\nSaveImage Node - Extra PNG Info:") - if extra_pnginfo: - print(json.dumps(extra_pnginfo, indent=2)) - else: - print("No extra PNG info available") + # Process each image + for i, image in enumerate(images): + # Convert the tensor image to numpy array + img = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + # Generate filename with counter if needed + if add_counter_to_filename: + filename = f"{filename_prefix}_{self.counter:05d}" + self.counter += 1 + else: + filename = f"{filename_prefix}" + + # Set file extension and prepare saving parameters + if file_format == "png": + filename += ".png" + file_extension = ".png" + save_kwargs = {"optimize": True, "compress_level": self.compress_level} + pnginfo = PngImagePlugin.PngInfo() + elif file_format == "jpeg": + filename += ".jpg" + file_extension = ".jpg" + save_kwargs = {"quality": quality, "optimize": True} + elif file_format == "webp": + filename += ".webp" + file_extension = ".webp" + save_kwargs = {"quality": quality, "lossless": lossless_webp} + + # Full save path + file_path = os.path.join(self.output_dir, filename) + + # Save the image with metadata + try: + if file_format == "png": + if metadata: + pnginfo.add_text("parameters", metadata) + if save_workflow_json and extra_pnginfo is not None: + workflow_json = json.dumps(extra_pnginfo) + 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: + print(f"Error adding EXIF data: {e}") + img.save(file_path, format="JPEG", **save_kwargs) + elif file_format == "webp": + # For WebP, also use piexif for metadata + 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: + print(f"Error adding EXIF data: {e}") + img.save(file_path, format="WEBP", **save_kwargs) + + results.append({ + "filename": filename, + "subfolder": "", + "type": self.type + }) + + # Notify UI about saved image + PromptServer.instance.send_sync("image", { + "filename": filename, + "subfolder": "", + "type": self.type, + }) + + except Exception as e: + print(f"Error saving image: {e}") - # Return the image unchanged - return (image,) + return results + + def process_image(self, image, filename_prefix="ComfyUI", file_format="png", prompt=None, extra_pnginfo=None, + lossless_webp=True, quality=100, save_workflow_json=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) + + # Convert single image to list for consistent processing + images = [image[0]] if len(image.shape) == 3 else [img for img in image] + + # Save all images + results = self.save_images( + images, + filename_prefix, + file_format, + prompt, + extra_pnginfo, + lossless_webp, + quality, + save_workflow_json, + add_counter_to_filename + ) + + # Return the first saved filename and the original image + filename = results[0]["filename"] if results else "" + return (image, filename) \ No newline at end of file diff --git a/refs/jpeg_civitai_exif_userComment_example b/refs/jpeg_civitai_exif_userComment_example index b9fc40ab..bf5935fa 100644 --- a/refs/jpeg_civitai_exif_userComment_example +++ b/refs/jpeg_civitai_exif_userComment_example @@ -2,13 +2,10 @@ a dynamic and dramatic digital artwork featuring a stylized anthropomorphic whit Negative prompt: Steps: 30, Sampler: Undefined, CFG scale: 3.5, Seed: 90300501, Size: 832x1216, Clip skip: 2, Created Date: 2025-03-05T13:51:18.1770234Z, Civitai resources: [{"type":"checkpoint","modelVersionId":691639,"modelName":"FLUX","modelVersionName":"Dev"},{"type":"lora","weight":0.4,"modelVersionId":1202162,"modelName":"Velvet\u0027s Mythic Fantasy Styles | Flux \u002B Pony \u002B illustrious","modelVersionName":"Flux Gothic Lines"},{"type":"lora","weight":0.8,"modelVersionId":1470588,"modelName":"Velvet\u0027s Mythic Fantasy Styles | Flux \u002B Pony \u002B illustrious","modelVersionName":"Flux Retro"},{"type":"lora","weight":0.75,"modelVersionId":746484,"modelName":"Elden Ring - Yoshitaka Amano","modelVersionName":"V1"},{"type":"lora","weight":0.2,"modelVersionId":914935,"modelName":"Ink-style","modelVersionName":"ink-dynamic"},{"type":"lora","weight":0.2,"modelVersionId":1189379,"modelName":"Painterly Fantasy by ChronoKnight - [FLUX \u0026 IL]","modelVersionName":"FLUX"},{"type":"lora","weight":0.2,"modelVersionId":757030,"modelName":"Mezzotint Artstyle for Flux - by Ethanar","modelVersionName":"V1"}], Civitai metadata: {} -, masterpiece, best quality, good quality, very aesthetic, absurdres, newest, 8K, depth of field, focused subject, dynamic angle, dutch angle, from below, epic half body portrait, gritty, wabi sabi, looking at viewer, woman is a geisha, parted lips, -holographic skin, holofoil glitter, faint, glowing, ethereal, neon hair, glowing hair, otherworldly glow, she is dangerous, - - - +holographic skin, holofoil glitter, faint, glowing, ethereal, neon hair, glowing hair, otherworldly glow, she is dangerous +, , , Negative prompt: score_6, score_5, score_4, bad quality, worst quality, worst detail, sketch, censorship, furry, window, headphones, Steps: 30, Sampler: Euler a, Schedule type: Simple, CFG scale: 7, Seed: 1405717592, Size: 832x1216, Model hash: 1ad6ca7f70, Model: waiNSFWIllustrious_v100, Denoising strength: 0.35, Hires CFG Scale: 5, Hires upscale: 1.3, Hires steps: 20, Hires upscaler: 4x-AnimeSharp, Lora hashes: "ck-shadow-circuit-IL: 88e247aa8c3d, ck-nc-cyberpunk-IL-000011: 935e6755554c, ck-neon-retrowave-IL: edafb9df7da1, ck-yoneyama-mai-IL-000014: 1b9305692a2e", Version: f2.0.1v1.10.1-1.10.1, Diffusion in Low Bits: Automatic (fp16 LoRA)