feat: enhance metadata handling with generation type and source image extraction

tests/test_metadata_utils.py

@@ -146,6 +146,8 @@ def test_build_imh_metadata_includes_workflow_and_prompt():
         "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"}}}
 
@@ -154,6 +156,8 @@ def test_build_imh_metadata_includes_workflow_and_prompt():
     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():

tests/test_workflow_extractor.py

@@ -68,3 +68,59 @@ def test_workflow_extractor_lora_detection():
 
     assert "loras" not in missing
     assert data["lora_list"] == [{"name": "detail.safetensors", "weight": 0.8}]
+
+
+
+def test_workflow_extractor_detects_img2img_lineage():
+    prompt = {
+        "1": {
+            "class_type": "LoadImage",
+            "inputs": {"image": "inputs/base.png"},
+        },
+        "2": {
+            "class_type": "VAEEncode",
+            "inputs": {"pixels": ["1", 0], "vae": ["5", 2]},
+        },
+        "3": {
+            "class_type": "KSampler",
+            "inputs": {
+                "seed": 123,
+                "steps": 20,
+                "cfg": 7.5,
+                "sampler_name": "euler",
+                "scheduler": "normal",
+                "denoise": 0.45,
+                "model": ["5", 0],
+                "positive": ["6", 0],
+                "negative": ["7", 0],
+                "latent_image": ["2", 0],
+            },
+        },
+        "4": {
+            "class_type": "VAEDecode",
+            "inputs": {"samples": ["3", 0], "vae": ["5", 2]},
+        },
+        "5": {
+            "class_type": "CheckpointLoaderSimple",
+            "inputs": {"ckpt_name": "model.safetensors"},
+        },
+        "6": {
+            "class_type": "CLIPTextEncode",
+            "inputs": {"text": "positive", "clip": ["5", 1]},
+        },
+        "7": {
+            "class_type": "CLIPTextEncode",
+            "inputs": {"text": "negative", "clip": ["5", 1]},
+        },
+        "8": {
+            "class_type": "MetaHubSaveNode",
+            "inputs": {"images": ["4", 0]},
+        },
+    }
+
+    extractor = WorkflowExtractor(prompt)
+    data, _missing = extractor.extract(save_node_id="8")
+
+    assert data["generation_type"] == "img2img"
+    assert data["source_image"]["fileName"] == "base.png"
+    assert data["source_image"]["relativePath"] == "inputs/base.png"