Add distributed list splitting/collection and branch routing

api/orchestration/prompt_transform.py

@@ -88,6 +88,35 @@ def _find_downstream_nodes(prompt_obj, start_ids):
     return connected
 
 
+def _find_downstream_of_output_slot(prompt_obj, node_id, slot_index):
+    """Return all nodes reachable from a specific output slot on one node."""
+    source_id = str(node_id)
+    try:
+        slot = int(slot_index)
+    except (TypeError, ValueError):
+        return set()
+
+    direct_consumers = set()
+    for candidate_id, candidate_node in _iter_prompt_nodes(prompt_obj):
+        inputs = candidate_node.get("inputs", {})
+        for value in inputs.values():
+            if not (isinstance(value, list) and len(value) == 2):
+                continue
+            if str(value[0]) != source_id:
+                continue
+            try:
+                input_slot = int(value[1])
+            except (TypeError, ValueError):
+                continue
+            if input_slot == slot:
+                direct_consumers.add(str(candidate_id))
+                break
+
+    if not direct_consumers:
+        return set()
+    return _find_downstream_nodes(prompt_obj, list(direct_consumers))
+
+
 def _create_numeric_id_generator(prompt_obj):
     """Return a closure that yields new numeric string IDs."""
     max_id = 0
@@ -125,15 +154,90 @@ def _find_upstream_nodes(prompt_obj, start_ids):
     return connected
 
 
+def _resolve_participants(enabled_worker_ids, delegate_master):
+    worker_ids = [str(worker_id) for worker_id in (enabled_worker_ids or [])]
+    if delegate_master:
+        return worker_ids
+    return ["master"] + worker_ids
+
+
+def _remove_dangling_input_refs(prompt_obj):
+    """Drop input links that point to nodes no longer present in prompt_obj."""
+    existing_ids = set(prompt_obj.keys())
+    for _node_id, node in _iter_prompt_nodes(prompt_obj):
+        inputs = node.get("inputs", {})
+        for input_name, input_value in list(inputs.items()):
+            if isinstance(input_value, list) and len(input_value) == 2:
+                source_id = str(input_value[0])
+                if source_id not in existing_ids:
+                    inputs.pop(input_name, None)
+
+
+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)
+
+    assigned_slots = set()
+    if isinstance(assigned_branch, (list, tuple, set)):
+        for value in assigned_branch:
+            try:
+                idx = int(value)
+            except (TypeError, ValueError):
+                continue
+            if idx >= 0:
+                assigned_slots.add(idx)
+    else:
+        try:
+            idx = int(assigned_branch)
+        except (TypeError, ValueError):
+            idx = -1
+        if idx >= 0:
+            assigned_slots.add(idx)
+
+    try:
+        total_slots = int(num_branches)
+    except (TypeError, ValueError):
+        total_slots = 2
+    total_slots = max(2, min(total_slots, 10))
+
+    downstream_by_slot = {}
+    for slot_idx in range(total_slots):
+        downstream = _find_downstream_of_output_slot(prompt_obj, branch_id, slot_idx)
+        downstream.discard(branch_id)
+        downstream_by_slot[slot_idx] = downstream
+
+    keep = {branch_id}
+    for slot_idx in assigned_slots:
+        keep.update(downstream_by_slot.get(slot_idx, set()))
+
+    remove = set()
+    for slot_idx in range(total_slots):
+        if slot_idx in assigned_slots:
+            continue
+        remove.update(downstream_by_slot.get(slot_idx, set()))
+
+    remove -= keep
+    for node_id in remove:
+        prompt_obj.pop(node_id, None)
+
+    _remove_dangling_input_refs(prompt_obj)
+    return prompt_obj
+
+
 def prune_prompt_for_worker(prompt_obj):
     """Prune worker prompt to distributed nodes and their upstream dependencies."""
     collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
+    list_collector_ids = find_nodes_by_class(prompt_obj, "DistributedListCollector")
     upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
-    distributed_ids = collector_ids + upscale_ids
+    branch_ids = find_nodes_by_class(prompt_obj, "DistributedBranch")
+    distributed_ids = collector_ids + list_collector_ids + upscale_ids + branch_ids
     if not distributed_ids:
         return prompt_obj
 
     connected = _find_upstream_nodes(prompt_obj, distributed_ids)
+    if branch_ids:
+        connected.update(_find_downstream_nodes(prompt_obj, branch_ids))
+
     pruned_prompt = {}
     for node_id in connected:
         node = prompt_obj.get(node_id)
@@ -146,7 +250,7 @@ def prune_prompt_for_worker(prompt_obj):
         if dist_id not in pruned_prompt:
             continue
         downstream = _find_downstream_nodes(prompt_obj, [dist_id])
