guard worker prompts from branch-collector downstream outputs

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

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()