Preserve singleton mapped image batches

README.md

@@ -112,7 +112,7 @@ Mapped execution currently supports:
 
 One `Modal Map Input` boundary is supported per remote component. Non-mapped boundary inputs are broadcast unchanged to every item. Mapped outputs are reassembled in item order, concatenating batchable tensors when possible and otherwise preserving an ordered list.
 
-Ordinary remote components without `Modal Map Input` still preserve ComfyUI's zipped batch behavior at the remote boundary. If a compatible batch reaches a primitive socket such as `seed: INT`, Modal-Sync itemizes it instead of injecting the whole list into the primitive widget input. If the target node declares `INPUT_IS_LIST`, Modal-Sync runs the component once as an ordinary subgraph.
+Ordinary remote components without `Modal Map Input` still preserve ComfyUI's zipped batch behavior at the remote boundary. If a compatible batch reaches a primitive socket such as `seed: INT`, Modal-Sync itemizes it instead of injecting the whole list into the primitive widget input. If the target node declares `INPUT_IS_LIST`, Modal-Sync runs the component once as an ordinary subgraph. Scheduler-list output contracts remain stable for singleton runs: an ordinary one-item `IMAGE` result stays wrapped as one `[B,H,W,C]` value instead of losing its batch dimension when it returns to ComfyUI.
 
 Mapped components can contain both one-time execute targets and per-item execute targets. For example, two remote samplers may share one upstream model loader while only one sampler fans out over latents. Modal-Sync keeps the invariant upstream work separate from the per-item work so the sibling branch still runs once.
 

docs/LEARNINGS.md

@@ -1,5 +1,10 @@
 # Learnings
 
+## 2026-08-15
+
+- ComfyUI `OUTPUT_IS_LIST` is a strict runtime container contract, not only planning metadata. If a proxy advertises a scheduler-list `IMAGE` but returns a raw `[1,H,W,C]` tensor, ComfyUI extends that tensor as the output list and passes `[H,W,C]` downstream; `PreviewImage` then treats height as batch size and saves `H` arrays shaped `[W,C]`. Normalize every scheduler-list proxy result to a Python list even when singleton execution skipped mapped aggregation.
+- Multi-item mapped tests do not cover the singleton contract. Their aggregation path already wraps a joined tensor, so regression coverage also needs an ordinary remote result with one `IMAGE` and must verify that ComfyUI list merging preserves the complete `[1,H,W,C]` tensor as one scheduler item.
+
 ## 2026-08-14
 
 - A costly local re-entry is a graph property, not just a single-edge pattern. The planner can identify every affected local node by intersecting two local-only closures: nodes reachable downstream from any remote node and nodes that can reach any remote node upstream. This catches both one-node and multi-node remote-to-local-to-remote paths while excluding local preview branches that never re-enter remote execution.

modal_executor_node.py

@@ -303,6 +303,38 @@ def _normalize_proxy_payload(payload: Any) -> Mapping[str, Any]:
     return payload
 
 
+def _normalize_scheduler_list_outputs(
+    payload: Mapping[str, Any],
+    outputs: Sequence[Any],
+) -> tuple[Any, ...]:
+    """Make remote output containers match the proxy's scheduler-list declarations."""
+    normalized_outputs = list(outputs)
+    boundary_outputs = payload.get("boundary_outputs")
+    if not isinstance(boundary_outputs, Sequence) or isinstance(
+        boundary_outputs,
+        str | bytes | bytearray,
+    ):
+        return tuple(normalized_outputs)
+
+    for output_index, boundary_output in enumerate(boundary_outputs):
+        if output_index >= len(normalized_outputs):
+            break
+        if not isinstance(boundary_output, Mapping):
+            continue
+        if not bool(boundary_output.get("scheduler_is_list", False)):
+            continue
+        if isinstance(normalized_outputs[output_index], list):
+            continue
+        normalized_outputs[output_index] = [normalized_outputs[output_index]]
+        logger.debug(
+            "Wrapped singleton remote output %d for component=%s to satisfy its scheduler-list contract.",
+            output_index,
+            payload.get("component_id"),
+        )
+
+    return tuple(normalized_outputs)
+
+
 def _normalize_prompt_id(value: Any) -> str | None:
     """Return one non-empty prompt id string when available."""
     if value is None:
@@ -551,12 +583,13 @@ def _build_proxy_node_class(
             )
 
             async with _modal_workflow_execution_slot(payload):
-                outputs = tuple(
+                outputs = _normalize_scheduler_list_outputs(
+                    payload,
                     await _execute_payload_async(
                         get_remote_executor_client(),
                         payload,
                         _normalize_proxy_kwargs(kwargs),
-                    )
+                    ),
                 )
             logger.debug(
                 "Remote execution completed for payload kind=%s with %d outputs.",

tests/test_modal_executor_node.py

@@ -651,6 +651,74 @@ def test_mapped_boundary_payload_requests_scheduler_list_output(
     assert api_intercept_module._proxy_boundary_output_is_list(payload) is True
 
 
+def test_scheduler_list_proxy_wraps_singleton_image_tensor(
+    modal_executor_module: Any,
+) -> None:
+    """A singleton mapped IMAGE must retain its batch dimension in ComfyUI."""
+    torch = pytest.importorskip("torch")
+
+    fake_nodes_module = types.SimpleNamespace(
+        NODE_CLASS_MAPPINGS={},
+        NODE_DISPLAY_NAME_MAPPINGS={},
+    )
+    proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
+        output_types=("IMAGE",),
+        output_names=("image",),
+        output_is_list=(True,),
+        nodes_module=fake_nodes_module,
+        is_output_node=False,
+    )
+    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
+    image = torch.zeros((1, 8, 6, 3), dtype=torch.float32)
+
+    class FakeClient:
+        """Return one ordinary IMAGE tensor from a singleton remote execution."""
+
+        async def execute_payload_async(
+            self,
+            payload: dict[str, Any],
+            kwargs: dict[str, Any],
+        ) -> tuple[Any, ...]:
+            """Return the unwrapped tensor produced by the ordinary subgraph path."""
+            assert payload["boundary_outputs"][0]["scheduler_is_list"] is True
+            assert kwargs == {}
+            return (image,)
+
+    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
+    try:
+        result = asyncio.run(
+            proxy_class.execute(
+                original_node_data={
+                    "payload_kind": "subgraph",
+                    "prompt_id": "singleton-image-prompt",
+                    "component_id": "singleton-image-component",
+                    "boundary_outputs": [
+                        {
+                            "proxy_output_name": "image",
+                            "node_id": "11",
+                            "output_index": 0,
+                            "io_type": "IMAGE",
+                            "is_list": False,
+                            "scheduler_is_list": True,
+                        }
+                    ],
+                },
+                unique_id="singleton-image-component",
+            )
+        )
+    finally:
+        modal_executor_module.set_remote_executor_client_factory(None)
+
+    assert isinstance(result.result[0], list)
+    assert len(result.result[0]) == 1
+    assert result.result[0][0] is image
+    scheduler_items: list[Any] = []
+    scheduler_items.extend(result.result[0])
+    assert len(scheduler_items) == 1
+    assert scheduler_items[0] is image
+    assert tuple(scheduler_items[0].shape) == (1, 8, 6, 3)
+
+
 def test_component_proxy_emits_completion_token_after_remote_execution(
     modal_executor_module: Any,
 ) -> None: