diff --git a/metadata_utils.py b/metadata_utils.py index b0aa175..0437956 100644 --- a/metadata_utils.py +++ b/metadata_utils.py @@ -3,7 +3,7 @@ Compatibility shim for MetaHub Save Node metadata helpers. The implementation is kept in metadata_utils_impl.py; this module exposes the same public API expected by the save nodes while pinning the published node -version to 1.2.1. +version to 1.2.2. """ try: @@ -11,7 +11,7 @@ try: except ImportError: import metadata_utils_impl as _impl -METAHUB_SAVE_NODE_VERSION = "1.2.1" +METAHUB_SAVE_NODE_VERSION = "1.2.2" _impl.METAHUB_SAVE_NODE_VERSION = METAHUB_SAVE_NODE_VERSION for _name in dir(_impl): diff --git a/metadata_utils_impl.py b/metadata_utils_impl.py index b77974f..50f042c 100644 --- a/metadata_utils_impl.py +++ b/metadata_utils_impl.py @@ -14,7 +14,7 @@ from xml.sax.saxutils import escape as xml_escape import numpy as np from PIL import Image, PngImagePlugin -METAHUB_SAVE_NODE_VERSION = "1.2.1" +METAHUB_SAVE_NODE_VERSION = "1.2.2" try: import piexif diff --git a/pyproject.toml b/pyproject.toml index fde97c3..48a62a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "imagemetahub-comfyui-save" description = "Official companion node for Image MetaHub. Saves A1111/Civitai-compatible parameters plus extended Image MetaHub metadata (workflow JSON, SHA256 model hashes, VRAM peak/total, GPU device, generation time, and steps/sec) for reproducibility and benchmarking." -version = "1.2.1" +version = "1.2.2" license = {file = "LICENSE"} # classifiers = [ # # For OS-independent nodes (works on all operating systems) diff --git a/tests/test_metadata_utils.py b/tests/test_metadata_utils.py index 658ebab..50765ed 100644 --- a/tests/test_metadata_utils.py +++ b/tests/test_metadata_utils.py @@ -334,7 +334,7 @@ def test_extract_workflow_attribution_from_node_properties(): assert attribution["token"] == "imhcrt_br_creator_workflow_v1_random" assert attribution["source"] == "metahub_save_node" - assert attribution["node_version"] == "1.2.1" + assert attribution["node_version"] == "1.2.2" def test_build_imh_metadata_includes_attribution_without_a1111_parameters(): diff --git a/tests/test_workflow_extractor.py b/tests/test_workflow_extractor.py index 7fb048c..5592277 100644 --- a/tests/test_workflow_extractor.py +++ b/tests/test_workflow_extractor.py @@ -1,6 +1,76 @@ from workflow_extractor import WorkflowExtractor +def test_qwen_image_21_uses_distinct_conditioning_outputs(): + prompt = { + "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "qwen_image_2.1.safetensors"}}, + "2": {"class_type": "String Literal", "inputs": {"string": "linked positive"}}, + "3": {"class_type": "TextEncodeQwenImage21", "inputs": { + "prompt": ["2", 0], "negative_prompt": "direct negative", + }}, + "4": {"class_type": "KSampler", "inputs": { + "model": ["1", 0], "positive": ["3", 0], "negative": ["3", 1], + "seed": 42, "steps": 25, "cfg": 1, + "sampler_name": "euler", "scheduler": "simple", "denoise": 1, + }}, + "5": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0]}}, + "6": {"class_type": "MetaHubSaveNode", "inputs": {"images": ["5", 0]}}, + } + + data, missing = WorkflowExtractor(prompt).extract(save_node_id="6") + + assert data["positive"] == "linked positive" + assert data["negative"] == "direct negative" + assert data["model_name"] == "qwen_image_2.1.safetensors" + assert (data["seed"], data["steps"], data["cfg"]) == (42, 25, 1.0) + assert (data["sampler_name"], data["scheduler"]) == ("euler", "simple") + assert "positive" not in missing and "negative" not in missing + + prompt["3"]["inputs"]["negative_prompt"] = "" + data, missing = WorkflowExtractor(prompt).extract(save_node_id="6") + assert data["negative"] == "" + assert "negative" not in missing + + prompt["3"]["inputs"]["negative_prompt"] = ["8", 0] + prompt["7"] = {"class_type": "String Literal", "inputs": {"string": ""}} + prompt["8"] = {"class_type": "ComfySwitchNode", "inputs": { + "switch": False, "on_false": ["7", 0], "on_true": ["9", 0], + }} + prompt["9"] = {"class_type": "PrimitiveStringMultiline", "inputs": { + "value": "inactive negative", + }} + data, missing = WorkflowExtractor(prompt).extract(save_node_id="6") + assert data["negative"] == "" + assert "negative" not in missing + + prompt["3"]["inputs"]["negative_prompt"] = ["7", 0] + data, missing = WorkflowExtractor(prompt).extract(save_node_id="6") + assert data["negative"] == "" + assert "negative" not in missing + + +def test_qwen_image_21_prompts_stay_separate_through_cfg_guider(): + prompt = { + "1": {"class_type": "TextEncodeQwenImage21", "inputs": { + "prompt": "positive text", "negative_prompt": "negative text", + }}, + "2": {"class_type": "CFGGuider", "inputs": { + "positive": ["1", 0], "negative": ["1", 1], "cfg": 1, + }}, + "3": {"class_type": "SamplerCustomAdvanced", "inputs": { + "guider": ["2", 0], "noise": ["4", 0], "sigmas": ["5", 0], + }}, + "4": {"class_type": "RandomNoise", "inputs": {"noise_seed": 42}}, + "5": {"class_type": "BasicScheduler", "inputs": {"steps": 25}}, + } + + data, missing = WorkflowExtractor(prompt).extract() + + assert data["positive"] == "positive text" + assert data["negative"] == "negative text" + assert "positive" not in missing and "negative" not in missing + + def test_workflow_extractor_basic_prompt(): prompt = { "1": { diff --git a/workflow_extractor.py b/workflow_extractor.py index 315bfa5..3dcb5be 100644 --- a/workflow_extractor.py +++ b/workflow_extractor.py @@ -48,6 +48,7 @@ class WorkflowExtractor: "CLIPTextEncodeSDXLPlus", "CLIPTextEncodeSDXLRefiner", "CLIPTextEncodeSD3", + "TextEncodeQwenImage21", ] VAE_DECODE_NODES = [ @@ -255,7 +256,15 @@ class WorkflowExtractor: negative_conn = inputs.get("negative") or inputs.get("negative_cond") positive = self._extract_text_from_connection(positive_conn) negative = self._extract_text_from_connection(negative_conn) - if positive is None: + guider_id = self._get_connection_node_id(inputs.get("guider")) + guider = self._get_node(guider_id) + guider_inputs = guider.get("inputs", {}) if guider else {} + if "positive" in guider_inputs or "negative" in guider_inputs: + if positive is None: + positive = self._extract_text_from_connection(guider_inputs.get("positive")) + if negative is None: + negative = self._extract_text_from_connection(guider_inputs.get("negative")) + elif positive is None: positive = self._extract_text_from_connection(inputs.get("guider")) if negative is None: negative = self._extract_text_from_connection(inputs.get("negative_conditioning")) @@ -518,24 +527,29 @@ class WorkflowExtractor: node_id = self._get_connection_node_id(conn) if not node_id: return None - return self._extract_text_from_node(node_id) + return self._extract_text_from_node(node_id, conn[1]) - def _extract_text_from_node(self, start_node_id: str) -> Optional[str]: + def _extract_text_from_node(self, start_node_id: str, output_slot: int = 0) -> Optional[str]: texts: List[str] = [] - queue = [start_node_id] - visited: Set[str] = set() + queue = [(start_node_id, output_slot)] + visited: Set[Tuple[str, int]] = set() while queue: - node_id = queue.pop(0) - if node_id in visited: + node_id, current_slot = queue.pop(0) + if (node_id, current_slot) in visited: continue - visited.add(node_id) + visited.add((node_id, current_slot)) node = self._get_node(node_id) if not node: continue class_type = self._class_type(node) if class_type == "ConditioningZeroOut": return "" + if class_type == "TextEncodeQwenImage21": + text = self._get_qwen_image_21_text(node, current_slot) + if text is not None: + texts.append(text) + continue if self._is_prompt_encoder(class_type): text = self._get_clip_text(node) if text is not None: @@ -543,13 +557,33 @@ class WorkflowExtractor: continue for input_val in node.get("inputs", {}).values(): conn_id = self._get_connection_node_id(input_val) - if conn_id and conn_id not in visited: - queue.append(conn_id) + if conn_id and (conn_id, input_val[1]) not in visited: + queue.append((conn_id, input_val[1])) if texts: return "\n".join(texts) return None + def _get_qwen_image_21_text(self, node: Dict[str, Any], output_slot: int) -> Optional[str]: + if output_slot not in (0, 1): + return None + key = "prompt" if output_slot == 0 else "negative_prompt" + inputs = node.get("inputs", {}) + value = self._get_literal_input(inputs, key) + if isinstance(value, str): + return value + conn = inputs.get(key) + source_id = self._get_connection_node_id(conn) + if source_id: + source = self._get_node(source_id) + if source: + for source_key in ("text", "string", "value", "prompt", "negative_prompt"): + source_value = self._get_literal_input(source.get("inputs", {}), source_key) + if isinstance(source_value, str): + return source_value + return self._resolve_string_from_connection(conn) + return None + def _get_clip_text(self, node: Dict[str, Any]) -> Optional[str]: inputs = node.get("inputs", {}) text = self._get_literal_input(inputs, "text") @@ -589,6 +623,12 @@ class WorkflowExtractor: if class_type == "ConditioningZeroOut": return "" + if class_type in ("PrimitiveStringMultiline", "String Literal"): + for key in ("value", "string", "text"): + value = self._get_literal_input(inputs, key) + if isinstance(value, str): + return value + if self._is_prompt_encoder(class_type): return self._get_clip_text(node)