fix: extract connected prompt and seed values

2026-06-01 15:33:06 -03:00
parent 1b2106b6af
commit ec1fc25637
2 changed files with 225 additions and 2 deletions
@@ -135,6 +135,98 @@ def test_workflow_extractor_detects_lora_alternate_keys():
    assert data["lora_list"] == [{"name": "style.safetensors", "weight": 0.65}]


+def test_workflow_extractor_resolves_connected_seed_and_prompt_strings():
+    prompt = {
+        "1": {
+            "class_type": "UNETLoader",
+            "inputs": {"unet_name": "zImageTurboNSFW_82BF16FP8.safetensors"},
+        },
+        "2": {
+            "class_type": "Lora Loader (LoraManager)",
+            "inputs": {
+                "text": "<lora:Z-Detail-Slider:0.95>",
+                "loras": {
+                    "__value__": [
+                        {
+                            "name": "Z-Detail-Slider",
+                            "strength": 0.75,
+                            "active": True,
+                            "clipStrength": 0.75,
+                        }
+                    ]
+                },
+                "model": ["1", 0],
+            },
+        },
+        "3": {
+            "class_type": "easy positive",
+            "inputs": {"positive": "rabbit"},
+        },
+        "4": {
+            "class_type": "JoinStrings",
+            "inputs": {"delimiter": " ", "string1": ["3", 0]},
+        },
+        "5": {
+            "class_type": "easy stylesSelector",
+            "inputs": {"styles": "fooocus_styles", "select_styles": [], "positive": ["4", 0]},
+        },
+        "6": {
+            "class_type": "CLIPTextEncode",
+            "inputs": {"text": ["5", 0], "clip": ["2", 1]},
+        },
+        "7": {
+            "class_type": "ConditioningZeroOut",
+            "inputs": {"conditioning": ["6", 0]},
+        },
+        "8": {
+            "class_type": "SeedGenerator",
+            "inputs": {"seed": 862877771869053},
+        },
+        "9": {
+            "class_type": "ModelSamplingAuraFlow",
+            "inputs": {"shift": 3.0, "model": ["2", 0]},
+        },
+        "10": {
+            "class_type": "KSampler",
+            "inputs": {
+                "seed": ["8", 0],
+                "steps": 9,
+                "cfg": 1.0,
+                "sampler_name": "euler",
+                "scheduler": "simple",
+                "denoise": 1.0,
+                "model": ["9", 0],
+                "positive": ["6", 0],
+                "negative": ["7", 0],
+            },
+        },
+        "11": {
+            "class_type": "VAEDecode",
+            "inputs": {"samples": ["10", 0], "vae": ["12", 0]},
+        },
+        "12": {
+            "class_type": "VAELoader",
+            "inputs": {"vae_name": "ae.safetensors"},
+        },
+        "13": {
+            "class_type": "MetaHubSaveNode",
+            "inputs": {"images": ["11", 0]},
+        },
+    }
+
+    extractor = WorkflowExtractor(prompt)
+    data, missing = extractor.extract(save_node_id="13")
+
+    assert data["seed"] == 862877771869053
+    assert data["positive"] == "rabbit"
+    assert data["negative"] == ""
+    assert data["model_name"] == "zImageTurboNSFW_82BF16FP8.safetensors"
+    assert data["vae_name"] == "ae.safetensors"
+    assert data["lora_list"] == [{"name": "Z-Detail-Slider", "weight": 0.75}]
+    assert "seed" not in missing
+    assert "positive" not in missing
+
+

 def test_workflow_extractor_detects_img2img_lineage():
    prompt = {
@@ -153,6 +153,10 @@ class WorkflowExtractor:
        inputs = node.get("inputs", {})

        seed = self._coerce_int(self._get_literal_input(inputs, "seed"))
+        if seed is None:
+            seed = self._coerce_int(
+                self._resolve_scalar_from_connection(inputs.get("seed"), ("seed", "value"))
+            )
        steps = self._coerce_int(self._get_literal_input(inputs, "steps"))
        cfg = self._coerce_float(self._get_first_literal(inputs, ("cfg", "cfg_scale")))
        sampler_name = self._get_first_literal(inputs, ("sampler_name", "sampler"))
@@ -230,6 +234,10 @@ class WorkflowExtractor:
                continue
            has_lora_nodes = True
            inputs = node.get("inputs", {})
+            stacked_loras = self._extract_lora_manager_entries(inputs)
+            if stacked_loras:
+                loras.extend(stacked_loras)
+                continue
            lora_name = self._get_first_literal(inputs, ("lora_name", "lora", "lora_name_1", "lora_1"))
            if not lora_name:
                continue
@@ -247,6 +255,34 @@ class WorkflowExtractor:

        return loras, has_lora_nodes

+    def _extract_lora_manager_entries(self, inputs: Dict[str, Any]) -> List[Dict[str, Any]]:
+        raw_loras = inputs.get("loras")
+        if isinstance(raw_loras, dict) and isinstance(raw_loras.get("__value__"), list):
+            raw_loras = raw_loras["__value__"]
+        if not isinstance(raw_loras, list):
+            return []
+
+        loras: List[Dict[str, Any]] = []
+        for raw_lora in raw_loras:
+            if not isinstance(raw_lora, dict):
+                continue
+            if raw_lora.get("active") is False:
+                continue
+            name = raw_lora.get("name") or raw_lora.get("lora_name")
+            if not name:
+                continue
+            weight = self._coerce_float(
+                raw_lora.get("strength")
+                if raw_lora.get("strength") is not None
+                else raw_lora.get("modelStrength")
+            )
+            if weight is None:
+                weight = self._coerce_float(raw_lora.get("clipStrength"))
+            if weight is None:
+                weight = 1.0
+            loras.append({"name": str(name), "weight": float(weight)})
+        return loras
+
    def extract_lineage(self, sampler_node_id: Optional[str]) -> Dict[str, Any]:
        if not sampler_node_id:
            return {}
@@ -379,9 +415,11 @@ class WorkflowExtractor:
            if not node:
                continue
            class_type = self._class_type(node)
+            if class_type == "ConditioningZeroOut":
+                return ""
            if class_type in self.CLIP_NODES:
                text = self._get_clip_text(node)
-                if text:
+                if text is not None:
                    texts.append(text)
                    continue
            for input_val in node.get("inputs", {}).values():
@@ -396,8 +434,11 @@ class WorkflowExtractor:
    def _get_clip_text(self, node: Dict[str, Any]) -> Optional[str]:
        inputs = node.get("inputs", {})
        text = self._get_literal_input(inputs, "text")
-        if text:
+        if text is not None and str(text).strip():
            return str(text)
+        text_from_conn = self._resolve_string_from_connection(inputs.get("text"))
+        if text_from_conn is not None:
+            return text_from_conn
        parts: List[str] = []
        for key in ("text_g", "text_l"):
            value = self._get_literal_input(inputs, key)
@@ -407,6 +448,96 @@ class WorkflowExtractor:
            return "\n".join(parts)
        return None

+    def _resolve_string_from_connection(self, conn: Any) -> Optional[str]:
+        node_id = self._get_connection_node_id(conn)
+        if not node_id:
+            return None
+        return self._resolve_string_from_node(node_id, set())
+
+    def _resolve_string_from_node(self, node_id: str, visited: Set[str]) -> Optional[str]:
+        if node_id in visited:
+            return None
+        visited.add(node_id)
+
+        node = self._get_node(node_id)
+        if not node:
+            return None
+
+        inputs = node.get("inputs", {})
+        class_type = self._class_type(node)
+        class_lower = class_type.lower()
+
+        if class_type == "ConditioningZeroOut":
+            return ""
+
+        if class_type in self.CLIP_NODES:
+            return self._get_clip_text(node)
+
+        if class_lower == "joinstrings" or "join strings" in class_lower:
+            delimiter = self._get_literal_input(inputs, "delimiter")
+            joiner = str(delimiter) if delimiter is not None else " "
+            parts: List[str] = []
+            for key in ("string1", "string2", "string3", "string4"):
+                value = self._get_literal_input(inputs, key)
+                if value is None:
+                    value = self._resolve_string_from_connection_with_visited(inputs.get(key), visited)
+                if value is not None and str(value).strip():
+                    parts.append(str(value).strip())
+            return joiner.join(parts) if parts else ""
+
+        for key in ("text", "positive", "negative", "prompt", "string", "value"):
+            value = self._get_literal_input(inputs, key)
+            if value is not None and str(value).strip():
+                return str(value)
+
+        for key in ("positive", "text", "string1", "string2", "conditioning"):
+            value = self._resolve_string_from_connection_with_visited(inputs.get(key), visited)
+            if value is not None and str(value).strip():
+                return value
+
+        for input_val in inputs.values():
+            value = self._resolve_string_from_connection_with_visited(input_val, visited)
+            if value is not None and str(value).strip():
+                return value
+
+        return None
+
+    def _resolve_string_from_connection_with_visited(self, conn: Any, visited: Set[str]) -> Optional[str]:
+        node_id = self._get_connection_node_id(conn)
+        if not node_id:
+            return None
+        return self._resolve_string_from_node(node_id, visited)
+
+    def _resolve_scalar_from_connection(self, conn: Any, keys: Tuple[str, ...]) -> Any:
+        node_id = self._get_connection_node_id(conn)
+        if not node_id:
+            return None
+
+        queue = [node_id]
+        visited: Set[str] = set()
+        while queue:
+            current_id = queue.pop(0)
+            if current_id in visited:
+                continue
+            visited.add(current_id)
+
+            node = self._get_node(current_id)
+            if not node:
+                continue
+
+            inputs = node.get("inputs", {})
+            for key in keys:
+                value = self._get_literal_input(inputs, key)
+                if value is not None:
+                    return value
+
+            for input_val in inputs.values():
+                conn_id = self._get_connection_node_id(input_val)
+                if conn_id and conn_id not in visited:
+                    queue.append(conn_id)
+
+        return None
+
    def _get_checkpoint_name(self, node: Dict[str, Any]) -> Optional[str]:
        inputs = node.get("inputs", {})
        for key in ("ckpt_name", "checkpoint", "model_name", "unet_name", "diffusion_model_name", "model"):