86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
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())
|