Fix collector handling of ComfyUI list inputs

.github/FUNDING.yml

@@ -1,4 +1,3 @@
 # These are supported funding model platforms
 
 github: robertvoy
-buy_me_a_coffee: robertvoy

.github/workflows/publish_action.yml

@@ -19,6 +19,6 @@ jobs:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@v1
+        uses: Comfy-Org/publish-node-action@main
         with:
           personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

api/job_routes.py

@@ -217,7 +217,7 @@ async def distributed_queue_endpoint(request):
         return await handle_api_error(request, exc, 400)
 
     try:
-        prompt_id, worker_count = await orchestrate_distributed_execution(
+        prompt_id, prompt_number, worker_count, node_errors = await orchestrate_distributed_execution(
             payload.prompt,
             payload.workflow_meta,
             payload.client_id,
@@ -227,6 +227,8 @@ async def distributed_queue_endpoint(request):
         )
         return web.json_response({
             "prompt_id": prompt_id,
+            "number": prompt_number,
+            "node_errors": node_errors,
             "worker_count": worker_count,
             "auto_prepare_supported": True,
         })

api/queue_orchestration.py

@@ -208,7 +208,7 @@ async def orchestrate_distributed_execution(
     """Core orchestration logic for the /distributed/queue endpoint.
 
     Returns:
-        tuple[str, int]: (prompt_id, worker_count)
+        tuple[str, int, int, dict]: (prompt_id, number, worker_count, node_errors)
     """
     ensure_distributed_state()
     execution_trace_id = trace_execution_id or _generate_execution_trace_id()
@@ -317,8 +317,18 @@ async def orchestrate_distributed_execution(
 
     if not job_id_map:
         trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
-        prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id)
-        return prompt_id, 0
+        queue_result = await queue_prompt_payload(
+            prompt_obj,
+            workflow_meta,
+            client_id,
+            include_queue_metadata=True,
+        )
+        return (
+            queue_result["prompt_id"],
+            queue_result["number"],
+            0,
+            queue_result.get("node_errors", {}),
+        )
 
     for job_id in job_id_map.values():
         await _ensure_distributed_queue(job_id)
@@ -392,9 +402,17 @@ async def orchestrate_distributed_execution(
             ]
         )
 
-    prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id)
+    queue_result = await queue_prompt_payload(
+        master_prompt,
+        workflow_meta,
+        client_id,
+        include_queue_metadata=True,
+    )
+    prompt_id = queue_result["prompt_id"]
+    prompt_number = queue_result["number"]
+    node_errors = queue_result.get("node_errors", {})
     trace_debug(
         execution_trace_id,
         f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
     )
-    return prompt_id, len(worker_payloads)
+    return prompt_id, prompt_number, len(worker_payloads), node_errors

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}
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "ComfyUI-Distributed"
 description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
-version = "1.4.1"
+version = "1.4.4"
 license = {file = "LICENSE"}
 dependencies  = []
 

tests/api/test_distributed_queue.py

@@ -162,7 +162,7 @@ def _load_job_routes_module():
     sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
 
     queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
-    queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 1))
+    queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 7, 1, {}))
     sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
 
     @dataclass(frozen=True)
@@ -221,7 +221,7 @@ job_routes = _load_job_routes_module()
 
 
 class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
-    async def test_distributed_queue_happy_path_returns_prompt_id(self):
+    async def test_distributed_queue_happy_path_returns_prompt_metadata(self):
         request = _FakeRequest(
             {
                 "prompt": {"1": {"class_type": "Node"}},
@@ -233,12 +233,14 @@ class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
         with patch.object(
             job_routes,
             "orchestrate_distributed_execution",
-            new=AsyncMock(return_value=("prompt_123", 2)),
+            new=AsyncMock(return_value=("prompt_123", 42, 2, {})),
         ):
             response = await job_routes.distributed_queue_endpoint(request)
 
         self.assertEqual(response.status, 200)
         self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
+        self.assertEqual(response.payload.get("number"), 42)
+        self.assertEqual(response.payload.get("node_errors"), {})
         self.assertTrue(response.payload.get("auto_prepare_supported"))
 
     async def test_distributed_queue_missing_prompt_returns_400(self):

tests/test_async_helpers.py

@@ -0,0 +1,91 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+class _PromptQueue:
+    def __init__(self):
+        self.items = []
+
+    def put(self, item):
+        self.items.append(item)
+
+
+def _load_async_helpers_module():
+    module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py"
+    package_name = "dist_async_helpers_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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    execution_module = types.ModuleType("execution")
+
+    async def _validate_prompt(prompt_id, prompt, partial_execution_targets):
+        return (True, None, ["9"], {})
+
+    execution_module.validate_prompt = _validate_prompt
+    execution_module.SENSITIVE_EXTRA_DATA_KEYS = []
+    sys.modules["execution"] = execution_module
+
+    prompt_server = types.SimpleNamespace(
+        trigger_on_prompt=lambda payload: payload,
+        number=12,
+        prompt_queue=_PromptQueue(),
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
+    sys.modules["server"] = server_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.get_server_loop = lambda: None
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.utils.async_helpers", 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, prompt_server
+
+
+async_helpers, prompt_server = _load_async_helpers_module()
+
+
+class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
+    async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
+        result = await async_helpers.queue_prompt_payload(
+            {"1": {"class_type": "Node"}},
+            workflow_meta={"id": "workflow-1"},
+            client_id="client-1",
+            include_queue_metadata=True,
+        )
+
+        self.assertIsInstance(result["prompt_id"], str)
+        self.assertTrue(result["prompt_id"])
+        self.assertEqual(result["number"], 12)
+        self.assertEqual(result["node_errors"], {})
+
+        self.assertEqual(prompt_server.number, 13)
+        self.assertEqual(len(prompt_server.prompt_queue.items), 1)
+        queued_item = prompt_server.prompt_queue.items[0]
+        self.assertEqual(queued_item[0], 12)
+        extra_data = queued_item[3]
+        self.assertEqual(extra_data["client_id"], "client-1")
+        self.assertIn("create_time", extra_data)
+        self.assertIsInstance(extra_data["create_time"], int)
+        self.assertGreater(extra_data["create_time"], 0)
+        self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
+
+
+if __name__ == "__main__":
+    unittest.main()

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"}

utils/async_helpers.py

@@ -3,6 +3,7 @@ Async helper utilities for ComfyUI-Distributed.
 """
 import asyncio
 import threading
+import time
 import uuid
 import execution
 import server
@@ -104,7 +105,12 @@ class PromptValidationError(RuntimeError):
         super().__init__(f"Invalid prompt: {merged}")
 
 
-async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
+async def queue_prompt_payload(
+    prompt_obj,
+    workflow_meta=None,
+    client_id=None,
+    include_queue_metadata=False,
+):
     """Validate and queue a prompt via ComfyUI's prompt queue."""
     payload = {"prompt": prompt_obj}
     payload = prompt_server.trigger_on_prompt(payload)
@@ -117,7 +123,7 @@ async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
         node_errors = valid[3] if len(valid) > 3 else {}
         raise PromptValidationError(error_payload, node_errors)
 
-    extra_data = {}
+    extra_data = {"create_time": int(time.time() * 1000)}
     if workflow_meta:
         extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
     if client_id:
@@ -132,4 +138,12 @@ async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
     prompt_server.number = number + 1
     prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
     prompt_server.prompt_queue.put(prompt_queue_item)
+
+    if include_queue_metadata:
+        return {
+            "prompt_id": prompt_id,
+            "number": number,
+            "node_errors": {},
+        }
+
     return prompt_id