diff --git a/metadata_utils.py b/metadata_utils.py index 2f1538e..b3a8aa6 100644 --- a/metadata_utils.py +++ b/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}") diff --git a/tests/test_metadata_utils.py b/tests/test_metadata_utils.py index c7a2cf1..f15e48e 100644 --- a/tests/test_metadata_utils.py +++ b/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"]