-        has_removed_downstream = any(node_id != dist_id for node_id in downstream)
+        has_removed_downstream = any(node_id != dist_id and node_id not in connected for node_id in downstream)
         if has_removed_downstream:
             preview_id = next_id()
             pruned_prompt[preview_id] = {
@@ -217,8 +321,11 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
 def generate_job_id_map(prompt_index, prefix):
     """Create stable per-node job IDs for distributed nodes."""
     job_map = {}
-    distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
-        "UltimateSDUpscaleDistributed"
+    distributed_nodes = (
+        prompt_index.nodes_for_class("DistributedCollector")
+        + prompt_index.nodes_for_class("DistributedListCollector")
+        + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
+        + prompt_index.nodes_for_class("DistributedBranch")
     )
     for node_id in distributed_nodes:
         job_map[node_id] = f"{prefix}_{node_id}"
@@ -313,6 +420,110 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id,
             inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
 
 
+def _override_list_splitter_nodes(prompt_copy, prompt_index, participant_id, enabled_worker_ids, delegate_master):
+    """Configure participant chunk mapping for DistributedListSplitter nodes."""
+    participants = _resolve_participants(enabled_worker_ids, delegate_master)
+    total_participants = max(1, len(participants))
+    participant_id = str(participant_id)
+
+    if participant_id in participants:
+        participant_index = participants.index(participant_id)
+    else:
+        participant_index = 0
+
+    for node_id in prompt_index.nodes_for_class("DistributedListSplitter"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+        inputs = node.setdefault("inputs", {})
+        inputs["participant_index"] = participant_index
+        inputs["total_participants"] = total_participants
+
+
+def _override_list_collector_nodes(
+    prompt_copy,
+    prompt_index,
+    is_master,
+    participant_id,
+    job_id_map,
+    master_url,
+    enabled_json,
+    delegate_master,
+):
+    """Configure DistributedListCollector nodes for master/worker role."""
+    for node_id in prompt_index.nodes_for_class("DistributedListCollector"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+
+        inputs = node.setdefault("inputs", {})
+        inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
+        inputs["is_worker"] = not is_master
+        inputs["enabled_worker_ids"] = enabled_json
+        if is_master:
+            inputs["delegate_only"] = bool(delegate_master)
+            inputs.pop("master_url", None)
+            inputs.pop("worker_id", None)
+        else:
+            inputs["master_url"] = master_url
+            inputs["worker_id"] = str(participant_id)
+            inputs["delegate_only"] = False
+
+
+def _override_branch_nodes(
+    prompt_copy,
+    prompt_index,
+    is_master,
+    participant_id,
+    enabled_worker_ids,
+    job_id_map,
+    delegate_master,
+):
+    """Assign branch slots per participant and prune non-assigned branch paths."""
+    participants = _resolve_participants(enabled_worker_ids, delegate_master)
+    participant_id = str(participant_id)
+    participant_count = len(participants)
+    participant_pos = participants.index(participant_id) if participant_id in participants else None
+
+    for node_id in prompt_index.nodes_for_class("DistributedBranch"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+
+        inputs = node.setdefault("inputs", {})
+        try:
+            num_branches = int(inputs.get("num_branches", 2))
+        except (TypeError, ValueError):
+            num_branches = 2
+        num_branches = max(2, min(num_branches, 10))
+
+        assigned_slots = []
+        if participant_pos is not None and participant_count > 0:
+            assigned_slots = [slot for slot in range(num_branches) if slot % participant_count == participant_pos]
+
+        if len(assigned_slots) == 1:
+            assigned_branch = assigned_slots[0]
+        else:
+            assigned_branch = -1
+
+        inputs["is_worker"] = not is_master
+        inputs["worker_id"] = "" if is_master else participant_id
+        inputs["assigned_branch"] = assigned_branch
+        inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
+
+        if is_master and delegate_master:
+            continue
+
+        prompt_copy = prune_prompt_for_branch_worker(
+            prompt_copy,
+            node_id,
+            assigned_slots,
+            num_branches,
+        )
+
+    return prompt_copy
+
+
 def apply_participant_overrides(
     prompt_copy,
     participant_id,
@@ -329,6 +540,13 @@ def apply_participant_overrides(
 
     _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
     _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
+    _override_list_splitter_nodes(
+        prompt_copy,
+        prompt_index,
+        participant_id,
+        enabled_worker_ids,
+        delegate_master,
+    )
     _override_collector_nodes(
         prompt_copy,
         prompt_index,
@@ -339,6 +557,16 @@ def apply_participant_overrides(
         enabled_json,
         delegate_master,
     )
+    _override_list_collector_nodes(
+        prompt_copy,
+        prompt_index,
+        is_master,
+        participant_id,
+        job_id_map,
+        master_url,
+        enabled_json,
+        delegate_master,
+    )
     _override_upscale_nodes(
         prompt_copy,
         prompt_index,
@@ -348,5 +576,14 @@ def apply_participant_overrides(
         master_url,
         enabled_json,
     )
+    prompt_copy = _override_branch_nodes(
+        prompt_copy,
+        prompt_index,
+        is_master,
+        participant_id,
+        enabled_worker_ids,
+        job_id_map,
+        delegate_master,
+    )
 
     return prompt_copy

api/queue_orchestration.py

@@ -314,8 +314,15 @@ async def orchestrate_distributed_execution(
         prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id)
         return prompt_id, 0
 
-    for job_id in job_id_map.values():
-        await _ensure_distributed_queue(job_id)
+    queue_node_ids = (
+        prompt_index.nodes_for_class("DistributedCollector")
+        + prompt_index.nodes_for_class("DistributedListCollector")
+        + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
+    )
+    for node_id in queue_node_ids:
+        job_id = job_id_map.get(node_id)
+        if job_id:
+            await _ensure_distributed_queue(job_id)
 
     master_prompt = prompt_index.copy_prompt()
     master_prompt = apply_participant_overrides(
@@ -329,7 +336,9 @@ async def orchestrate_distributed_execution(
     )
 
     if delegate_master:
-        collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
+        collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector") + find_nodes_by_class(
+            master_prompt, "DistributedListCollector"
+        )
         upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
         if upscale_nodes:
             debug_log(

distributed.py

@@ -23,6 +23,9 @@ from .nodes import (
     DistributedValue,
     AudioBatchDivider,
     DistributedEmptyImage,
+    DistributedListSplitter,
+    DistributedListCollector,
+    DistributedBranch,
     AnyType,
     ByPassTypeTuple,
     any_type,

nodes/__init__.py

@@ -10,9 +10,15 @@ from .utilities import (
     any_type,
 )
 from .collector import DistributedCollectorNode
+from .list_splitter import DistributedListSplitter
+from .list_collector import DistributedListCollector
+from .branch import DistributedBranch
 
 NODE_CLASS_MAPPINGS = {
     "DistributedCollector": DistributedCollectorNode,
+    "DistributedListSplitter": DistributedListSplitter,
+    "DistributedListCollector": DistributedListCollector,
+    "DistributedBranch": DistributedBranch,
     "DistributedSeed": DistributedSeed,
     "DistributedModelName": DistributedModelName,
     "DistributedValue": DistributedValue,
@@ -22,6 +28,9 @@ NODE_CLASS_MAPPINGS = {
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "DistributedCollector": "Distributed Collector",
+    "DistributedListSplitter": "Distributed List Splitter",
+    "DistributedListCollector": "Distributed List Collector",
+    "DistributedBranch": "Distributed Branch",
     "DistributedSeed": "Distributed Seed",
     "DistributedModelName": "Distributed Model Name",
     "DistributedValue": "Distributed Value",

nodes/branch.py

@@ -0,0 +1,26 @@
+from .utilities import ByPassTypeTuple, any_type
+
+
+class DistributedBranch:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input": (any_type,),
+                "num_branches": ("INT", {"default": 2, "min": 2, "max": 10, "step": 1}),
+            },
+            "hidden": {
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "worker_id": ("STRING", {"default": ""}),
+                "assigned_branch": ("INT", {"default": -1, "min": -1, "max": 9}),
+                "multi_job_id": ("STRING", {"default": ""}),
+            },
+        }
+
+    RETURN_TYPES = ByPassTypeTuple((any_type,))
+    RETURN_NAMES = ByPassTypeTuple(tuple([f"branch_{idx + 1}" for idx in range(10)]))
+    FUNCTION = "branch"
+    CATEGORY = "utils"
+
+    def branch(self, input, num_branches=2, is_worker=False, worker_id="", assigned_branch=-1, multi_job_id=""):
+        return tuple([input] * 10)

nodes/list_collector.py

@@ -0,0 +1,266 @@
+import asyncio
+import base64
+import io
+import json
+import time
+
+import aiohttp
+import torch
+
+from ..utils.async_helpers import run_async_in_server_loop
+from ..utils.config import get_worker_timeout_seconds, is_master_delegate_only
+from ..utils.constants import HEARTBEAT_INTERVAL
+from ..utils.image import ensure_contiguous, tensor_to_pil
+from ..utils.logging import debug_log, log
+from ..utils.network import get_client_session
+
+
+def _get_prompt_server_instance():
+    import server as _server
+
+    return _server.PromptServer.instance
+
+
+def _throw_if_processing_interrupted():
+    try:
+        import comfy.model_management as model_management
+
+        model_management.throw_exception_if_processing_interrupted()
+    except Exception:
+        # Tests and non-Comfy contexts may not have interruption support.
+        return
+
+
+class DistributedListCollector:
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True,)
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+            },
+            "hidden": {
+                "multi_job_id": ("STRING", {"default": ""}),
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "master_url": ("STRING", {"default": ""}),
+                "enabled_worker_ids": ("STRING", {"default": "[]"}),
+                "worker_id": ("STRING", {"default": ""}),
+                "delegate_only": ("BOOLEAN", {"default": False}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("images",)
+    FUNCTION = "run"
+    CATEGORY = "image"
+
+    def run(
+        self,
+        images,
+        multi_job_id="",
+        is_worker=False,
+        master_url="",
+        enabled_worker_ids="[]",
+        worker_id="",
+        delegate_only=False,
+    ):
+        image_list = self._normalize_image_list(images)
+        if not multi_job_id:
+            return (image_list,)
+
+        return run_async_in_server_loop(
+            self.execute(
+                image_list,
+                multi_job_id=multi_job_id,
+                is_worker=is_worker,
+                master_url=master_url,
+                enabled_worker_ids=enabled_worker_ids,
+                worker_id=worker_id,
+                delegate_only=delegate_only,
+            )
+        )
+
+    def _normalize_image_list(self, images):
+        if images is None:
+            return []
+        if isinstance(images, list):
+            return list(images)
+        return [images]
+
+    async def send_list_to_master(self, image_list, multi_job_id, master_url, worker_id):
+        if not image_list:
+            return
+
+        sendable = []
+        for idx, image in enumerate(image_list):
+            if not isinstance(image, torch.Tensor):
+                continue
+            batch = image
+            if batch.ndim == 3:
+                batch = batch.unsqueeze(0)
+            if batch.ndim != 4 or batch.shape[0] <= 0:
+                continue
+            sendable.append((idx, batch))
+
+        if not sendable:
+            return
+
+        session = await get_client_session()
+        url = f"{master_url}/distributed/job_complete"
+
+        for send_pos, (image_index, image_batch) in enumerate(sendable):
+            img = tensor_to_pil(image_batch, 0)
+            byte_io = io.BytesIO()
+            img.save(byte_io, format="PNG", compress_level=0)
+            encoded_image = base64.b64encode(byte_io.getvalue()).decode("utf-8")
+            payload = {
+                "job_id": str(multi_job_id),
+                "worker_id": str(worker_id),
+                "batch_idx": int(image_index),
+                "image": f"data:image/png;base64,{encoded_image}",
+                "is_last": bool(send_pos == len(sendable) - 1),
+            }
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    timeout=aiohttp.ClientTimeout(total=60),
+                ) as response:
+                    response.raise_for_status()
+            except Exception as exc:
+                log(f"Worker - Failed to send list item to master: {exc}")
+                debug_log(f"Worker - Full error details: URL={url}")
+                raise
+
+    def _store_worker_result(self, worker_images, item):
+        worker_id = str(item.get("worker_id", ""))
+        tensor = item.get("tensor")
+        image_index = item.get("image_index")
+        if not worker_id or tensor is None or image_index is None:
+            return 0
+
+        worker_images.setdefault(worker_id, {})
+        worker_images[worker_id][int(image_index)] = tensor
+        return 1
+
+    def _parse_enabled_workers(self, enabled_worker_ids):
+        try:
+            raw = json.loads(enabled_worker_ids)
+        except Exception:
+            raw = []
+
+        workers = []
+        seen = set()
+        for worker_id in raw if isinstance(raw, list) else []:
+            worker_id_str = str(worker_id)
+            if worker_id_str in seen:
+                continue
+            seen.add(worker_id_str)
+            workers.append(worker_id_str)
+        return workers
+
+    def _reorder_and_combine_list(self, worker_images, worker_order, master_images, delegate_mode):
+        ordered = []
+        if not delegate_mode:
+            ordered.extend(master_images)
+
+        ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
+        seen = set()
+
+        for worker_id in ordered_worker_ids:
+            seen.add(worker_id)
+            worker_entries = worker_images.get(worker_id, {})
+            for image_index in sorted(worker_entries.keys()):
+                ordered.append(worker_entries[image_index])
+
+        for worker_id in sorted(worker_images.keys()):
+            if worker_id in seen:
+                continue
+            worker_entries = worker_images.get(worker_id, {})
+            for image_index in sorted(worker_entries.keys()):
+                ordered.append(worker_entries[image_index])
+
+        combined = []
+        for image in ordered:
+            if isinstance(image, torch.Tensor):
+                if image.is_cuda:
+                    image = image.cpu()
+                image = ensure_contiguous(image)
+            combined.append(image)
+
+        return combined
+
+    async def execute(
+        self,
+        images,
+        multi_job_id="",
+        is_worker=False,
+        master_url="",
+        enabled_worker_ids="[]",
+        worker_id="",
+        delegate_only=False,
+    ):
+        image_list = self._normalize_image_list(images)
+
+        if is_worker:
+            debug_log(f"Worker - Job {multi_job_id} complete. Sending {len(image_list)} image list item(s) to master")
+            await self.send_list_to_master(image_list, multi_job_id, master_url, worker_id)
+            return (image_list,)
+
+        delegate_mode = bool(delegate_only or is_master_delegate_only())
+        enabled_workers = self._parse_enabled_workers(enabled_worker_ids)
+        expected_workers = set(enabled_workers)
+        if not expected_workers:
+            return (image_list,)
+
+        prompt_server = _get_prompt_server_instance()
+
+        async with prompt_server.distributed_jobs_lock:
+            if multi_job_id not in prompt_server.distributed_pending_jobs:
+                prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
+                debug_log(f"Master - Initialized queue early for list collector job {multi_job_id}")
+
+        master_images = [] if delegate_mode else image_list
+
+        worker_images = {}
+        workers_done = set()
+        base_timeout = float(get_worker_timeout_seconds())
+        slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
+        last_activity = time.time()
+
+        try:
+            while len(workers_done) < len(expected_workers):
+                _throw_if_processing_interrupted()
+                try:
+                    async with prompt_server.distributed_jobs_lock:
+                        queue = prompt_server.distributed_pending_jobs[multi_job_id]
+                    result = await asyncio.wait_for(queue.get(), timeout=slice_timeout)
+                except asyncio.TimeoutError:
+                    if (time.time() - last_activity) < base_timeout:
+                        continue
+                    missing_workers = sorted(expected_workers - workers_done)
+                    log(
+                        "Master - List collector heartbeat timeout. "
+                        f"Still waiting for workers: {missing_workers}"
+                    )
+                    break
+
+                worker_id_value = str(result.get("worker_id", ""))
+                is_last = bool(result.get("is_last", False))
+                self._store_worker_result(worker_images, result)
+                last_activity = time.time()
+                base_timeout = float(get_worker_timeout_seconds())
+                if is_last and worker_id_value in expected_workers:
+                    workers_done.add(worker_id_value)
+
+        finally:
+            async with prompt_server.distributed_jobs_lock:
+                prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
+
+        combined = self._reorder_and_combine_list(worker_images, enabled_workers, master_images, delegate_mode)
+        debug_log(
+            f"Master - List collector job {multi_job_id} complete. Combined {len(combined)} image list item(s)."
+        )
+        return (combined,)

nodes/list_splitter.py

@@ -0,0 +1,43 @@
+from .utilities import _chunk_bounds
+
+
+class DistributedListSplitter:
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True,)
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+            },
+            "hidden": {
+                "participant_index": ("INT", {"default": 0, "min": 0, "max": 1024}),
+                "total_participants": ("INT", {"default": 1, "min": 1, "max": 1024}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("images",)
+    FUNCTION = "split"
+    CATEGORY = "image"
+
+    def split(self, images, participant_index=0, total_participants=1):
+        if images is None:
+            return ([],)
+
+        image_list = images if isinstance(images, list) else [images]
+
+        split_count = max(1, int(total_participants or 1))
+        if split_count <= 1:
+            return (list(image_list),)
+
+        bounds = _chunk_bounds(len(image_list), split_count)
+        index = int(participant_index or 0)
+        if index < 0:
+            index = 0
+        if index >= split_count:
+            index = split_count - 1
+
+        start, end = bounds[index]
+        return (list(image_list[start:end]),)

tests/test_branch.py

@@ -0,0 +1,76 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    nodes_pkg = types.ModuleType(f"{package_name}.nodes")
+    nodes_pkg.__path__ = []
+    sys.modules[f"{package_name}.nodes"] = nodes_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[1] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_branch_module():
+    package_name = "dist_branch_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "nodes/utilities.py", "nodes.utilities")
+    return _load_module(package_name, "nodes/branch.py", "nodes.branch")
+
+
+branch_module = _load_branch_module()
+
+
+class DistributedBranchTests(unittest.TestCase):
+    def test_no_distribution_outputs_all_branches(self):
+        node = branch_module.DistributedBranch()
+        token = {"k": "v"}
+
+        outputs = node.branch(token, num_branches=2, assigned_branch=-1)
+
+        self.assertEqual(len(outputs), 10)
+        self.assertTrue(all(output is token for output in outputs))
+
+    def test_assigned_branch_still_passes_input(self):
+        node = branch_module.DistributedBranch()
+        token = "payload"
+
+        outputs = node.branch(token, num_branches=3, is_worker=True, worker_id="worker-a", assigned_branch=1)
+
+        self.assertEqual(outputs[0], "payload")
+        self.assertEqual(outputs[1], "payload")
+        self.assertEqual(outputs[9], "payload")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_list_collector.py

@@ -0,0 +1,206 @@
+import asyncio
+import importlib.util
+import json
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    nodes_pkg = types.ModuleType(f"{package_name}.nodes")
+    nodes_pkg.__path__ = []
+    sys.modules[f"{package_name}.nodes"] = nodes_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.get_worker_timeout_seconds = lambda: 0.2
+    config_module.is_master_delegate_only = lambda: False
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.HEARTBEAT_INTERVAL = 1.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+
+    def ensure_contiguous(tensor):
+        return tensor.contiguous()
+
+    def tensor_to_pil(image_batch, index):
+        tensor = image_batch[index].detach().cpu().clamp(0, 1)
+        array = (tensor.numpy() * 255.0).astype("uint8")
+        return Image.fromarray(array)
+
+    image_module.ensure_contiguous = ensure_contiguous
+    image_module.tensor_to_pil = tensor_to_pil
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def get_client_session():
+        raise RuntimeError("test should patch get_client_session")
+
+    network_module.get_client_session = get_client_session
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[1] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_list_collector_module():
+    package_name = "dist_list_collector_testpkg"
+    _bootstrap_package(package_name)
+    return _load_module(package_name, "nodes/list_collector.py", "nodes.list_collector")
+
+
+collector_module = _load_list_collector_module()
+
+
+class _FakeResponseCtx:
+    def __init__(self, recorder, payload, url):
+        self._recorder = recorder
+        self._payload = payload
+        self._url = url
+
+    async def __aenter__(self):
+        self._recorder.append({"url": self._url, "payload": self._payload})
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+    def raise_for_status(self):
+        return None
+
+
+class _FakeSession:
+    def __init__(self):
+        self.calls = []
+
+    def post(self, url, json=None, timeout=None):
+        return _FakeResponseCtx(self.calls, json, url)
+
+
+class _FakePromptServer:
+    def __init__(self):
+        self.distributed_pending_jobs = {}
+        self.distributed_jobs_lock = asyncio.Lock()
+
+
+class DistributedListCollectorTests(unittest.IsolatedAsyncioTestCase):
+    def setUp(self):
+        self.node = collector_module.DistributedListCollector()
+
+    async def test_worker_mode_sends_images_to_master(self):
+        fake_session = _FakeSession()
+
+        async def _get_session():
+            return fake_session
+
+        collector_module.get_client_session = _get_session
+
+        images = [
+            torch.zeros((1, 4, 4, 3), dtype=torch.float32),
+            torch.ones((1, 4, 4, 3), dtype=torch.float32),
+        ]
+
+        result = await self.node.execute(
+            images,
+            multi_job_id="job-list",
+            is_worker=True,
+            master_url="http://master.local:8188",
+            worker_id="worker-a",
+        )
+
+        self.assertEqual(len(result[0]), 2)
+        self.assertEqual(len(fake_session.calls), 2)
+        self.assertEqual(fake_session.calls[0]["url"], "http://master.local:8188/distributed/job_complete")
+        self.assertEqual(fake_session.calls[0]["payload"]["batch_idx"], 0)
+        self.assertFalse(fake_session.calls[0]["payload"]["is_last"])
+        self.assertEqual(fake_session.calls[1]["payload"]["batch_idx"], 1)
+        self.assertTrue(fake_session.calls[1]["payload"]["is_last"])
+
+    async def test_master_mode_collects_results_as_list(self):
+        fake_server = _FakePromptServer()
+        collector_module._get_prompt_server_instance = lambda: fake_server
+
+        queue = asyncio.Queue()
+        worker_a = torch.full((1, 4, 4, 3), 0.25, dtype=torch.float32)
+        worker_b = torch.full((1, 4, 4, 3), 0.75, dtype=torch.float32)
+        await queue.put({"tensor": worker_a, "worker_id": "worker-a", "image_index": 0, "is_last": True})
+        await queue.put({"tensor": worker_b, "worker_id": "worker-b", "image_index": 0, "is_last": True})
+        fake_server.distributed_pending_jobs["job-list"] = queue
+
+        master_images = [torch.zeros((1, 4, 4, 3), dtype=torch.float32)]
+
+        result = await self.node.execute(
+            master_images,
+            multi_job_id="job-list",
+            is_worker=False,
+            enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
+        )
+
+        self.assertEqual(len(result[0]), 3)
+        self.assertTrue(torch.equal(result[0][0], master_images[0]))
+        self.assertTrue(torch.equal(result[0][1], worker_a))
+        self.assertTrue(torch.equal(result[0][2], worker_b))
+
+    async def test_delegate_only_mode_returns_worker_images_only(self):
+        fake_server = _FakePromptServer()
+        collector_module._get_prompt_server_instance = lambda: fake_server
+
+        queue = asyncio.Queue()
+        worker_img = torch.full((1, 4, 4, 3), 0.5, dtype=torch.float32)
+        await queue.put({"tensor": worker_img, "worker_id": "worker-a", "image_index": 0, "is_last": True})
+        fake_server.distributed_pending_jobs["job-list"] = queue
+
+        master_images = [torch.zeros((1, 4, 4, 3), dtype=torch.float32)]
+
+        result = await self.node.execute(
+            master_images,
+            multi_job_id="job-list",
+            is_worker=False,
+            enabled_worker_ids=json.dumps(["worker-a"]),
+            delegate_only=True,
+        )
+
+        self.assertEqual(len(result[0]), 1)
+        self.assertTrue(torch.equal(result[0][0], worker_img))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_list_splitter.py

@@ -0,0 +1,108 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    nodes_pkg = types.ModuleType(f"{package_name}.nodes")
+    nodes_pkg.__path__ = []
+    sys.modules[f"{package_name}.nodes"] = nodes_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[1] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_list_splitter_module():
+    package_name = "dist_list_splitter_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "nodes/utilities.py", "nodes.utilities")
+    return _load_module(package_name, "nodes/list_splitter.py", "nodes.list_splitter")
+
+
+splitter_module = _load_list_splitter_module()
+
+
+class DistributedListSplitterTests(unittest.TestCase):
+    def setUp(self):
+        self.node = splitter_module.DistributedListSplitter()
+
+    def _images(self, count):
+        return [torch.full((1, 4, 4, 3), float(idx), dtype=torch.float32) for idx in range(count)]
+
+    def test_single_participant_returns_full_list(self):
+        images = self._images(5)
+        out = self.node.split(images, participant_index=0, total_participants=1)[0]
+        self.assertEqual(len(out), 5)
+        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in out], [0, 1, 2, 3, 4])
+
+    def test_two_participants_split_six_images_evenly(self):
+        images = self._images(6)
+        p0 = self.node.split(images, participant_index=0, total_participants=2)[0]
+        p1 = self.node.split(images, participant_index=1, total_participants=2)[0]
+        self.assertEqual(len(p0), 3)
+        self.assertEqual(len(p1), 3)
+        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p0], [0, 1, 2])
+        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p1], [3, 4, 5])
+
+    def test_three_participants_split_seven_images(self):
+        images = self._images(7)
+        p0 = self.node.split(images, participant_index=0, total_participants=3)[0]
+        p1 = self.node.split(images, participant_index=1, total_participants=3)[0]
+        p2 = self.node.split(images, participant_index=2, total_participants=3)[0]
+        self.assertEqual([len(p0), len(p1), len(p2)], [3, 2, 2])
+        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p0], [0, 1, 2])
+        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p1], [3, 4])
+        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p2], [5, 6])
+
+    def test_more_participants_than_images_yields_empty_slices(self):
+        images = self._images(2)
+        outputs = [self.node.split(images, participant_index=i, total_participants=5)[0] for i in range(5)]
+        self.assertEqual([len(chunk) for chunk in outputs], [1, 1, 0, 0, 0])
+
+    def test_mixed_dimensions_are_preserved(self):
+        images = [
+            torch.zeros((1, 4, 4, 3), dtype=torch.float32),
+            torch.zeros((1, 8, 6, 3), dtype=torch.float32),
+            torch.zeros((1, 3, 9, 3), dtype=torch.float32),
+        ]
+
+        p0 = self.node.split(images, participant_index=0, total_participants=2)[0]
+        p1 = self.node.split(images, participant_index=1, total_participants=2)[0]
+
+        self.assertEqual([tuple(t.shape) for t in p0], [(1, 4, 4, 3), (1, 8, 6, 3)])
+        self.assertEqual([tuple(t.shape) for t in p1], [(1, 3, 9, 3)])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_prompt_transform.py

