fix: extract Qwen Image 2.1 prompts from both conditioning outputs

tests/test_workflow_extractor.py

@@ -1,6 +1,43 @@
 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"] = ["7", 0]
+    prompt["7"] = {"class_type": "String Literal", "inputs": {"string": ""}}
+    data, missing = WorkflowExtractor(prompt).extract(save_node_id="6")
+    assert data["negative"] == ""
+    assert "negative" not in missing
+
+
 def test_workflow_extractor_basic_prompt():
     prompt = {
         "1": {

workflow_extractor.py

@@ -48,6 +48,7 @@ class WorkflowExtractor:
         "CLIPTextEncodeSDXLPlus",
         "CLIPTextEncodeSDXLRefiner",
         "CLIPTextEncodeSD3",
+        "TextEncodeQwenImage21",
     ]
 
     VAE_DECODE_NODES = [
@@ -518,24 +519,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 +549,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")