fix: harden per-GPU preview outputs after review

- Typed serverless pools validate the unpruned workflow, so a preview on
  any new slot failed the whole run; both serverless compilers now accept
  every Collector output slot (COLLECTOR_OUTPUT_COUNT).
- Serverless pools are config workers: their serverless:<pool>:<job>
  participants now share the pool's worker_N slot instead of overflowing.
- A worker returning only empty frames no longer raises in the preview
  split (the combined batch already skipped them).
- Worker previews are views into the combined batch, not a second copy,
  and a frame-count guard skips them rather than misattribute frames.
- Frontend: one shared config fetch per load; unnamed slots drop their
  saved label instead of storing a redundant one.
- Tooltip notes that the per-GPU outputs are unavailable downstream of
  Ultimate SD Upscale Distributed.

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

api/orchestration/serverless_media_prompt.py

@@ -16,6 +16,7 @@ from graphlib import CycleError, TopologicalSorter
 from ...serverless_bridge.contract import SOCKET_KINDS, VALUE_SOCKETS, identifier, validate_bindings
 from ...serverless_bridge.image_transfer import DEFAULT_FORMAT, validate_options
 from .serverless_prompt import (
+    COLLECTOR_OUTPUT_COUNT,
     _PRIMITIVE_NODES,
     _REMOVED_SCALAR_NODES,
     _is_link,
@@ -99,7 +100,7 @@ def _validate_prompt(prepared_prompt, *, require_collectors=True):
             if source_class in _REMOVED_SCALAR_NODES or source_class in _PRIMITIVE_NODES:
                 if slot != 0:
                     raise ValueError(f"invalid output slot {slot} for scalar node {source_id!r}")
-            elif source_class == _COLLECTOR and slot >= len(COLLECTOR_SOCKETS):
+            elif source_class == _COLLECTOR and slot >= COLLECTOR_OUTPUT_COUNT:
                 raise ValueError(f"invalid output slot {slot} for collector {source_id!r}")
             parents.add(source_id)
         dependencies[node_id] = parents

api/orchestration/serverless_prompt.py

@@ -29,8 +29,12 @@ _PRIMITIVE_NODES = {
     "PrimitiveInt": 0,
     "PrimitiveString": "",
 }
+# DistributedCollector outputs: images, audio, then the per-GPU previews
+# (PER_GPU_OUTPUT_NAMES in nodes/collector.py). Previews are valid link
+# sources in a workflow even though a pool only ever returns images and audio.
+COLLECTOR_OUTPUT_COUNT = 2 + 9
 _REMOVED_NODE_OUTPUTS = {
-    "DistributedCollector": {0, 1},
+    "DistributedCollector": set(range(COLLECTOR_OUTPUT_COUNT)),
     "DistributedSeed": {0},
     "DistributedValue": {0},
     "DistributedModelName": {0},

nodes/collector.py

@@ -169,6 +169,9 @@ class DistributedCollectorNode:
                 # RGB is opaque; never discard alpha from another participant.
                 tensor = torch.cat((tensor, torch.ones_like(tensor[..., :1])), dim=-1)
             normalized.append(ensure_contiguous(tensor))
+        if len(normalized) == 1:
+            # torch.cat would copy a single batch for nothing.
+            return normalized[0]
         return ensure_contiguous(torch.cat(normalized, dim=0))
 
     def _normalize_images_input(self, images):
@@ -668,6 +671,28 @@ class DistributedCollectorNode:
             worker_audio[worker_id] = self._normalize_audio_input(media["audio"])
         return stored
 
+    @staticmethod
+    def _ordered_worker_frames(worker_images, worker_order):
+        """Yield (worker_id, image_index, tensor) in combined-batch order.
+
+        Expected workers come in ``worker_order``; any unexpected worker ids
+        follow in sorted order.
+        """
+        seen = set()
+        for worker_id_str in [str(worker_id) for worker_id in (worker_order or [])]:
+            seen.add(worker_id_str)
+            frames = worker_images.get(worker_id_str)
+            if not frames:
+                continue
+            for idx in sorted(frames):
+                yield worker_id_str, idx, frames[idx]
+        for worker_id_str in sorted(worker_images):
+            if worker_id_str in seen:
+                continue
+            frames = worker_images[worker_id_str]
+            for idx in sorted(frames):
+                yield worker_id_str, idx, frames[idx]
+
     def _reorder_and_combine_tensors(
         self,
         worker_images: dict,
@@ -685,23 +710,9 @@ class DistributedCollectorNode:
                 ordered_tensors.append(images_on_cpu[i:i+1])
                 sources.append(f"master image {i}")
 
-        ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
-        seen = set()
-        for worker_id_str in ordered_worker_ids:
-            seen.add(worker_id_str)
-            if worker_id_str not in worker_images:
-                continue
-            for idx in sorted(worker_images[worker_id_str].keys()):
-                ordered_tensors.append(worker_images[worker_id_str][idx])
-                sources.append(f"worker {worker_id_str!r} image {idx}")
-
-        # Append any unexpected worker ids deterministically.
-        for worker_id_str in sorted(worker_images.keys()):
-            if worker_id_str in seen:
-                continue
-            for idx in sorted(worker_images[worker_id_str].keys()):
-                ordered_tensors.append(worker_images[worker_id_str][idx])
-                sources.append(f"worker {worker_id_str!r} image {idx}")
+        for worker_id_str, idx, tensor in self._ordered_worker_frames(worker_images, worker_order):
+            ordered_tensors.append(tensor)
+            sources.append(f"worker {worker_id_str!r} image {idx}")
 
         cpu_tensors = []
         for t in ordered_tensors:
@@ -717,11 +728,22 @@ class DistributedCollectorNode:
         return None
 
     @staticmethod
-    def _worker_output_slots(participant_ids):
-        """Map worker ids to worker_N preview slots.
+    def _config_worker_id(participant_id):
+        """Config entry a participant belongs to.
 
-        Configured workers keep their config position; other participants
-        (e.g. serverless) take the free slots after the configured workers.
+        Serverless pools run as ``serverless:<pool id>:<job>`` participants
+        (see queue_orchestration) and all preview on their pool's slot.
+        """
+        if participant_id.startswith("serverless:"):
+            return participant_id[len("serverless:"):].rsplit(":", 1)[0]
+        return participant_id
+
+    @staticmethod
+    def _worker_output_slots(participant_ids):
+        """Map participant ids to worker_N preview slots (many-to-one).
+
+        Configured workers keep their config position; participants outside
+        the config take the free slots after the configured workers.
         """
         try:
             config_ids = [str(worker.get("id")) for worker in load_config().get("workers", [])]
@@ -729,34 +751,53 @@ class DistributedCollectorNode:
             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)
+        overflow = {}
+        for participant_id in participant_ids:
+            config_id = DistributedCollectorNode._config_worker_id(participant_id)
+            if config_id in config_ids:
+                index = config_ids.index(config_id)
             else:
-                index, next_free = next_free, next_free + 1
+                index = overflow.setdefault(config_id, len(config_ids) + len(overflow))
             if index < MAX_WORKER_OUTPUTS:
-                slots[worker_id] = PER_GPU_OUTPUT_NAMES[index + 1]
+                slots[participant_id] = PER_GPU_OUTPUT_NAMES[index + 1]
             else:
-                debug_log(f"Collector - No preview output for worker {worker_id}; only in the combined images")
+                debug_log(f"Collector - No preview output for {participant_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."""
+    def _split_participant_images(self, worker_images, worker_order, master_images, combined, master_frames):
+        """Return each participant's own images keyed by preview slot.
+
+        ``combined`` holds the master's ``master_frames`` frames and then the
+        worker frames in ``_ordered_worker_frames`` order, so worker previews
+        are views into it rather than a second copy of every frame.
+        """
         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:
+        frames = list(self._ordered_worker_frames(worker_images, worker_order))
+        slots = self._worker_output_slots(list(dict.fromkeys(worker_id for worker_id, _, _ in frames)))
+        segments = {}  # slot -> [start, end) frame runs within combined
+        offset = master_frames
+        for worker_id, _idx, tensor in frames:
+            count = tensor.shape[0]
+            # Empty placeholders were dropped from the combined batch too.
+            if count == 0:
                 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)]
-            )
+            slot = slots.get(worker_id)
+            if slot is not None:
+                runs = segments.setdefault(slot, [])
+                if runs and runs[-1][1] == offset:
+                    runs[-1][1] = offset + count
+                else:
+                    runs.append([offset, offset + count])
+            offset += count
+        if offset != combined.shape[0]:
+            # Never attribute frames to the wrong GPU; the combined output is unaffected.
+            log(f"Master - Per-GPU previews skipped: expected {offset} combined frames, got {combined.shape[0]}")
+            return participants
+        for slot, runs in segments.items():
+            views = [combined[start:end] for start, end in runs]
+            participants[slot] = views[0] if len(views) == 1 else torch.cat(views, dim=0)
         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):
