106 lines
2.7 KiB
Python
106 lines
2.7 KiB
Python
"""
|
|
Helpers for constructing and embedding PNG metadata.
|
|
|
|
Provides:
|
|
- build_iteration_metadata(): creates the metadata dict for one iteration
|
|
- embed_in_pnginfo(): writes metadata into PIL PngInfo object
|
|
- format_a1111_params(): formats params in A1111-compatible style
|
|
"""
|
|
|
|
import json
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
|
|
def build_iteration_metadata(
|
|
batch_name: str,
|
|
batch_id: str,
|
|
index: int,
|
|
total: int,
|
|
params: dict,
|
|
mode: str = "matrix",
|
|
) -> dict:
|
|
"""Build a complete metadata dictionary for one iteration."""
|
|
return {
|
|
"batch_name": batch_name,
|
|
"batch_id": batch_id,
|
|
"index": index,
|
|
"total": total,
|
|
"params": params,
|
|
"mode": mode,
|
|
}
|
|
|
|
|
|
def embed_in_pnginfo(
|
|
pnginfo: PngInfo,
|
|
iteration_meta: dict,
|
|
workflow: dict | None = None,
|
|
prompt: dict | None = None,
|
|
) -> PngInfo:
|
|
"""
|
|
Embed metadata into a PngInfo object.
|
|
|
|
Adds:
|
|
- "workflow" key (standard ComfyUI workflow JSON)
|
|
- "prompt" key (standard ComfyUI prompt JSON)
|
|
- "wi_iteration" key (our custom iteration metadata)
|
|
- "parameters" key (A1111-compatible parameter string)
|
|
"""
|
|
if workflow is not None:
|
|
pnginfo.add_text("workflow", json.dumps(workflow))
|
|
|
|
if prompt is not None:
|
|
pnginfo.add_text("prompt", json.dumps(prompt))
|
|
|
|
pnginfo.add_text("wi_iteration", json.dumps(iteration_meta))
|
|
|
|
params = iteration_meta.get("params", {})
|
|
if params:
|
|
pnginfo.add_text("parameters", format_a1111_params(params))
|
|
|
|
return pnginfo
|
|
|
|
|
|
def format_a1111_params(params: dict) -> str:
|
|
"""
|
|
Format parameters in A1111-compatible string format.
|
|
|
|
Example output: "Steps: 20, CFG scale: 7.5, Sampler: euler"
|
|
"""
|
|
# Map common parameter names to A1111-style names
|
|
name_map = {
|
|
"steps": "Steps",
|
|
"cfg": "CFG scale",
|
|
"sampler": "Sampler",
|
|
"sampler_name": "Sampler",
|
|
"scheduler": "Scheduler",
|
|
"seed": "Seed",
|
|
"denoise": "Denoising strength",
|
|
"width": "Size-1",
|
|
"height": "Size-2",
|
|
}
|
|
|
|
parts = []
|
|
for key, value in params.items():
|
|
display_name = name_map.get(key.lower(), key.replace("_", " ").title())
|
|
parts.append(f"{display_name}: {value}")
|
|
|
|
return ", ".join(parts)
|
|
|
|
|
|
def extract_iteration_metadata(png_info_dict: dict) -> dict | None:
|
|
"""
|
|
Extract wi_iteration metadata from a PNG info dictionary.
|
|
Returns None if not found.
|
|
"""
|
|
wi_data = png_info_dict.get("wi_iteration")
|
|
if wi_data is None:
|
|
return None
|
|
|
|
if isinstance(wi_data, str):
|
|
try:
|
|
return json.loads(wi_data)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
|
|
return wi_data
|