nodes/LuqP2-ImageMetaHub-ComfyUI-Save
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add workflow auto-extraction and multi-format metadata saving

Browse Source 
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.
This commit is contained in:
LuqP2
2025-12-20 21:16:27 -03:00
parent 6923df5009
commit f6bdd59fdc
5 changed files with 1034 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+47 -25
View File
@@ -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
{
metadata_utils.py
+290 -29
View File
@@ -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
}
}
metahub_save_node.py
+276 -81
View File
@@ -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}
requirements.txt
+1
View File
@@ -1,2 +1,3 @@
Pillow>=10.0.0
numpy>=1.24.0
piexif>=1.1.3
workflow_extractor.py
+420
View File
@@ -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()