diff --git a/tests/smoke_workflow.py b/tests/smoke_workflow.py new file mode 100644 index 0000000..45115cd --- /dev/null +++ b/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()) diff --git a/tests/test_metadata_utils.py b/tests/test_metadata_utils.py index bb38b0b..3c559db 100644 --- a/tests/test_metadata_utils.py +++ b/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"