diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 9e8d02c..e2064b7 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -251,6 +251,22 @@ def _ensure_worker_output_node(prompt_obj): return prompt_obj +def _prune_worker_downstream_of_branch_collectors(prompt_obj): + """Drop worker-side nodes downstream of branch collectors.""" + collector_ids = find_nodes_by_class(prompt_obj, "DistributedBranchCollector") + if not collector_ids: + return prompt_obj + + downstream = _find_downstream_nodes(prompt_obj, collector_ids) + for node_id in downstream: + if node_id in collector_ids: + continue + prompt_obj.pop(node_id, None) + + _remove_dangling_input_refs(prompt_obj) + return prompt_obj + + def prune_prompt_for_branch_worker(prompt_obj, branch_node_id, assigned_branch, num_branches): """Prune non-assigned branch paths while keeping shared downstream nodes.""" branch_id = str(branch_node_id) @@ -775,6 +791,7 @@ def apply_participant_overrides( ) if not is_master: + prompt_copy = _prune_worker_downstream_of_branch_collectors(prompt_copy) prompt_copy = _ensure_worker_output_node(prompt_copy) return prompt_copy diff --git a/tests/test_prompt_transform.py b/tests/test_prompt_transform.py index c56a793..9b8ea82 100644 --- a/tests/test_prompt_transform.py +++ b/tests/test_prompt_transform.py @@ -762,5 +762,24 @@ class ApplyOverridesBranchCollectorTests(unittest.TestCase): self.assertEqual(master_result["3"]["inputs"]["multi_job_id"], "run_2") self.assertEqual(worker_result["3"]["inputs"]["multi_job_id"], "run_2") + def test_worker_prunes_nodes_downstream_of_branch_collector(self): + prompt = { + "1": {"class_type": "KSampler", "inputs": {}}, + "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}}, + "3": {"class_type": "Blur", "inputs": {"image": ["2", 1]}}, + "4": { + "class_type": "DistributedBranchCollector", + "inputs": {"branch_2": ["3", 0], "num_branches": 2}, + }, + "5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}}, + } + + worker_result = _apply(prompt, "worker-a", enabled_worker_ids=["worker-a"]) + + class_types = {node.get("class_type") for node in worker_result.values()} + self.assertIn("DistributedBranchCollector", class_types) + self.assertNotIn("SaveImage", class_types) + self.assertIn("PreviewImage", class_types) + if __name__ == "__main__": unittest.main()