diff --git a/metadata_utils.py b/metadata_utils.py
index 69e2072..fd40b85 100644
--- a/metadata_utils.py
+++ b/metadata_utils.py
@@ -8,7 +8,8 @@ import json
import os
from datetime import datetime
from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
+from xml.sax.saxutils import escape as xml_escape
import numpy as np
from PIL import Image, PngImagePlugin
@@ -427,6 +428,52 @@ def build_imh_metadata(params: dict, workflow_json: dict) -> dict:
}
+def _serialize_metadata_json(value: Any) -> Optional[str]:
+ if value is None:
+ return None
+ if isinstance(value, str):
+ return value
+ try:
+ return json.dumps(value, ensure_ascii=False)
+ except Exception:
+ return None
+
+
+def _get_workflow_prompt_texts(imh_metadata: dict) -> Tuple[Optional[str], Optional[str]]:
+ if not imh_metadata:
+ return None, None
+ workflow_text = _serialize_metadata_json(imh_metadata.get("workflow"))
+ prompt_text = _serialize_metadata_json(
+ imh_metadata.get("prompt_api") or imh_metadata.get("prompt")
+ )
+ return workflow_text, prompt_text
+
+
+def build_comfyui_xmp_packet(imh_metadata: dict) -> Optional[bytes]:
+ workflow_text, prompt_text = _get_workflow_prompt_texts(imh_metadata)
+ if not workflow_text and not prompt_text:
+ return None
+
+ fields = []
+ if workflow_text:
+ fields.append(f"{xml_escape(workflow_text)}")
+ if prompt_text:
+ fields.append(f"{xml_escape(prompt_text)}")
+
+ xmp_payload = (
+ ''
+ ''
+ ''
+ ''
+ f"{''.join(fields)}"
+ ''
+ ''
+ ''
+ ''
+ )
+ return xmp_payload.encode("utf-8")
+
+
# ============================================================================
# WORKFLOW PARSING
# ============================================================================
@@ -545,6 +592,11 @@ def save_png_with_metadata(image: Image.Image, image_path: str, a1111_metadata:
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)
+ workflow_text, prompt_text = _get_workflow_prompt_texts(imh_metadata)
+ if workflow_text:
+ png_info.add_itxt("workflow", workflow_text)
+ if prompt_text:
+ png_info.add_itxt("prompt", prompt_text)
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}")
@@ -582,12 +634,22 @@ def save_jpeg_with_metadata(
if image.mode in ("RGBA", "LA", "P"):
image = image.convert("RGB")
save_kwargs: Dict[str, Any] = {"format": "JPEG", "quality": quality}
+ xmp_bytes = build_comfyui_xmp_packet(imh_metadata)
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)
+ if xmp_bytes:
+ save_kwargs["xmp"] = xmp_bytes
+ try:
+ image.save(image_path, **save_kwargs)
+ except TypeError:
+ if "xmp" in save_kwargs:
+ save_kwargs.pop("xmp", None)
+ image.save(image_path, **save_kwargs)
+ else:
+ raise
except Exception as e:
print(f"[MetaHub] Warning: JPEG metadata save failed for {image_path}: {e}")
image.save(image_path, "JPEG", quality=quality)
@@ -605,12 +667,22 @@ def save_webp_with_metadata(
"""
try:
save_kwargs: Dict[str, Any] = {"format": "WEBP", "quality": quality}
+ xmp_bytes = build_comfyui_xmp_packet(imh_metadata)
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)
+ if xmp_bytes:
+ save_kwargs["xmp"] = xmp_bytes
+ try:
+ image.save(image_path, **save_kwargs)
+ except TypeError:
+ if "xmp" in save_kwargs:
+ save_kwargs.pop("xmp", None)
+ image.save(image_path, **save_kwargs)
+ else:
+ raise
except Exception as e:
print(f"[MetaHub] Warning: WebP metadata save failed for {image_path}: {e}")
image.save(image_path, "WEBP", quality=quality)