279 lines
8.1 KiB
Python
279 lines
8.1 KiB
Python
import json
|
|
import math
|
|
from pathlib import Path
|
|
import copy
|
|
|
|
from metadata_utils import (
|
|
build_a1111_metadata,
|
|
build_imh_metadata,
|
|
ensure_metahub_save_node,
|
|
ensure_prompt_in_workflow,
|
|
get_workflow_json,
|
|
save_png_with_metadata,
|
|
)
|
|
from PIL import Image
|
|
|
|
|
|
def _read_png_text_chunks(path: Path) -> dict:
|
|
data = path.read_bytes()
|
|
if data[:8] != b"\x89PNG\r\n\x1a\n":
|
|
raise AssertionError("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 _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": {}}}
|
|
|
|
result = ensure_prompt_in_workflow(workflow, prompt)
|
|
|
|
assert result["prompt"] == prompt
|
|
|
|
|
|
def test_ensure_prompt_in_workflow_preserves_existing_prompt():
|
|
existing = {"9": {"class_type": "SaveImage", "inputs": {}}}
|
|
workflow = {"prompt": existing}
|
|
prompt = {"1": {"class_type": "KSampler", "inputs": {}}}
|
|
|
|
result = ensure_prompt_in_workflow(workflow, prompt)
|
|
|
|
assert result["prompt"] is existing
|
|
|
|
|
|
def test_ensure_metahub_save_node_updates_prompt_by_id():
|
|
workflow = {"prompt": {"7": {"class_type": "SaveImage", "inputs": {}}}}
|
|
|
|
ensure_metahub_save_node(workflow, "7")
|
|
|
|
assert workflow["prompt"]["7"]["class_type"] == "MetaHubSaveNode"
|
|
|
|
|
|
def test_ensure_metahub_save_node_updates_single_prompt_saveimage():
|
|
workflow = {"prompt": {"1": {"class_type": "SaveImage", "inputs": {}}}}
|
|
|
|
ensure_metahub_save_node(workflow, None)
|
|
|
|
assert workflow["prompt"]["1"]["class_type"] == "MetaHubSaveNode"
|
|
|
|
|
|
def test_ensure_metahub_save_node_skips_multiple_prompt_saveimage():
|
|
workflow = {
|
|
"prompt": {
|
|
"1": {"class_type": "SaveImage", "inputs": {}},
|
|
"2": {"class_type": "SaveImage", "inputs": {}},
|
|
}
|
|
}
|
|
|
|
ensure_metahub_save_node(workflow, None)
|
|
|
|
assert workflow["prompt"]["1"]["class_type"] == "SaveImage"
|
|
assert workflow["prompt"]["2"]["class_type"] == "SaveImage"
|
|
|
|
|
|
def test_ensure_metahub_save_node_updates_workflow_node_by_id():
|
|
workflow = {
|
|
"workflow": {
|
|
"nodes": [
|
|
{
|
|
"id": 5,
|
|
"type": "SaveImage",
|
|
"title": "Save Image",
|
|
"properties": {"node_name": "Save Image"},
|
|
}
|
|
]
|
|
}
|
|
}
|
|
|
|
ensure_metahub_save_node(workflow, "5")
|
|
|
|
node = workflow["workflow"]["nodes"][0]
|
|
assert node["type"] == "MetaHubSaveNode"
|
|
assert node["title"] == "MetaHub Save Image"
|
|
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_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",
|
|
"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": [],
|
|
"generation_type": "img2img",
|
|
"source_image": {"fileName": "base.png", "relativePath": "inputs/base.png"},
|
|
}
|
|
workflow_json = {"workflow": {"nodes": []}, "prompt": {"1": {"class_type": "KSampler"}}}
|
|
|
|
imh = build_imh_metadata(params, workflow_json)
|
|
|
|
assert imh["generator"] == "ComfyUI"
|
|
assert imh["workflow"] == workflow_json["workflow"]
|
|
assert imh["prompt_api"] == workflow_json["prompt"]
|
|
assert imh["generation_type"] == "img2img"
|
|
assert imh["source_image"]["fileName"] == "base.png"
|
|
|
|
|
|
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",
|
|
"negative": "neg",
|
|
"steps": 20,
|
|
"sampler": "euler",
|
|
"cfg": 7.5,
|
|
"seed": 123,
|
|
"width": 512,
|
|
"height": 768,
|
|
"model_name": "model.safetensors",
|
|
"model_hash": "abcdef1234",
|
|
"lora_hashes": {},
|
|
}
|
|
|
|
text = build_a1111_metadata(params)
|
|
|
|
assert "Negative prompt: neg" in text
|
|
assert "Steps: 20" in text
|
|
assert "Sampler: euler" in text
|
|
assert "CFG scale: 7.5" in text
|
|
assert "Seed: 123" in 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"]
|
|
|
|
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"
|