feat: metadata tests

tests/smoke_workflow.py

@@ -0,0 +1,85 @@
+import json
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from PIL import Image
+
+from metadata_utils import build_imh_metadata, ensure_metahub_save_node, save_png_with_metadata
+
+
+def _read_png_text_chunks(path: Path) -> dict:
+    data = path.read_bytes()
+    if data[:8] != b"\x89PNG\r\n\x1a\n":
+        raise RuntimeError("Not a PNG file")
+    offset = 8
+    chunks = {}
+    while offset + 8 <= len(data):
+        length = int.from_bytes(data[offset:offset + 4], "big")
+        chunk_type = data[offset + 4:offset + 8]
+        chunk_data = data[offset + 8:offset + 8 + length]
+        offset += 12 + length
+        if chunk_type == b"tEXt":
+            if b"\x00" in chunk_data:
+                keyword, text = chunk_data.split(b"\x00", 1)
+                chunks[keyword.decode("latin-1")] = text.decode("latin-1")
+    return chunks
+
+
+def main() -> int:
+    output = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("smoke_workflow.png")
+
+    params = {
+        "positive": "smoke",
+        "negative": "",
+        "seed": 123,
+        "steps": 20,
+        "cfg": 7.5,
+        "sampler": "euler",
+        "scheduler": "normal",
+        "model_name": "model.safetensors",
+        "model_hash": "abcdef1234",
+        "width": 256,
+        "height": 256,
+        "lora_list": [],
+    }
+    workflow_json = {
+        "workflow": {"nodes": [{"id": 7, "type": "SaveImage", "title": "Save Image"}], "links": []},
+        "prompt": {"7": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}}},
+    }
+
+    ensure_metahub_save_node(workflow_json, "7")
+    imh_metadata = build_imh_metadata(params, workflow_json)
+
+    image = Image.new("RGB", (2, 2), color=(0, 0, 0))
+    save_png_with_metadata(image, str(output), "params", imh_metadata)
+
+    chunks = _read_png_text_chunks(output)
+    if "workflow" not in chunks or "prompt" not in chunks:
+        print("[FAIL] Missing tEXt workflow/prompt chunks")
+        return 1
+
+    try:
+        workflow = json.loads(chunks["workflow"])
+        prompt = json.loads(chunks["prompt"])
+    except json.JSONDecodeError as exc:
+        print(f"[FAIL] Invalid JSON in chunks: {exc}")
+        return 1
+
+    if workflow.get("nodes", [{}])[0].get("type") != "MetaHubSaveNode":
+        print("[FAIL] workflow save node is not MetaHubSaveNode")
+        return 1
+    if prompt.get("7", {}).get("class_type") != "MetaHubSaveNode":
+        print("[FAIL] prompt save node is not MetaHubSaveNode")
+        return 1
+
+    print("[OK] tEXt workflow/prompt present and MetaHubSaveNode preserved")
+    print(f"[OK] Wrote {output}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/test_metadata_utils.py

@@ -1,4 +1,5 @@
 import json
+import math
 from pathlib import Path
 import copy
 
@@ -31,6 +32,11 @@ def _read_png_text_chunks(path: Path) -> dict:
     return chunks
 
 
+def _save_simple_png(path: Path) -> None:
+    image = Image.new("RGB", (2, 2), color=(0, 0, 0))
+    image.save(path, "PNG")
+
+
 def test_ensure_prompt_in_workflow_adds_prompt():
     workflow = {"workflow": {"nodes": []}}
     prompt = {"1": {"class_type": "KSampler", "inputs": {}}}
@@ -110,6 +116,22 @@ def test_get_workflow_json_wraps_raw_workflow():
     assert result == {"workflow": extra_pnginfo}
 
 
+def test_get_workflow_json_accepts_extra_pnginfo_envelope():
+    extra_pnginfo = {"extra_pnginfo": {"workflow": {"nodes": []}}}
+
+    result = get_workflow_json(extra_pnginfo)
+
+    assert result == {"workflow": {"nodes": []}}
+
+
+def test_get_workflow_json_accepts_prompt_dict():
+    prompt = {"1": {"class_type": "KSampler", "inputs": {}}}
+
+    result = get_workflow_json(prompt)
+
+    assert result == {"prompt": prompt}
+
+
 def test_build_imh_metadata_includes_workflow_and_prompt():
     params = {
         "positive": "pos",
@@ -134,6 +156,32 @@ def test_build_imh_metadata_includes_workflow_and_prompt():
     assert imh["prompt_api"] == workflow_json["prompt"]
 
 
+def test_build_imh_metadata_sanitizes_nan():
+    params = {
+        "positive": "pos",
+        "negative": "neg",
+        "seed": 123,
+        "steps": 20,
+        "cfg": 7.5,
+        "sampler": "euler",
+        "scheduler": "normal",
+        "model_name": "model.safetensors",
+        "model_hash": "abcdef1234",
+        "width": 512,
+        "height": 768,
+        "lora_list": [],
+    }
+    workflow_json = {
+        "workflow": {"nodes": [{"id": 1, "widgets_values": [float("nan")]}]},
+        "prompt": {"1": {"class_type": "KSampler", "inputs": {"seed": float("nan")}}},
+    }
+
+    imh = build_imh_metadata(params, workflow_json)
+
+    assert imh["workflow"]["nodes"][0]["widgets_values"][0] is None
+    assert imh["prompt_api"]["1"]["inputs"]["seed"] is None
+
+
 def test_build_a1111_metadata_formats_text():
     params = {
         "positive": "pos",
@@ -185,3 +233,42 @@ def test_save_png_with_metadata_writes_workflow_prompt(tmp_path):
     chunk_text = _read_png_text_chunks(file_path)
     assert "workflow" in chunk_text
     assert "prompt" in chunk_text
+
+
+def test_save_png_with_metadata_roundtrip_for_comfyui(tmp_path):
+    params = {
+        "positive": "pos",
+        "negative": "neg",
+        "seed": 123,
+        "steps": 20,
+        "cfg": 7.5,
+        "sampler": "euler",
+        "scheduler": "normal",
+        "model_name": "model.safetensors",
+        "model_hash": "abcdef1234",
+        "width": 512,
+        "height": 768,
+        "lora_list": [],
+    }
+    workflow_json = {
+        "workflow": {"nodes": [{"id": 7, "type": "SaveImage", "title": "Save Image"}], "links": []},
+        "prompt": {"7": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}}},
+    }
+
+    ensure_metahub_save_node(workflow_json, "7")
+    imh_metadata = build_imh_metadata(params, workflow_json)
+
+    file_path = tmp_path / "roundtrip.png"
+    _save_simple_png(file_path)
+    image = Image.open(file_path)
+    save_png_with_metadata(image, str(file_path), "params", imh_metadata)
+
+    text_chunks = _read_png_text_chunks(file_path)
+    assert "workflow" in text_chunks
+    assert "prompt" in text_chunks
+
+    workflow_loaded = json.loads(text_chunks["workflow"])
+    prompt_loaded = json.loads(text_chunks["prompt"])
+
+    assert workflow_loaded["nodes"][0]["type"] == "MetaHubSaveNode"
+    assert prompt_loaded["7"]["class_type"] == "MetaHubSaveNode"