From d29704693cffd4d45171ded19ac0695be4bbcf71 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Sun, 19 Apr 2026 14:57:53 -0700 Subject: [PATCH] Handle heterogeneous mapped latents --- README.md | 2 +- docs/LEARNINGS.md | 1 + serialization.py | 28 ++++++++++++++++--- tests/test_modal_executor_node.py | 45 +++++++++++++++++++++++++++++++ tests/test_serialization.py | 24 +++++++++++++++++ 5 files changed, 96 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 47cbea3..96f0b93 100644 --- a/README.md +++ b/README.md @@ -255,7 +255,7 @@ Current mapped-execution rules: - ordinary remote components without `Modal Map Input` still preserve ComfyUI's usual zipped batch semantics at the remote boundary: if multiple boundary inputs arrive as compatible batches, Modal-Sync fans that component out per item instead of injecting the raw list into primitive widget sockets like `seed: INT` - that implicit batching path now partitions one-time and per-item execute nodes too, so if two remote samplers share the same upstream `MODEL` but only one sampler receives a multi-item batch input, the sibling sampler still runs once instead of repeating once per batch item - non-mapped boundary inputs are broadcast unchanged to every per-item execution -- mapped outputs are reassembled in item order, concatenating batchable tensors back together when possible +- mapped outputs are reassembled in item order, concatenating batchable tensors back together when possible and otherwise preserving the ordered per-item list when shapes differ - per-item remote node status updates are suppressed, and streamed UI events from mapped item runs are filtered to the nodes that actually belong to that per-item payload so static sibling branches do not repaint the UI on every item - aggregate mapped-batch completion now shades the global Modal status pill like a progress bar and shows counts such as `3/16`, while node-local progress bars remain reserved for real streamed node progress only - the representative remote node also shows a small non-bar batch badge like `3/16` during mapped execution so you can see which remote island owns the batch without implying that the node itself emits sampler-style progress diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 011cc01..a1fcb8a 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -2,6 +2,7 @@ ## 2026-04-19 +- Mapped remote outputs cannot always be reassembled by concatenating tensors along batch dimension `0`. A flow like `KSampler -> UpscaleLatentBy` with per-item upscale factors can produce latents with different spatial sizes, so the safe aggregation rule is "concatenate when shapes match, otherwise preserve the ordered list of per-item outputs." - The frontend cannot treat mapped execution like a prompt-wide single active node. Per-lane progress is the durable source of truth there: ordinary non-mapped runs can still use `activeNodeId`, but mapped remote styling should derive purple/active state from whichever nodes currently hold live lane or numeric progress, and lane ownership needs to move when one worker advances from node A to node B. - The test bootstrap should not hard-code `~/git/ComfyUI` as its only source of ComfyUI imports. Reading `COMFYUI_ROOT` first makes it easy to validate the extension against a disposable checkout in `/tmp` while preserving the old local-default behavior. - A usable `uv` test environment for this repo needs more than `pytest`. Importing current ComfyUI APIs during tests also requires `aiohttp`, `Pillow`, `torch`, `safetensors`, `numpy`, and `tqdm`, so those belong in the dedicated `test` dependency group instead of being installed ad hoc on one machine. diff --git a/serialization.py b/serialization.py index 5f209d9..d4178f1 100644 --- a/serialization.py +++ b/serialization.py @@ -264,6 +264,11 @@ def _join_latent_batches(values: Sequence[Any]) -> Any: return merged +def _join_mapped_values_as_list(values: Sequence[Any]) -> list[Any]: + """Return mapped outputs as an ordered Python list.""" + return list(values) + + def join_mapped_values(values: Sequence[Any], io_type: str, is_list: bool) -> Any: """Reassemble ordered per-item mapped outputs into one proxy output value.""" if not values: @@ -280,7 +285,15 @@ def join_mapped_values(values: Sequence[Any], io_type: str, is_list: bool) -> An normalized_io_type = str(io_type) if normalized_io_type == "LATENT": - return _join_latent_batches(values) + try: + return _join_latent_batches(values) + except RuntimeError as exc: + logger.info( + "Falling back to list aggregation for mapped LATENT outputs because batch concatenation " + "failed: %s", + exc, + ) + return _join_mapped_values_as_list(values) try: torch = _import_torch() @@ -288,6 +301,15 @@ def join_mapped_values(values: Sequence[Any], io_type: str, is_list: bool) -> An torch = None if torch is not None and all(isinstance(value, torch.Tensor) for value in values): - return torch.cat(list(values), dim=0) + try: + return torch.cat(list(values), dim=0) + except RuntimeError as exc: + logger.info( + "Falling back to list aggregation for mapped %s outputs because tensor concatenation " + "failed: %s", + normalized_io_type, + exc, + ) + return _join_mapped_values_as_list(values) - return list(values) + return _join_mapped_values_as_list(values) diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index d951d04..2a8b3d6 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -2288,6 +2288,51 @@ def test_invoke_mapped_remote_engine_async_executes_static_branch_once( ] +def test_invoke_mapped_remote_engine_async_keeps_heterogeneous_latents_as_list( + remote_modal_app_module: Any, + serialization_module: Any, + monkeypatch: Any, +) -> None: + """Mapped LATENT outputs should remain a list when per-item spatial shapes differ.""" + torch = pytest.importorskip("torch") + + async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes: + hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload) + scale = int(hydrated_inputs["remote_input_0"]) + latent = { + "samples": torch.zeros((1, 4, scale, scale), dtype=torch.float32), + "batch_index": [scale], + } + return serialization_module.serialize_node_outputs((latent,)) + + monkeypatch.setattr( + remote_modal_app_module, + "invoke_remote_engine_async", + fake_invoke_remote_engine_async, + ) + + payload = { + "payload_kind": "mapped_subgraph", + "component_id": "7", + "prompt_id": "prompt-1", + "mapped_input": {"proxy_input_name": "remote_input_0", "io_type": "INT"}, + "boundary_outputs": [{"io_type": "LATENT", "is_list": False}], + "extra_data": {"client_id": "client-1"}, + } + + response = asyncio.run( + remote_modal_app_module._invoke_mapped_remote_engine_async( + payload, + serialization_module.serialize_node_inputs({"remote_input_0": [32, 35]}), + ) + ) + + outputs = serialization_module.deserialize_node_outputs(response) + assert len(outputs) == 1 + assert isinstance(outputs[0], list) + assert [item["samples"].shape for item in outputs[0]] == [(1, 4, 32, 32), (1, 4, 35, 35)] + + def test_invoke_implicitly_mapped_subgraph_async_zips_batched_boundary_inputs( remote_modal_app_module: Any, serialization_module: Any, diff --git a/tests/test_serialization.py b/tests/test_serialization.py index 7b21417..2a57b5b 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -101,3 +101,27 @@ def test_split_and_join_latent_batch_for_mapped_execution(serialization_module: assert all(item["samples"].shape[0] == 1 for item in split_items) assert torch.equal(rejoined["samples"], latent["samples"]) assert rejoined["batch_index"] == [0, 1, 2] + + +def test_join_mapped_latents_falls_back_to_list_when_shapes_differ( + serialization_module: Any, +) -> None: + """Mapped LATENT outputs should stay ordered as a list when batch concatenation is impossible.""" + torch = pytest.importorskip("torch") + latents = [ + { + "samples": torch.zeros((1, 4, 32, 32), dtype=torch.float32), + "batch_index": [0], + }, + { + "samples": torch.zeros((1, 4, 35, 35), dtype=torch.float32), + "batch_index": [1], + }, + ] + + rejoined = serialization_module.join_mapped_values(latents, "LATENT", is_list=False) + + assert isinstance(rejoined, list) + assert len(rejoined) == 2 + assert rejoined[0]["samples"].shape == (1, 4, 32, 32) + assert rejoined[1]["samples"].shape == (1, 4, 35, 35)