@@ -1079,9 +1120,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})")
+                    master_frames = 0 if delegate_mode or images_on_cpu is None else 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,
+                        worker_images, enabled_workers, images_on_cpu if master_frames else None,
+                        combined, master_frames,
                     )
                 return (combined, combined_audio)
             except Exception as e:

nodes/v3.py

@@ -136,7 +136,9 @@ class DistributedCollector(io.ComfyNode):
                 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.')
+                                        'Nodes attached to a GPU that did not take part are skipped. '
+                                        'After Ultimate SD Upscale Distributed every GPU contributes '
+                                        'tiles, so only the combined images output is available.')
                 for name in PER_GPU_OUTPUT_NAMES],
             is_input_list=True, accept_all_inputs=True,
         )

tests/test_collector_image_transport.py

@@ -283,3 +283,68 @@ def test_undistributed_run_previews_everything_as_master(media_modules):
     pass_through = module.DistributedCollectorNode()
     pass_through.run(images=[images], multi_job_id=["job"], pass_through=[True])
     assert pass_through.participant_images == {}
+
+
+def test_empty_worker_batch_keeps_master_preview_and_combined_result(media_modules):
+    module, _ = media_modules
+    module.load_config = lambda: {"workers": [{"id": "worker-a"}]}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    queue.put_nowait({"worker_id": "worker-a", "image_index": 0,
+                      "tensor": torch.empty(0, 8, 12, 3), "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"]'))
+
+    assert combined.shape[0] == 1
+    assert list(collector.participant_images) == ["master"]
+
+
+def test_serverless_pool_jobs_share_their_pools_preview_slot(media_modules):
+    module, _ = media_modules
+    module.load_config = lambda: {"workers": [{"id": "gpu-a"}, {"id": "pool", "type": "serverless"}]}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    for job, value in ((2, 0.75), (1, 0.5)):
+        queue.put_nowait({"worker_id": f"serverless:pool:{job}", "image_index": 0,
+                          "tensor": torch.full((1, 8, 12, 3), value), "is_last": True})
+    module.prompt_server.distributed_pending_jobs = {"image-test": queue}
+    collector = module.DistributedCollectorNode()
+
+    asyncio.run(collector.execute(
+        torch.zeros(1, 8, 12, 3), None, multi_job_id="image-test",
+        enabled_worker_ids='["serverless:pool:1", "serverless:pool:2"]'))
+
+    previews = collector.participant_images
+    assert set(previews) == {"master", "worker_2"}
+    # Pool jobs appear in their enabled order, not delivery order.
+    assert previews["worker_2"][:, 0, 0, 0].tolist() == [0.5, 0.75]
+
+
+def test_serverless_validators_accept_every_collector_output(media_modules):
+    module, _ = media_modules
+    legacy = _load_file("serverless_prompt_slot_check", ROOT / "api/orchestration/serverless_prompt.py")
+    assert legacy.COLLECTOR_OUTPUT_COUNT == 2 + len(module.PER_GPU_OUTPUT_NAMES)
+    assert legacy._REMOVED_NODE_OUTPUTS["DistributedCollector"] == set(range(legacy.COLLECTOR_OUTPUT_COUNT))
+
+
+def test_worker_previews_are_views_of_the_combined_batch(media_modules):
+    module, _ = media_modules
+    module.load_config = lambda: {"workers": [{"id": "worker-a"}]}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    for index in range(2):
+        queue.put_nowait({"worker_id": "worker-a", "image_index": index,
+                          "tensor": torch.full((1, 8, 12, 3), 0.5), "is_last": index == 1})
+    module.prompt_server.distributed_pending_jobs = {"image-test": queue}
+    collector = module.DistributedCollectorNode()
+
+    combined, _ = asyncio.run(collector.execute(
+        torch.zeros(1, 8, 12, 3), None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]'))
+
+    preview = collector.participant_images["worker_1"]
+    assert preview.shape == (2, 8, 12, 3) and preview.is_contiguous()
+    assert preview.data_ptr() == combined[1:].data_ptr()

tests/test_serverless_media_prompt.py

@@ -424,3 +424,21 @@ assert graph[bindings[0]["node_id"]]["class_type"] == "MediaExportV1"
 '''
     result = subprocess.run([sys.executable, "-c", script, str(ROOT)], text=True, capture_output=True)
     assert result.returncode == 0, result.stderr
+
+
+def test_links_from_collector_preview_outputs_are_pruned_not_rejected():
+    # Per-GPU preview sockets follow images/audio; a pool never returns them,
+    # but a workflow previewing them must still compile for typed pools.
+    prompt = _prompt()
+    prompt["preview:1"] = _node("PreviewImage", images=["collector:1", 2])
+    prompt["preview:2"] = _node("PreviewImage", images=["collector:1", legacy.COLLECTOR_OUTPUT_COUNT - 1])
+    graph, bindings = _run(prompt)
+    assert [binding["socket"] for binding in bindings] == ["images"]
+    assert not any(node["class_type"] == "PreviewImage" for node in graph.values())
+
+
+def test_rejects_collector_output_slots_beyond_the_schema():
+    prompt = _prompt()
+    prompt["preview:1"] = _node("PreviewImage", images=["collector:1", legacy.COLLECTOR_OUTPUT_COUNT])
+    with pytest.raises(ValueError, match="invalid output slot"):
+        _run(prompt)

web/collectorOutputs.js

@@ -7,6 +7,7 @@ const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
 
 const trackedNodes = new Set();
 let workers = [];
+let workersPromise = null;
 let listenerAttached = false;
 
 async function fetchWorkers() {
@@ -20,13 +21,30 @@ async function fetchWorkers() {
     }
 }
 
+// Collectors created together (a workflow load) share one config fetch.
+function loadWorkers() {
+    if (!workersPromise) {
+        workersPromise = fetchWorkers().then((list) => {
+            workers = list;
+            if (list.length === 0) workersPromise = null;
+            return list;
+        });
+    }
+    return workersPromise;
+}
+
 // Only the displayed label changes; slot names stay fixed so links survive.
+// Unnamed slots drop the label so saved workflows carry no redundant copies.
 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;
+        if (worker?.name) {
+            output.label = `${output.name} (${worker.name})`;
+        } else {
+            delete output.label;
+        }
     }
     node.setDirtyCanvas?.(true, true);
 }
@@ -67,7 +85,7 @@ app.registerExtension({
                 return originalOnRemoved ? originalOnRemoved.call(this) : undefined;
             };
 
-            if (workers.length === 0) workers = await fetchWorkers();
+            if (workers.length === 0) await loadWorkers();
             labelWorkerOutputs(node, workers);
         } catch (error) {
             console.error("Error in DistributedCollector outputs extension:", error);

web/tests/collectorOutputs.test.js

@@ -5,14 +5,15 @@ 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) {
+function setup(workers, onFetch = () => {}) {
     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 }) }),
+        fetch: async () => { onFetch(); return { ok: true, json: async () => ({ workers }) }; },
+        Promise,
         window: { addEventListener: (name, fn) => listeners.set(name, fn) },
         parseInt,
         Array,
@@ -32,10 +33,26 @@ describe("DistributedCollector per-GPU outputs", () => {
         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"]);
+            ["worker_1 (GPU A)", "worker_2 (GPU B)", undefined]);
         expect(node.outputs[2].label).toBeUndefined();
     });
 
+    it("shares one config fetch across collectors created together", async () => {
+        let fetches = 0;
+        const { extension, node } = setup([{ id: "a", name: "GPU A" }], () => { fetches += 1; });
+        const second = { ...node, outputs: node.outputs.map((o) => ({ ...o })) };
+        await Promise.all([extension.nodeCreated(node), extension.nodeCreated(second)]);
+        expect(fetches).toBe(1);
+        expect(second.outputs[3].label).toBe("worker_1 (GPU A)");
+    });
+
+    it("clears a stale saved label when the slot has no configured worker", async () => {
+        const { extension, node } = setup([]);
+        node.outputs[3].label = "worker_1 (Old PC)";
+        await extension.nodeCreated(node);
+        expect("label" in node.outputs[3]).toBe(false);
+    });
+
     it("relabels when the worker list changes", async () => {
         const { extension, listeners, node } = setup([{ id: "a", name: "GPU A" }]);
         await extension.nodeCreated(node);