From 213d491f06852f368e3fef4be45f870c7b2fd6dd Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 10 Oct 2026 05:59:52 +0000 Subject: [PATCH] fix: preserve image alpha in distributed collection --- api/job_routes.py | 4 +- nodes/collector.py | 57 +++++- tests/test_collector_image_transport.py | 225 ++++++++++++++++++++++++ 3 files changed, 280 insertions(+), 6 deletions(-) create mode 100644 tests/test_collector_image_transport.py diff --git a/api/job_routes.py b/api/job_routes.py index b35655b..ce25eca 100644 --- a/api/job_routes.py +++ b/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: diff --git a/nodes/collector.py b/nodes/collector.py index d577b0f..9934092 100644 --- a/nodes/collector.py +++ b/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 diff --git a/tests/test_collector_image_transport.py b/tests/test_collector_image_transport.py new file mode 100644 index 0000000..6994167 --- /dev/null +++ b/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))