From dba3a6abd9aa467d8bbf0eba13683f176c97fa41 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Fri, 6 Nov 2026 14:09:02 -0700 Subject: [PATCH] Add Remote Progressive Preview node and run every mapped item (#9) * Expect resolved STRING type for ModalMapInput boundaries 69a33e3 made _remote_output_io_type resolve wildcard "*" outputs through their linked input, so a ModalMapInput boundary fed by a STRING list now reports STRING, matching the mapped_input type the test already expects. Co-Authored-By: Claude Opus 5.5 * Add Remote Progressive Preview node A local PreviewImage-style node that shows one empty slot per mapped item when remote work starts and fills each slot as soon as that item's images stream back, instead of waiting for ComfyUI to deliver the whole result. - Planning records direct consumers as progressive_preview_target_node_ids and exempts them from local bridge materializers. - The remote shim streams boundary outputs for those targets; mapped item payloads carry map_item_count and dispatch emits a start event. - The host saves streamed images to temp, records them per prompt/node, and sends modal_progressive_preview websocket events. - The node reuses streamed slots on final execution and only materializes bridge refs that were never streamed. - web/progressive_preview.js renders the slot grid and hands off to the native preview on the final executed event. Co-Authored-By: Claude Opus 5.5 * Execute every mapped item instead of reusing shared-session outputs Mapped items share one remote session. The restored-session short-circuit, meant for identical per-lane seed phases, looked outputs up by node id only, so after the first item on a container every later item skipped sampling and returned that earlier item's image under its own bridge key. Payloads carrying map_item_index now always execute. Co-Authored-By: Claude Opus 5.5 * Keep finished progressive previews handed off to the native preview Hide the panel with an inline display style so the frontend's DOM-widget wrapper cannot override the hidden attribute, and ignore replayed events for a prompt whose final images were already delivered. Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Claude Opus 5.5 --- .gitignore | 3 + README.md | 13 + __init__.py | 4 + cloud_prompt_execution.py | 19 + cloud_prompt_server_shims.py | 62 ++- component_planning.py | 77 ++- docs/LEARNINGS.md | 21 + progressive_preview.py | 201 +++++++ progressive_preview_node.py | 133 +++++ remote/local_execution.py | 19 + remote/local_ui_events.py | 107 ++++ remote/mapped_execution.py | 70 ++- remote/payload_stream.py | 49 +- remote_plan_types.py | 1 + tests/progressive_preview_frontend.test.mjs | 168 ++++++ tests/test_mapped_session_short_circuit.py | 113 ++++ tests/test_progressive_preview.py | 565 ++++++++++++++++++++ tests/test_progressive_preview_frontend.py | 17 + tests/test_prompt_rewrite_mapped.py | 4 +- web/modal_toggle.js | 2 + web/progressive_preview.js | 403 ++++++++++++++ 21 files changed, 1998 insertions(+), 53 deletions(-) create mode 100644 progressive_preview.py create mode 100644 progressive_preview_node.py create mode 100644 tests/progressive_preview_frontend.test.mjs create mode 100644 tests/test_mapped_session_short_circuit.py create mode 100644 tests/test_progressive_preview.py create mode 100644 tests/test_progressive_preview_frontend.py create mode 100644 web/progressive_preview.js diff --git a/.gitignore b/.gitignore index c500369..9b8d737 100644 --- a/.gitignore +++ b/.gitignore @@ -66,6 +66,8 @@ !/prompt_affinity_planning.py !/prompt_interception.py !/prompt_rewrite.py +!/progressive_preview.py +!/progressive_preview_node.py !/proxy_node_factory.py !/proxy_payloads.py !/queue_bridge.py @@ -136,6 +138,7 @@ !/tests/ /tests/* !/tests/*.py +!/tests/*.mjs !/tests/fixtures/ /tests/fixtures/* !/tests/fixtures/*.json diff --git a/README.md b/README.md index f5bb828..ca076f5 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,7 @@ The back-end layer is pluggable. Each provider implements only host discovery, p - [Canvas State](#canvas-state) - [Remote Execution Configurator Panel](#remote-execution-configurator-panel) - [Batched And Mapped Workflows](#batched-and-mapped-workflows) + - [Remote Progressive Preview](#remote-progressive-preview) - [LLM Nodes](#llm-nodes) - [Modal LLM](#modal-llm) - [Modal Endpoint Chat](#modal-endpoint-chat) @@ -342,6 +343,18 @@ Ordinary components without `Modal Map Input` preserve ComfyUI's normal zipped-b Mapped progress is summarized on the global status pill and a representative node with counts such as `3/16`. Fan-out is clamped to `COMFY_MODAL_MAX_INFLIGHT_CALLS`. +### Remote Progressive Preview + +`Remote Progressive Preview` (category `Modal`) is a drop-in replacement for `PreviewImage` on the output of remote work. Connect it directly to a remote `IMAGE` output and keep it local. + +- When a mapped component starts, the node shows one empty slot per item. +- As each remote item finishes, its images are streamed straight to the host and fill that item's slot. Slots fill in completion order, not item order. This does not wait for ComfyUI's scheduler, which only hands a node its input after every item has finished. +- When ComfyUI finally runs the node, it reuses the images it already received instead of downloading them again, and the node switches to the normal ComfyUI image preview. + +Use it in place of `PreviewImage` when the same remote output also feeds later remote nodes. That output stays in remote session storage, and the progressive node handles the session references itself, so no separate local download step is added. Only a slot that was never streamed is downloaded at the end. On SSH and Vast.ai back-ends, the host cannot read remote session storage, so the final value is delivered through ComfyUI as usual. Streamed slots still fill in progressively there. + +A non-mapped remote component fills a single slot when it finishes. The node must be wired directly to the remote output; slots are not streamed through intermediate local nodes. + ## LLM Nodes ### Modal LLM diff --git a/__init__.py b/__init__.py index 7896a53..7723210 100644 --- a/__init__.py +++ b/__init__.py @@ -52,6 +52,7 @@ try: ModalMapInput, ModalUniversalExecutor, ) + from .progressive_preview_node import RemoteProgressivePreview from .remote_configuration_nodes import ( ModalConfiguration, R2StorageConfiguration, @@ -68,6 +69,7 @@ except ModuleNotFoundError: # pragma: no cover - used during local non-Comfy im ModalEndpointChat = None # type: ignore[assignment] ModalLLM = None # type: ignore[assignment] ModalUniversalExecutor = None # type: ignore[assignment] + RemoteProgressivePreview = None # type: ignore[assignment] ModalConfiguration = None # type: ignore[assignment] R2StorageConfiguration = None # type: ignore[assignment] RemoteExecutionConfigurator = None # type: ignore[assignment] @@ -96,6 +98,7 @@ else: ModalUniversalExecutor, ModalMapInput, ModalArtifactFinalizer, + RemoteProgressivePreview, ModalEndpointChat, ModalLLM, ModalConfiguration, @@ -121,6 +124,7 @@ __all__ = [ "ModalConfiguration", "ModalUniversalExecutor", "RemoteExecutionConfigurator", + "RemoteProgressivePreview", "R2StorageConfiguration", "SshConfiguration", "VastConfiguration", diff --git a/cloud_prompt_execution.py b/cloud_prompt_execution.py index ed3be9e..c036ea3 100644 --- a/cloud_prompt_execution.py +++ b/cloud_prompt_execution.py @@ -1278,6 +1278,23 @@ def _execute_subgraph_prompt( ) +def _payload_is_mapped_item(payload: Mapping[str, Any]) -> bool: + """Return whether a payload runs one item of a mapped component. + + Mapped items share one remote session, so a session output already stored + under the same node id belongs to a different item and must never be + reused in place of executing this item. + """ + if payload.get("map_item_index") is None: + return False + logger.info( + "Executing mapped item component=%s instead of restoring session outputs " + "that earlier items stored in the shared session.", + payload.get("component_id", "modal-subgraph"), + ) + return True + + def _short_circuit_restored_session_output_subgraph( *, payload: dict[str, Any], @@ -1289,6 +1306,8 @@ def _short_circuit_restored_session_output_subgraph( boundary_outputs = list(payload.get("boundary_outputs", [])) if session_handle is None or not boundary_outputs: return None + if _payload_is_mapped_item(payload): + return None if resolution_stats.input_ref_count <= 0 or resolution_stats.replay_count > 0: return None if any( diff --git a/cloud_prompt_server_shims.py b/cloud_prompt_server_shims.py index 7dea1d3..f38a356 100644 --- a/cloud_prompt_server_shims.py +++ b/cloud_prompt_server_shims.py @@ -42,6 +42,11 @@ def _collapse_cache_slot(slot_values: Any, is_list: bool) -> Any: return _prompt_server_hooks().collapse_cache_slot(slot_values, is_list) +def _boundary_target_node_ids(boundary_output: dict[str, Any], key: str) -> list[str]: + """Return the non-empty local target node ids listed under one boundary key.""" + return [str(node_id) for node_id in boundary_output.get(key, []) or [] if str(node_id)] + + def _emit_cloud_info(message: str, *args: Any) -> None: """Delegate timestamped cloud logging to the stable entrypoint.""" _prompt_server_hooks().emit_cloud_info(message, *args) @@ -295,12 +300,13 @@ class _TracingPromptServer(_NullPromptServer): """Configure streamed remote boundary-output publication for this execution.""" outputs_by_node_id: dict[str, list[dict[str, Any]]] = {} for boundary_output in boundary_outputs: - preview_target_node_ids = [ - str(node_id) - for node_id in boundary_output.get("preview_target_node_ids", []) - if str(node_id) - ] - if not preview_target_node_ids: + preview_target_node_ids = _boundary_target_node_ids( + boundary_output, "preview_target_node_ids" + ) + progressive_target_node_ids = _boundary_target_node_ids( + boundary_output, "progressive_preview_target_node_ids" + ) + if not preview_target_node_ids and not progressive_target_node_ids: continue if str(boundary_output.get("io_type")) != "IMAGE": continue @@ -340,28 +346,32 @@ class _TracingPromptServer(_NullPromptServer): if output_index >= len(cache_outputs): continue - preview_target_node_ids = [ - str(target_node_id) - for target_node_id in boundary_output.get("preview_target_node_ids", []) - if str(target_node_id) - ] - if not preview_target_node_ids: + preview_target_node_ids = _boundary_target_node_ids( + boundary_output, "preview_target_node_ids" + ) + progressive_target_node_ids = _boundary_target_node_ids( + boundary_output, "progressive_preview_target_node_ids" + ) + if not preview_target_node_ids and not progressive_target_node_ids: continue - self._status_callback( - { - "event_type": "boundary_output", - "node_id": str(node_id), - "output_index": output_index, - "io_type": str(boundary_output.get("io_type", "")), - "is_list": bool(boundary_output.get("is_list", False)), - "preview_target_node_ids": preview_target_node_ids, - "value": _collapse_cache_slot( - slot_values=cache_outputs[output_index], - is_list=bool(boundary_output.get("is_list", False)), - ), - } - ) + boundary_event: dict[str, Any] = { + "event_type": "boundary_output", + "node_id": str(node_id), + "output_index": output_index, + "io_type": str(boundary_output.get("io_type", "")), + "is_list": bool(boundary_output.get("is_list", False)), + "preview_target_node_ids": preview_target_node_ids, + "value": _collapse_cache_slot( + slot_values=cache_outputs[output_index], + is_list=bool(boundary_output.get("is_list", False)), + ), + } + if progressive_target_node_ids: + boundary_event["progressive_preview_target_node_ids"] = ( + progressive_target_node_ids + ) + self._status_callback(boundary_event) self._published_boundary_outputs.add(publication_key) def send_sync( diff --git a/component_planning.py b/component_planning.py index e7da7f0..c1028ca 100644 --- a/component_planning.py +++ b/component_planning.py @@ -11,6 +11,10 @@ from typing import Any, Iterable, Mapping if __package__: from .execution_environments import ExecutionAssignment, ExecutionProvider from .modal_executor_node import MODAL_MAP_INPUT_NODE_ID + from .progressive_preview import ( + PROGRESSIVE_PREVIEW_NODE_ID, + PROGRESSIVE_PREVIEW_TARGETS_KEY, + ) from .remote_graph_analysis import ( _build_consumer_map, _build_remote_components, @@ -36,6 +40,10 @@ if __package__: else: # pragma: no cover - flat import inside the Modal container. from execution_environments import ExecutionAssignment, ExecutionProvider from modal_executor_node import MODAL_MAP_INPUT_NODE_ID + from progressive_preview import ( + PROGRESSIVE_PREVIEW_NODE_ID, + PROGRESSIVE_PREVIEW_TARGETS_KEY, + ) from remote_graph_analysis import ( _build_consumer_map, _build_remote_components, @@ -193,6 +201,12 @@ def _build_component_plan( prompt=prompt, local_consumers=local_consumers, ), + progressive_preview_target_node_ids=( + _progressive_preview_target_node_ids( + prompt=prompt, + local_consumers=local_consumers, + ) + ), ) if not has_downstream_consumer: @@ -416,21 +430,57 @@ def _subset_component_prompt( } +def _local_consumer_node_ids_of_class( + *, + prompt: dict[str, Any], + local_consumers: list[InputTarget], + class_type: str, +) -> list[str]: + """Return direct local consumers of one boundary output with a given class.""" + node_ids: set[str] = set() + for local_consumer in local_consumers: + consumer_prompt_node = prompt.get(local_consumer.node_id) + if consumer_prompt_node is None: + continue + if str(consumer_prompt_node.get("class_type")) != class_type: + continue + node_ids.add(str(local_consumer.node_id)) + return sorted(node_ids) + + def _preview_target_node_ids( *, prompt: dict[str, Any], local_consumers: list[InputTarget], ) -> list[str]: """Return direct local PreviewImage consumers for one remote boundary output.""" - preview_target_node_ids: set[str] = set() - for local_consumer in local_consumers: - consumer_prompt_node = prompt.get(local_consumer.node_id) - if consumer_prompt_node is None: - continue - if str(consumer_prompt_node.get("class_type")) != "PreviewImage": - continue - preview_target_node_ids.add(str(local_consumer.node_id)) - return sorted(preview_target_node_ids) + return _local_consumer_node_ids_of_class( + prompt=prompt, local_consumers=local_consumers, class_type="PreviewImage" + ) + + +def _progressive_preview_target_node_ids( + *, + prompt: dict[str, Any], + local_consumers: list[InputTarget], +) -> list[str]: + """Return direct local progressive preview consumers for one boundary output.""" + return _local_consumer_node_ids_of_class( + prompt=prompt, + local_consumers=local_consumers, + class_type=PROGRESSIVE_PREVIEW_NODE_ID, + ) + + +def _is_progressive_preview_consumer( + prompt: Mapping[str, Any], + consumer: InputTarget, +) -> bool: + """Return whether one consumer materializes remote bridge refs by itself.""" + prompt_node = prompt.get(consumer.node_id) + return isinstance(prompt_node, Mapping) and ( + str(prompt_node.get("class_type")) == PROGRESSIVE_PREVIEW_NODE_ID + ) def _component_downstream_closure( @@ -1093,6 +1143,11 @@ def _mark_remote_to_remote_session_boundaries( boundary_output.session_consumer_node_ids = sorted( {consumer.node_id for consumer in remote_consumers} ) + non_returning_local_consumers = [ + consumer + for consumer in non_returning_local_consumers + if not _is_progressive_preview_consumer(prompt, consumer) + ] if non_returning_local_consumers: materializer_node_id = ( f"__ModalLocalBridgeMaterializer__" @@ -1171,6 +1226,10 @@ def _boundary_output_payload( "is_list": boundary_output.is_list, "preview_target_node_ids": list(boundary_output.preview_target_node_ids), } + if boundary_output.progressive_preview_target_node_ids: + payload[PROGRESSIVE_PREVIEW_TARGETS_KEY] = list( + boundary_output.progressive_preview_target_node_ids + ) if boundary_output.session_output: payload["session_output"] = True if mapped_output is not None: diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index dfbb5c5..bae8ead 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -904,3 +904,24 @@ is still listed. `COMFY_MODAL_EXECUTION_MODE` used to default to `local`. That mode runs rewritten components inside the local ComfyUI process, and users who forgot to set the variable kept landing in it. They saw confusing failures, such as `/assets/...` model paths that only exist in remote storage. `settings.configured_execution_mode()` now owns the default (`remote`), and every host-side reader goes through it. Modal containers never set the variable, so their code relied on the old default to behave as local in-process workers: custom-node sync off, no nested Modal dispatch. The Modal image environment now pins `COMFY_MODAL_EXECUTION_MODE=local` explicitly, as SSH and Vast workers already did. Any new worker runtime must do the same, or it inherits the host default and may try to dispatch remotely from inside a worker. + + +## Progressive previews bypass ComfyUI's list scheduling (2026-11-06) + +ComfyUI only runs a node once the upstream node has produced its entire output. A mapped proxy finishes only after every item, so no downstream node can show per-item results while the mapped run is in progress. Progressive display must ride the remote stream, not ComfyUI's data flow. The remote shim already published `boundary_output` stream events as each node finished, and mapped item payloads already carried `map_item_index`. `RemoteProgressivePreview` builds on that: + +- Planning records direct consumers as `progressive_preview_target_node_ids` on the boundary output. The key is serialized only when non-empty, so payload fingerprints and exact-payload tests are unchanged for workflows that do not use the node. +- Mapped dispatch sends a `start` event with the item count before any item runs. Item payloads carry `map_item_count`. +- The host saves each streamed IMAGE to temp, records it in `PROGRESSIVE_PREVIEW_REGISTRY` keyed by `(prompt_id, node_id)`, and sends a `modal_progressive_preview` websocket event. +- The node is `is_input_list=True`, so one call sees every item. It gets the prompt id from `comfy_execution.utils.get_executing_context()`. ComfyUI v3 hidden inputs expose the prompt dict, not its id. + +Related bug found while investigating (not yet fixed): `_rewrite_hybrid_split_proxies` ignores `local_materializer_node_id` and session routes. A plain `PreviewImage` beside a remote continuation of a hybrid static/mapped component therefore receives bridge-ref dicts and fails in `save_images` with `KeyError: 0`. The progressive node avoids this because the planner exempts it from materializers and it materializes refs itself. + + +## Mapped items silently reused one item's output (2026-11-06) + +`_short_circuit_restored_session_output_subgraph` (cloud and host fallback) skips execution when the remote session already holds a value for every session-backed boundary output. It looks that value up by `(session_id, node_id, output_index)`. It was added for per-lane static seed phases, whose outputs really are identical. Mapped items, however, share one remote session by design. Once one item on a container had produced node 9's IMAGE, every later item on that container found that value, skipped sampling (about 2 s instead of 30 s, with zero progress events) and returned the earlier item's image under its own, distinct bridge key. Downstream remote consumers, such as per-item SaveImage, and local previews all received duplicates. It surfaced only because the progressive preview showed streamed slots that differed from the final images. + +Symptoms to look for: per-item stream summaries with `progress_event_count=0` and short durations, and the container log line `Skipping prompt_executor_execute ... session-backed outputs were restored`. The fix: payloads that carry `map_item_index` never take the shortcut. Seed payloads carry no item index and keep it. + +Separately, a DOM widget's `hidden` attribute can lose to inline styles set by the frontend's widget wrapper. The progressive panel now hides itself with an inline `display: none`, and it ignores replayed events for a prompt it has already finalized. diff --git a/progressive_preview.py b/progressive_preview.py new file mode 100644 index 0000000..ebb5c84 --- /dev/null +++ b/progressive_preview.py @@ -0,0 +1,201 @@ +"""Shared state for progressive remote image previews. + +Remote components stream completed IMAGE boundary values (one per mapped item) +straight to the local host. Each streamed value is saved as temporary preview +images and recorded in a per-prompt slot registry, so the local +``RemoteProgressivePreview`` node can show slots filling in while the remote +work is still running and can reuse those saved images when ComfyUI finally +executes it. +""" + +from __future__ import annotations + +import logging +import threading +from collections import OrderedDict +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +logger = logging.getLogger(__name__) + +PROGRESSIVE_PREVIEW_NODE_ID = "RemoteProgressivePreview" +PROGRESSIVE_PREVIEW_EVENT = "modal_progressive_preview" +PROGRESSIVE_PREVIEW_TARGETS_KEY = "progressive_preview_target_node_ids" +_REGISTRY_ENTRY_LIMIT = 128 + + +@dataclass(frozen=True) +class PreviewImageRef: + """Identify one saved preview image the ComfyUI ``/view`` route can serve.""" + + filename: str + subfolder: str + type: str + + @classmethod + def from_ui(cls, ui_image: Mapping[str, Any]) -> PreviewImageRef: + """Build a reference from one ComfyUI ``ui.images`` entry.""" + return cls( + filename=str(ui_image["filename"]), + subfolder=str(ui_image.get("subfolder", "")), + type=str(ui_image.get("type", "temp")), + ) + + def as_ui(self) -> dict[str, str]: + """Return the ComfyUI ``ui.images`` entry for this image.""" + return {"filename": self.filename, "subfolder": self.subfolder, "type": self.type} + + +@dataclass +class ProgressivePreviewState: + """Track the filled slots of one progressive preview node for one prompt.""" + + slot_count: int | None = None + slots: dict[int, list[PreviewImageRef]] = field(default_factory=dict) + + def effective_slot_count(self) -> int: + """Return the declared slot count, grown to cover every filled slot.""" + highest_filled = max(self.slots, default=-1) + 1 + return max(self.slot_count or 0, highest_filled) + + def is_complete(self) -> bool: + """Return whether every declared slot has received its images.""" + if self.slot_count is None or self.slot_count <= 0: + return False + return all(index in self.slots for index in range(self.slot_count)) + + def ordered_images(self) -> list[PreviewImageRef]: + """Return every filled image in slot order.""" + return [image for index in sorted(self.slots) for image in self.slots[index]] + + +class ProgressivePreviewRegistry: + """Thread-safe registry of progressive preview slots keyed by prompt and node.""" + + def __init__(self, entry_limit: int = _REGISTRY_ENTRY_LIMIT) -> None: + """Create an empty registry that keeps at most ``entry_limit`` entries.""" + self._lock = threading.Lock() + self._entries: OrderedDict[tuple[str, str], ProgressivePreviewState] = ( + OrderedDict() + ) + self._entry_limit = entry_limit + + def begin(self, prompt_id: str, node_id: str, slot_count: int) -> None: + """Reset one node's slots for a new prompt with a known slot count.""" + with self._lock: + self._drop_other_prompts_locked(prompt_id, node_id) + self._entries[(prompt_id, node_id)] = ProgressivePreviewState( + slot_count=max(int(slot_count), 0) + ) + self._entries.move_to_end((prompt_id, node_id)) + self._evict_locked() + + def fill( + self, + prompt_id: str, + node_id: str, + slot_index: int, + images: Sequence[PreviewImageRef], + slot_count: int | None = None, + ) -> ProgressivePreviewState: + """Store one slot's images and return a snapshot of the node's state.""" + with self._lock: + key = (prompt_id, node_id) + state = self._entries.get(key) + if state is None: + self._drop_other_prompts_locked(prompt_id, node_id) + state = ProgressivePreviewState() + self._entries[key] = state + if slot_count is not None: + state.slot_count = max(state.slot_count or 0, int(slot_count)) + state.slots[int(slot_index)] = list(images) + self._entries.move_to_end(key) + self._evict_locked() + return _copy_state(state) + + def get(self, prompt_id: str, node_id: str) -> ProgressivePreviewState | None: + """Return a snapshot of one node's slots, if any were recorded.""" + with self._lock: + state = self._entries.get((prompt_id, node_id)) + return _copy_state(state) if state is not None else None + + def pop(self, prompt_id: str, node_id: str) -> ProgressivePreviewState | None: + """Remove and return one node's slots, if any were recorded.""" + with self._lock: + return self._entries.pop((prompt_id, node_id), None) + + def clear(self) -> None: + """Forget every recorded slot.""" + with self._lock: + self._entries.clear() + + def _drop_other_prompts_locked(self, prompt_id: str, node_id: str) -> None: + """Discard stale slots the same node recorded for earlier prompts.""" + for key in [ + key + for key in self._entries + if key[1] == node_id and key[0] != prompt_id + ]: + del self._entries[key] + + def _evict_locked(self) -> None: + """Drop the oldest entries beyond the configured limit.""" + while len(self._entries) > self._entry_limit: + evicted_key, _state = self._entries.popitem(last=False) + logger.debug( + "Evicted progressive preview slots prompt=%s node=%s.", + evicted_key[0], + evicted_key[1], + ) + + +def _copy_state(state: ProgressivePreviewState) -> ProgressivePreviewState: + """Return a detached copy of one registry state.""" + return ProgressivePreviewState( + slot_count=state.slot_count, + slots={index: list(images) for index, images in state.slots.items()}, + ) + + +PROGRESSIVE_PREVIEW_REGISTRY = ProgressivePreviewRegistry() + + +def progressive_preview_targets(boundary_output: Mapping[str, Any]) -> list[str]: + """Return the progressive preview node ids attached to one boundary payload.""" + return [ + str(node_id) + for node_id in boundary_output.get(PROGRESSIVE_PREVIEW_TARGETS_KEY, []) or [] + if str(node_id) + ] + + +def iter_image_batches(value: Any) -> Iterator[Any]: + """Yield every IMAGE tensor batch contained in a streamed boundary value.""" + if isinstance(value, (list, tuple)): + for item in value: + yield from iter_image_batches(item) + return + if hasattr(value, "shape") and hasattr(value, "cpu"): + if len(value.shape) == 3: + yield value.unsqueeze(0) + return + yield value + return + logger.warning( + "Ignoring non-image value of type %s in a progressive preview stream.", + type(value).__name__, + ) + + +def save_temp_preview_images(value: Any) -> list[PreviewImageRef]: + """Save every IMAGE batch in ``value`` as temporary ComfyUI preview images.""" + import nodes + + preview_node = nodes.PreviewImage() + saved: list[PreviewImageRef] = [] + for image_batch in iter_image_batches(value): + result = preview_node.save_images(images=image_batch) + ui_images = ((result or {}).get("ui") or {}).get("images") or [] + saved.extend(PreviewImageRef.from_ui(ui_image) for ui_image in ui_images) + return saved diff --git a/progressive_preview_node.py b/progressive_preview_node.py new file mode 100644 index 0000000..eb16498 --- /dev/null +++ b/progressive_preview_node.py @@ -0,0 +1,133 @@ +"""Local ComfyUI node that shows remote IMAGE results as they stream in.""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Mapping +from typing import Any + +from comfy_api.latest import _io as io + +if __package__: + from .progressive_preview import ( + PROGRESSIVE_PREVIEW_NODE_ID, + PROGRESSIVE_PREVIEW_REGISTRY, + PreviewImageRef, + ProgressivePreviewState, + save_temp_preview_images, + ) + from .session_state import is_remote_session_bridge_ref_payload +else: # pragma: no cover - flat import inside the Modal container. + from progressive_preview import ( + PROGRESSIVE_PREVIEW_NODE_ID, + PROGRESSIVE_PREVIEW_REGISTRY, + PreviewImageRef, + ProgressivePreviewState, + save_temp_preview_images, + ) + from session_state import is_remote_session_bridge_ref_payload + +logger = logging.getLogger(__name__) + + +def _executing_prompt_and_node() -> tuple[str | None, str | None]: + """Return the prompt id and node id ComfyUI is currently executing.""" + from comfy_execution.utils import get_executing_context + + context = get_executing_context() + if context is None: + return None, None + return str(context.prompt_id), str(context.node_id) + + +def _materialize_bridge_ref(bridge_ref: Mapping[str, Any]) -> Any: + """Download one remote session bridge value for local preview.""" + if __package__: + from .remote.modal_app import materialize_remote_session_bridge_ref_locally + else: # pragma: no cover - flat import inside the Modal container. + from remote.modal_app import materialize_remote_session_bridge_ref_locally + + return materialize_remote_session_bridge_ref_locally(bridge_ref) + + +async def _resolve_item_images( + index: int, + item: Any, + state: ProgressivePreviewState | None, +) -> list[PreviewImageRef]: + """Return saved preview images for one input item, reusing streamed slots.""" + if not is_remote_session_bridge_ref_payload(item): + return save_temp_preview_images(item) + streamed_images = state.slots.get(index) if state is not None else None + if streamed_images: + return list(streamed_images) + logger.info( + "Materializing progressive preview slot %d locally because it was not streamed.", + index, + ) + materialized = await asyncio.to_thread(_materialize_bridge_ref, item) + return save_temp_preview_images(materialized) + + +async def resolve_progressive_preview_images( + items: list[Any], + state: ProgressivePreviewState | None, +) -> list[PreviewImageRef]: + """Return the final preview images, preferring images already streamed.""" + if state is not None and state.is_complete(): + logger.info( + "Reusing %d streamed progressive preview slot(s) without re-downloading.", + state.slot_count, + ) + return state.ordered_images() + resolved: list[PreviewImageRef] = [] + for index, item in enumerate(items): + resolved.extend(await _resolve_item_images(index, item, state)) + return resolved + + +class RemoteProgressivePreview(io.ComfyNode): + """Preview remote images, filling one slot per item as remote work finishes.""" + + @classmethod + def define_schema(cls) -> io.Schema: + """Declare a PreviewImage-style output node that accepts the whole list.""" + return io.Schema( + node_id=PROGRESSIVE_PREVIEW_NODE_ID, + display_name="Remote Progressive Preview", + category="Modal", + description=( + "Preview images produced by remote execution. When remote work " + "starts, the node shows one empty slot per mapped item and fills " + "each slot as soon as that item finishes, without waiting for " + "ComfyUI to receive the whole result." + ), + inputs=[ + io.Image.Input( + "images", + tooltip="Images from a remote node, usually a mapped component.", + ) + ], + outputs=[], + is_output_node=True, + is_input_list=True, + ) + + @classmethod + async def execute(cls, images: list[Any]) -> io.NodeOutput: + """Finalize the preview from streamed slots or the delivered images.""" + prompt_id, node_id = _executing_prompt_and_node() + state = ( + PROGRESSIVE_PREVIEW_REGISTRY.pop(prompt_id, node_id) + if prompt_id and node_id + else None + ) + saved_images = await resolve_progressive_preview_images(list(images), state) + logger.info( + "Finalized progressive preview prompt=%s node=%s images=%d.", + prompt_id, + node_id, + len(saved_images), + ) + return io.NodeOutput(ui={"images": [image.as_ui() for image in saved_images]}) diff --git a/remote/local_execution.py b/remote/local_execution.py index 3c0f4af..238908d 100644 --- a/remote/local_execution.py +++ b/remote/local_execution.py @@ -1268,6 +1268,23 @@ def _execute_subgraph_prompt( return tuple(outputs) +def _payload_is_mapped_item(payload: Mapping[str, Any]) -> bool: + """Return whether a payload runs one item of a mapped component. + + Mapped items share one remote session, so a session output already stored + under the same node id belongs to a different item and must never be + reused in place of executing this item. + """ + if payload.get("map_item_index") is None: + return False + logger.info( + "Executing mapped item component=%s instead of restoring session outputs " + "that earlier items stored in the shared session.", + payload.get("component_id", "modal-subgraph"), + ) + return True + + def _short_circuit_restored_session_output_subgraph( *, payload: dict[str, Any], @@ -1279,6 +1296,8 @@ def _short_circuit_restored_session_output_subgraph( boundary_outputs = list(payload.get("boundary_outputs", [])) if session_handle is None or not boundary_outputs: return None + if _payload_is_mapped_item(payload): + return None if resolution_stats.input_ref_count <= 0 or resolution_stats.replay_count > 0: return None if any( diff --git a/remote/local_ui_events.py b/remote/local_ui_events.py index 0f65e7c..74b2bb9 100644 --- a/remote/local_ui_events.py +++ b/remote/local_ui_events.py @@ -396,6 +396,113 @@ def _emit_local_preview_boundary_output( ) +def _progressive_preview_module() -> Any: + """Return the shared progressive preview module for package or flat imports.""" + try: + from .. import progressive_preview + except ImportError: + import progressive_preview # type: ignore[no-redef] + return progressive_preview + + +def _send_progressive_preview_event( + *, + client_id: str | None, + payload: dict[str, Any], +) -> None: + """Send and record one progressive preview websocket event.""" + if client_id is None: + return + prompt_server = _lookup_local_prompt_server() + if prompt_server is None: + return + event_name = _progressive_preview_module().PROGRESSIVE_PREVIEW_EVENT + prompt_server.send_sync(event_name, payload, client_id) + _record_local_modal_ui_event(event_name, payload, client_id) + + +def _emit_local_progressive_preview_start( + *, + prompt_id: str | None, + client_id: str | None, + node_ids: list[str], + slot_count: int, +) -> None: + """Announce the number of empty slots each progressive preview should show.""" + if not prompt_id or not node_ids: + return + registry = _progressive_preview_module().PROGRESSIVE_PREVIEW_REGISTRY + for node_id in node_ids: + registry.begin(prompt_id, node_id, slot_count) + logger.info( + "Initialized progressive preview prompt=%s node=%s slots=%d.", + prompt_id, + node_id, + slot_count, + ) + _send_progressive_preview_event( + client_id=client_id, + payload={ + "phase": "start", + "prompt_id": prompt_id, + "node": node_id, + "slot_count": int(slot_count), + }, + ) + + +def _emit_local_progressive_preview_slot( + *, + prompt_id: str | None, + client_id: str | None, + node_ids: list[str], + slot_index: int, + slot_count: int | None, + image_value: Any, +) -> None: + """Save one streamed IMAGE value and fill its slot on each progressive preview.""" + if not prompt_id or not node_ids: + return + preview_module = _progressive_preview_module() + try: + images = preview_module.save_temp_preview_images(image_value) + except ModuleNotFoundError: + logger.warning( + "Progressive preview is unavailable because ComfyUI nodes could not be imported." + ) + return + if not images: + logger.warning( + "Progressive preview slot %d for prompt=%s carried no images.", + slot_index, + prompt_id, + ) + return + for node_id in node_ids: + state = preview_module.PROGRESSIVE_PREVIEW_REGISTRY.fill( + prompt_id, node_id, slot_index, images, slot_count + ) + logger.info( + "Filled progressive preview prompt=%s node=%s slot=%d/%d images=%d.", + prompt_id, + node_id, + slot_index + 1, + state.effective_slot_count(), + len(images), + ) + _send_progressive_preview_event( + client_id=client_id, + payload={ + "phase": "slot", + "prompt_id": prompt_id, + "node": node_id, + "slot_index": int(slot_index), + "slot_count": state.effective_slot_count(), + "images": [image.as_ui() for image in images], + }, + ) + + def _allowed_suppressed_stream_node_ids(payload: dict[str, Any]) -> set[str]: """Return the node ids that may surface UI events for a suppressed mapped/static stream.""" allowed_node_ids = { diff --git a/remote/mapped_execution.py b/remote/mapped_execution.py index 27cef80..daaa4a7 100644 --- a/remote/mapped_execution.py +++ b/remote/mapped_execution.py @@ -31,7 +31,10 @@ from .local_execution import ( _node_input_type_map, _resolve_required_subgraph_nodes, ) -from .local_ui_events import _emit_local_modal_progress +from .local_ui_events import ( + _emit_local_modal_progress, + _emit_local_progressive_preview_start, +) from .modal_deployment import ( ModalRemoteInvocationError, _lookup_deployed_remote_engine, @@ -168,6 +171,7 @@ def _build_phase_subgraph_payload( suppress_status_stream: bool = False, lane_id: str | None = None, item_index: int | None = None, + item_count: int | None = None, ) -> dict[str, Any]: """Return one explicit static or mapped subgraph payload.""" phase_definition = _mapped_phase_definition(payload, phase_key) @@ -201,6 +205,8 @@ def _build_phase_subgraph_payload( ) if item_index is not None: phase_payload["map_item_index"] = int(item_index) + if item_count is not None: + phase_payload["map_item_count"] = int(item_count) return phase_payload @@ -290,6 +296,7 @@ def _execute_mapped_subgraph_payload( total_items = len(mapped_items) _emit_local_mapped_progress(payload, 0, total_items) + _emit_mapped_progressive_preview_start(payload, total_items) per_item_outputs: list[tuple[Any, ...]] = [] for item_index, item_value in enumerate(mapped_items): if _local_processing_interrupted(): @@ -302,9 +309,12 @@ def _execute_mapped_subgraph_payload( suppress_status_stream=True, lane_id="0", item_index=item_index, + item_count=total_items, ) else: - item_payload = _build_mapped_item_payload(payload, item_index, 0) + item_payload = _build_mapped_item_payload( + payload, item_index, 0, item_count=total_items + ) item_inputs = dict(broadcast_inputs) item_inputs[mapped_input_name] = item_value logger.info( @@ -356,6 +366,7 @@ def _build_mapped_item_payload( payload: dict[str, Any], item_index: int, lane_index: int, + item_count: int | None = None, ) -> dict[str, Any]: """Return one per-item subgraph payload derived from a mapped remote component payload.""" if _mapped_phase_definition(payload, "mapped_phase") is not None: @@ -366,6 +377,7 @@ def _build_mapped_item_payload( suppress_status_stream=True, lane_id=str(lane_index), item_index=item_index, + item_count=item_count, ) item_payload = copy.deepcopy(payload) item_payload["payload_kind"] = "subgraph" @@ -375,6 +387,8 @@ def _build_mapped_item_payload( item_payload["mapped_input"] = None item_payload["suppress_status_stream"] = True item_payload["map_item_index"] = item_index + if item_count is not None: + item_payload["map_item_count"] = int(item_count) item_payload["mapped_progress_lane_id"] = str(lane_index) item_payload["mapped_progress_display_node_id"] = str( payload.get("component_id", "modal-subgraph") @@ -558,6 +572,52 @@ def _emit_local_mapped_progress( ) +def _mapped_progressive_preview_targets(payload: dict[str, Any]) -> list[str]: + """Return progressive preview nodes fed by this payload's per-item outputs.""" + try: + from ..progressive_preview import progressive_preview_targets + except ImportError: + from progressive_preview import progressive_preview_targets # type: ignore[no-redef] + + candidate_outputs = [ + boundary_output + for boundary_output in payload.get("boundary_outputs", []) or [] + if _is_mapped_boundary_output(boundary_output, payload) + ] + mapped_phase = _mapped_phase_definition(payload, "mapped_phase") + if mapped_phase is not None: + candidate_outputs.extend(mapped_phase.get("boundary_outputs", []) or []) + target_node_ids: list[str] = [] + for boundary_output in candidate_outputs: + for node_id in progressive_preview_targets(boundary_output): + if node_id not in target_node_ids: + target_node_ids.append(node_id) + return target_node_ids + + +def _emit_mapped_progressive_preview_start( + payload: dict[str, Any], + total_items: int, +) -> None: + """Lay out one empty progressive preview slot per mapped item before dispatch.""" + target_node_ids = _mapped_progressive_preview_targets(payload) + if not target_node_ids: + return + extra_data = payload.get("extra_data") or {} + _emit_local_progressive_preview_start( + prompt_id=( + str(payload["prompt_id"]) if payload.get("prompt_id") is not None else None + ), + client_id=( + str(extra_data["client_id"]) + if extra_data.get("client_id") is not None + else None + ), + node_ids=target_node_ids, + slot_count=total_items, + ) + + def _mapped_progress_owner_component_id(payload: dict[str, Any]) -> str | None: """Return the stable component id used to track one mapped worker lane locally.""" owner_component_id = payload.get( @@ -1068,6 +1128,7 @@ async def _invoke_implicitly_mapped_subgraph_async( payload, mapped_execute_node_ids, ) + _emit_mapped_progressive_preview_start(hybrid_payload, total_items) cleanup_payload = _build_remote_session_cleanup_payload(hybrid_payload) execution_mode = configured_execution_mode() use_seeded_remote_lanes = ( @@ -1167,7 +1228,10 @@ async def _invoke_implicitly_mapped_subgraph_async( raise_if_local_interrupted() try: item_payload = _build_mapped_item_payload( - hybrid_payload, item_index, lane_index + hybrid_payload, + item_index, + lane_index, + item_count=total_items, ) item_inputs = dict(broadcast_inputs) for input_name, items in split_inputs.items(): diff --git a/remote/payload_stream.py b/remote/payload_stream.py index d9d16f5..13be75f 100644 --- a/remote/payload_stream.py +++ b/remote/payload_stream.py @@ -18,6 +18,7 @@ from .local_ui_events import ( _emit_local_modal_status, _emit_local_preview_boundary_output, _emit_local_preview_image, + _emit_local_progressive_preview_slot, _progress_stream_event_metadata, _remote_execution_destination, _remote_execution_identity, @@ -348,25 +349,47 @@ def _handle_boundary_output_event( stream_event: dict[str, Any], ) -> None: """Forward one streamed boundary value to its local preview targets.""" - target_node_ids = [ - str(node_id) - for node_id in stream_event.get("preview_target_node_ids", []) - if str(node_id) - ] - if not target_node_ids: + target_node_ids = _stream_event_node_ids(stream_event, "preview_target_node_ids") + progressive_node_ids = _stream_event_node_ids( + stream_event, "progressive_preview_target_node_ids" + ) + if not target_node_ids and not progressive_node_ids: return logger.debug( - "Forwarding streamed Modal boundary output previews for component=%s source_node=%s targets=%s.", + "Forwarding streamed Modal boundary output previews for component=%s source_node=%s targets=%s progressive_targets=%s.", state.payload.get("component_id"), stream_event.get("node_id"), target_node_ids, + progressive_node_ids, ) - _emit_local_preview_boundary_output( - prompt_id=state.prompt_id, - client_id=state.client_id, - preview_target_node_ids=target_node_ids, - image_value=deserialize_value(stream_event.get("value")), - ) + image_value = deserialize_value(stream_event.get("value")) + if progressive_node_ids: + _emit_local_progressive_preview_slot( + prompt_id=state.prompt_id, + client_id=state.client_id, + node_ids=progressive_node_ids, + slot_index=_optional_payload_int(state.payload, "map_item_index") or 0, + slot_count=_optional_payload_int(state.payload, "map_item_count"), + image_value=image_value, + ) + if target_node_ids: + _emit_local_preview_boundary_output( + prompt_id=state.prompt_id, + client_id=state.client_id, + preview_target_node_ids=target_node_ids, + image_value=image_value, + ) + + +def _stream_event_node_ids(stream_event: dict[str, Any], key: str) -> list[str]: + """Return the non-empty node ids listed under one stream event key.""" + return [str(node_id) for node_id in stream_event.get(key, []) or [] if str(node_id)] + + +def _optional_payload_int(payload: dict[str, Any], key: str) -> int | None: + """Return one optional integer payload field.""" + value = payload.get(key) + return int(value) if value is not None else None def _handle_cached_node_event( diff --git a/remote_plan_types.py b/remote_plan_types.py index c6af946..ad32c32 100644 --- a/remote_plan_types.py +++ b/remote_plan_types.py @@ -64,6 +64,7 @@ class BoundaryOutputSpec: io_type: str is_list: bool preview_target_node_ids: list[str] = field(default_factory=list) + progressive_preview_target_node_ids: list[str] = field(default_factory=list) session_output: bool = False session_consumer_node_ids: list[str] = field(default_factory=list) local_materializer_node_id: str | None = None diff --git a/tests/progressive_preview_frontend.test.mjs b/tests/progressive_preview_frontend.test.mjs new file mode 100644 index 0000000..991d104 --- /dev/null +++ b/tests/progressive_preview_frontend.test.mjs @@ -0,0 +1,168 @@ +import assert from "node:assert/strict"; +import { readFile } from "node:fs/promises"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; +import { runInThisContext } from "node:vm"; + +const repoRoot = path.resolve(path.dirname(fileURLToPath(import.meta.url)), ".."); + +class FakeElement { + constructor(tagName) { + this.tagName = tagName; + this.children = []; + this.dataset = {}; + this.style = {}; + this.hidden = false; + this.className = ""; + this.textContent = ""; + this.title = ""; + this.src = ""; + } + + appendChild(child) { + this.children.push(child); + return child; + } + + replaceChildren(...children) { + this.children = [...children]; + } +} + +globalThis.document = { + head: new FakeElement("head"), + createElement: (tagName) => new FakeElement(tagName), + getElementById: () => null, +}; + +const listeners = new Map(); +const extensions = []; +const graphNodes = new Map(); +globalThis.__progressiveAppStub = { + registerExtension: (extension) => extensions.push(extension), + rootGraph: { getNodeById: (id) => graphNodes.get(String(id)) ?? null }, + graph: { setDirtyCanvas() {} }, +}; +globalThis.__progressiveApiStub = { + addEventListener: (name, handler) => listeners.set(name, handler), + apiURL: (route) => `/api${route}`, +}; + +const sourcePath = path.join(repoRoot, "web", "progressive_preview.js"); +const originalSource = await readFile(sourcePath, "utf8"); +runInThisContext( + [ + "const app = globalThis.__progressiveAppStub;", + "const api = globalThis.__progressiveApiStub;", + originalSource.replace(/^import .*?;\n/gm, ""), + "globalThis.__progressiveExports = {", + " applyProgressivePreviewEvent, createSlotState, slotGridColumns, filledSlotCount,", + " progressivePreviewHeaderText, progressivePreviewImageUrl, resolveProgressivePreviewNode,", + " handleProgressivePreviewEvent, finishProgressivePreview, progressivePreviewPanels,", + "};", + ].join("\n"), + { filename: sourcePath }, +); +const preview = globalThis.__progressiveExports; + +// Slot state: start lays out empty slots, slots fill out of order, a new prompt resets. +const state = preview.createSlotState(); +assert.equal(preview.applyProgressivePreviewEvent(state, { phase: "start", prompt_id: "p1", slot_count: 4 }), true); +assert.equal(state.slotCount, 4); +assert.equal(preview.progressivePreviewHeaderText(state), "Remote preview · 0/4 received"); +preview.applyProgressivePreviewEvent(state, { + phase: "slot", prompt_id: "p1", slot_index: 2, slot_count: 4, images: [{ filename: "c.png" }], +}); +assert.equal(preview.filledSlotCount(state), 1); +assert.equal(preview.progressivePreviewHeaderText(state), "Remote preview · 1/4 received"); +preview.applyProgressivePreviewEvent(state, { + phase: "slot", prompt_id: "p2", slot_index: 0, images: [{ filename: "a.png" }], +}); +assert.equal(state.promptId, "p2"); +assert.equal(state.slotCount, 1); +assert.equal(preview.progressivePreviewHeaderText(state), "Remote preview · all 1 received"); +assert.equal(preview.applyProgressivePreviewEvent(state, { phase: "bogus", prompt_id: "p2" }), false); + +// Layout and URLs. +assert.equal(preview.slotGridColumns(0), 1); +assert.equal(preview.slotGridColumns(4), 2); +assert.equal(preview.slotGridColumns(5), 3); +assert.equal( + preview.progressivePreviewImageUrl({ filename: "a b.png", subfolder: "", type: "temp" }), + "/api/view?filename=a+b.png&subfolder=&type=temp", +); + +// Subgraph node ids resolve through each subgraph node. +const innerNode = { id: 18, comfyClass: "RemoteProgressivePreview" }; +graphNodes.set("12", { id: 12, subgraph: { getNodeById: (id) => (String(id) === "18" ? innerNode : null) } }); +assert.equal(preview.resolveProgressivePreviewNode("12:18"), innerNode); +assert.equal(preview.resolveProgressivePreviewNode("99"), null); + +// End to end: events mount a panel with placeholders, fill it, and hand back on executed. +const widgets = []; +const node = { + id: 7, + comfyClass: "RemoteProgressivePreview", + size: [100, 100], + imgs: ["stale"], + addDOMWidget: (name, type, element, options) => { + const widget = { name, type, element, options }; + widgets.push(widget); + return widget; + }, + setSize(size) { + this.size = size; + }, + setDirtyCanvas() {}, +}; +graphNodes.set("7", node); +for (const extension of extensions) { + extension.setup?.(); +} +const handler = listeners.get("modal_progressive_preview"); +assert.equal(typeof handler, "function"); +assert.equal(handler({ detail: { phase: "start", prompt_id: "p3", node: "7", slot_count: 3 } }), true); +const panel = preview.progressivePreviewPanels.get(node); +assert.equal(widgets.length, 1); +assert.equal(panel.visible, true); +assert.equal(node.imgs, undefined); +assert.equal(widgets[0].options.getMinHeight() > 0, true); +assert.equal(panel.grid.children.length, 3); +assert.equal(panel.grid.children.every((cell) => cell.className.includes("is-empty")), true); +handler({ + detail: { phase: "slot", prompt_id: "p3", node: "7", slot_index: 1, slot_count: 3, images: [{ filename: "b.png", type: "temp" }, { filename: "b2.png", type: "temp" }] }, +}); +const filledCell = panel.grid.children[1]; +assert.equal(filledCell.className.includes("is-empty"), false); +assert.equal(filledCell.children[0].src, "/api/view?filename=b.png&subfolder=&type=temp"); +assert.equal(filledCell.children[1].textContent, "×2"); +assert.equal(panel.header.textContent, "Remote preview · 1/3 received"); +assert.equal(handler({ detail: { phase: "start", prompt_id: "p3", node: "99", slot_count: 3 } }), false); + +const nodeType = { prototype: {} }; +let originalCalled = false; +nodeType.prototype.onExecuted = () => { + originalCalled = true; +}; +await extensions[0].beforeRegisterNodeDef(nodeType, { name: "RemoteProgressivePreview" }); +nodeType.prototype.onExecuted.call(node, { images: [] }); +assert.equal(originalCalled, true); +assert.equal(panel.visible, false); +assert.equal(panel.root.style.display, "none"); +assert.equal(widgets[0].options.getMinHeight(), 0); + +// Replayed events for the finished prompt (browser refocus) must not resurrect the panel. +assert.equal(handler({ detail: { phase: "start", prompt_id: "p3", node: "7", slot_count: 3 } }), false); +assert.equal( + handler({ detail: { phase: "slot", prompt_id: "p3", node: "7", slot_index: 0, images: [{ filename: "a.png" }] } }), + false, +); +assert.equal(panel.visible, false); + +// The next prompt shows the panel again. +assert.equal(handler({ detail: { phase: "start", prompt_id: "p4", node: "7", slot_count: 2 } }), true); +assert.equal(panel.visible, true); +assert.equal(panel.root.style.display, ""); +assert.equal(panel.grid.children.length, 2); + +console.log("progressive preview frontend checks passed"); diff --git a/tests/test_mapped_session_short_circuit.py b/tests/test_mapped_session_short_circuit.py new file mode 100644 index 0000000..448bee6 --- /dev/null +++ b/tests/test_mapped_session_short_circuit.py @@ -0,0 +1,113 @@ +"""Mapped items must execute instead of reusing outputs from a shared session.""" + +from __future__ import annotations + +import importlib +from types import SimpleNamespace +from typing import Any + +import pytest + + +class _FakeSessionStore: + """Session store that already holds an earlier item's node output.""" + + def __init__(self) -> None: + """Start with no recorded lookups.""" + self.lookups: list[Any] = [] + + def get_output(self, ref: Any) -> str: + """Return the stale value an earlier item stored for this node.""" + self.lookups.append(ref) + return "earlier-item-output" + + def put_bridge_output(self, session_handle: Any, **kwargs: Any) -> Any: + """Pretend to publish a bridge output into the session.""" + return SimpleNamespace(session_id=session_handle.session_id) + + +@pytest.fixture +def cloud_prompt_execution_module(modal_cloud_module: Any) -> Any: + """Return the flat-imported cloud prompt execution module.""" + return importlib.import_module("cloud_prompt_execution") + + +def _payload(map_item_index: int | None) -> dict[str, Any]: + """Return a payload whose only boundary output is session-backed.""" + payload: dict[str, Any] = { + "component_id": "1__mapped::item:5", + "boundary_outputs": [ + {"node_id": "9", "output_index": 0, "io_type": "IMAGE", "session_output": True} + ], + } + if map_item_index is not None: + payload["map_item_index"] = map_item_index + return payload + + +def _short_circuit(module: Any, store: _FakeSessionStore, payload: dict[str, Any]) -> Any: + """Invoke the short-circuit with every other precondition satisfied.""" + stats = module._RemoteSessionBridgeResolutionStats() + stats.input_ref_count = 5 + return module._short_circuit_restored_session_output_subgraph( + payload=payload, + hydrated_inputs={}, + session_handle=SimpleNamespace(session_id="shared-session"), + resolution_stats=stats, + ) + + +def _install_cloud_doubles(module: Any, store: _FakeSessionStore, monkeypatch: Any) -> None: + """Route the cloud module's session and bridge helpers to test doubles.""" + monkeypatch.setattr(module, "remote_session_store", lambda: store) + monkeypatch.setattr( + module, + "_build_remote_session_bridge_record", + lambda **kwargs: SimpleNamespace(bridge_key="RSB_test"), + ) + monkeypatch.setattr(module, "_store_remote_session_bridge_record", lambda record: None) + monkeypatch.setattr(module, "_store_remote_session_bridge_value", lambda key, value: None) + + +def test_cloud_mapped_item_never_reuses_shared_session_output( + cloud_prompt_execution_module: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A mapped item must run even when the shared session already holds its node output.""" + store = _FakeSessionStore() + _install_cloud_doubles(cloud_prompt_execution_module, store, monkeypatch) + + result = _short_circuit(cloud_prompt_execution_module, store, _payload(map_item_index=5)) + + assert result is None + assert store.lookups == [] + + +def test_cloud_seed_phase_still_short_circuits( + cloud_prompt_execution_module: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Static seed phases keep reusing outputs that are identical for every lane.""" + store = _FakeSessionStore() + _install_cloud_doubles(cloud_prompt_execution_module, store, monkeypatch) + + result = _short_circuit(cloud_prompt_execution_module, store, _payload(map_item_index=None)) + + assert result is not None and len(result) == 1 + assert len(store.lookups) == 1 + + +def test_local_mapped_item_never_reuses_shared_session_output( + local_execution_module: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The host-side fallback applies the same rule as the cloud executor.""" + store = _FakeSessionStore() + monkeypatch.setattr( + local_execution_module._host_session_bridge, "_REMOTE_SESSION_STORE", store + ) + + result = _short_circuit(local_execution_module, store, _payload(map_item_index=0)) + + assert result is None + assert store.lookups == [] diff --git a/tests/test_progressive_preview.py b/tests/test_progressive_preview.py new file mode 100644 index 0000000..f131aa3 --- /dev/null +++ b/tests/test_progressive_preview.py @@ -0,0 +1,565 @@ +"""Tests for progressive remote image previews.""" + +from __future__ import annotations + +import asyncio +import importlib +import types +from pathlib import Path +from typing import Any + +import pytest + +from conftest import PACKAGE_NAME +from modal_executor_test_support import _FakeRewritePreviewImageNode +from modal_executor_test_support import _FakeRewriteRemoteImageNode +from modal_executor_test_support import _FakeRewriteRemoteTextNode +from modal_executor_test_support import _FakeRewriteLocalSinkNode + + +@pytest.fixture(scope="session") +def progressive_preview_module(extension_package: object) -> Any: + """Return the shared progressive preview state module.""" + return importlib.import_module(f"{PACKAGE_NAME}.progressive_preview") + + +@pytest.fixture(scope="session") +def progressive_preview_node_module(extension_package: object) -> Any: + """Return the RemoteProgressivePreview node module.""" + return importlib.import_module(f"{PACKAGE_NAME}.progressive_preview_node") + + +@pytest.fixture(autouse=True) +def _clear_progressive_registry(progressive_preview_module: Any) -> None: + """Start every test with an empty global registry.""" + progressive_preview_module.PROGRESSIVE_PREVIEW_REGISTRY.clear() + + +def _image_ref(progressive_preview_module: Any, name: str) -> Any: + """Return one temp preview image reference.""" + return progressive_preview_module.PreviewImageRef( + filename=name, subfolder="", type="temp" + ) + + +def _bridge_ref(item: int) -> dict[str, Any]: + """Return a minimal payload recognized as a remote session bridge ref.""" + return {"__comfy_modal_remote_session_bridge_ref__": True, "item": item} + + +def test_registry_tracks_slots_until_complete(progressive_preview_module: Any) -> None: + """Slots fill out of order, report completeness, and return images in slot order.""" + registry = progressive_preview_module.ProgressivePreviewRegistry() + registry.begin("prompt-1", "18", 3) + registry.fill("prompt-1", "18", 2, [_image_ref(progressive_preview_module, "c.png")]) + state = registry.fill( + "prompt-1", "18", 0, [_image_ref(progressive_preview_module, "a.png")] + ) + + assert not state.is_complete() + assert state.effective_slot_count() == 3 + + registry.fill("prompt-1", "18", 1, [_image_ref(progressive_preview_module, "b.png")]) + final_state = registry.pop("prompt-1", "18") + + assert final_state is not None and final_state.is_complete() + assert [image.filename for image in final_state.ordered_images()] == [ + "a.png", + "b.png", + "c.png", + ] + assert registry.get("prompt-1", "18") is None + + +def test_registry_fill_without_begin_grows_slot_count( + progressive_preview_module: Any, +) -> None: + """A non-mapped stream fills slot 0 and never claims to know the full count.""" + registry = progressive_preview_module.ProgressivePreviewRegistry() + state = registry.fill( + "prompt-1", "18", 0, [_image_ref(progressive_preview_module, "a.png")] + ) + + assert state.slot_count is None + assert state.effective_slot_count() == 1 + assert not state.is_complete() + + +def test_registry_drops_earlier_prompts_for_the_same_node( + progressive_preview_module: Any, +) -> None: + """A new prompt replaces a node's stale slots and the registry stays bounded.""" + registry = progressive_preview_module.ProgressivePreviewRegistry(entry_limit=2) + registry.begin("prompt-1", "18", 2) + registry.begin("prompt-2", "18", 4) + registry.begin("prompt-2", "19", 1) + registry.begin("prompt-2", "20", 1) + + assert registry.get("prompt-1", "18") is None + assert registry.get("prompt-2", "18") is None + assert registry.get("prompt-2", "19") is not None + assert registry.get("prompt-2", "20") is not None + + +def test_iter_image_batches_flattens_lists_and_adds_batch_dimension( + progressive_preview_module: Any, +) -> None: + """Lists of IMAGE tensors flatten and single frames gain a batch dimension.""" + torch = pytest.importorskip("torch") + batches = list( + progressive_preview_module.iter_image_batches( + [torch.zeros(2, 4, 4, 3), [torch.zeros(4, 4, 3)], "not-an-image"] + ) + ) + + assert [tuple(batch.shape) for batch in batches] == [(2, 4, 4, 3), (1, 4, 4, 3)] + + +def test_resolve_reuses_complete_streamed_slots( + progressive_preview_module: Any, + progressive_preview_node_module: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A fully streamed preview never saves or downloads the delivered value again.""" + monkeypatch.setattr( + progressive_preview_node_module, + "save_temp_preview_images", + lambda value: pytest.fail("streamed slots should be reused"), + ) + state = progressive_preview_module.ProgressivePreviewState( + slot_count=2, + slots={ + 1: [_image_ref(progressive_preview_module, "b.png")], + 0: [_image_ref(progressive_preview_module, "a.png")], + }, + ) + + images = asyncio.run( + progressive_preview_node_module.resolve_progressive_preview_images( + [_bridge_ref(0), _bridge_ref(1)], state + ) + ) + + assert [image.filename for image in images] == ["a.png", "b.png"] + + +def test_resolve_materializes_only_missing_bridge_slots( + progressive_preview_module: Any, + progressive_preview_node_module: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Partially streamed refs reuse streamed slots and download only the gaps.""" + materialized: list[dict[str, Any]] = [] + monkeypatch.setattr( + progressive_preview_node_module, + "_materialize_bridge_ref", + lambda ref: materialized.append(ref) or f"tensor-{ref['item']}", + ) + monkeypatch.setattr( + progressive_preview_node_module, + "save_temp_preview_images", + lambda value: [_image_ref(progressive_preview_module, f"{value}.png")], + ) + state = progressive_preview_module.ProgressivePreviewState( + slot_count=2, slots={0: [_image_ref(progressive_preview_module, "streamed.png")]} + ) + + images = asyncio.run( + progressive_preview_node_module.resolve_progressive_preview_images( + [_bridge_ref(0), _bridge_ref(1)], state + ) + ) + + assert [image.filename for image in images] == ["streamed.png", "tensor-1.png"] + assert materialized == [_bridge_ref(1)] + + +def test_node_schema_is_list_input_output_node( + progressive_preview_node_module: Any, +) -> None: + """The node behaves like PreviewImage but receives the whole mapped list at once.""" + schema = progressive_preview_node_module.RemoteProgressivePreview.define_schema() + + assert schema.node_id == "RemoteProgressivePreview" + assert schema.is_output_node is True + assert schema.is_input_list is True + assert schema.outputs == [] + + +def test_node_execute_pops_registry_and_returns_ui_images( + progressive_preview_module: Any, + progressive_preview_node_module: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Final execution returns the streamed images and clears the registry entry.""" + registry = progressive_preview_module.PROGRESSIVE_PREVIEW_REGISTRY + registry.begin("prompt-1", "18", 1) + registry.fill("prompt-1", "18", 0, [_image_ref(progressive_preview_module, "a.png")]) + monkeypatch.setattr( + progressive_preview_node_module, + "_executing_prompt_and_node", + lambda: ("prompt-1", "18"), + ) + + output = asyncio.run( + progressive_preview_node_module.RemoteProgressivePreview.execute( + [_bridge_ref(0)] + ) + ) + + assert output.ui == { + "images": [{"filename": "a.png", "subfolder": "", "type": "temp"}] + } + assert registry.get("prompt-1", "18") is None + + +class _FakePromptServer: + """Collect websocket events sent by local UI emitters.""" + + def __init__(self) -> None: + """Create an empty event log.""" + self.events: list[tuple[str, dict[str, Any], str | None]] = [] + + def send_sync(self, event: str, data: dict[str, Any], client_id: str | None) -> None: + """Record one websocket event.""" + self.events.append((event, data, client_id)) + + +def test_local_emitters_send_start_and_slot_events( + progressive_preview_module: Any, + local_ui_events_module: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Start lays out empty slots; each slot event carries its saved images.""" + prompt_server = _FakePromptServer() + monkeypatch.setattr( + local_ui_events_module, "_lookup_local_prompt_server", lambda: prompt_server + ) + monkeypatch.setattr( + local_ui_events_module, "_record_local_modal_ui_event", lambda *args: None + ) + monkeypatch.setattr( + progressive_preview_module, + "save_temp_preview_images", + lambda value: [_image_ref(progressive_preview_module, f"{value}.png")], + ) + + local_ui_events_module._emit_local_progressive_preview_start( + prompt_id="prompt-1", client_id="client-1", node_ids=["18"], slot_count=3 + ) + local_ui_events_module._emit_local_progressive_preview_slot( + prompt_id="prompt-1", + client_id="client-1", + node_ids=["18"], + slot_index=1, + slot_count=3, + image_value="item-1", + ) + + assert prompt_server.events == [ + ( + "modal_progressive_preview", + {"phase": "start", "prompt_id": "prompt-1", "node": "18", "slot_count": 3}, + "client-1", + ), + ( + "modal_progressive_preview", + { + "phase": "slot", + "prompt_id": "prompt-1", + "node": "18", + "slot_index": 1, + "slot_count": 3, + "images": [{"filename": "item-1.png", "subfolder": "", "type": "temp"}], + }, + "client-1", + ), + ] + state = progressive_preview_module.PROGRESSIVE_PREVIEW_REGISTRY.get("prompt-1", "18") + assert state is not None and sorted(state.slots) == [1] + + +def test_boundary_stream_event_fills_the_mapped_item_slot( + payload_stream_module: Any, + serialization_module: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A streamed boundary output fills the slot named by the item payload.""" + slot_calls: list[dict[str, Any]] = [] + preview_calls: list[dict[str, Any]] = [] + monkeypatch.setattr( + payload_stream_module, + "_emit_local_progressive_preview_slot", + lambda **kwargs: slot_calls.append(kwargs), + ) + monkeypatch.setattr( + payload_stream_module, + "_emit_local_preview_boundary_output", + lambda **kwargs: preview_calls.append(kwargs), + ) + state = types.SimpleNamespace( + prompt_id="prompt-1", + client_id="client-1", + payload={"component_id": "1__mapped::item:2", "map_item_index": 2, "map_item_count": 4}, + ) + + payload_stream_module._handle_boundary_output_event( + state, + { + "event_type": "boundary_output", + "node_id": "9", + "preview_target_node_ids": [], + "progressive_preview_target_node_ids": ["18"], + "value": serialization_module.serialize_value(["image-2"]), + }, + ) + + assert slot_calls == [ + { + "prompt_id": "prompt-1", + "client_id": "client-1", + "node_ids": ["18"], + "slot_index": 2, + "slot_count": 4, + "image_value": ["image-2"], + } + ] + assert preview_calls == [] + + +def test_mapped_item_payloads_carry_item_count(mapped_execution_module: Any) -> None: + """Per-item payloads tell the host how many slots the mapped run fills.""" + payload = { + "component_id": "1__mapped", + "mapped_input": {"proxy_input_name": "remote_input_0"}, + "boundary_outputs": [ + { + "node_id": "9", + "output_index": 0, + "io_type": "IMAGE", + "progressive_preview_target_node_ids": ["18"], + } + ], + "execute_node_ids": ["9"], + } + + item_payload = mapped_execution_module._build_mapped_item_payload( + payload, 1, 0, item_count=3 + ) + + assert item_payload["map_item_index"] == 1 + assert item_payload["map_item_count"] == 3 + assert mapped_execution_module._mapped_progressive_preview_targets(payload) == ["18"] + + +def test_mapped_start_lays_out_slots_for_progressive_targets( + mapped_execution_module: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Mapped dispatch announces one slot per item before any item runs.""" + start_calls: list[dict[str, Any]] = [] + monkeypatch.setattr( + mapped_execution_module, + "_emit_local_progressive_preview_start", + lambda **kwargs: start_calls.append(kwargs), + ) + payload = { + "component_id": "1__mapped", + "prompt_id": "prompt-1", + "extra_data": {"client_id": "client-1"}, + "mapped_input": {"proxy_input_name": "remote_input_0"}, + "boundary_outputs": [ + {"node_id": "9", "progressive_preview_target_node_ids": ["18"]} + ], + } + + mapped_execution_module._emit_mapped_progressive_preview_start(payload, 5) + + assert start_calls == [ + { + "prompt_id": "prompt-1", + "client_id": "client-1", + "node_ids": ["18"], + "slot_count": 5, + } + ] + + +class _FakeProgressivePreviewNode(_FakeRewritePreviewImageNode): + """Fake local RemoteProgressivePreview output node.""" + + +def test_progressive_preview_consumer_gets_bridge_without_materializer( + api_intercept_module: Any, + modal_executor_module: Any, + settings_module: Any, + sync_engine_module: Any, + tmp_path: Path, +) -> None: + """A progressive preview beside a remote continuation is streamed, not materialized.""" + settings = settings_module.ModalSyncSettings( + app_name="app", + auto_deploy=True, + allow_ephemeral_fallback=False, + enable_memory_snapshot=True, + enable_gpu_memory_snapshot=False, + execution_mode="local", + sync_custom_nodes=False, + volume_name="volume", + route_path="/modal/queue_prompt", + marker_property="is_modal_remote", + local_storage_root=tmp_path / "storage", + remote_storage_root="/storage", + custom_nodes_archive_name="custom_nodes_bundle.zip", + comfyui_root=None, + custom_nodes_dir=tmp_path / "custom_nodes", + ) + settings.custom_nodes_dir.mkdir() + sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) + fake_nodes_module = types.SimpleNamespace( + NODE_CLASS_MAPPINGS={ + "RemoteImage": _FakeRewriteRemoteImageNode, + "RemoteText": _FakeRewriteRemoteTextNode, + "RemoteProgressivePreview": _FakeProgressivePreviewNode, + "LocalSink": _FakeRewriteLocalSinkNode, + }, + NODE_DISPLAY_NAME_MAPPINGS={}, + ) + workflow = { + "nodes": [ + {"id": 1, "properties": {"is_modal_remote": True}}, + {"id": 2, "properties": {"is_modal_remote": True}}, + {"id": 3, "properties": {"is_modal_remote": False}}, + {"id": 4, "properties": {"is_modal_remote": False}}, + ] + } + prompt = { + "1": {"class_type": "RemoteImage", "inputs": {}}, + "2": {"class_type": "RemoteText", "inputs": {"image": ["1", 0]}}, + "3": {"class_type": "RemoteProgressivePreview", "inputs": {"images": ["1", 0]}}, + "4": {"class_type": "LocalSink", "inputs": {"text": ["2", 0]}}, + } + + rewritten_prompt, summary = api_intercept_module.rewrite_prompt_for_modal( + prompt=prompt, + workflow=workflow, + sync_engine=sync_engine, + settings=settings, + nodes_module=fake_nodes_module, + extra_data={"prompt_id": "prompt-progressive"}, + ) + + producer_payload = modal_executor_module._rehydrate_proxy_payload( + rewritten_prompt["1"]["inputs"]["original_node_data"], + unique_id="1", + ) + boundary_output = producer_payload["boundary_outputs"][0] + assert boundary_output["session_output"] is True + assert boundary_output["progressive_preview_target_node_ids"] == ["3"] + assert not any( + prompt_node["class_type"] + == api_intercept_module.MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID + for prompt_node in rewritten_prompt.values() + ) + dispatch_gate_node_id = summary.parallel_local_branch_node_ids[0] + assert rewritten_prompt["3"]["inputs"]["images"] == [dispatch_gate_node_id, 0] + assert rewritten_prompt[dispatch_gate_node_id]["inputs"]["value"] == ["1", 0] + + +class _FakeRemoteSaveNode: + """Fake remote output node consuming a mapped sampler result.""" + + RETURN_TYPES: tuple[str, ...] = () + RETURN_NAMES: tuple[str, ...] = () + OUTPUT_IS_LIST: tuple[bool, ...] = () + OUTPUT_NODE = True + + +def test_hybrid_mapped_output_streams_to_progressive_preview( + api_intercept_module: Any, + modal_executor_module: Any, + settings_module: Any, + sync_engine_module: Any, + tmp_path: Path, +) -> None: + """The logged failure shape: a mapped output feeding a remote save and a local preview.""" + import json + + from test_modal_test_workflow import ( + _FakeRegressionLatentSourceNode, + _FakeRegressionLocalSinkNode, + _FakeRegressionModalMapInputNode, + _FakeRegressionRemoteModelNode, + _FakeRegressionRemoteSamplerNode, + ) + + artifact = json.loads( + (Path(__file__).parent / "fixtures" / "modal_hybrid_split_regression.json").read_text() + ) + prompt = artifact["prompt"] + prompt["8"] = {"class_type": "RemoteProgressivePreview", "inputs": {"images": ["7", 0]}} + prompt["9"] = {"class_type": "RemoteSave", "inputs": {"images": ["7", 0]}} + workflow = artifact["workflow"] + workflow["nodes"].append({"id": 9, "properties": {"is_modal_remote": True}}) + settings = settings_module.ModalSyncSettings( + app_name="app", + auto_deploy=True, + allow_ephemeral_fallback=False, + enable_memory_snapshot=True, + enable_gpu_memory_snapshot=False, + execution_mode="local", + sync_custom_nodes=False, + volume_name="volume", + route_path="/modal/queue_prompt", + marker_property="is_modal_remote", + local_storage_root=tmp_path / "storage", + remote_storage_root="/storage", + custom_nodes_archive_name="custom_nodes_bundle.zip", + comfyui_root=None, + custom_nodes_dir=tmp_path / "custom_nodes", + ) + settings.custom_nodes_dir.mkdir() + sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) + fake_nodes_module = types.SimpleNamespace( + NODE_CLASS_MAPPINGS={ + "RemoteModel": _FakeRegressionRemoteModelNode, + "RemoteSampler": _FakeRegressionRemoteSamplerNode, + "LatentSource": _FakeRegressionLatentSourceNode, + "ModalMapInput": _FakeRegressionModalMapInputNode, + "LocalSink": _FakeRegressionLocalSinkNode, + "RemoteProgressivePreview": _FakeProgressivePreviewNode, + "RemoteSave": _FakeRemoteSaveNode, + }, + NODE_DISPLAY_NAME_MAPPINGS={}, + ) + + rewritten_prompt, summary = api_intercept_module.rewrite_prompt_for_modal( + prompt=prompt, + workflow=workflow, + sync_engine=sync_engine, + settings=settings, + nodes_module=fake_nodes_module, + extra_data={"prompt_id": "prompt-hybrid-progressive"}, + ) + + preview_source = rewritten_prompt["8"]["inputs"]["images"] + if preview_source[0] in summary.parallel_local_branch_node_ids: + preview_source = rewritten_prompt[preview_source[0]]["inputs"]["value"] + producer_id = str(preview_source[0]) + producer_payload = modal_executor_module._rehydrate_proxy_payload( + rewritten_prompt[producer_id]["inputs"]["original_node_data"], + unique_id=producer_id, + ) + progressive_outputs = [ + output + for output in producer_payload["boundary_outputs"] + if output.get("progressive_preview_target_node_ids") == ["8"] + ] + assert producer_id.endswith("__mapped") + assert progressive_outputs, producer_payload["boundary_outputs"] + assert progressive_outputs[0]["session_output"] is True + assert not any( + prompt_node["class_type"] + == api_intercept_module.MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID + for prompt_node in rewritten_prompt.values() + ) diff --git a/tests/test_progressive_preview_frontend.py b/tests/test_progressive_preview_frontend.py new file mode 100644 index 0000000..1fea697 --- /dev/null +++ b/tests/test_progressive_preview_frontend.py @@ -0,0 +1,17 @@ +"""Run the progressive preview frontend checks through pytest.""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def test_progressive_preview_frontend() -> None: + """Verify slot layout, filling, and native-preview handoff in the browser extension.""" + subprocess.run( + ["node", "tests/progressive_preview_frontend.test.mjs"], + cwd=REPO_ROOT, + check=True, + ) diff --git a/tests/test_prompt_rewrite_mapped.py b/tests/test_prompt_rewrite_mapped.py index e7194ec..8fc32b1 100644 --- a/tests/test_prompt_rewrite_mapped.py +++ b/tests/test_prompt_rewrite_mapped.py @@ -419,7 +419,7 @@ def test_rewrite_marks_local_modal_map_source_as_mapped_subgraph( assert payload["boundary_inputs"] == [ { "proxy_input_name": "remote_input_0", - "io_type": "*", + "io_type": "STRING", "targets": [{"node_id": "3", "input_name": "text"}], } ] @@ -431,7 +431,7 @@ def test_rewrite_marks_local_modal_map_source_as_mapped_subgraph( "boundary_inputs": [ { "proxy_input_name": "remote_input_0", - "io_type": "*", + "io_type": "STRING", "targets": [{"node_id": "3", "input_name": "text"}], } ], diff --git a/web/modal_toggle.js b/web/modal_toggle.js index 8515cda..1f9f434 100644 --- a/web/modal_toggle.js +++ b/web/modal_toggle.js @@ -1747,6 +1747,8 @@ function replayModalUiEvent(eventRecord) { handleModalStatus({ detail: payload }); } else if (eventName === "modal_progress") { handleModalProgress({ detail: payload }); + } else if (eventName === "modal_progressive_preview") { + api.dispatchEvent?.(new CustomEvent(eventName, { detail: payload })); } } finally { modalReplayedEventUpdatedAtMs = null; diff --git a/web/progressive_preview.js b/web/progressive_preview.js new file mode 100644 index 0000000..4fa4f62 --- /dev/null +++ b/web/progressive_preview.js @@ -0,0 +1,403 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +const PROGRESSIVE_PREVIEW_NODE_CLASS = "RemoteProgressivePreview"; +const PROGRESSIVE_PREVIEW_EVENT = "modal_progressive_preview"; +const PROGRESSIVE_PREVIEW_WIDGET_NAME = "progressive_preview"; +const PROGRESSIVE_PREVIEW_STYLE_ID = "comfy-progressive-preview-style"; +const PROGRESSIVE_PREVIEW_MIN_HEIGHT = 160; +const PROGRESSIVE_PREVIEW_MIN_WIDTH = 280; + +/** @type {WeakMap} */ +const progressivePreviewPanels = new WeakMap(); + +/** + * @typedef {{filename: string, subfolder?: string, type?: string}} PreviewImage + * @typedef {{promptId: string|null, slotCount: number, slots: Map}} SlotState + * @typedef {{node: object, root: HTMLElement, header: HTMLElement, grid: HTMLElement, + * cells: HTMLElement[], state: SlotState, visible: boolean, widget: object|null, + * finishedPromptId: string|null}} ProgressivePreviewPanel + */ + +/** + * Create an empty progressive preview slot state. + * @returns {SlotState} + */ +function createSlotState() { + return { promptId: null, slotCount: 0, slots: new Map() }; +} + +/** + * Apply one progressive preview websocket payload to a slot state. + * @param {SlotState} state + * @param {Record} detail + * @returns {boolean} Whether the payload was applied. + */ +function applyProgressivePreviewEvent(state, detail) { + if (!detail || typeof detail !== "object") { + return false; + } + const promptId = detail.prompt_id == null ? null : String(detail.prompt_id); + if (promptId !== state.promptId) { + state.promptId = promptId; + state.slotCount = 0; + state.slots = new Map(); + } + const declaredCount = Math.max(0, Number(detail.slot_count) || 0); + if (detail.phase === "start") { + state.slotCount = declaredCount; + state.slots = new Map(); + return true; + } + if (detail.phase !== "slot") { + return false; + } + const slotIndex = Math.max(0, Number(detail.slot_index) || 0); + const images = Array.isArray(detail.images) ? detail.images : []; + state.slots.set(slotIndex, images); + state.slotCount = Math.max(state.slotCount, declaredCount, slotIndex + 1); + return true; +} + +/** + * Return a near-square column count for a slot grid. + * @param {number} slotCount + * @returns {number} + */ +function slotGridColumns(slotCount) { + return Math.max(1, Math.ceil(Math.sqrt(Math.max(1, slotCount)))); +} + +/** + * Return how many slots have received images. + * @param {SlotState} state + * @returns {number} + */ +function filledSlotCount(state) { + let filled = 0; + for (let index = 0; index < state.slotCount; index += 1) { + if ((state.slots.get(index) ?? []).length > 0) { + filled += 1; + } + } + return filled; +} + +/** + * Return the header text for a slot state. + * @param {SlotState} state + * @returns {string} + */ +function progressivePreviewHeaderText(state) { + if (state.slotCount <= 0) { + return "Waiting for remote images…"; + } + const filled = filledSlotCount(state); + if (filled >= state.slotCount) { + return `Remote preview · all ${state.slotCount} received`; + } + return `Remote preview · ${filled}/${state.slotCount} received`; +} + +/** + * Build the ComfyUI /view URL for one saved preview image. + * @param {PreviewImage} image + * @returns {string} + */ +function progressivePreviewImageUrl(image) { + const params = new URLSearchParams({ + filename: String(image.filename ?? ""), + subfolder: String(image.subfolder ?? ""), + type: String(image.type ?? "temp"), + }); + const path = `/view?${params.toString()}`; + return typeof api?.apiURL === "function" ? api.apiURL(path) : path; +} + +/** + * Look up one node in a graph by string or numeric id. + * @param {object|null|undefined} graph + * @param {string} id + * @returns {object|null} + */ +function graphNodeById(graph, id) { + if (!graph) { + return null; + } + const direct = graph.getNodeById?.(id) ?? null; + if (direct) { + return direct; + } + const numericId = Number(id); + return Number.isFinite(numericId) ? graph.getNodeById?.(numericId) ?? null : null; +} + +/** + * Resolve a prompt node id, including "subgraph:node" paths, to a live graph node. + * @param {string|number} nodeIdValue + * @returns {object|null} + */ +function resolveProgressivePreviewNode(nodeIdValue) { + const segments = String(nodeIdValue ?? "").split(":").filter(Boolean); + let graph = app.rootGraph ?? app.graph; + let node = null; + for (let index = 0; index < segments.length; index += 1) { + node = graphNodeById(graph, segments[index]); + if (!node) { + return null; + } + if (index < segments.length - 1) { + graph = node.subgraph ?? null; + } + } + return node; +} + +/** + * Return whether a graph node is a progressive preview node. + * @param {object|null} node + * @returns {boolean} + */ +function isProgressivePreviewNode(node) { + return (node?.comfyClass ?? node?.type) === PROGRESSIVE_PREVIEW_NODE_CLASS; +} + +/** Install the panel stylesheet once. */ +function installProgressivePreviewStyles() { + if (typeof document === "undefined" || document.getElementById?.(PROGRESSIVE_PREVIEW_STYLE_ID)) { + return; + } + const style = document.createElement("style"); + style.id = PROGRESSIVE_PREVIEW_STYLE_ID; + style.textContent = ` + .comfy-progressive-preview { display: flex; flex-direction: column; gap: 6px; + width: 100%; height: 100%; box-sizing: border-box; padding: 4px; overflow: hidden; } + .comfy-progressive-preview[hidden] { display: none; } + .comfy-progressive-preview-header { font-size: 11px; opacity: 0.75; + color: var(--input-text, #ddd); white-space: nowrap; } + .comfy-progressive-preview-grid { display: grid; gap: 4px; flex: 1; min-height: 0; } + .comfy-progressive-preview-cell { position: relative; aspect-ratio: 1; min-height: 0; + border-radius: 4px; overflow: hidden; background: var(--comfy-input-bg, #222); + display: flex; align-items: center; justify-content: center; } + .comfy-progressive-preview-cell.is-empty { border: 1px dashed var(--border-color, #555); + animation: comfy-progressive-preview-pulse 1.6s ease-in-out infinite; } + .comfy-progressive-preview-cell img { width: 100%; height: 100%; object-fit: contain; } + .comfy-progressive-preview-label { position: absolute; left: 4px; bottom: 2px; + font-size: 10px; color: var(--input-text, #ddd); opacity: 0.7; } + .comfy-progressive-preview-badge { position: absolute; right: 4px; top: 3px; font-size: 10px; + padding: 0 4px; border-radius: 6px; background: rgba(0, 0, 0, 0.6); color: #fff; } + @keyframes comfy-progressive-preview-pulse { 50% { opacity: 0.45; } } + `; + document.head?.appendChild(style); +} + +/** + * Create the DOM panel and widget for one progressive preview node. + * @param {object} node + * @returns {ProgressivePreviewPanel|null} + */ +function ensureProgressivePreviewPanel(node) { + const existing = progressivePreviewPanels.get(node); + if (existing) { + return existing; + } + if (typeof document === "undefined" || typeof node?.addDOMWidget !== "function") { + return null; + } + const root = document.createElement("div"); + root.className = "comfy-progressive-preview"; + root.hidden = true; + root.style.display = "none"; + const header = document.createElement("div"); + header.className = "comfy-progressive-preview-header"; + const grid = document.createElement("div"); + grid.className = "comfy-progressive-preview-grid"; + root.appendChild(header); + root.appendChild(grid); + /** @type {ProgressivePreviewPanel} */ + const panel = { + node, + root, + header, + grid, + cells: [], + state: createSlotState(), + visible: false, + widget: null, + finishedPromptId: null, + }; + try { + panel.widget = node.addDOMWidget(PROGRESSIVE_PREVIEW_WIDGET_NAME, "progressive-preview", root, { + serialize: false, + hideOnZoom: false, + getMinHeight: () => (panel.visible ? PROGRESSIVE_PREVIEW_MIN_HEIGHT : 0), + }); + } catch (error) { + console.warn("Unable to mount the remote progressive preview panel.", error); + return null; + } + if (panel.widget) { + panel.widget.serialize = false; + } + progressivePreviewPanels.set(node, panel); + return panel; +} + +/** + * Return the cell element for one slot, creating it when needed. + * @param {ProgressivePreviewPanel} panel + * @param {number} index + * @returns {HTMLElement} + */ +function progressivePreviewCell(panel, index) { + let cell = panel.cells[index]; + if (!cell) { + cell = document.createElement("div"); + panel.cells[index] = cell; + } + return cell; +} + +/** + * Render one slot cell as an empty placeholder or its first image. + * @param {HTMLElement} cell + * @param {number} index + * @param {PreviewImage[]} images + */ +function renderProgressivePreviewCell(cell, index, images) { + const firstImage = images[0]; + const imageKey = firstImage ? progressivePreviewImageUrl(firstImage) : ""; + if (cell.dataset.imageKey === imageKey && cell.dataset.imageCount === String(images.length)) { + return; + } + cell.dataset.imageKey = imageKey; + cell.dataset.imageCount = String(images.length); + cell.replaceChildren?.(); + cell.className = `comfy-progressive-preview-cell${firstImage ? "" : " is-empty"}`; + cell.title = firstImage ? `Item ${index + 1}` : `Item ${index + 1} · waiting for remote result`; + if (firstImage) { + const image = document.createElement("img"); + image.src = imageKey; + image.alt = `Remote preview item ${index + 1}`; + cell.appendChild(image); + } + if (images.length > 1) { + const badge = document.createElement("span"); + badge.className = "comfy-progressive-preview-badge"; + badge.textContent = `×${images.length}`; + cell.appendChild(badge); + } + const label = document.createElement("span"); + label.className = "comfy-progressive-preview-label"; + label.textContent = String(index + 1); + cell.appendChild(label); +} + +/** + * Render a panel's header and slot grid from its state. + * @param {ProgressivePreviewPanel} panel + */ +function renderProgressivePreviewPanel(panel) { + const { state } = panel; + panel.header.textContent = progressivePreviewHeaderText(state); + panel.grid.style.gridTemplateColumns = `repeat(${slotGridColumns(state.slotCount)}, minmax(0, 1fr))`; + const cells = []; + for (let index = 0; index < state.slotCount; index += 1) { + const cell = progressivePreviewCell(panel, index); + renderProgressivePreviewCell(cell, index, state.slots.get(index) ?? []); + cells.push(cell); + } + panel.cells.length = state.slotCount; + panel.grid.replaceChildren?.(...cells); +} + +/** + * Show or hide the progressive panel, handing the node back to the native preview when hidden. + * @param {ProgressivePreviewPanel} panel + * @param {boolean} visible + */ +function setProgressivePreviewVisible(panel, visible) { + if (panel.visible === visible) { + return; + } + panel.visible = visible; + panel.root.hidden = !visible; + panel.root.style.display = visible ? "" : "none"; + const node = panel.node; + if (visible) { + node.imgs = undefined; + node.imageIndex = null; + node.setSize?.([ + Math.max(Number(node.size?.[0]) || 0, PROGRESSIVE_PREVIEW_MIN_WIDTH), + Math.max(Number(node.size?.[1]) || 0, PROGRESSIVE_PREVIEW_MIN_HEIGHT + 40), + ]); + } + node.setDirtyCanvas?.(true, true); +} + +/** + * Handle one progressive preview websocket event. + * @param {{detail?: Record}} event + * @returns {boolean} Whether a node panel was updated. + */ +function handleProgressivePreviewEvent(event) { + const detail = event?.detail; + const node = resolveProgressivePreviewNode(detail?.node); + if (!isProgressivePreviewNode(node)) { + return false; + } + const panel = ensureProgressivePreviewPanel(node); + if (!panel || isFinishedProgressivePreviewEvent(panel, detail)) { + return false; + } + if (!applyProgressivePreviewEvent(panel.state, detail)) { + return false; + } + renderProgressivePreviewPanel(panel); + setProgressivePreviewVisible(panel, true); + app.graph?.setDirtyCanvas?.(true, true); + return true; +} + +/** + * Return whether an event belongs to a prompt whose final images were already shown. + * Replayed events after a browser refocus must not resurrect a finished panel. + * @param {ProgressivePreviewPanel} panel + * @param {Record} detail + * @returns {boolean} + */ +function isFinishedProgressivePreviewEvent(panel, detail) { + const promptId = detail?.prompt_id == null ? null : String(detail.prompt_id); + return promptId !== null && promptId === panel.finishedPromptId; +} + +/** + * Hide a node's progressive panel once ComfyUI delivers its final images. + * @param {object} node + */ +function finishProgressivePreview(node) { + const panel = progressivePreviewPanels.get(node); + if (panel) { + panel.finishedPromptId = panel.state.promptId; + setProgressivePreviewVisible(panel, false); + } +} + +app.registerExtension({ + name: "comfy.modal.progressivePreview", + setup() { + installProgressivePreviewStyles(); + api.addEventListener(PROGRESSIVE_PREVIEW_EVENT, handleProgressivePreviewEvent); + }, + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name !== PROGRESSIVE_PREVIEW_NODE_CLASS) { + return; + } + const originalOnExecuted = nodeType.prototype.onExecuted; + nodeType.prototype.onExecuted = function onProgressivePreviewExecuted(...args) { + const result = originalOnExecuted?.apply(this, args); + finishProgressivePreview(this); + return result; + }; + }, +});