feat: per-GPU preview outputs on the Distributed Collector

Append master and worker_1..worker_8 IMAGE outputs after images/audio so
each participant's own batch can be previewed. worker_N follows the Nth
configured worker so a slot keeps its GPU when others are toggled; slots
with no images return ExecutionBlocker(None) so attached nodes are
skipped silently. The frontend labels worker_N with the worker's name.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

nodes/collector.py

@@ -21,6 +21,11 @@ from ..utils import transfer_codec
 
 prompt_server = _server.PromptServer.instance
 
+# Per-GPU preview outputs follow the combined images/audio outputs. worker_N is
+# the Nth worker in the config, so a slot keeps its GPU when others are toggled.
+MAX_WORKER_OUTPUTS = 8
+PER_GPU_OUTPUT_NAMES = ("master",) + tuple(f"worker_{i + 1}" for i in range(MAX_WORKER_OUTPUTS))
+
 
 def _encode_all_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes):
     """Materialize every transfer chunk (runs in an executor)."""
@@ -113,7 +118,12 @@ class DistributedCollectorNode:
     RETURN_NAMES = ("images", "audio")
     FUNCTION = "run"
     CATEGORY = "image"
-    
+
+    def __init__(self):
+        # Image batch each participant produced in this run, keyed by
+        # PER_GPU_OUTPUT_NAMES slot. Runtime instances are per-execution.
+        self.participant_images = {}
+
     @staticmethod
     def _unwrap_list_input(value):
         """Unwrap scalar inputs when ComfyUI passes them via INPUT_IS_LIST."""
@@ -232,6 +242,9 @@ class DistributedCollectorNode:
         if not multi_job_id or pass_through:
             if pass_through:
                 debug_log("Collector: pass-through mode enabled, returning images unchanged")
+            elif images is not None:
+                # Not distributed: everything was produced on this GPU.
+                self.participant_images = {"master": images}
             return (images, audio if audio is not None else empty_audio)
 
         # Use async helper to run in server loop
@@ -703,6 +716,49 @@ class DistributedCollectorNode:
             return self._combine_image_tensors([fallback_images], ["master fallback"])
         return None
 
+    @staticmethod
+    def _worker_output_slots(participant_ids):
+        """Map worker ids to worker_N preview slots.
+
+        Configured workers keep their config position; other participants
+        (e.g. serverless) take the free slots after the configured workers.
+        """
+        try:
+            config_ids = [str(worker.get("id")) for worker in load_config().get("workers", [])]
+        except Exception as e:
+            debug_log(f"Collector - Could not read worker order for preview outputs: {e}")
+            config_ids = []
+        slots = {}
+        next_free = len(config_ids)
+        for worker_id in participant_ids:
+            if worker_id in config_ids:
+                index = config_ids.index(worker_id)
+            else:
+                index, next_free = next_free, next_free + 1
+            if index < MAX_WORKER_OUTPUTS:
+                slots[worker_id] = PER_GPU_OUTPUT_NAMES[index + 1]
+            else:
+                debug_log(f"Collector - No preview output for worker {worker_id}; only in the combined images")
+        return slots
+
+    def _split_participant_images(self, worker_images, worker_order, master_images):
+        """Return each participant's own image batch keyed by preview slot."""
+        participants = {}
+        if master_images is not None:
+            participants["master"] = master_images
+        ordered = [str(w) for w in (worker_order or [])]
+        ordered += sorted(str(w) for w in worker_images if str(w) not in ordered)
+        images_by_id = {str(w): frames for w, frames in worker_images.items()}
+        for worker_id, slot in self._worker_output_slots(ordered).items():
+            frames = images_by_id.get(worker_id)
+            if not frames:
+                continue
+            tensors = [ensure_contiguous(frames[idx].cpu()) for idx in sorted(frames)]
+            participants[slot] = self._combine_image_tensors(
+                tensors, [f"worker {worker_id!r} image {idx}" for idx in sorted(frames)]
+            )
+        return participants
+
     async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
         if is_worker:
             # Worker mode: send images and audio to master in a single batch
@@ -781,6 +837,8 @@ class DistributedCollectorNode:
             exact_requested = False
         num_workers = len(expected_workers)
         if num_workers == 0:
