fix: preserve image alpha in distributed collection

api/job_routes.py

@@ -125,7 +125,9 @@ def _decode_canonical_png_tensor(image_payload):
 
     try:
         with Image.open(io.BytesIO(png_bytes)) as img:
-            img = img.convert("RGB")
+            # PNG callbacks must retain alpha just like local IMAGE tensors.
+            has_alpha = "A" in img.getbands() or "transparency" in img.info
+            img = img.convert("RGBA" if has_alpha else "RGB")
             tensor = pil_to_tensor(img)
         return ensure_contiguous(tensor)
     except Exception as exc:

nodes/collector.py

@@ -65,6 +65,46 @@ class DistributedCollectorNode:
             return value[0]
         return value
 
+    @staticmethod
+    def _combine_image_tensors(tensors, sources):
+        """Validate BHWC batches, retaining alpha and deterministic input order."""
+        nonempty = []
+        reference_shape = None
+        reference_source = None
+        for tensor, source in zip(tensors, sources):
+            if not isinstance(tensor, torch.Tensor):
+                raise ValueError(f"Collector expected an IMAGE tensor from {source}")
+            shape = tuple(tensor.shape)
+            if tensor.ndim != 4 or shape[-1] not in (1, 3, 4) or min(shape[1:3]) <= 0:
+                raise ValueError(
+                    f"Collector received invalid IMAGE shape {shape} from {source}; "
+                    "expected [B, H, W, C] with 1, 3, or 4 channels"
+                )
+            # Empty master placeholders must not determine the worker image grid.
+            if shape[0] == 0:
+                continue
+            if reference_shape is None:
+                reference_shape, reference_source = shape, source
+            elif shape[1:3] != reference_shape[1:3]:
+                raise ValueError(
+                    "Collector cannot combine images with different spatial dimensions: "
+                    f"{reference_source} has shape {reference_shape}; {source} has shape {shape}"
+                )
+            nonempty.append(tensor)
+        if not nonempty:
+            raise ValueError("Collector cannot return an empty image batch")
+
+        channels = max(tensor.shape[-1] for tensor in nonempty)
+        normalized = []
+        for tensor in nonempty:
+            if tensor.shape[-1] == 1 and channels > 1:
+                tensor = tensor.expand(-1, -1, -1, 3)
+            if tensor.shape[-1] == 3 and channels == 4:
+                # 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))
+        return ensure_contiguous(torch.cat(normalized, dim=0))
+
     def _normalize_images_input(self, images):
         """Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
         if isinstance(images, (list, tuple)):
@@ -74,7 +114,9 @@ class DistributedCollectorNode:
                 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 self._combine_image_tensors(
+                images, [f"IMAGE list item {index}" for index in range(len(images))]
+            )
         return ensure_contiguous(images)
 
     def _normalize_audio_input(self, audio):
@@ -286,9 +328,11 @@ class DistributedCollectorNode:
     ):
         """Assemble final tensor, or return None when the job contains only audio."""
         ordered_tensors = []
+        sources = []
         if not delegate_mode and images_on_cpu is not None:
             for i in range(master_batch_size):
                 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()
@@ -298,6 +342,7 @@ class DistributedCollectorNode:
                 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()):
@@ -305,6 +350,7 @@ class DistributedCollectorNode:
                 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}")
 
         cpu_tensors = []
         for t in ordered_tensors:
@@ -314,9 +360,9 @@ class DistributedCollectorNode:
             cpu_tensors.append(t)
 
         if cpu_tensors:
-            return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
+            return self._combine_image_tensors(cpu_tensors, sources)
         if fallback_images is not None:
-            return ensure_contiguous(fallback_images)
+            return self._combine_image_tensors([fallback_images], ["master fallback"])
         return None
 
     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):
@@ -552,5 +598,6 @@ class DistributedCollectorNode:
                 return (combined, combined_audio)
             except Exception as e:
                 log(f"Master - Error combining images: {e}")
-                # Preserve collected audio even when image assembly fails.
-                return (images, combined_audio)
+                # A fallback can silently lose worker results or feed an empty
+                # batch into SaveImage. Surface the original collector error.
+                raise

tests/test_collector_image_transport.py

@@ -0,0 +1,225 @@
+"""Real PNG/tensor regressions; only server/config/HTTP plumbing is stubbed."""
+import asyncio
+import importlib.util
+import io
+import base64
+import sys
+import types
+from pathlib import Path
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+import torch
+from aiohttp import web
+from PIL import Image
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def _load_file(name, path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture
+def media_modules():
+    # Existing harnesses provide ComfyUI/config doubles. Restore their module
+    # replacements so this file can coexist with the rest of the test suite.
+    with patch.dict(sys.modules):
+        image_utils = _load_file("real_image_utils", ROOT / "utils/image.py")
+        collector_tests = _load_file("collector_harness", ROOT / "tests/test_collector_list_inputs.py")
+        collector = collector_tests._load_collector_module()
+        sys.modules["aiohttp"].web = web
+        route_tests = _load_file("route_harness", ROOT / "tests/api/test_distributed_queue.py")
+        routes = route_tests._load_job_routes_module()
+        for module in (collector, routes):
+            module.pil_to_tensor = image_utils.pil_to_tensor
+            module.ensure_contiguous = image_utils.ensure_contiguous
+        collector.tensor_to_pil = image_utils.tensor_to_pil
+        routes.prompt_server = collector.prompt_server
+        yield collector, routes
+
+
+def _png_payload(image):
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode("ascii")
+
+
+@pytest.mark.parametrize("mode", ["RGB", "RGBA", "L", "LA", "palette", "palette_alpha"])
+def test_png_decoder_preserves_alpha_and_normalizes_opaque_modes(media_modules, mode):
+    _, routes = media_modules
+    if mode.startswith("palette"):
+        image = Image.new("P", (12, 8), 0)
+        image.putpalette([16, 32, 64] + [0] * 765)
+        if mode == "palette_alpha":
+            image.info["transparency"] = 0
+    else:
+        color = {"RGB": (16, 32, 64), "RGBA": (16, 32, 64, 96), "L": 32, "LA": (32, 96)}[mode]
+        image = Image.new(mode, (12, 8), color)
+    alpha = mode in {"RGBA", "LA", "palette_alpha"}
+    expected = np.array(image.convert("RGBA" if alpha else "RGB"), dtype=np.float32) / 255
+    decoded = routes._decode_canonical_png_tensor(_png_payload(image))
+    assert decoded.shape == (1, 8, 12, 4 if alpha else 3)
+    assert decoded.is_contiguous()
+    np.testing.assert_array_equal(decoded[0].numpy(), expected)
+
+
+async def _master_collect(module, local, worker_images, delegate=False):
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    for index, tensor in enumerate(worker_images):
+        queue.put_nowait({"worker_id": "worker-a", "image_index": index,
+                          "tensor": tensor, "is_last": index == len(worker_images) - 1})
+    module.prompt_server.distributed_pending_jobs = {"image-test": queue}
+    return await module.DistributedCollectorNode().execute(
+        local, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]',
+        delegate_only=delegate,
+    )
+
+
+@pytest.mark.parametrize("local_channels,worker_channels", [(3, 3), (4, 4), (3, 4), (4, 3), (None, 4)])
+def test_worker_png_callback_and_collector_preserve_pixels_and_alpha(
+    media_modules, local_channels, worker_channels
+):
+    module, routes = media_modules
+    local = None if local_channels is None else torch.full((1, 8, 12, local_channels), 0.5)
+    worker = torch.full((2, 8, 12, worker_channels), 0.25)
+    if worker_channels == 4:
+        worker[..., 3] = torch.linspace(0, 1, 12).reshape(1, 1, 12)
+    responses = []
+
+    async def scenario():
+        module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+        module.prompt_server.distributed_pending_jobs = {"image-test": asyncio.Queue()}
+
+        class Response:
+            def __init__(self, payload):
+                self.payload = payload
+
+            async def __aenter__(self):
+                async def request_json():
+                    return self.payload
+                response = await routes.job_complete_endpoint(types.SimpleNamespace(json=request_json))
+                responses.append(response.status)
+                assert response.status == 200, response.text
+                return self
+
+            async def __aexit__(self, *_args):
+                return False
+
+            def raise_for_status(self):
+                pass
+
+        class Session:
+            def post(self, _url, json, timeout):
+                return Response(json)
+
+        async def get_session():
+            return Session()
+
+        module.get_client_session = get_session
+        module.encode_audio_payload = lambda _audio: None
+        collector = module.DistributedCollectorNode()
+        await collector.send_batch_to_master(worker, None, "image-test", "http://master", "worker-a")
+        result, _ = await collector.execute(
+            local, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]',
+            delegate_only=local is None,
+        )
+        assert "image-test" not in module.prompt_server.distributed_pending_jobs
+        return result
+
+    result = asyncio.run(scenario())
+    channels = max(local_channels or 3, worker_channels)
+    offset = 0 if local is None else 1
+    assert responses == [200, 200]
+    assert result.shape == (offset + 2, 8, 12, channels)
+    assert result.is_contiguous()
+    if local is not None:
+        assert torch.equal(result[:1, ..., :local_channels], local)
+        if local_channels == 3 and channels == 4:
+            assert torch.all(result[:1, ..., 3] == 1)
+    assert torch.allclose(result[offset:, ..., :worker_channels], worker, atol=1 / 255)
+    if worker_channels == 3 and channels == 4:
+        assert torch.all(result[offset:, ..., 3] == 1)
+
+
+def test_collector_mixed_rgb_rgba_list_promotes_rgb_to_opaque_rgba(media_modules):
+    module, _ = media_modules
+    rgb = torch.zeros(1, 8, 12, 3)
+    rgba = torch.full((1, 8, 12, 4), 0.5)
+    result, _ = module.DistributedCollectorNode().run(images=[rgb, rgba])
+    assert result.shape == (2, 8, 12, 4)
+    assert torch.equal(result[:1, ..., :3], rgb)
+    assert torch.all(result[:1, ..., 3] == 1)
+    assert torch.equal(result[1:], rgba)
+
+
+@pytest.mark.parametrize("other_channels", [1, 3, 4])
+def test_collector_preserves_grayscale_or_promotes_it_for_color_batches(media_modules, other_channels):
+    module, _ = media_modules
+    gray = torch.full((1, 8, 12, 1), 0.25)
+    other = torch.full((1, 8, 12, other_channels), 0.5)
+    result, _ = module.DistributedCollectorNode().run(images=[gray, other])
+    assert result.shape == (2, 8, 12, other_channels)
+    assert torch.all(result[:1, ..., :min(other_channels, 3)] == 0.25)
+    if other_channels == 4:
+        assert torch.all(result[:1, ..., 3] == 1)
+    assert torch.equal(result[1:], other)
+
+
+def test_mixed_channel_merge_keeps_configured_worker_and_image_order(media_modules):
+    module, _ = media_modules
+    collector = module.DistributedCollectorNode()
+    master = torch.zeros(1, 8, 12, 3)
+    worker_a = torch.full((1, 8, 12, 4), 0.25)
+    worker_b = torch.full((1, 8, 12, 3), 0.5)
+    unexpected = torch.full((1, 8, 12, 3), 0.75)
+    result = collector._reorder_and_combine_tensors(
+        {"worker-a": {0: worker_a}, "worker-b": {2: worker_b, 0: master},
+         "unexpected": {0: unexpected}},
+        ["worker-b", "worker-a"], 1, master, False, master,
+    )
+    assert result.shape == (5, 8, 12, 4)
+    assert result[:, 0, 0, 0].tolist() == [0, 0, 0.5, 0.25, 0.75]
+    assert result[:, 0, 0, 3].tolist() == [1, 1, 1, 0.25, 1]
+
+
+def test_collector_spatial_mismatch_raises_with_shapes_not_successful_fallback(media_modules):
+    module, _ = media_modules
+    local = torch.zeros(1, 8, 12, 3)
+    worker = torch.zeros(1, 8, 13, 3)
+    with pytest.raises(ValueError, match="spatial dimensions") as error:
+        asyncio.run(_master_collect(module, local, [worker]))
+    message = str(error.value)
+    assert "master image 0" in message
+    assert "worker 'worker-a' image 0" in message
+    assert "(1, 8, 12, 3)" in message and "(1, 8, 13, 3)" in message
+    assert "image-test" not in module.prompt_server.distributed_pending_jobs
+
+
+@pytest.mark.parametrize("shape", [(1, 8, 12), (1, 8, 12, 5)])
+def test_collector_rejects_invalid_worker_image_shape(media_modules, shape):
+    module, _ = media_modules
+    with pytest.raises(ValueError, match="worker 'worker-a' image 0"):
+        asyncio.run(_master_collect(module, None, [torch.zeros(shape)], delegate=True))
+
+
+def test_delegate_master_ignores_empty_placeholder_shape(media_modules):
+    module, _ = media_modules
+    placeholder = torch.empty(0, 64, 64, 3)
+    worker = torch.full((1, 8, 12, 4), 0.5)
+    result, _ = asyncio.run(_master_collect(module, placeholder, [worker], delegate=True))
+    assert torch.equal(result, worker)
+
+
+def test_empty_image_results_raise_instead_of_reaching_save_image(media_modules):
+    module, _ = 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))