diff --git a/metadata_utils.py b/metadata_utils.py index 4b6b739..7d4eae7 100644 --- a/metadata_utils.py +++ b/metadata_utils.py @@ -397,6 +397,7 @@ def build_imh_metadata(params: dict, workflow_json: dict) -> dict: "vae": params.get('vae_name', ''), "denoise": params.get('denoise', 1.0), "generation_type": params.get('generation_type'), + "parent_image": _sanitize(params.get('parent_image')), "source_image": _sanitize(params.get('source_image')), "width": params['width'], "height": params['height'], diff --git a/metahub_save_node.py b/metahub_save_node.py index d1027a0..0560bc7 100644 --- a/metahub_save_node.py +++ b/metahub_save_node.py @@ -198,6 +198,9 @@ class MetaHubSaveNode: prompt_data = prompt if isinstance(prompt, dict) else workflow_json.get("prompt", {}) if not isinstance(prompt_data, dict): prompt_data = {} + parent_image = None + if isinstance(extra_pnginfo, dict): + parent_image = extra_pnginfo.get("parent_image") workflow_json = utils.ensure_prompt_in_workflow(workflow_json, prompt_data) save_node_id = str(unique_id) if unique_id is not None else None @@ -309,6 +312,7 @@ class MetaHubSaveNode: "project_name": project_name, "lora_list": lora_list, "lora_hashes": lora_hashes, + "parent_image": parent_image, "generation_type": extracted.get("generation_type"), "source_image": extracted.get("source_image"), # Performance metrics (Tier 1, 2, 3) diff --git a/tests/test_metadata_utils.py b/tests/test_metadata_utils.py index 6218bf2..80eb766 100644 --- a/tests/test_metadata_utils.py +++ b/tests/test_metadata_utils.py @@ -147,6 +147,7 @@ def test_build_imh_metadata_includes_workflow_and_prompt(): "height": 768, "lora_list": [], "generation_type": "img2img", + "parent_image": {"fileName": "selected.png", "relativePath": "library/selected.png"}, "source_image": {"fileName": "base.png", "relativePath": "inputs/base.png"}, } workflow_json = {"workflow": {"nodes": []}, "prompt": {"1": {"class_type": "KSampler"}}} @@ -157,6 +158,7 @@ def test_build_imh_metadata_includes_workflow_and_prompt(): assert imh["workflow"] == workflow_json["workflow"] assert imh["prompt_api"] == workflow_json["prompt"] assert imh["generation_type"] == "img2img" + assert imh["parent_image"]["fileName"] == "selected.png" assert imh["source_image"]["fileName"] == "base.png"