From 2d5a56a95ea4c02cdebde42167a41b5033283c76 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 2 Mar 2026 07:06:17 +0000 Subject: [PATCH] Prevent prompt_no_outputs on pruned worker branch prompts --- api/orchestration/prompt_transform.py | 81 +++++++++++++++++++++++++++ tests/test_prompt_transform.py | 35 ++++++++++-- web/ui.js | 6 ++ 3 files changed, 118 insertions(+), 4 deletions(-) diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 993bbb6..53dea3d 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -173,6 +173,84 @@ def _remove_dangling_input_refs(prompt_obj): inputs.pop(input_name, None) +def _has_terminal_output_nodes(prompt_obj): + """Return True when the prompt already has at least one terminal output node.""" + for _node_id, node in _iter_prompt_nodes(prompt_obj): + if node.get("class_type") in {"PreviewImage", "SaveImage"}: + return True + return False + + +def _add_preview_node(prompt_obj, next_id_fn, source_node_id, slot_index, title_suffix): + """Attach a PreviewImage output node to the given source output slot.""" + preview_id = next_id_fn() + prompt_obj[preview_id] = { + "inputs": { + "images": [str(source_node_id), int(slot_index)], + }, + "class_type": "PreviewImage", + "_meta": { + "title": f"Preview Image ({title_suffix})", + }, + } + return preview_id + + +def _ensure_worker_output_node(prompt_obj): + """Guarantee worker prompts retain at least one output node after pruning. + + Branch pruning can remove all explicit output nodes for worker prompts, which + causes ComfyUI validation to reject the prompt (`prompt_no_outputs`). + """ + if _has_terminal_output_nodes(prompt_obj): + return prompt_obj + + next_id = _create_numeric_id_generator(prompt_obj) + + # Prefer a DistributedJoin output tied to the assigned branch for this worker. + for node_id, node in _iter_prompt_nodes(prompt_obj): + if node.get("class_type") != "DistributedJoin": + continue + inputs = node.get("inputs", {}) + try: + assigned_branch = int(inputs.get("assigned_branch", -1)) + except (TypeError, ValueError): + assigned_branch = -1 + if assigned_branch >= 0: + _add_preview_node(prompt_obj, next_id, node_id, assigned_branch, "auto-added worker output") + return prompt_obj + + # Otherwise, anchor to the assigned DistributedBranch slot if available. + for node_id, node in _iter_prompt_nodes(prompt_obj): + if node.get("class_type") != "DistributedBranch": + continue + inputs = node.get("inputs", {}) + try: + assigned_branch = int(inputs.get("assigned_branch", -1)) + except (TypeError, ValueError): + assigned_branch = -1 + if assigned_branch >= 0: + _add_preview_node(prompt_obj, next_id, node_id, assigned_branch, "auto-added worker output") + return prompt_obj + + # Idle participants (assigned_branch=-1) still need a terminal node to pass + # validation; keep this cheap by emitting a tiny synthetic image preview. + empty_id = next_id() + prompt_obj[empty_id] = { + "class_type": "DistributedEmptyImage", + "inputs": { + "height": 64, + "width": 64, + "channels": 3, + }, + "_meta": { + "title": "Distributed Empty Image (auto-added worker output)", + }, + } + _add_preview_node(prompt_obj, next_id, empty_id, 0, "auto-added worker output") + 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) @@ -695,4 +773,7 @@ def apply_participant_overrides( delegate_master, ) + if not is_master: + 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 5407ebf..fc9ad93 100644 --- a/tests/test_prompt_transform.py +++ b/tests/test_prompt_transform.py @@ -684,10 +684,37 @@ class ApplyOverridesBranchTests(unittest.TestCase): enabled_worker_ids=["worker-a", "worker-b", "worker-c"], ) self.assertEqual(worker_b_result["2"]["inputs"]["assigned_branch"], -1) - self.assertNotIn("3", worker_b_result) - self.assertNotIn("4", worker_b_result) - self.assertNotIn("5", worker_b_result) - self.assertNotIn("6", worker_b_result) + class_types = {node.get("class_type") for node in worker_b_result.values()} + self.assertNotIn("Blur", class_types) + self.assertNotIn("Sharpen", class_types) + + def test_worker_with_pruned_outputs_gets_auto_preview_for_assigned_branch(self): + prompt = { + "1": {"class_type": "KSampler", "inputs": {}}, + "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 3}}, + "3": {"class_type": "Blur", "inputs": {"image": ["2", 0]}}, + "4": {"class_type": "PreviewImage", "inputs": {"images": ["3", 0]}}, + } + + worker_result = _apply(prompt, "worker-a", enabled_worker_ids=["worker-a", "worker-b"]) + preview_nodes = [node for node in worker_result.values() if node.get("class_type") == "PreviewImage"] + + self.assertTrue(preview_nodes) + self.assertIn(["2", 1], [node.get("inputs", {}).get("images") for node in preview_nodes]) + + def test_unassigned_worker_gets_idle_fallback_output(self): + prompt = _branch_prompt() + worker_result = _apply( + prompt, + "worker-c", + enabled_worker_ids=["worker-a", "worker-b", "worker-c"], + ) + + preview_nodes = [node for node in worker_result.values() if node.get("class_type") == "PreviewImage"] + empty_nodes = [node for node in worker_result.values() if node.get("class_type") == "DistributedEmptyImage"] + + self.assertTrue(preview_nodes) + self.assertTrue(empty_nodes) class ApplyOverridesJoinTests(unittest.TestCase): diff --git a/web/ui.js b/web/ui.js index 2f992a1..71e13ab 100644 --- a/web/ui.js +++ b/web/ui.js @@ -602,6 +602,12 @@ export class DistributedUI { logBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl; logBtn.classList.add("btn--log"); logBtn.title = "View remote worker log"; + // Keep disabled workers from showing a log button before state sync. + if (!data.enabled) { + logBtn.classList.add("is-hidden"); + logBtn.style.display = "none"; + logBtn.disabled = true; + } controlsWrapper.appendChild(logBtn); } else { const controls = this.createWorkerControls(data.id, {