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

metadata_utils.py

@@ -396,6 +396,8 @@ def build_imh_metadata(params: dict, workflow_json: dict) -> dict:
         "model_hash": params['model_hash'],
         "vae": params.get('vae_name', ''),
         "denoise": params.get('denoise', 1.0),
+        "generation_type": params.get('generation_type'),
+        "source_image": _sanitize(params.get('source_image')),
         "width": params['width'],
         "height": params['height'],
 

metahub_save_node.py

@@ -309,6 +309,8 @@ class MetaHubSaveNode:
                 "project_name": project_name,
                 "lora_list": lora_list,
                 "lora_hashes": lora_hashes,
+                "generation_type": extracted.get("generation_type"),
+                "source_image": extracted.get("source_image"),
                 # Performance metrics (Tier 1, 2, 3)
                 "vram_peak_mb": vram_peak_mb if vram_peak_mb is not None else gpu_metrics.get("vram_peak_mb"),
                 "gpu_device": gpu_device_override if gpu_device_override else gpu_metrics.get("gpu_device"),

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"

workflow_extractor.py

@@ -54,6 +54,11 @@ class WorkflowExtractor:
         "VAEDecode (Tiled)",
     ]
 
+    LOAD_IMAGE_NODES = [
+        "LoadImage",
+        "LoadImageMask",
+    ]
+
     def __init__(self, prompt: Optional[Dict[str, Any]]):
         self.prompt: Dict[str, Any] = {}
         if isinstance(prompt, dict):
@@ -116,6 +121,10 @@ class WorkflowExtractor:
         elif has_lora_nodes:
             missing.add("loras")
 
+        lineage = self.extract_lineage(sampler_node_id)
+        if lineage:
+            data.update(lineage)
+
         return data, missing
 
     def find_sampler_for_save_node(self, save_node_id: Optional[str]) -> Optional[str]:
@@ -236,6 +245,73 @@ class WorkflowExtractor:
 
         return loras, has_lora_nodes
 
+    def extract_lineage(self, sampler_node_id: Optional[str]) -> Dict[str, Any]:
+        if not sampler_node_id:
+            return {}
+
+        node = self._get_node(sampler_node_id)
+        if not node:
+            return {}
+
+        inputs = node.get("inputs", {})
+        latent_conn = inputs.get("latent_image") or inputs.get("latent") or inputs.get("samples")
+        start_node_id = self._get_connection_node_id(latent_conn)
+        if not start_node_id:
+            return {}
+
+        queue = [start_node_id]
+        visited: Set[str] = set()
+        source_image: Optional[Dict[str, Any]] = None
+        has_load_image = False
+        has_inpaint = False
+        has_outpaint = False
+
+        while queue:
+            current_id = queue.pop(0)
+            if current_id in visited:
+                continue
+            visited.add(current_id)
+
+            current = self._get_node(current_id)
+            if not current:
+                continue
+
+            class_type = self._class_type(current)
+            class_lower = class_type.lower()
+
+            if class_type in self.LOAD_IMAGE_NODES or class_lower in {"loadimage", "loadimagemask"}:
+                has_load_image = True
+                if not source_image:
+                    image_value = current.get("inputs", {}).get("image")
+                    if isinstance(image_value, str) and image_value.strip():
+                        normalized = image_value.replace("\\", "/")
+                        source_image = {
+                            "fileName": normalized.split("/")[-1],
+                            "relativePath": normalized,
+                            "nodeId": current_id,
+                            "nodeType": class_type,
+                        }
+
+            if "outpaint" in class_lower or class_lower in {"padforoutpaint", "imagepadforoutpaint"}:
+                has_outpaint = True
+
+            if "inpaint" in class_lower or class_lower == "setlatentnoisemask" or "mask" in class_lower:
+                has_inpaint = True
+
+            for input_value in current.get("inputs", {}).values():
+                upstream_id = self._get_connection_node_id(input_value)
+                if upstream_id and upstream_id not in visited:
+                    queue.append(upstream_id)
+
+        if not has_load_image:
+            return {}
+
+        generation_type = "outpaint" if has_outpaint else "inpaint" if has_inpaint else "img2img"
+        return {
+            "generation_type": generation_type,
+            "source_image": source_image,
+        }
+
     def _find_sampler_from_images_source(self, start_node_id: str) -> Optional[str]:
         start_node = self._get_node(start_node_id)
         if not start_node: