feat: add JSON serialization for metadata and enhance workflow handling in save functions

metadata_utils.py

@@ -449,6 +449,17 @@ def _get_workflow_prompt_texts(imh_metadata: dict) -> Tuple[Optional[str], Optio
     return workflow_text, prompt_text
 
 
+def _serialize_metadata_json_ascii(value: Any) -> Optional[str]:
+    if value is None:
+        return None
+    if isinstance(value, str):
+        return value
+    try:
+        return json.dumps(value, ensure_ascii=True)
+    except Exception:
+        return None
+
+
 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:
@@ -547,10 +558,25 @@ def get_workflow_json(extra_pnginfo: Optional[dict]) -> dict:
 
     # extra_pnginfo can be a dict or list
     if isinstance(extra_pnginfo, list) and len(extra_pnginfo) > 0:
-        return extra_pnginfo[0]
-    elif isinstance(extra_pnginfo, dict):
+        extra_pnginfo = extra_pnginfo[0]
+    if not isinstance(extra_pnginfo, dict):
+        return {}
+
+    if isinstance(extra_pnginfo.get("extra_pnginfo"), dict):
+        extra_pnginfo = extra_pnginfo["extra_pnginfo"]
+
+    if "workflow" in extra_pnginfo or "prompt" in extra_pnginfo:
         return extra_pnginfo
 
+    if "nodes" in extra_pnginfo and "links" in extra_pnginfo:
+        return {"workflow": extra_pnginfo}
+
+    if extra_pnginfo and all(
+        isinstance(value, dict) and "class_type" in value
+        for value in extra_pnginfo.values()
+    ):
+        return {"prompt": extra_pnginfo}
+
     return {}
 
 
@@ -579,6 +605,13 @@ def ensure_metahub_save_node(workflow_json: dict, save_node_id: Optional[str]) -
     target_id = str(save_node_id) if save_node_id is not None else None
 
     prompt = workflow_json.get("prompt")
+    if isinstance(prompt, str):
+        try:
+            prompt = json.loads(prompt)
+        except Exception:
+            prompt = None
+        else:
+            workflow_json["prompt"] = prompt
     if isinstance(prompt, dict):
         if target_id and target_id in prompt:
             node = prompt.get(target_id)
@@ -594,6 +627,13 @@ def ensure_metahub_save_node(workflow_json: dict, save_node_id: Optional[str]) -
                 prompt[save_nodes[0]]["class_type"] = "MetaHubSaveNode"
 
     workflow = workflow_json.get("workflow")
+    if isinstance(workflow, str):
+        try:
+            workflow = json.loads(workflow)
+        except Exception:
+            workflow = None
+        else:
+            workflow_json["workflow"] = workflow
     if not isinstance(workflow, dict):
         return
     nodes = workflow.get("nodes")
@@ -668,11 +708,14 @@ 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)
+        workflow_text = _serialize_metadata_json_ascii(imh_metadata.get("workflow"))
+        prompt_text = _serialize_metadata_json_ascii(
+            imh_metadata.get("prompt_api") or imh_metadata.get("prompt")
+        )
         if workflow_text:
-            png_info.add_itxt("workflow", workflow_text)
+            png_info.add_text("workflow", workflow_text)
         if prompt_text:
-            png_info.add_itxt("prompt", prompt_text)
+            png_info.add_text("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}")

tests/test_metadata_utils.py

@@ -1,3 +1,4 @@
+import json
 import copy
 
 from metadata_utils import (
@@ -5,7 +6,10 @@ from metadata_utils import (
     build_imh_metadata,
     ensure_metahub_save_node,
     ensure_prompt_in_workflow,
+    get_workflow_json,
+    save_png_with_metadata,
 )
+from PIL import Image
 
 
 def test_ensure_prompt_in_workflow_adds_prompt():
@@ -79,6 +83,14 @@ def test_ensure_metahub_save_node_updates_workflow_node_by_id():
     assert node["properties"]["node_name"] == "MetaHub Save Image"
 
 
+def test_get_workflow_json_wraps_raw_workflow():
+    extra_pnginfo = {"nodes": [], "links": []}
+
+    result = get_workflow_json(extra_pnginfo)
+
+    assert result == {"workflow": extra_pnginfo}
+
+
 def test_build_imh_metadata_includes_workflow_and_prompt():
     params = {
         "positive": "pos",
@@ -128,3 +140,25 @@ def test_build_a1111_metadata_formats_text():
     assert "Size: 512x768" in text
     assert "Model: model.safetensors" in text
     assert "Model hash: abcdef1234" in text
+
+
+def test_save_png_with_metadata_writes_workflow_prompt(tmp_path):
+    imh_metadata = {
+        "workflow": {"nodes": []},
+        "prompt_api": {"1": {"class_type": "KSampler"}},
+    }
+    image = Image.new("RGB", (2, 2), color=(0, 0, 0))
+    file_path = tmp_path / "test.png"
+
+    save_png_with_metadata(image, str(file_path), "params", imh_metadata)
+
+    loaded = Image.open(file_path)
+    text = getattr(loaded, "text", {})
+
+    workflow_text = text.get("workflow")
+    prompt_text = text.get("prompt")
+
+    assert workflow_text
+    assert prompt_text
+    assert json.loads(workflow_text) == imh_metadata["workflow"]
+    assert json.loads(prompt_text) == imh_metadata["prompt_api"]