From ec1fc25637d2e349ed2345c2415e74c49b98bfe3 Mon Sep 17 00:00:00 2001 From: LuqP2 Date: Mon, 1 Jun 2026 15:33:06 -0300 Subject: [PATCH] fix: extract connected prompt and seed values --- tests/test_workflow_extractor.py | 92 +++++++++++++++++++++ workflow_extractor.py | 135 ++++++++++++++++++++++++++++++- 2 files changed, 225 insertions(+), 2 deletions(-) diff --git a/tests/test_workflow_extractor.py b/tests/test_workflow_extractor.py index e25f5c2..8108742 100644 --- a/tests/test_workflow_extractor.py +++ b/tests/test_workflow_extractor.py @@ -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": "", + "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 = { diff --git a/workflow_extractor.py b/workflow_extractor.py index 5924b87..bdd1170 100644 --- a/workflow_extractor.py +++ b/workflow_extractor.py @@ -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"):