Handle heterogeneous mapped latents

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

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.

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)

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,

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)