diff --git a/__pycache__/metahub_save_video_node.cpython-310.pyc b/__pycache__/metahub_save_video_node.cpython-310.pyc index 8382beb..14ca554 100644 Binary files a/__pycache__/metahub_save_video_node.cpython-310.pyc and b/__pycache__/metahub_save_video_node.cpython-310.pyc differ diff --git a/__pycache__/video_metadata_utils.cpython-310.pyc b/__pycache__/video_metadata_utils.cpython-310.pyc index ddf9644..c7c2e2b 100644 Binary files a/__pycache__/video_metadata_utils.cpython-310.pyc and b/__pycache__/video_metadata_utils.cpython-310.pyc differ diff --git a/metadata_utils.py b/metadata_utils.py index 2e933d9..4b6b739 100644 --- a/metadata_utils.py +++ b/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'], diff --git a/metahub_save_node.py b/metahub_save_node.py index f3cc327..d1027a0 100644 --- a/metahub_save_node.py +++ b/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"), diff --git a/tests/test_metadata_utils.py b/tests/test_metadata_utils.py index 3c559db..6218bf2 100644 --- a/tests/test_metadata_utils.py +++ b/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(): diff --git a/tests/test_workflow_extractor.py b/tests/test_workflow_extractor.py index 9a5470b..51e0827 100644 --- a/tests/test_workflow_extractor.py +++ b/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" diff --git a/workflow_extractor.py b/workflow_extractor.py index 26d66f2..bbd12e8 100644 --- a/workflow_extractor.py +++ b/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: