From e2125c7529f7603b2421e1a2df7e4838fab16052 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 26 May 2026 23:02:26 +0000 Subject: [PATCH] Fix collector handling of ComfyUI list inputs --- nodes/collector.py | 59 ++++++++ tests/test_collector_list_inputs.py | 205 ++++++++++++++++++++++++++++ 2 files changed, 264 insertions(+) create mode 100644 tests/test_collector_list_inputs.py diff --git a/nodes/collector.py b/nodes/collector.py index f32a8ca..4528985 100644 --- a/nodes/collector.py +++ b/nodes/collector.py @@ -22,6 +22,7 @@ prompt_server = _server.PromptServer.instance class DistributedCollectorNode: + INPUT_IS_LIST = True EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100} @classmethod @@ -55,7 +56,65 @@ class DistributedCollectorNode: FUNCTION = "run" CATEGORY = "image" + @staticmethod + def _unwrap_list_input(value): + """Unwrap scalar inputs when ComfyUI passes them via INPUT_IS_LIST.""" + if isinstance(value, (list, tuple)) and len(value) == 1: + return value[0] + return value + + def _normalize_images_input(self, images): + """Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor.""" + if isinstance(images, (list, tuple)): + if not images: + raise ValueError("Collector received an empty image list") + if not all(isinstance(image, torch.Tensor) for image in images): + raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances") + if len(images) == 1: + return ensure_contiguous(images[0]) + return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0)) + return ensure_contiguous(images) + + def _normalize_audio_input(self, audio): + """Collapse ComfyUI list AUDIO inputs into a single AUDIO payload when present.""" + if not isinstance(audio, (list, tuple)): + return audio + + audio_items = [item for item in audio if item is not None] + if not audio_items: + return None + if len(audio_items) == 1: + return audio_items[0] + + waveforms = [] + sample_rate = 44100 + for item in audio_items: + if not isinstance(item, dict): + raise TypeError("Collector expected AUDIO list items to be dictionaries") + waveform = item.get("waveform") + if waveform is None or waveform.numel() == 0: + continue + waveforms.append(waveform) + if sample_rate == 44100: + sample_rate = item.get("sample_rate", 44100) + + if not waveforms: + return None + return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate} + def run(self, images, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False): + images = self._normalize_images_input(images) + audio = self._normalize_audio_input(audio) + load_balance = self._unwrap_list_input(load_balance) + multi_job_id = self._unwrap_list_input(multi_job_id) + is_worker = self._unwrap_list_input(is_worker) + master_url = self._unwrap_list_input(master_url) + enabled_worker_ids = self._unwrap_list_input(enabled_worker_ids) + worker_batch_size = self._unwrap_list_input(worker_batch_size) + worker_id = self._unwrap_list_input(worker_id) + pass_through = self._unwrap_list_input(pass_through) + delegate_only = self._unwrap_list_input(delegate_only) + # Create empty audio if not provided empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100} diff --git a/tests/test_collector_list_inputs.py b/tests/test_collector_list_inputs.py new file mode 100644 index 0000000..e8de4f1 --- /dev/null +++ b/tests/test_collector_list_inputs.py @@ -0,0 +1,205 @@ +import importlib.util +import sys +import types +import asyncio +from pathlib import Path + +import torch + + +def _load_collector_module(): + module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py" + package_name = "dist_collector_list_testpkg" + + 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 + + class _Routes: + def post(self, _path): + return lambda fn: fn + + def get(self, _path): + return lambda fn: fn + + prompt_server = types.SimpleNamespace( + routes=_Routes(), + distributed_jobs_lock=None, + distributed_pending_jobs={}, + ) + server_module = types.ModuleType("server") + server_module.PromptServer = types.SimpleNamespace(instance=prompt_server) + sys.modules["server"] = server_module + + comfy_module = types.ModuleType("comfy") + model_management = types.ModuleType("comfy.model_management") + + class InterruptProcessingException(Exception): + pass + + model_management.InterruptProcessingException = InterruptProcessingException + model_management.throw_exception_if_processing_interrupted = lambda: None + comfy_module.model_management = model_management + + comfy_utils = types.ModuleType("comfy.utils") + + class ProgressBar: + def __init__(self, _total): + self.total = _total + self.updates = [] + + def update(self, value): + self.updates.append(value) + + comfy_utils.ProgressBar = ProgressBar + comfy_module.utils = comfy_utils + sys.modules["comfy"] = comfy_module + sys.modules["comfy.model_management"] = model_management + sys.modules["comfy.utils"] = comfy_utils + + aiohttp_module = types.ModuleType("aiohttp") + aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total) + sys.modules["aiohttp"] = aiohttp_module + + 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.1 + config_module.load_config = lambda: {"workers": []} + 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() if hasattr(tensor, "contiguous") else tensor + + image_module.ensure_contiguous = _ensure_contiguous + image_module.tensor_to_pil = lambda *_args, **_kwargs: None + image_module.pil_to_tensor = lambda value: value + sys.modules[f"{package_name}.utils.image"] = image_module + + network_module = types.ModuleType(f"{package_name}.utils.network") + network_module.build_worker_url = lambda worker: "http://worker" + network_module.get_client_session = lambda: None + network_module.probe_worker = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.utils.network"] = network_module + + audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload") + audio_payload_module.encode_audio_payload = lambda audio: audio + sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_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 + + spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.collector", module_path) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_collector_opts_into_comfyui_list_inputs(): + collector = _load_collector_module().DistributedCollectorNode + + assert collector.INPUT_IS_LIST is True + + +def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs(): + collector = _load_collector_module().DistributedCollectorNode() + first = torch.zeros(1, 2, 2, 3) + second = torch.ones(1, 2, 2, 3) + + images, audio = collector.run( + images=[first, second], + load_balance=[False], + audio=[None], + multi_job_id=[""], + is_worker=[False], + master_url=[""], + enabled_worker_ids=["[]"], + worker_batch_size=[1], + worker_id=[""], + pass_through=[False], + delegate_only=[False], + ) + + assert tuple(images.shape) == (2, 2, 2, 3) + assert torch.equal(images[0:1], first) + assert torch.equal(images[1:2], second) + assert tuple(audio["waveform"].shape) == (1, 2, 1) + + +def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final_item(): + module = _load_collector_module() + collector = module.DistributedCollectorNode() + first = torch.zeros(1, 2, 2, 3) + second = torch.ones(1, 2, 2, 3) + posted_payloads = [] + + class _FakeImage: + def save(self, fp, format=None, compress_level=None): + fp.write(b"png-bytes") + + class _FakeResponse: + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + def raise_for_status(self): + return None + + class _FakeSession: + def post(self, url, json, timeout): + posted_payloads.append(json) + return _FakeResponse() + + async def _fake_get_client_session(): + return _FakeSession() + + module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage() + module.get_client_session = _fake_get_client_session + module.encode_audio_payload = lambda _audio: None + + images, audio = collector.run( + images=[first, second], + load_balance=[False], + audio=[None], + multi_job_id=["job-list-1"], + is_worker=[True], + master_url=["http://master"], + enabled_worker_ids=["[]"], + worker_batch_size=[1], + worker_id=["worker-a"], + pass_through=[False], + delegate_only=[False], + ) + + assert tuple(images.shape) == (2, 2, 2, 3) + assert tuple(audio["waveform"].shape) == (1, 2, 1) + assert len(posted_payloads) == 2 + assert [payload["batch_idx"] for payload in posted_payloads] == [0, 1] + assert [payload["is_last"] for payload in posted_payloads] == [False, True] + assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"} + assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}