From f6bdd59fdccaa4059e4cefb8bc93f6f80ec35055 Mon Sep 17 00:00:00 2001 From: LuqP2 Date: Sat, 20 Dec 2025 21:16:27 -0300 Subject: [PATCH] Add workflow auto-extraction and multi-format metadata saving Auto-detect sampler/model/vae/prompts/loras with concrete node lists and multi-sampler heuristic Switch to None sentinels for overrides and add filename placeholders + sanitization Add JPEG/WebP metadata via EXIF (piexif) and update README/docs Implement UI preview structure for ComfyUI frontend, matching Save Image behavior. --- README.md | 72 +++++--- metadata_utils.py | 319 +++++++++++++++++++++++++++++--- metahub_save_node.py | 357 +++++++++++++++++++++++++++-------- requirements.txt | 1 + workflow_extractor.py | 420 ++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 1034 insertions(+), 135 deletions(-) create mode 100644 workflow_extractor.py diff --git a/README.md b/README.md index 2f82010..9f4ab16 100644 --- a/README.md +++ b/README.md @@ -4,9 +4,11 @@ Advanced image saving node for ComfyUI with dual metadata support. ## Features +- **Auto-Extraction** - Detects sampler params, prompts, model/VAE, and LoRAs directly from your workflow +- **Multi-Format** - PNG, JPEG, and WebP with metadata injection +- **Filename Patterns** - Placeholder-based filenames with sanitization - **A1111/Civitai Compatible** - Saves metadata in tEXt chunk ("parameters") recognized by Automatic1111, Civitai, and most SD tools -- **Image MetaHub Compatible** - Saves extended metadata in iTXt chunk ("imagemetahub_data") with full workflow JSON -- **Auto-Detection** - Automatically extracts LoRAs and their weights from your workflow +- **Image MetaHub Compatible** - Saves extended metadata in iTXt chunk ("imagemetahub_data") with full workflow JSON - **Model Hashes** - Calculates SHA256 hashes (AutoV2 format) for models and LoRAs - **IMH Pro Fields** - Support for user tags, notes, and project names - **Performance** - Hash caching and graceful degradation ensure fast generation @@ -38,43 +40,63 @@ Advanced image saving node for ComfyUI with dual metadata support. 1. Download this repository as ZIP 2. Extract to `ComfyUI/custom_nodes/ImageMetaHub-ComfyUI-Save` -3. Install dependencies: `pip install Pillow>=10.0.0 numpy>=1.24.0` +3. Install dependencies: `pip install Pillow>=10.0.0 numpy>=1.24.0 piexif>=1.1.3` 4. Restart ComfyUI ## Usage -### Basic Usage +### Quick Start 1. Add "MetaHub Save Image" node to your workflow -2. Connect the `images` output from your sampler -3. Fill in the generation parameters: - - **positive**: Your positive prompt - - **negative**: Your negative prompt - - **seed**, **steps**, **cfg**: Generation settings - - **sampler_name**, **scheduler**: Sampler configuration - - **model_name**: Your checkpoint filename (e.g., `model.safetensors`) +2. Connect the `images` output from your decoder/sampler +3. Generate! The node auto-detects and saves metadata. -4. Generate! Images are saved with full metadata. +Auto-detected fields: +- seed, steps, cfg +- sampler and scheduler +- model name + hash +- positive/negative prompts +- LoRAs + weights +- VAE name -### Advanced Features +### Override (Optional) +Connect any override input to replace auto-detected values (seed, steps, cfg, sampler_name, scheduler, model_name, positive, negative, denoise, vae_name). -#### Auto-Detection of LoRAs -The node automatically scans your workflow for LoRA Loader nodes and extracts: -- LoRA filename -- Weight/strength value +### Filename Pattern +Use `filename_pattern` to customize names (default `ComfyUI_%counter%`). Invalid filename characters are sanitized and missing values become `unknown`. `filename_prefix` is deprecated but still supported. -No manual input needed! +| Placeholder | Description | Example | +| --- | --- | --- | +| `%counter%` | Auto-increment (00001, 00002...) | 00042 | +| `%seed%` | Seed used | 1234567890 | +| `%date%` | Date YYYY-MM-DD | 2025-01-15 | +| `%time%` | Time HH-MM-SS | 14-30-22 | +| `%datetime%` | Date + time | 2025-01-15_14-30-22 | +| `%model%` | Model name (no extension) | dreamshaper_v8 | +| `%sampler%` | Sampler name | euler_ancestral | +| `%steps%` | Steps | 20 | +| `%cfg%` | CFG scale | 7.5 | +| `%width%` | Image width | 1024 | +| `%height%` | Image height | 768 | -#### IMH Pro Fields (Optional) +### File Formats +- `file_format`: PNG, JPEG, or WebP +- `quality`: JPEG/WebP quality (1-100) +- `output_path`: custom output directory (empty = ComfyUI default) +- PNG uses tEXt "parameters" and iTXt "imagemetahub_data" +- JPEG/WebP use EXIF UserComment (A1111) and ImageDescription (IMH JSON) + +If structured metadata cannot be written, the image is still saved. + +### Multi-Sampler Heuristic +When multiple samplers exist, the node prefers the sampler that feeds the latent used by the VAEDecode connected to the saved images; it falls back to the first sampler found. + +### IMH Pro Fields (Optional) Organize your generations with Image MetaHub Pro features: - **user_tags**: Tag your images (e.g., "portrait, fantasy, character") - **notes**: Add generation notes or experiments - **project_name**: Group images by project -#### Custom Output -- **filename_prefix**: Change output filename prefix (default: "ComfyUI") -- **output_path**: Save to custom directory (empty = ComfyUI default) - ## Configuration (Optional) If your models are in non-standard locations, set environment variables: @@ -93,7 +115,7 @@ export COMFYUI_VAE_PATH="/path/to/vae" ## Metadata Formats -### A1111/Civitai Format (tEXt chunk: "parameters") +### A1111/Civitai Format (PNG: tEXt "parameters", JPEG/WebP: EXIF UserComment) ``` masterpiece, best quality, 1girl, portrait @@ -101,7 +123,7 @@ Negative prompt: ugly, blurry Steps: 20, Sampler: euler, CFG scale: 7.0, Seed: 12345, Size: 512x768, Model: mymodel, Model hash: abc1234567, Lora hashes: "detail: def9876543" ``` -### Image MetaHub Format (iTXt chunk: "imagemetahub_data") +### Image MetaHub Format (PNG: iTXt "imagemetahub_data", JPEG/WebP: EXIF ImageDescription) ```json { diff --git a/metadata_utils.py b/metadata_utils.py index 5843c40..1d5af3c 100644 --- a/metadata_utils.py +++ b/metadata_utils.py @@ -1,16 +1,24 @@ """ Metadata Utils for MetaHub Save Node -Handles hash calculation, metadata formatting, and PNG chunk injection. +Handles hash calculation, metadata formatting, and metadata injection. """ import hashlib import json import os + from datetime import datetime from pathlib import Path -from typing import List, Dict, Optional +from typing import Any, Dict, List, Optional import numpy as np from PIL import Image, PngImagePlugin +try: + import piexif + from piexif import helper as piexif_helper +except Exception: + piexif = None + piexif_helper = None + # ============================================================================ # HASH CALCULATION SYSTEM @@ -327,7 +335,9 @@ def extract_loras_from_workflow(workflow_json: dict) -> List[dict]: inputs = node_data.get('inputs', {}) lora_name = inputs.get('lora_name', '') if lora_name: - weight = inputs.get('strength_model', inputs.get('strength', 1.0)) + weight = inputs.get('strength_model') + if weight is None: + weight = inputs.get('strength', inputs.get('strength_clip', 1.0)) loras.append({ "name": lora_name, "weight": float(weight) @@ -382,31 +392,171 @@ def inject_metadata_chunks(image_path: str, a1111_metadata: str, imh_metadata: d imh_metadata: IMH metadata dict (will be JSON serialized) """ try: - # Load existing image img = Image.open(image_path) - - # Create PngInfo object - png_info = PngImagePlugin.PngInfo() - - # Add tEXt chunk for A1111 compatibility - png_info.add_text("parameters", a1111_metadata) - - # Add iTXt chunk for IMH data (supports UTF-8) - imh_json = json.dumps(imh_metadata, ensure_ascii=False) - png_info.add_itxt("imagemetahub_data", imh_json) - - # Save with metadata (Pillow optimizes to avoid full re-encode) - img.save(image_path, "PNG", pnginfo=png_info, compress_level=4) + save_png_with_metadata(img, image_path, a1111_metadata, imh_metadata) except Exception as e: # Log error but don't raise - image is already saved print(f"[MetaHub] Warning: Could not inject metadata into {image_path}: {e}") +def save_png_with_metadata(image: Image.Image, image_path: str, a1111_metadata: str, imh_metadata: dict) -> None: + """ + Saves PNG image with A1111 tEXt and IMH iTXt metadata. + """ + try: + png_info = PngImagePlugin.PngInfo() + png_info.add_text("parameters", a1111_metadata or "") + imh_json = json.dumps(imh_metadata or {}, ensure_ascii=False) + png_info.add_itxt("imagemetahub_data", imh_json) + image.save(image_path, "PNG", pnginfo=png_info, compress_level=4) + except Exception as e: + print(f"[MetaHub] Warning: PNG metadata save failed for {image_path}: {e}") + image.save(image_path, "PNG", compress_level=4) + + +def _build_exif_bytes(a1111_metadata: str, imh_metadata: dict) -> Optional[bytes]: + if piexif is None or piexif_helper is None: + return None + try: + exif_dict = {"0th": {}, "Exif": {}, "GPS": {}, "Interop": {}, "1st": {}, "thumbnail": None} + imh_json = json.dumps(imh_metadata or {}, ensure_ascii=False) + if imh_json: + exif_dict["0th"][piexif.ImageIFD.ImageDescription] = imh_json.encode("utf-8", errors="replace") + if a1111_metadata: + exif_dict["Exif"][piexif.ExifIFD.UserComment] = piexif_helper.UserComment.dump( + a1111_metadata, encoding="unicode" + ) + return piexif.dump(exif_dict) + except Exception: + return None + + +def save_jpeg_with_metadata( + image: Image.Image, + image_path: str, + quality: int, + a1111_metadata: str, + imh_metadata: dict, +) -> None: + """ + Saves JPEG image with EXIF metadata (UserComment + ImageDescription). + """ + try: + if image.mode in ("RGBA", "LA", "P"): + image = image.convert("RGB") + save_kwargs: Dict[str, Any] = {"format": "JPEG", "quality": quality} + exif_bytes = _build_exif_bytes(a1111_metadata, imh_metadata) + if exif_bytes: + save_kwargs["exif"] = exif_bytes + elif a1111_metadata: + save_kwargs["comment"] = a1111_metadata.encode("utf-8", errors="replace") + image.save(image_path, **save_kwargs) + except Exception as e: + print(f"[MetaHub] Warning: JPEG metadata save failed for {image_path}: {e}") + image.save(image_path, "JPEG", quality=quality) + + +def save_webp_with_metadata( + image: Image.Image, + image_path: str, + quality: int, + a1111_metadata: str, + imh_metadata: dict, +) -> None: + """ + Saves WebP image with EXIF metadata. + """ + try: + save_kwargs: Dict[str, Any] = {"format": "WEBP", "quality": quality} + exif_bytes = _build_exif_bytes(a1111_metadata, imh_metadata) + if exif_bytes: + save_kwargs["exif"] = exif_bytes + elif a1111_metadata: + save_kwargs["comment"] = a1111_metadata + image.save(image_path, **save_kwargs) + except Exception as e: + print(f"[MetaHub] Warning: WebP metadata save failed for {image_path}: {e}") + image.save(image_path, "WEBP", quality=quality) + + # ============================================================================ # FILE MANAGEMENT # ============================================================================ +_INVALID_FILENAME_CHARS = '<>:"/\\|?*' +_KNOWN_EXTENSIONS = (".png", ".jpg", ".jpeg", ".webp") + + +def _normalize_file_format(file_format: str) -> str: + fmt = (file_format or "PNG").strip().upper() + if fmt == "JPG": + fmt = "JPEG" + if fmt not in ("PNG", "JPEG", "WEBP"): + fmt = "PNG" + return fmt + + +def _get_extension(file_format: str) -> str: + fmt = _normalize_file_format(file_format) + if fmt == "JPEG": + return ".jpg" + if fmt == "WEBP": + return ".webp" + return ".png" + + +def _strip_known_extension(name: str) -> str: + lower = name.lower() + for ext in _KNOWN_EXTENSIONS: + if lower.endswith(ext): + return name[: -len(ext)] + return name + + +def sanitize_filename(name: str) -> str: + if not name: + return "unknown" + cleaned = str(name) + for ch in _INVALID_FILENAME_CHARS: + cleaned = cleaned.replace(ch, "_") + cleaned = cleaned.strip().strip(".") + return cleaned or "unknown" + + +def _format_placeholder_value(value: Any) -> str: + if value is None: + return "unknown" + if isinstance(value, str) and not value.strip(): + return "unknown" + return str(value) + + +def resolve_placeholders(pattern: str, params: dict, counter: int) -> str: + now = datetime.now() + model_name = params.get("model_name") or "" + model_base = Path(model_name).stem if model_name else "" + + replacements = { + "%counter%": f"{counter:05d}", + "%seed%": _format_placeholder_value(params.get("seed")), + "%date%": now.strftime("%Y-%m-%d"), + "%time%": now.strftime("%H-%M-%S"), + "%datetime%": now.strftime("%Y-%m-%d_%H-%M-%S"), + "%model%": _format_placeholder_value(model_base), + "%sampler%": _format_placeholder_value(params.get("sampler")), + "%steps%": _format_placeholder_value(params.get("steps")), + "%cfg%": _format_placeholder_value(params.get("cfg")), + "%width%": _format_placeholder_value(params.get("width")), + "%height%": _format_placeholder_value(params.get("height")), + } + + resolved = pattern or "ComfyUI_%counter%" + for placeholder, value in replacements.items(): + resolved = resolved.replace(placeholder, value) + return resolved + + def get_output_directory(output_path: str) -> Path: """ Determines output directory with automatic creation (mkdir -p). @@ -439,23 +589,30 @@ def get_output_directory(output_path: str) -> Path: return output_dir -def get_next_filename(output_dir: Path, prefix: str) -> Path: +def get_next_filename(output_dir: Path, filename_pattern: str, params: dict, file_format: str) -> Path: """ Generates next available filename with auto-increment counter. - Format: {prefix}_{counter:05d}.png + Format: resolved pattern + extension Example: ComfyUI_00001.png, ComfyUI_00002.png Args: output_dir: Directory to save file - prefix: Filename prefix + filename_pattern: Filename pattern with placeholders + params: Dict with placeholder values + file_format: PNG, JPEG, or WEBP Returns: Full path to next available filename """ counter = 1 + uses_counter = "%counter%" in (filename_pattern or "") + extension = _get_extension(file_format) while True: - filename = f"{prefix}_{counter:05d}.png" + resolved = resolve_placeholders(filename_pattern, params, counter) + if not uses_counter and counter > 1: + resolved = f"{resolved}_{counter:05d}" + filename = sanitize_filename(_strip_known_extension(resolved)) + extension full_path = output_dir / filename if not full_path.exists(): return full_path @@ -501,14 +658,47 @@ def tensor_to_pil(image_tensor: np.ndarray) -> Image.Image: return pil_image -def save_image_batch(images: np.ndarray, output_dir: Path, prefix: str) -> List[Path]: +def save_image_with_metadata( + image: Image.Image, + file_path: Path, + file_format: str, + quality: int, + a1111_metadata: str, + imh_metadata: dict, +) -> None: + fmt = _normalize_file_format(file_format) + if fmt == "PNG": + save_png_with_metadata(image, str(file_path), a1111_metadata, imh_metadata) + elif fmt == "JPEG": + save_jpeg_with_metadata(image, str(file_path), quality, a1111_metadata, imh_metadata) + elif fmt == "WEBP": + save_webp_with_metadata(image, str(file_path), quality, a1111_metadata, imh_metadata) + else: + save_png_with_metadata(image, str(file_path), a1111_metadata, imh_metadata) + + +def save_image_batch( + images: np.ndarray, + output_dir: Path, + filename_pattern: str, + params: dict, + file_format: str = "PNG", + quality: int = 95, + a1111_metadata: str = "", + imh_metadata: Optional[dict] = None, +) -> List[Path]: """ - Saves batch of images with auto-increment filenames. + Saves batch of images with auto-increment filenames and metadata. Args: images: Batch of images from ComfyUI (shape: [B, H, W, C]) output_dir: Output directory - prefix: Filename prefix + filename_pattern: Filename pattern with placeholders + params: Dict with placeholder values + file_format: PNG, JPEG, or WEBP + quality: JPEG/WebP quality (1-100) + a1111_metadata: A1111 metadata string + imh_metadata: IMH metadata dict Returns: List of saved file paths @@ -520,11 +710,82 @@ def save_image_batch(images: np.ndarray, output_dir: Path, prefix: str) -> List[ pil_image = tensor_to_pil(image_tensor) # Get next filename - file_path = get_next_filename(output_dir, prefix) + file_path = get_next_filename(output_dir, filename_pattern, params, file_format) - # Save image (metadata will be added later) - pil_image.save(file_path, "PNG", compress_level=4) - - saved_paths.append(file_path) + # Save image with metadata + try: + save_image_with_metadata(pil_image, file_path, file_format, quality, a1111_metadata, imh_metadata or {}) + saved_paths.append(file_path) + except Exception as e: + print(f"[MetaHub] Warning: Could not save image {file_path}: {e}") + try: + fmt = _normalize_file_format(file_format) + if fmt == "JPEG": + if pil_image.mode in ("RGBA", "LA", "P"): + pil_image = pil_image.convert("RGB") + pil_image.save(file_path, "JPEG", quality=quality) + elif fmt == "WEBP": + pil_image.save(file_path, "WEBP", quality=quality) + else: + pil_image.save(file_path, "PNG", compress_level=4) + saved_paths.append(file_path) + except Exception as save_error: + print(f"[MetaHub] Warning: Fallback save failed for {file_path}: {save_error}") return saved_paths + + +def build_ui_preview(saved_paths: List[Path], output_base: Path) -> dict: + """ + Builds UI preview structure for ComfyUI frontend. + + ComfyUI frontend expects relative paths from output directory with format: + { + "ui": { + "images": [ + {"filename": "file.png", "subfolder": "subdir", "type": "output"} + ] + } + } + + Args: + saved_paths: List of saved file paths + output_base: Base output directory (ComfyUI default output/) + + Returns: + Dict with {"ui": {"images": [...]}} structure for preview + """ + ui_images = [] + + for file_path in saved_paths: + try: + # Calculate relative path from output directory + relative_path = file_path.relative_to(output_base) + + # Split into subfolder + filename + parts = relative_path.parts + if len(parts) > 1: + # Has subfolder(s): output/project/img.png -> subfolder="project" + subfolder = str(Path(*parts[:-1])) + filename = parts[-1] + else: + # No subfolder: output/img.png -> subfolder="" + subfolder = "" + filename = str(relative_path.name) + + ui_images.append({ + "filename": filename, + "subfolder": subfolder, + "type": "output" + }) + + except ValueError: + # Path is outside output_base (e.g., absolute path like D:/MyImages) + # Preview won't work - frontend can only serve from output/ or temp/ + print(f"[MetaHub] Warning: Preview unavailable for {file_path.name} (outside output directory)") + + return { + "ui": { + "images": ui_images + } + } diff --git a/metahub_save_node.py b/metahub_save_node.py index 3891ab9..4df0ad3 100644 --- a/metahub_save_node.py +++ b/metahub_save_node.py @@ -3,8 +3,11 @@ MetaHub Save Image Node for ComfyUI Advanced image saving with dual metadata support (A1111/Civitai + Image MetaHub) """ -from typing import Dict, Any +from pathlib import Path from . import metadata_utils as utils +from .workflow_extractor import WorkflowExtractor + +_HINT_SHOWN = False class MetaHubSaveNode: @@ -14,7 +17,7 @@ class MetaHubSaveNode: Features: - A1111/Civitai compatible metadata (tEXt chunk "parameters") - Image MetaHub metadata (iTXt chunk "imagemetahub_data") - - Auto-detection of LoRAs from workflow + - Auto-detection of workflow parameters (sampler, prompts, model, VAE, LoRAs) - SHA256 model hashes (AutoV2/Civitai format) - Graceful degradation (never interrupts generation) """ @@ -24,68 +27,77 @@ class MetaHubSaveNode: return { "required": { "images": ("IMAGE",), + }, + "optional": { + "filename_pattern": ("STRING", { + "default": "ComfyUI_%counter%", + "tooltip": "Filename pattern with placeholders" + }), + "file_format": (["PNG", "JPEG", "WebP"], { + "default": "PNG", + "tooltip": "File format" + }), + "quality": ("INT", { + "default": 95, + "min": 1, + "max": 100, + "tooltip": "JPEG/WebP quality" + }), + "output_path": ("STRING", { + "default": "", + "tooltip": "Custom output directory (empty = ComfyUI default)" + }), + "seed": ("INT", { + "default": None, + "forceInput": True, + "tooltip": "Override seed" + }), + "steps": ("INT", { + "default": None, + "forceInput": True, + "tooltip": "Override steps" + }), + "cfg": ("FLOAT", { + "default": None, + "forceInput": True, + "tooltip": "Override CFG scale" + }), + "sampler_name": ("STRING", { + "default": None, + "forceInput": True, + "tooltip": "Override sampler name" + }), + "scheduler": ("STRING", { + "default": None, + "forceInput": True, + "tooltip": "Override scheduler" + }), + "model_name": ("STRING", { + "default": None, + "forceInput": True, + "tooltip": "Override model filename" + }), "positive": ("STRING", { "multiline": True, - "default": "", - "tooltip": "Positive prompt" + "default": None, + "forceInput": True, + "tooltip": "Override positive prompt" }), "negative": ("STRING", { "multiline": True, - "default": "", - "tooltip": "Negative prompt" - }), - "seed": ("INT", { - "default": 0, - "min": 0, - "max": 0xffffffffffffffff, - "tooltip": "Generation seed" - }), - "steps": ("INT", { - "default": 20, - "min": 1, - "max": 10000, - "tooltip": "Sampling steps" - }), - "cfg": ("FLOAT", { - "default": 7.0, - "min": 0.0, - "max": 100.0, - "step": 0.1, - "tooltip": "CFG scale" - }), - "sampler_name": ("STRING", { - "default": "euler", - "tooltip": "Sampler name" - }), - "scheduler": ("STRING", { - "default": "normal", - "tooltip": "Scheduler" - }), - "model_name": ("STRING", { - "default": "", - "tooltip": "Model filename (e.g., model.safetensors)" - }), - }, - "optional": { - "vae_name": ("STRING", { - "default": "", - "tooltip": "VAE filename" + "default": None, + "forceInput": True, + "tooltip": "Override negative prompt" }), "denoise": ("FLOAT", { - "default": 1.0, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "tooltip": "Denoise strength" + "default": None, + "forceInput": True, + "tooltip": "Override denoise strength" }), - "upscale_model": ("STRING", { - "default": "", - "tooltip": "Upscale model name" - }), - "generation_time": ("FLOAT", { - "default": 0.0, - "min": 0.0, - "tooltip": "Generation time in seconds" + "vae_name": ("STRING", { + "default": None, + "forceInput": True, + "tooltip": "Override VAE filename" }), "user_tags": ("STRING", { "default": "", @@ -100,18 +112,24 @@ class MetaHubSaveNode: "default": "", "tooltip": "IMH Pro: Project name" }), + "generation_time": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "tooltip": "Generation time in seconds" + }), "filename_prefix": ("STRING", { "default": "ComfyUI", - "tooltip": "Filename prefix" + "tooltip": "Deprecated: use filename_pattern" }), - "output_path": ("STRING", { + "upscale_model": ("STRING", { "default": "", - "tooltip": "Custom output directory (empty = ComfyUI default)" + "tooltip": "Upscale model name" }), }, "hidden": { "prompt": "PROMPT", - "extra_pnginfo": "EXTRA_PNGINFO" + "extra_pnginfo": "EXTRA_PNGINFO", + "unique_id": "UNIQUE_ID" }, } @@ -121,29 +139,206 @@ class MetaHubSaveNode: CATEGORY = "image/save" DESCRIPTION = "Save images with A1111/Civitai and Image MetaHub metadata" - def save_images(self, images, positive, negative, seed, steps, cfg, sampler_name, scheduler, model_name, vae_name="", denoise=1.0, upscale_model="", generation_time=0.0, user_tags="", notes="", project_name="", filename_prefix="ComfyUI", output_path="", prompt=None, extra_pnginfo=None): - print(f"[MetaHub] Starting save_images for batch of {len(images)} image(s)") - workflow_json = utils.get_workflow_json(extra_pnginfo) - lora_list = utils.extract_loras_from_workflow(workflow_json) - print(f"[MetaHub] Detected {len(lora_list)} LoRA(s) from workflow") - print(f"[MetaHub] Calculating hash for model: {model_name}") - model_hash = utils.calculate_model_hash(model_name, model_type="checkpoint") - print(f"[MetaHub] Calculating hashes for LoRAs...") - lora_hashes = utils.calculate_lora_hashes(lora_list) - height, width = images[0].shape[0], images[0].shape[1] - print(f"[MetaHub] Image dimensions: {width}x{height}") - output_dir = utils.get_output_directory(output_path) - print(f"[MetaHub] Saving to: {output_dir}") - saved_paths = utils.save_image_batch(images, output_dir, filename_prefix) - print(f"[MetaHub] Saved {len(saved_paths)} image(s)") - for idx, file_path in enumerate(saved_paths): - print(f"[MetaHub] Injecting metadata into: {file_path.name}") - params = {'positive': positive, 'negative': negative, 'steps': steps, 'sampler': sampler_name, 'scheduler': scheduler, 'cfg': cfg, 'seed': seed, 'width': width, 'height': height, 'model_name': model_name, 'model_hash': model_hash, 'vae_name': vae_name, 'denoise': denoise, 'upscale_model': upscale_model, 'generation_time': generation_time, 'user_tags': user_tags, 'notes': notes, 'project_name': project_name, 'lora_list': lora_list, 'lora_hashes': lora_hashes} + def save_images( + self, + images, + filename_pattern="ComfyUI_%counter%", + file_format="PNG", + quality=95, + output_path="", + seed=None, + steps=None, + cfg=None, + sampler_name=None, + scheduler=None, + model_name=None, + positive=None, + negative=None, + denoise=None, + vae_name=None, + user_tags="", + notes="", + project_name="", + generation_time=0.0, + filename_prefix="ComfyUI", + upscale_model="", + prompt=None, + extra_pnginfo=None, + unique_id=None, + ): + global _HINT_SHOWN + + try: + workflow_json = utils.get_workflow_json(extra_pnginfo) + prompt_data = prompt if isinstance(prompt, dict) else workflow_json.get("prompt", {}) + if not isinstance(prompt_data, dict): + prompt_data = {} + + extractor = WorkflowExtractor(prompt_data) + extracted, missing_fields = extractor.extract( + save_node_id=str(unique_id) if unique_id is not None else None + ) + lora_list = extracted.get("lora_list") or utils.extract_loras_from_workflow(workflow_json) + + def resolve_value(manual_value, extracted_value, default_value): + if manual_value is not None: + return manual_value + if extracted_value is not None: + return extracted_value + return default_value + + def normalize_int(value, default_value): + try: + return int(value) + except (TypeError, ValueError): + return default_value + + def normalize_float(value, default_value): + try: + return float(value) + except (TypeError, ValueError): + return default_value + + seed_value = normalize_int(resolve_value(seed, extracted.get("seed"), 0), 0) + steps_value = normalize_int(resolve_value(steps, extracted.get("steps"), 20), 20) + cfg_value = normalize_float(resolve_value(cfg, extracted.get("cfg"), 7.0), 7.0) + sampler_value = resolve_value(sampler_name, extracted.get("sampler_name"), "euler") + scheduler_value = resolve_value(scheduler, extracted.get("scheduler"), "normal") + model_name_value = resolve_value(model_name, extracted.get("model_name"), "") + positive_value = resolve_value(positive, extracted.get("positive"), "") + negative_value = resolve_value(negative, extracted.get("negative"), "") + denoise_value = normalize_float(resolve_value(denoise, extracted.get("denoise"), 1.0), 1.0) + vae_name_value = resolve_value(vae_name, extracted.get("vae_name"), "") + + quality_value = normalize_int(quality if quality is not None else 95, 95) + quality_value = max(1, min(100, quality_value)) + + default_pattern = "ComfyUI_%counter%" + pattern_value = (filename_pattern or "").strip() + if pattern_value and pattern_value != default_pattern: + final_pattern = pattern_value + elif filename_prefix and str(filename_prefix).strip(): + final_pattern = f"{filename_prefix}_%counter%" + else: + final_pattern = default_pattern + + if model_name_value: + model_hash = utils.calculate_model_hash(model_name_value, model_type="checkpoint") + else: + model_hash = "0000000000" + lora_hashes = utils.calculate_lora_hashes(lora_list) if lora_list else {} + + height, width = images[0].shape[0], images[0].shape[1] + + params = { + "positive": positive_value, + "negative": negative_value, + "steps": steps_value, + "sampler": sampler_value, + "scheduler": scheduler_value, + "cfg": cfg_value, + "seed": seed_value, + "width": width, + "height": height, + "model_name": model_name_value, + "model_hash": model_hash, + "vae_name": vae_name_value, + "denoise": denoise_value, + "upscale_model": upscale_model, + "generation_time": generation_time, + "user_tags": user_tags, + "notes": notes, + "project_name": project_name, + "lora_list": lora_list, + "lora_hashes": lora_hashes, + } + a1111_metadata = utils.build_a1111_metadata(params) imh_metadata = utils.build_imh_metadata(params, workflow_json) - utils.inject_metadata_chunks(str(file_path), a1111_metadata, imh_metadata) - print(f"[MetaHub] ✓ All images saved successfully with metadata") - return {} + + output_dir = utils.get_output_directory(output_path) + saved_paths = utils.save_image_batch( + images, + output_dir, + final_pattern, + params, + file_format, + quality_value, + a1111_metadata, + imh_metadata, + ) + + warn_fields = { + "seed", + "steps", + "cfg", + "sampler_name", + "scheduler", + "model_name", + "positive", + "negative", + "loras", + } + label_map = { + "sampler_name": "sampler", + "model_name": "model", + } + manual_inputs = { + "seed": seed, + "steps": steps, + "cfg": cfg, + "sampler_name": sampler_name, + "scheduler": scheduler, + "model_name": model_name, + "positive": positive, + "negative": negative, + } + missing_warn = [] + for field in missing_fields: + if field not in warn_fields: + continue + if field == "loras": + if lora_list: + continue + missing_warn.append("loras") + continue + if manual_inputs.get(field) is not None: + continue + missing_warn.append(label_map.get(field, field)) + if missing_warn: + missing_warn = sorted(set(missing_warn)) + print( + "[ImageMetaHub-Save] ⚠ Some params not detected " + f"({', '.join(missing_warn)}), image saved with available metadata" + ) + + auto_summary_parts = [] + if seed is None and extracted.get("seed") is not None: + auto_summary_parts.append(f"seed={seed_value}") + if steps is None and extracted.get("steps") is not None: + auto_summary_parts.append(f"steps={steps_value}") + if model_name is None and extracted.get("model_name"): + model_display = Path(str(model_name_value)).stem or str(model_name_value) + auto_summary_parts.append(f"model={model_display}") + auto_summary = ", ".join(auto_summary_parts) + + for file_path in saved_paths: + suffix = "" + if not _HINT_SHOWN and auto_summary: + suffix = f" (auto-detected: {auto_summary})" + print(f"[ImageMetaHub-Save] ✓ Saved: {file_path.name}{suffix}") + if not _HINT_SHOWN: + print("💡 Organize your AI images → github.com/LuqP2/ImageMetaHub") + _HINT_SHOWN = True + + # Build preview structure for ComfyUI UI + output_base = utils.get_output_directory("") # Get default ComfyUI output + return utils.build_ui_preview(saved_paths, output_base) + + except Exception as e: + print(f"[ImageMetaHub-Save] Warning: Save failed: {e}") + # Return empty UI structure on error + return {"ui": {"images": []}} NODE_CLASS_MAPPINGS = {"MetaHubSaveNode": MetaHubSaveNode} diff --git a/requirements.txt b/requirements.txt index 4d5361c..d108f24 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ Pillow>=10.0.0 numpy>=1.24.0 +piexif>=1.1.3 diff --git a/workflow_extractor.py b/workflow_extractor.py new file mode 100644 index 0000000..26d66f2 --- /dev/null +++ b/workflow_extractor.py @@ -0,0 +1,420 @@ +""" +Workflow Extractor for MetaHub Save Node +Auto-extracts sampler params, prompts, model, VAE, and LoRAs from ComfyUI prompt data. +""" + +from typing import Any, Dict, List, Optional, Set, Tuple + + +class WorkflowExtractor: + SAMPLER_NODES = [ + "KSampler", + "KSamplerAdvanced", + "SamplerCustom", + "SamplerCustomAdvanced", + "KSamplerSelect", + "KSampler (Efficient)", + "ImpactKSamplerBasicPipe", + "ImpactKSamplerAdvancedPipe", + "ImpactKSampler", + ] + + CHECKPOINT_NODES = [ + "CheckpointLoaderSimple", + "CheckpointLoader", + "UNETLoader", + "UnetLoaderGGUF", + ] + + VAE_NODES = [ + "VAELoader", + "CheckpointLoaderSimple", + "CheckpointLoader", + ] + + LORA_NODES = [ + "LoraLoader", + "LoraLoaderModelOnly", + "LoraLoaderModelOnlyAdvanced", + "LoraLoaderAdvanced", + "LoraLoaderAny", + ] + + CLIP_NODES = [ + "CLIPTextEncode", + "CLIPTextEncodeSDXL", + "CLIPTextEncodeSDXLPlus", + "CLIPTextEncodeSDXLRefiner", + "CLIPTextEncodeSD3", + ] + + VAE_DECODE_NODES = [ + "VAEDecode", + "VAEDecodeTiled", + "VAEDecode (Tiled)", + ] + + def __init__(self, prompt: Optional[Dict[str, Any]]): + self.prompt: Dict[str, Any] = {} + if isinstance(prompt, dict): + for node_id, node_data in prompt.items(): + self.prompt[str(node_id)] = node_data + + def extract(self, save_node_id: Optional[str] = None) -> Tuple[Dict[str, Any], Set[str]]: + data: Dict[str, Any] = {"lora_list": []} + missing: Set[str] = set() + + sampler_node_id = self.find_sampler_for_save_node(save_node_id) + if sampler_node_id: + sampler_params = self.extract_sampler_params(sampler_node_id) + for key in ("seed", "steps", "cfg", "sampler_name", "scheduler", "denoise"): + value = sampler_params.get(key) + if value is None: + missing.add(key) + else: + data[key] = value + + positive, negative = self.extract_prompts(sampler_node_id) + if positive is None: + missing.add("positive") + else: + data["positive"] = positive + if negative is None: + missing.add("negative") + else: + data["negative"] = negative + + model_name = self.extract_model_name(sampler_node_id) + if model_name is None: + missing.add("model_name") + else: + data["model_name"] = model_name + else: + missing.update( + { + "seed", + "steps", + "cfg", + "sampler_name", + "scheduler", + "denoise", + "positive", + "negative", + "model_name", + } + ) + + vae_name = self.extract_vae_name(save_node_id) + if vae_name is None: + missing.add("vae_name") + else: + data["vae_name"] = vae_name + + lora_list, has_lora_nodes = self.extract_loras() + if lora_list: + data["lora_list"] = lora_list + elif has_lora_nodes: + missing.add("loras") + + return data, missing + + def find_sampler_for_save_node(self, save_node_id: Optional[str]) -> Optional[str]: + sampler_nodes = self._find_nodes_by_type(self.SAMPLER_NODES) + if not sampler_nodes: + return None + + if save_node_id: + save_node = self._get_node(save_node_id) + if save_node: + images_conn = save_node.get("inputs", {}).get("images") + start_node_id = self._get_connection_node_id(images_conn) + if start_node_id: + sampler_id = self._find_sampler_from_images_source(start_node_id) + if sampler_id: + return sampler_id + + return sampler_nodes[0] + + def extract_sampler_params(self, sampler_node_id: str) -> Dict[str, Any]: + node = self._get_node(sampler_node_id) + if not node: + return {} + inputs = node.get("inputs", {}) + + seed = self._coerce_int(self._get_literal_input(inputs, "seed")) + steps = self._coerce_int(self._get_literal_input(inputs, "steps")) + cfg = self._coerce_float(self._get_first_literal(inputs, ("cfg", "cfg_scale"))) + sampler_name = self._get_first_literal(inputs, ("sampler_name", "sampler")) + scheduler = self._get_first_literal(inputs, ("scheduler", "scheduler_name")) + denoise = self._coerce_float(self._get_first_literal(inputs, ("denoise", "denoise_strength"))) + + return { + "seed": seed, + "steps": steps, + "cfg": cfg, + "sampler_name": sampler_name, + "scheduler": scheduler, + "denoise": denoise, + } + + def extract_model_name(self, sampler_node_id: str) -> Optional[str]: + node = self._get_node(sampler_node_id) + if not node: + return None + model_conn = node.get("inputs", {}).get("model") + start_node_id = self._get_connection_node_id(model_conn) + if not start_node_id: + return None + checkpoint_node_id = self._bfs_upstream( + [start_node_id], + lambda n: self._class_type(n) in self.CHECKPOINT_NODES, + ) + if not checkpoint_node_id: + return None + checkpoint_node = self._get_node(checkpoint_node_id) + if not checkpoint_node: + return None + return self._get_checkpoint_name(checkpoint_node) + + def extract_vae_name(self, save_node_id: Optional[str]) -> Optional[str]: + vae_decode_id = self._find_vae_decode_node(save_node_id) + if not vae_decode_id: + return None + decode_node = self._get_node(vae_decode_id) + if not decode_node: + return None + vae_conn = decode_node.get("inputs", {}).get("vae") + start_node_id = self._get_connection_node_id(vae_conn) + if not start_node_id: + return None + vae_node_id = self._bfs_upstream( + [start_node_id], + lambda n: self._class_type(n) in self.VAE_NODES, + ) + if not vae_node_id: + return None + vae_node = self._get_node(vae_node_id) + if not vae_node: + return None + return self._get_vae_name(vae_node) + + def extract_prompts(self, sampler_node_id: str) -> Tuple[Optional[str], Optional[str]]: + node = self._get_node(sampler_node_id) + if not node: + return None, None + inputs = node.get("inputs", {}) + positive_conn = inputs.get("positive") or inputs.get("positive_cond") + negative_conn = inputs.get("negative") or inputs.get("negative_cond") + positive = self._extract_text_from_connection(positive_conn) + negative = self._extract_text_from_connection(negative_conn) + return positive, negative + + def extract_loras(self) -> Tuple[List[Dict[str, Any]], bool]: + loras: List[Dict[str, Any]] = [] + has_lora_nodes = False + + for _, node in self.prompt.items(): + class_type = self._class_type(node) + if not self._is_lora_node(class_type): + continue + has_lora_nodes = True + inputs = node.get("inputs", {}) + lora_name = self._get_first_literal(inputs, ("lora_name", "lora")) + if not lora_name: + continue + strength_model = self._get_first_literal( + inputs, + ("strength_model", "strength", "strength_unet"), + ) + weight = self._coerce_float(strength_model) + if weight is None: + strength_clip = self._get_first_literal(inputs, ("strength_clip",)) + weight = self._coerce_float(strength_clip) + if weight is None: + weight = 1.0 + loras.append({"name": lora_name, "weight": float(weight)}) + + return loras, has_lora_nodes + + def _find_sampler_from_images_source(self, start_node_id: str) -> Optional[str]: + start_node = self._get_node(start_node_id) + if not start_node: + return None + if self._class_type(start_node) in self.SAMPLER_NODES: + return start_node_id + + if self._class_type(start_node) in self.VAE_DECODE_NODES: + samples_conn = start_node.get("inputs", {}).get("samples") + samples_node_id = self._get_connection_node_id(samples_conn) + if samples_node_id: + sampler_id = self._bfs_upstream( + [samples_node_id], + lambda n: self._class_type(n) in self.SAMPLER_NODES, + ) + if sampler_id: + return sampler_id + + return self._bfs_upstream( + [start_node_id], + lambda n: self._class_type(n) in self.SAMPLER_NODES, + ) + + def _find_vae_decode_node(self, save_node_id: Optional[str]) -> Optional[str]: + if not save_node_id: + return None + save_node = self._get_node(save_node_id) + if not save_node: + return None + images_conn = save_node.get("inputs", {}).get("images") + start_node_id = self._get_connection_node_id(images_conn) + if not start_node_id: + return None + return self._bfs_upstream( + [start_node_id], + lambda n: self._class_type(n) in self.VAE_DECODE_NODES, + ) + + def _find_nodes_by_type(self, types: List[str]) -> List[str]: + matched = [] + for node_id, node in self.prompt.items(): + if self._class_type(node) in types: + matched.append(node_id) + return matched + + def _extract_text_from_connection(self, conn: Any) -> Optional[str]: + node_id = self._get_connection_node_id(conn) + if not node_id: + return None + return self._extract_text_from_node(node_id) + + def _extract_text_from_node(self, start_node_id: str) -> Optional[str]: + texts: List[str] = [] + queue = [start_node_id] + visited: Set[str] = set() + + while queue: + node_id = queue.pop(0) + if node_id in visited: + continue + visited.add(node_id) + node = self._get_node(node_id) + if not node: + continue + class_type = self._class_type(node) + if class_type in self.CLIP_NODES: + text = self._get_clip_text(node) + if text: + texts.append(text) + continue + for input_val in node.get("inputs", {}).values(): + conn_id = self._get_connection_node_id(input_val) + if conn_id and conn_id not in visited: + queue.append(conn_id) + + if texts: + return "\n".join(texts) + return None + + def _get_clip_text(self, node: Dict[str, Any]) -> Optional[str]: + inputs = node.get("inputs", {}) + text = self._get_literal_input(inputs, "text") + if text: + return str(text) + parts: List[str] = [] + for key in ("text_g", "text_l"): + value = self._get_literal_input(inputs, key) + if value: + parts.append(str(value)) + if parts: + return "\n".join(parts) + return None + + def _get_checkpoint_name(self, node: Dict[str, Any]) -> Optional[str]: + inputs = node.get("inputs", {}) + for key in ("ckpt_name", "checkpoint", "model_name", "unet_name"): + value = self._get_literal_input(inputs, key) + if value: + return str(value) + return None + + def _get_vae_name(self, node: Dict[str, Any]) -> Optional[str]: + inputs = node.get("inputs", {}) + for key in ("vae_name", "ckpt_name", "model_name"): + value = self._get_literal_input(inputs, key) + if value: + return str(value) + return None + + def _get_node(self, node_id: Optional[str]) -> Optional[Dict[str, Any]]: + if node_id is None: + return None + return self.prompt.get(str(node_id)) + + @staticmethod + def _class_type(node: Dict[str, Any]) -> str: + return str(node.get("class_type", "")) + + @staticmethod + def _is_connection(value: Any) -> bool: + return isinstance(value, (list, tuple)) and len(value) >= 2 + + def _get_connection_node_id(self, value: Any) -> Optional[str]: + if not self._is_connection(value): + return None + return str(value[0]) + + def _get_literal_input(self, inputs: Dict[str, Any], key: str) -> Any: + value = inputs.get(key) + if self._is_connection(value): + return None + return value + + def _get_first_literal(self, inputs: Dict[str, Any], keys: Tuple[str, ...]) -> Any: + for key in keys: + value = self._get_literal_input(inputs, key) + if value is not None: + return value + return None + + def _bfs_upstream(self, start_ids: List[str], match_fn) -> Optional[str]: + queue = list(start_ids) + visited: Set[str] = set() + while queue: + node_id = queue.pop(0) + if node_id in visited: + continue + visited.add(node_id) + node = self._get_node(node_id) + if not node: + continue + if match_fn(node): + return node_id + for input_val in node.get("inputs", {}).values(): + conn_id = self._get_connection_node_id(input_val) + if conn_id and conn_id not in visited: + queue.append(conn_id) + return None + + @staticmethod + def _coerce_int(value: Any) -> Optional[int]: + if value is None: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + @staticmethod + def _coerce_float(value: Any) -> Optional[float]: + if value is None: + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + @classmethod + def _is_lora_node(cls, class_type: str) -> bool: + if class_type in cls.LORA_NODES: + return True + return "lora" in class_type.lower()