@@ -81,13 +81,36 @@ def _delegate_prompt():
     }
 
 
+def _list_splitter_prompt():
+    """1(List source) → 2(DistributedListSplitter) → 3(DistributedListCollector) → 4(SaveImage)."""
+    return {
+        "1": {"class_type": "ImageListSource", "inputs": {}},
+        "2": {"class_type": "DistributedListSplitter", "inputs": {"images": ["1", 0]}},
+        "3": {"class_type": "DistributedListCollector", "inputs": {"images": ["2", 0]}},
+        "4": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}},
+    }
+
+
+def _branch_prompt():
+    """1 → 2(DistributedBranch) -> slot0:3->4, slot1:5->6."""
+    return {
+        "1": {"class_type": "KSampler", "inputs": {}},
+        "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+        "3": {"class_type": "Blur", "inputs": {"image": ["2", 0]}},
+        "4": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}},
+        "5": {"class_type": "Sharpen", "inputs": {"image": ["2", 1]}},
+        "6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
+    }
+
+
 def _apply(prompt, participant_id, enabled_worker_ids=None, delegate_master=False):
     if enabled_worker_ids is None:
         enabled_worker_ids = ["worker-a", "worker-b"]
-    idx = pt.PromptIndex(prompt)
+    prompt_copy = json.loads(json.dumps(prompt))
+    idx = pt.PromptIndex(prompt_copy)
     job_id_map = pt.generate_job_id_map(idx, "run")
     return pt.apply_participant_overrides(
-        prompt,
+        prompt_copy,
         participant_id=participant_id,
         enabled_worker_ids=enabled_worker_ids,
         job_id_map=job_id_map,
@@ -250,6 +273,23 @@ class PrunePromptForWorkerTests(unittest.TestCase):
         self.assertIn("2", result)
         self.assertNotIn("3", result)
 
+    def test_list_collector_is_treated_as_distributed_anchor(self):
+        prompt = _list_splitter_prompt()
+        result = pt.prune_prompt_for_worker(prompt)
+        self.assertIn("1", result)
+        self.assertIn("2", result)
+        self.assertIn("3", result)
+        self.assertNotIn("4", result)
+
+    def test_branch_anchor_keeps_downstream_for_later_branch_pruning(self):
+        prompt = _branch_prompt()
+        result = pt.prune_prompt_for_worker(prompt)
+        self.assertIn("2", result)
+        self.assertIn("3", result)
+        self.assertIn("4", result)
+        self.assertIn("5", result)
+        self.assertIn("6", result)
+
 
 # ---------------------------------------------------------------------------
 # prepare_delegate_master_prompt
@@ -327,6 +367,22 @@ class GenerateJobIdMapTests(unittest.TestCase):
         job_map = pt.generate_job_id_map(idx, "run")
         self.assertEqual(job_map["5"], "run_5")
 
+    def test_maps_list_collector_nodes(self):
+        prompt = {
+            "8": {"class_type": "DistributedListCollector", "inputs": {}},
+        }
+        idx = pt.PromptIndex(prompt)
+        job_map = pt.generate_job_id_map(idx, "run")
+        self.assertEqual(job_map["8"], "run_8")
+
+    def test_maps_branch_nodes(self):
+        prompt = {
+            "9": {"class_type": "DistributedBranch", "inputs": {"num_branches": 2}},
+        }
+        idx = pt.PromptIndex(prompt)
+        job_map = pt.generate_job_id_map(idx, "run")
+        self.assertEqual(job_map["9"], "run_9")
+
     def test_empty_prompt_returns_empty_map(self):
         idx = pt.PromptIndex({})
         self.assertEqual(pt.generate_job_id_map(idx, "prefix"), {})
@@ -489,5 +545,108 @@ class ApplyOverridesValueTests(unittest.TestCase):
         self.assertEqual(result["1"]["inputs"]["worker_id"], "")
 
 
+class ApplyOverridesListSplitterTests(unittest.TestCase):
+    def _splitter_prompt(self):
+        return {"1": {"class_type": "DistributedListSplitter", "inputs": {}}}
+
+    def test_master_gets_participant_index_zero(self):
+        result = _apply(self._splitter_prompt(), "master", enabled_worker_ids=["worker-a", "worker-b"])
+        self.assertEqual(result["1"]["inputs"]["participant_index"], 0)
+        self.assertEqual(result["1"]["inputs"]["total_participants"], 3)
+
+    def test_worker_gets_incremented_participant_index(self):
+        result = _apply(self._splitter_prompt(), "worker-a", enabled_worker_ids=["worker-a", "worker-b"])
+        self.assertEqual(result["1"]["inputs"]["participant_index"], 1)
+        self.assertEqual(result["1"]["inputs"]["total_participants"], 3)
+
+    def test_delegate_mode_shifts_worker_to_index_zero(self):
+        result = _apply(
+            self._splitter_prompt(),
+            "worker-a",
+            enabled_worker_ids=["worker-a", "worker-b"],
+            delegate_master=True,
+        )
+        self.assertEqual(result["1"]["inputs"]["participant_index"], 0)
+        self.assertEqual(result["1"]["inputs"]["total_participants"], 2)
+
+
+class ApplyOverridesListCollectorTests(unittest.TestCase):
+    def _collector_prompt(self):
+        return {"1": {"class_type": "DistributedListCollector", "inputs": {}}}
+
+    def test_worker_sets_master_url_and_worker_id(self):
+        result = _apply(self._collector_prompt(), "worker-a", enabled_worker_ids=["worker-a"])
+        self.assertTrue(result["1"]["inputs"]["is_worker"])
+        self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com")
+        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a")
+
+    def test_master_sets_delegate_only_and_clears_worker_fields(self):
+        prompt = {
+            "1": {
+                "class_type": "DistributedListCollector",
+                "inputs": {"master_url": "stale", "worker_id": "stale"},
+            }
+        }
+        result = _apply(prompt, "master", enabled_worker_ids=["worker-a"], delegate_master=True)
+        self.assertFalse(result["1"]["inputs"]["is_worker"])
+        self.assertTrue(result["1"]["inputs"]["delegate_only"])
+        self.assertNotIn("master_url", result["1"]["inputs"])
+        self.assertNotIn("worker_id", result["1"]["inputs"])
+
+
+class ApplyOverridesBranchTests(unittest.TestCase):
+    def test_master_gets_branch_zero_worker_gets_branch_one(self):
+        prompt = _branch_prompt()
+        master_result = _apply(prompt, "master", enabled_worker_ids=["worker-a"])
+        worker_result = _apply(prompt, "worker-a", enabled_worker_ids=["worker-a"])
+
+        self.assertEqual(master_result["2"]["inputs"]["assigned_branch"], 0)
+        self.assertEqual(worker_result["2"]["inputs"]["assigned_branch"], 1)
+        self.assertIn("3", master_result)
+        self.assertIn("4", master_result)
+        self.assertNotIn("5", master_result)
+        self.assertNotIn("6", master_result)
+        self.assertIn("5", worker_result)
+        self.assertIn("6", worker_result)
+        self.assertNotIn("3", worker_result)
+        self.assertNotIn("4", worker_result)
+
+    def test_delegate_mode_assigns_worker_a_to_branch_zero(self):
+        prompt = _branch_prompt()
+        worker_result = _apply(
+            prompt,
+            "worker-a",
+            enabled_worker_ids=["worker-a", "worker-b"],
+            delegate_master=True,
+        )
+        self.assertEqual(worker_result["2"]["inputs"]["assigned_branch"], 0)
+
+    def test_pruning_keeps_shared_nodes_between_branches(self):
+        prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+            "3": {"class_type": "NodeA", "inputs": {"image": ["2", 0]}},
+            "4": {"class_type": "NodeB", "inputs": {"image": ["2", 1]}},
+            "5": {"class_type": "SaveImage", "inputs": {"images": ["3", 0], "aux": ["4", 0]}},
+        }
+        worker_result = _apply(prompt, "master", enabled_worker_ids=["worker-a"])
+        self.assertIn("3", worker_result)
+        self.assertNotIn("4", worker_result)
+        self.assertIn("5", worker_result)
+        self.assertNotIn("aux", worker_result["5"]["inputs"])
+
+    def test_more_participants_than_branches_sets_unassigned_to_minus_one(self):
+        prompt = _branch_prompt()
+        worker_b_result = _apply(
+            prompt,
+            "worker-b",
+            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)
+
 if __name__ == "__main__":
     unittest.main()

