diff --git a/metadata_utils.py b/metadata_utils.py index fd40b85..2f1538e 100644 --- a/metadata_utils.py +++ b/metadata_utils.py @@ -554,6 +554,82 @@ def get_workflow_json(extra_pnginfo: Optional[dict]) -> dict: return {} +def ensure_prompt_in_workflow(workflow_json: dict, prompt_data: Optional[dict]) -> dict: + """ + Ensures workflow JSON includes prompt data when extra_pnginfo is missing it. + """ + if not isinstance(workflow_json, dict): + workflow_json = {} + + if isinstance(prompt_data, dict) and prompt_data: + existing_prompt = workflow_json.get("prompt") + if not isinstance(existing_prompt, dict) or not existing_prompt: + workflow_json["prompt"] = prompt_data + + return workflow_json + + +def ensure_metahub_save_node(workflow_json: dict, save_node_id: Optional[str]) -> None: + """ + Ensures saved workflow keeps MetaHubSaveNode instead of SaveImage. + """ + if not isinstance(workflow_json, dict): + return + + target_id = str(save_node_id) if save_node_id is not None else None + + prompt = workflow_json.get("prompt") + if isinstance(prompt, dict): + if target_id and target_id in prompt: + node = prompt.get(target_id) + if isinstance(node, dict): + node["class_type"] = "MetaHubSaveNode" + else: + save_nodes = [ + node_id + for node_id, node in prompt.items() + if isinstance(node, dict) and node.get("class_type") == "SaveImage" + ] + if len(save_nodes) == 1: + prompt[save_nodes[0]]["class_type"] = "MetaHubSaveNode" + + workflow = workflow_json.get("workflow") + if not isinstance(workflow, dict): + return + nodes = workflow.get("nodes") + if not isinstance(nodes, list): + return + + def update_workflow_node(node: dict) -> None: + if "type" in node: + node["type"] = "MetaHubSaveNode" + if "class_type" in node: + node["class_type"] = "MetaHubSaveNode" + if node.get("title") in ("Save Image", "SaveImage"): + node["title"] = "MetaHub Save Image" + props = node.get("properties") + if isinstance(props, dict) and props.get("node_name") in ("SaveImage", "Save Image"): + props["node_name"] = "MetaHub Save Image" + + if target_id: + for node in nodes: + if not isinstance(node, dict): + continue + node_id = node.get("id") + if node_id is not None and str(node_id) == target_id: + update_workflow_node(node) + return + else: + save_nodes = [ + node + for node in nodes + if isinstance(node, dict) + and (node.get("type") == "SaveImage" or node.get("class_type") == "SaveImage") + ] + if len(save_nodes) == 1: + update_workflow_node(save_nodes[0]) + + # ============================================================================ # PNG CHUNK INJECTION # ============================================================================ diff --git a/metahub_save_node.py b/metahub_save_node.py index 8fb3def..e14e5c5 100644 --- a/metahub_save_node.py +++ b/metahub_save_node.py @@ -194,9 +194,13 @@ class MetaHubSaveNode: if not isinstance(prompt_data, dict): prompt_data = {} + 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 + utils.ensure_metahub_save_node(workflow_json, save_node_id) + extractor = WorkflowExtractor(prompt_data) extracted, missing_fields = extractor.extract( - save_node_id=str(unique_id) if unique_id is not None else None + save_node_id=save_node_id ) lora_list = extracted.get("lora_list") or utils.extract_loras_from_workflow(workflow_json)