From 910c260857f188cd7af748f747b3c84d1f80a65e Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sun, 1 Mar 2026 22:48:08 +0000 Subject: [PATCH] Add distributed list splitting/collection and branch routing --- api/orchestration/prompt_transform.py | 245 +++++++++++++++++++++++- api/queue_orchestration.py | 15 +- distributed.py | 3 + nodes/__init__.py | 9 + nodes/branch.py | 26 +++ nodes/list_collector.py | 266 ++++++++++++++++++++++++++ nodes/list_splitter.py | 43 +++++ tests/test_branch.py | 76 ++++++++ tests/test_list_collector.py | 206 ++++++++++++++++++++ tests/test_list_splitter.py | 108 +++++++++++ tests/test_prompt_transform.py | 163 +++++++++++++++- web/constants.js | 3 + web/executionUtils.js | 5 +- 13 files changed, 1158 insertions(+), 10 deletions(-) create mode 100644 nodes/branch.py create mode 100644 nodes/list_collector.py create mode 100644 nodes/list_splitter.py create mode 100644 tests/test_branch.py create mode 100644 tests/test_list_collector.py create mode 100644 tests/test_list_splitter.py diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 244bd6c..d07c666 100644 --- a/api/orchestration/prompt_transform.py +++ b/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 diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 06f3a7f..3e063f5 100644 --- a/api/queue_orchestration.py +++ b/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( diff --git a/distributed.py b/distributed.py index f44bc7c..2f77856 100644 --- a/distributed.py +++ b/distributed.py @@ -23,6 +23,9 @@ from .nodes import ( DistributedValue, AudioBatchDivider, DistributedEmptyImage, + DistributedListSplitter, + DistributedListCollector, + DistributedBranch, AnyType, ByPassTypeTuple, any_type, diff --git a/nodes/__init__.py b/nodes/__init__.py index d7327b3..33e10eb 100644 --- a/nodes/__init__.py +++ b/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", diff --git a/nodes/branch.py b/nodes/branch.py new file mode 100644 index 0000000..44263b5 --- /dev/null +++ b/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) diff --git a/nodes/list_collector.py b/nodes/list_collector.py new file mode 100644 index 0000000..031cc2b --- /dev/null +++ b/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,) diff --git a/nodes/list_splitter.py b/nodes/list_splitter.py new file mode 100644 index 0000000..aed4991 --- /dev/null +++ b/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]),) diff --git a/tests/test_branch.py b/tests/test_branch.py new file mode 100644 index 0000000..7b52898 --- /dev/null +++ b/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() diff --git a/tests/test_list_collector.py b/tests/test_list_collector.py new file mode 100644 index 0000000..6772583 --- /dev/null +++ b/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() diff --git a/tests/test_list_splitter.py b/tests/test_list_splitter.py new file mode 100644 index 0000000..f38a0c8 --- /dev/null +++ b/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() diff --git a/tests/test_prompt_transform.py b/tests/test_prompt_transform.py index 0b5bf6f..cd6be7b 100644 --- a/tests/test_prompt_transform.py +++ b/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() diff --git a/web/constants.js b/web/constants.js index af2400b..896f187 100644 --- a/web/constants.js +++ b/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', diff --git a/web/executionUtils.js b/web/executionUtils.js index 703ee4a..c9438b6 100644 --- a/web/executionUtils.js +++ b/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);