web/constants.js

@@ -224,6 +224,9 @@ export const ENDPOINTS = {
 
 export const NODE_CLASSES = {
     DISTRIBUTED_COLLECTOR: 'DistributedCollector',
+    DISTRIBUTED_LIST_SPLITTER: 'DistributedListSplitter',
+    DISTRIBUTED_LIST_COLLECTOR: 'DistributedListCollector',
+    DISTRIBUTED_BRANCH: 'DistributedBranch',
     DISTRIBUTED_SEED: 'DistributedSeed',
     DISTRIBUTED_EMPTY_IMAGE: 'DistributedEmptyImage',
     UPSCALE_DISTRIBUTED: 'UltimateSDUpscaleDistributed',

web/executionUtils.js

@@ -7,9 +7,12 @@ export function setupInterceptor(extension) {
     api.queuePrompt = async (number, prompt) => {
         if (extension.isEnabled) {
             const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0;
+            const hasListSplitter = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_LIST_SPLITTER).length > 0;
+            const hasListCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_LIST_COLLECTOR).length > 0;
+            const hasBranch = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_BRANCH).length > 0;
             const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0;
 
-            if (hasCollector || hasDistUpscale) {
+            if (hasCollector || hasListSplitter || hasListCollector || hasBranch || hasDistUpscale) {
                 const result = await executeParallelDistributed(extension, prompt);
                 // Immediate status check for instant feedback
                 checkAllWorkerStatuses(extension);