+            if images is not None:
+                self.participant_images = {"master": images}
             return (images, audio if audio is not None else self.EMPTY_AUDIO)
 
         # Create the queue before any expensive local work to avoid job_complete race.
@@ -1021,6 +1079,10 @@ class DistributedCollectorNode:
                 else:
                     debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
                               f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
+                    self.participant_images = self._split_participant_images(
+                        worker_images, enabled_workers,
+                        None if delegate_mode or master_batch_size == 0 else images_on_cpu,
+                    )
                 return (combined, combined_audio)
             except Exception as e:
                 if serverless_results or exact_requested:

nodes/v3.py

@@ -5,9 +5,10 @@ This avoids sharing mutable instance state through V3's sanitized class clones.
 """
 import comfy.samplers
 from comfy_api.v0_0_2 import io
+from comfy_execution.graph_utils import ExecutionBlocker
 
 from . import utilities as _utilities
-from .collector import DistributedCollectorNode as _CollectorRuntime
+from .collector import DistributedCollectorNode as _CollectorRuntime, PER_GPU_OUTPUT_NAMES
 from ..utils import transfer_codec as _transfer_codec
 from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
 from .image_list import DistributedImageListEnd, DistributedImageListStart, DistributedImageListGetValue
@@ -130,7 +131,13 @@ class DistributedCollector(io.ComfyNode):
                     io.Int.Input('transfer_quality', default=_transfer_codec.DEFAULT_TRANSFER_QUALITY,
                                  min=1, max=100, step=1, optional=True,
                                  tooltip=_transfer_codec.TRANSFER_QUALITY_TOOLTIP)],
-            outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
+            # Per-GPU outputs are appended so saved links to images/audio keep their slots.
+            outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')] + [
+                io.Image.Output(id=name, display_name=name,
+                                tooltip='Images produced by this GPU only, for previewing. '
+                                        'worker_N is the Nth worker in the Distributed panel. '
+                                        'Nodes attached to a GPU that did not take part are skipped.')
+                for name in PER_GPU_OUTPUT_NAMES],
             is_input_list=True, accept_all_inputs=True,
         )
 
@@ -140,12 +147,16 @@ class DistributedCollector(io.ComfyNode):
                 worker_id='', pass_through=False, delegate_only=False,
                 transfer_format=_transfer_codec.DEFAULT_TRANSFER_FORMAT,
                 transfer_quality=_transfer_codec.DEFAULT_TRANSFER_QUALITY):
-        return io.NodeOutput(*_CollectorRuntime().run(
+        runtime = _CollectorRuntime()
+        result = runtime.run(
             images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
             is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
             worker_batch_size=worker_batch_size, worker_id=worker_id,
             pass_through=pass_through, delegate_only=delegate_only,
-            transfer_format=transfer_format, transfer_quality=transfer_quality))
+            transfer_format=transfer_format, transfer_quality=transfer_quality)
+        # A silent blocker skips previews wired to a GPU that produced nothing.
+        per_gpu = [runtime.participant_images.get(name, ExecutionBlocker(None)) for name in PER_GPU_OUTPUT_NAMES]
+        return io.NodeOutput(*result, *per_gpu)
 
 
 class UltimateSDUpscaleDistributed(io.ComfyNode):

tests/helpers/v3_runtime_check.py

@@ -43,6 +43,11 @@ ADDED_OPTIONAL_INPUTS = {
     },
 }
 
+# Outputs appended after the V1 baseline outputs, so saved links keep their slots.
+ADDED_OUTPUTS = {
+    'DistributedCollector': ['master'] + [f'worker_{index + 1}' for index in range(8)],
+}
+
 # Nodes added after the V1 baseline; not part of saved-workflow contracts.
 # DistributedWorkerSink is internal (dev-only) and only inserted by the master.
 ADDED_NODES = {'DistributedWorkerSink', 'DistributedImageListStart', 'DistributedImageListEnd',
@@ -374,7 +379,8 @@ async def main():
         assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id]
     assert 'DistributedMediaCollector' not in mapping
     collector_info = mapping['DistributedCollector'].GET_NODE_INFO_V1()
-    assert collector_info['output'] == ['IMAGE', 'AUDIO'], collector_info
+    assert collector_info['output'] == ['IMAGE', 'AUDIO'] + ['IMAGE'] * 9, collector_info
+    assert collector_info['output_name'][:2] == ['images', 'audio'], collector_info
     assert not any(collector_info['output_is_list']), collector_info
     assert list(collector_info['input']['optional']) == [
         'images', 'audio', 'transfer_format', 'transfer_quality']
@@ -393,6 +399,11 @@ async def main():
             assert old['output'] == ['*'] * 10
             old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10
         new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
+        added_outputs = ADDED_OUTPUTS.get(node_id, [])
+        if added_outputs:
+            assert new['output_name'][len(old['output']):] == added_outputs, (node_id, new['output_name'])
+            for key in ('output', 'output_name', 'output_is_list'):
+                new[key] = new[key][:len(old['output'])]
         for group in ('required', 'optional'):
             old_inputs = old['input'].get(group, {})
             new_inputs = dict(new['input'].get(group, {}))

tests/test_collector_image_transport.py

@@ -223,3 +223,63 @@ def test_empty_image_results_raise_instead_of_reaching_save_image(media_modules)
     with pytest.raises(ValueError, match="empty image batch"):
         asyncio.run(_master_collect(module, torch.empty(0, 64, 64, 3),
                                     [torch.empty(0, 8, 12, 3)], delegate=True))
+
+
+def test_master_collection_records_each_participant_in_its_config_slot(media_modules):
+    module, _ = media_modules
+    module.load_config = lambda: {"workers": [{"id": "worker-off"}, {"id": "worker-a"}]}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    worker_frames = [torch.full((1, 8, 12, 3), value) for value in (0.25, 0.5)]
+    # Out-of-order delivery must still preview in image order.
+    for index in (1, 0):
+        queue.put_nowait({"worker_id": "worker-a", "image_index": index,
+                          "tensor": worker_frames[index], "is_last": index == 0})
+    queue.put_nowait({"worker_id": "serverless:pool", "image_index": 0,
+                      "tensor": torch.full((1, 8, 12, 3), 0.75), "is_last": True})
+    module.prompt_server.distributed_pending_jobs = {"image-test": queue}
+    master = torch.zeros(1, 8, 12, 3)
+    collector = module.DistributedCollectorNode()
+
+    combined, _ = asyncio.run(collector.execute(
+        master, None, multi_job_id="image-test", enabled_worker_ids='["worker-a", "serverless:pool"]'))
+
+    previews = collector.participant_images
+    assert combined.shape[0] == 4
+    # worker-a is the second configured worker; the unconfigured serverless
+    # participant takes the first slot after the configured workers.
+    assert set(previews) == {"master", "worker_2", "worker_3"}
+    assert torch.equal(previews["master"], master)
+    assert previews["worker_2"][:, 0, 0, 0].tolist() == [0.25, 0.5]
+    assert previews["worker_3"][:, 0, 0, 0].tolist() == [0.75]
+
+
+def test_delegate_master_has_no_master_preview(media_modules):
+    module, _ = media_modules
+    module.load_config = lambda: {"workers": [{"id": "worker-a"}]}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    worker = torch.full((1, 8, 12, 3), 0.5)
+    queue.put_nowait({"worker_id": "worker-a", "image_index": 0, "tensor": worker, "is_last": True})
+    module.prompt_server.distributed_pending_jobs = {"image-test": queue}
+    collector = module.DistributedCollectorNode()
+
+    asyncio.run(collector.execute(
+        torch.empty(0, 64, 64, 3), None, multi_job_id="image-test",
+        enabled_worker_ids='["worker-a"]', delegate_only=True))
+
+    assert list(collector.participant_images) == ["worker_1"]
+    assert torch.equal(collector.participant_images["worker_1"], worker)
+
+
+def test_undistributed_run_previews_everything_as_master(media_modules):
+    module, _ = media_modules
+    images = torch.zeros(2, 8, 12, 3)
+    collector = module.DistributedCollectorNode()
+    collector.run(images=[images])
+    assert list(collector.participant_images) == ["master"]
+    assert torch.equal(collector.participant_images["master"], images)
+
+    pass_through = module.DistributedCollectorNode()
+    pass_through.run(images=[images], multi_job_id=["job"], pass_through=[True])
+    assert pass_through.participant_images == {}

tests/test_v3_collector_transfer.py

@@ -13,6 +13,7 @@ import pytest
 
 ROOT = Path(__file__).resolve().parents[1]
 PKG = "dist_v3_transfer_testpkg"
+PER_GPU_OUTPUT_NAMES = ("master", "worker_1", "worker_2")
 
 
 class _Field:
@@ -51,10 +52,20 @@ def v3(monkeypatch):
     calls = []
 
     class CollectorRuntime:
+        participant_images_for_run = {}
+
+        def __init__(self):
+            self.participant_images = {}
+
         def run(self, **kwargs):
             calls.append(kwargs)
+            self.participant_images = dict(self.participant_images_for_run)
             return (kwargs["images"], kwargs["audio"])
 
+    class ExecutionBlocker:
+        def __init__(self, message):
+            self.message = message
+
     def module(name, **attrs):
         mod = types.ModuleType(name)
         mod.__dict__.update(attrs)
@@ -68,13 +79,16 @@ def v3(monkeypatch):
     module(f"{PKG}.utils", __path__=[])
     module(f"{PKG}.nodes.utilities")
     module(f"{PKG}.nodes.distributed_upscale", UltimateSDUpscaleDistributed=object)
-    module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime)
+    module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime,
+           PER_GPU_OUTPUT_NAMES=PER_GPU_OUTPUT_NAMES)
     module(f"{PKG}.nodes.image_list", DistributedImageListStart=object, DistributedImageListEnd=object,
            DistributedImageListGetValue=object)
     module("comfy", __path__=[])
     module("comfy.samplers", KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["normal"]))
     module("comfy_api", __path__=[])
     module("comfy_api.v0_0_2", io=_fake_io())
+    module("comfy_execution", __path__=[])
+    module("comfy_execution.graph_utils", ExecutionBlocker=ExecutionBlocker)
 
     codec_spec = importlib.util.spec_from_file_location(f"{PKG}.utils.transfer_codec", ROOT / "utils" / "transfer_codec.py")
     codec = importlib.util.module_from_spec(codec_spec)
@@ -85,6 +99,7 @@ def v3(monkeypatch):
     mod = importlib.util.module_from_spec(spec)
     monkeypatch.setitem(sys.modules, f"{PKG}.nodes.v3", mod)
     spec.loader.exec_module(mod)
+    mod._TestRuntime = CollectorRuntime
     return mod, codec, calls
 
 
@@ -127,3 +142,24 @@ def test_registered_collector_defaults_to_legacy_transfer(v3):
 
     assert calls[-1]["transfer_format"] == "legacy_png"
     assert calls[-1]["transfer_quality"] == 90
+
+
+def test_registered_collector_appends_per_gpu_outputs_after_combined_outputs(v3):
+    mod, _, _ = v3
+    schema = mod.DistributedCollector.define_schema()
+
+    assert [field.kind for field in schema.outputs] == ["Image", "Audio", "Image", "Image", "Image"]
+    assert [field.kwargs["display_name"] for field in schema.outputs] == [
+        "images", "audio", "master", "worker_1", "worker_2"]
+
+
+def test_registered_collector_blocks_per_gpu_outputs_without_images(v3):
+    mod, _, _ = v3
+    mod._TestRuntime.participant_images_for_run = {"master": "m", "worker_2": "w2"}
+
+    _, values, _ = mod.DistributedCollector.execute(images=["img"], audio=["aud"])
+
+    assert values[:2] == (["img"], ["aud"])
+    assert values[2] == "m"
+    assert type(values[3]).__name__ == "ExecutionBlocker" and values[3].message is None
+    assert values[4] == "w2"

web/collectorOutputs.js

@@ -0,0 +1,76 @@
+import { app } from "/scripts/app.js";
+import { ENDPOINTS, NODE_CLASSES } from "./constants.js";
+
+// worker_N previews the Nth configured worker (see nodes/collector.py).
+const WORKER_OUTPUT_PREFIX = "worker_";
+const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
+
+const trackedNodes = new Set();
+let workers = [];
+let listenerAttached = false;
+
+async function fetchWorkers() {
+    try {
+        const resp = await fetch(ENDPOINTS.CONFIG);
+        if (!resp.ok) return [];
+        const config = await resp.json();
+        return Array.isArray(config.workers) ? config.workers : [];
+    } catch {
+        return [];
+    }
+}
+
+// Only the displayed label changes; slot names stay fixed so links survive.
+function labelWorkerOutputs(node, configWorkers) {
+    for (const output of node.outputs || []) {
+        if (!output?.name?.startsWith(WORKER_OUTPUT_PREFIX)) continue;
+        const index = parseInt(output.name.slice(WORKER_OUTPUT_PREFIX.length), 10) - 1;
+        const worker = configWorkers[index];
+        output.label = worker?.name ? `${output.name} (${worker.name})` : output.name;
+    }
+    node.setDirtyCanvas?.(true, true);
+}
+
+function relabelAll() {
+    for (const node of trackedNodes) labelWorkerOutputs(node, workers);
+}
+
+function attachListener() {
+    if (listenerAttached || typeof window === "undefined") return;
+    window.addEventListener(WORKERS_CHANGED_EVENT, (event) => {
+        if (Array.isArray(event?.detail?.workers)) {
+            workers = event.detail.workers;
+            relabelAll();
+        }
+    });
+    listenerAttached = true;
+}
+
+app.registerExtension({
+    name: "Distributed.CollectorOutputs",
+    async nodeCreated(node) {
+        if (node.comfyClass !== NODE_CLASSES.DISTRIBUTED_COLLECTOR) return;
+        try {
+            attachListener();
+            trackedNodes.add(node);
+
+            const originalConfigure = node.configure;
+            node.configure = function (data) {
+                const result = originalConfigure ? originalConfigure.call(this, data) : undefined;
+                labelWorkerOutputs(this, workers);
+                return result;
+            };
+
+            const originalOnRemoved = node.onRemoved;
+            node.onRemoved = function () {
+                trackedNodes.delete(this);
+                return originalOnRemoved ? originalOnRemoved.call(this) : undefined;
+            };
+
+            if (workers.length === 0) workers = await fetchWorkers();
+            labelWorkerOutputs(node, workers);
+        } catch (error) {
+            console.error("Error in DistributedCollector outputs extension:", error);
+        }
+    },
+});

web/tests/collectorOutputs.test.js

@@ -0,0 +1,45 @@
+import { readFileSync } from "node:fs";
+import vm from "node:vm";
+import { describe, expect, it } from "vitest";
+
+const source = readFileSync(new URL("../collectorOutputs.js", import.meta.url), "utf8")
+    .replace(/^import .*;\r?\n/gm, "");
+
+function setup(workers) {
+    let extension;
+    const listeners = new Map();
+    vm.runInNewContext(source, {
+        app: { registerExtension: (ext) => { extension = ext; } },
+        ENDPOINTS: { CONFIG: "/distributed/config" },
+        NODE_CLASSES: { DISTRIBUTED_COLLECTOR: "DistributedCollector" },
+        fetch: async () => ({ ok: true, json: async () => ({ workers }) }),
+        window: { addEventListener: (name, fn) => listeners.set(name, fn) },
+        parseInt,
+        Array,
+        console,
+    });
+    const node = {
+        comfyClass: "DistributedCollector",
+        outputs: ["images", "audio", "master", "worker_1", "worker_2", "worker_3"].map((name) => ({ name })),
+    };
+    return { extension, listeners, node };
+}
+
+describe("DistributedCollector per-GPU outputs", () => {
+    it("labels worker_N with the Nth configured worker and keeps slot names", async () => {
+        const { extension, node } = setup([{ id: "a", name: "GPU A", enabled: false }, { id: "b", name: "GPU B" }]);
+        await extension.nodeCreated(node);
+        expect(node.outputs.map((o) => o.name)).toEqual(
+            ["images", "audio", "master", "worker_1", "worker_2", "worker_3"]);
+        expect(node.outputs.slice(3).map((o) => o.label)).toEqual(
+            ["worker_1 (GPU A)", "worker_2 (GPU B)", "worker_3"]);
+        expect(node.outputs[2].label).toBeUndefined();
+    });
+
+    it("relabels when the worker list changes", async () => {
+        const { extension, listeners, node } = setup([{ id: "a", name: "GPU A" }]);
+        await extension.nodeCreated(node);
+        listeners.get("distributed:workers-changed")({ detail: { workers: [{ id: "c", name: "Renamed" }] } });
+        expect(node.outputs[3].label).toBe("worker_1 (Renamed)");
+    });
+});