diff --git a/README.md b/README.md index f6c7754..76552ed 100644 --- a/README.md +++ b/README.md @@ -164,6 +164,8 @@ The frontend shows remote state directly on the canvas: - steady green border: finished for the current run - red border: queue-time or execution failure +When a remote node reports numeric progress through ComfyUI's progress hooks, the active Modal-marked node now also shows an in-node progress bar and percentage locally while the remote work is running. That is especially useful for sampler-style nodes such as `KSampler`. + The extension also keeps a global activity badge visible during queue-time sync and remote execution, including the period before the prompt is formally queued. ### 6. Asset sync expectations diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index 9f6a85c..336796c 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -82,6 +82,7 @@ class _NullPromptServer: """Initialize the no-op prompt server state.""" self.client_id: str | None = None self.last_node_id: str | None = None + self.last_prompt_id: str | None = None def send_sync(self, event: str, data: dict[str, Any], client_id: str | None) -> None: """Discard PromptExecutor progress and status events.""" @@ -131,6 +132,7 @@ class _TracingPromptServer(_NullPromptServer): self._status_callback = status_callback self._active_node_id: str | None = None self._active_node_started_at: float | None = None + self.last_prompt_id = prompt_id def _classify_node_role(self, class_type: str) -> str: """Return a coarse role name for a node class.""" @@ -185,6 +187,7 @@ class _TracingPromptServer(_NullPromptServer): ) self._active_node_id = str(next_node_id) self._active_node_started_at = time.perf_counter() + self.last_node_id = self._active_node_id _emit_cloud_info( "Remote node %s class_type=%s role=%s started", self._active_node_id, @@ -193,6 +196,49 @@ class _TracingPromptServer(_NullPromptServer): ) return + if event == "progress_state": + if self._status_callback is None: + return + + nodes_payload = data.get("nodes") + if not isinstance(nodes_payload, dict): + return + + tracked_node_id = self._active_node_id + tracked_node_state: dict[str, Any] | None = None + if tracked_node_id is not None: + candidate_state = nodes_payload.get(tracked_node_id) + if isinstance(candidate_state, dict): + tracked_node_state = candidate_state + + if tracked_node_state is None: + for node_state in nodes_payload.values(): + if isinstance(node_state, dict) and node_state.get("state") == "running": + tracked_node_state = node_state + break + + if tracked_node_state is None: + return + + display_node_id = tracked_node_state.get("display_node_id") + real_node_id = tracked_node_state.get("real_node_id") + reported_node_id = display_node_id or real_node_id or tracked_node_state.get("node_id") + if reported_node_id is None: + return + + self._status_callback( + { + "event_type": "node_progress", + "node_id": str(reported_node_id), + "display_node_id": ( + str(display_node_id) if display_node_id is not None else None + ), + "value": float(tracked_node_state.get("value", 0.0)), + "max": float(tracked_node_state.get("max", 1.0)), + } + ) + return + if event == "executed": executed_node_id = data.get("node") if executed_node_id is not None and str(executed_node_id) == self._active_node_id: @@ -427,6 +473,44 @@ def _ensure_headless_prompt_server_instance() -> None: logger.info("Installed headless PromptServer.instance for remote custom-node initialization.") +@contextmanager +def _temporary_progress_hook(prompt_server: _NullPromptServer) -> Iterator[None]: + """Install a ComfyUI progress hook so remote samplers emit numeric progress updates.""" + import comfy.utils + from comfy_execution.progress import get_progress_state + from comfy_execution.utils import get_executing_context + + previous_hook = comfy.utils.PROGRESS_BAR_HOOK + + def hook( + value: float, + total: float, + preview_image: Any, + prompt_id: str | None = None, + node_id: str | None = None, + ) -> None: + """Mirror ComfyUI progress-bar updates into the headless progress registry.""" + executing_context = get_executing_context() + if prompt_id is None and executing_context is not None: + prompt_id = executing_context.prompt_id + if node_id is None and executing_context is not None: + node_id = executing_context.node_id + if prompt_id is None: + prompt_id = prompt_server.last_prompt_id + if node_id is None: + node_id = prompt_server.last_node_id + if node_id is None: + return + + get_progress_state().update_progress(str(node_id), value, total, preview_image) + + comfy.utils.set_progress_bar_global_hook(hook) + try: + yield + finally: + comfy.utils.set_progress_bar_global_hook(previous_hook) + + def _ensure_default_custom_nodes_dir() -> Path | None: """Create the default ComfyUI custom_nodes directory when the image omits its contents.""" comfyui_root = _active_comfyui_root() @@ -848,12 +932,15 @@ def _execute_subgraph_prompt( execution = _load_execution_module() cache_type, cache_args = _prompt_executor_cache_config(execution) - with _temporary_node_mapping(None), _patched_folder_paths_absolute_lookup(): - prompt_server = _TracingPromptServer( + with ( + _temporary_node_mapping(None), + _patched_folder_paths_absolute_lookup(), + _temporary_progress_hook(prompt_server := _TracingPromptServer( component_id, prompt, status_callback=status_callback, - ) + )), + ): with _timed_phase("create_prompt_executor", component=component_id): executor_state = _get_or_create_prompt_executor_state( execution=execution, diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index b7bf5c8..01aa874 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -5,6 +5,7 @@ - The README had drifted into release-note territory, especially in the old `Important distinction:` block. For this project, the maintainable structure is: short overview, explicit local-vs-remote execution sections, concise operational caveats, and a separate learnings log for historical implementation details. - ComfyUI registry setup is mostly metadata, not code. The repo needs a real `[tool.comfy]` block in `pyproject.toml`, a stable immutable node-pack id, repository URLs, and an optional GitHub publish workflow; choosing a software license should stay explicit instead of being guessed during packaging work. - `comfy node pack` is not git-aware by default; it uses `.gitignore` as its practical packaging filter. For a clean registry artifact, this repo needs an allowlist-style `.gitignore` so untracked local files like logs, caches, virtualenvs, and scratch docs cannot leak into `node.zip`. +- Remote per-node phase events are not enough to mirror sampler progress in the local UI. The Modal worker also needs a headless equivalent of ComfyUI's global progress hook so `ProgressBar` updates feed `progress_state`, and the local relay should forward those numeric updates as a separate high-frequency UI event rather than overloading coarse execution-phase messages. ## 2026-04-17 diff --git a/remote/modal_app.py b/remote/modal_app.py index efd3bc0..356c63a 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -523,6 +523,34 @@ def _emit_local_modal_status( prompt_server.send_sync("modal_status", payload, client_id) +def _emit_local_modal_progress( + *, + prompt_id: str | None, + client_id: str | None, + node_id: str, + value: float, + max_value: float, + display_node_id: str | None = None, +) -> None: + """Forward remote numeric node progress into the local ComfyUI websocket stream.""" + if client_id is None: + return + + prompt_server = _lookup_local_prompt_server() + if prompt_server is None: + return + + payload: dict[str, Any] = { + "prompt_id": prompt_id, + "node_id": node_id, + "value": float(value), + "max": float(max_value), + } + if display_node_id is not None: + payload["display_node_id"] = display_node_id + prompt_server.send_sync("modal_progress", payload, client_id) + + def _should_stream_remote_progress(payload: dict[str, Any]) -> bool: """Return whether the local client has enough context to mirror remote node progress.""" extra_data = payload.get("extra_data") or {} @@ -551,6 +579,29 @@ def _consume_remote_payload_stream( for stream_event in stream_events: event_kind = str(stream_event.get("kind", "")) if event_kind == "progress": + if stream_event.get("event_type") == "node_progress": + reported_node_id = stream_event.get("node_id") + if reported_node_id is not None: + logger.debug( + "Forwarding streamed Modal node progress for component=%s node_id=%s value=%s max=%s.", + payload.get("component_id"), + reported_node_id, + stream_event.get("value"), + stream_event.get("max"), + ) + _emit_local_modal_progress( + prompt_id=prompt_id, + client_id=client_id, + node_id=str(reported_node_id), + value=float(stream_event.get("value", 0.0)), + max_value=float(stream_event.get("max", 1.0)), + display_node_id=( + str(stream_event["display_node_id"]) + if stream_event.get("display_node_id") is not None + else None + ), + ) + continue logger.info( "Forwarding streamed Modal progress for component=%s phase=%s active_node_id=%s.", payload.get("component_id"), diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index d3de8dd..b4fb2ab 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -266,6 +266,15 @@ def test_modal_cloud_streams_progress_and_result_events( "active_node_role": "model_load", } ) + status_callback( + { + "event_type": "node_progress", + "node_id": "7", + "display_node_id": "7", + "value": 3, + "max": 10, + } + ) status_callback({"phase": "execution_success"}) progress_callbacks.append({"component_id": payload["component_id"], "kwargs": kwargs_payload}) return b"serialized-outputs" @@ -288,6 +297,14 @@ def test_modal_cloud_streams_progress_and_result_events( "active_node_class_type": "UNETLoader", "active_node_role": "model_load", }, + { + "kind": "progress", + "event_type": "node_progress", + "node_id": "7", + "display_node_id": "7", + "value": 3, + "max": 10, + }, { "kind": "progress", "phase": "execution_success", @@ -748,6 +765,14 @@ def test_remote_modal_consumes_streamed_progress_and_result( "active_node_class_type": "UNETLoader", "active_node_role": "model_load", }, + { + "kind": "progress", + "event_type": "node_progress", + "node_id": "7", + "display_node_id": "7", + "value": 4, + "max": 20, + }, { "kind": "progress", "phase": "execution_success", @@ -774,6 +799,17 @@ def test_remote_modal_consumes_streamed_progress_and_result( }, "client-1", ), + ( + "modal_progress", + { + "prompt_id": "prompt-1", + "node_id": "7", + "display_node_id": "7", + "value": 4.0, + "max": 20.0, + }, + "client-1", + ), ( "modal_status", { @@ -786,6 +822,46 @@ def test_remote_modal_consumes_streamed_progress_and_result( ] +def test_modal_cloud_tracing_prompt_server_emits_numeric_node_progress( + modal_cloud_module: Any, +) -> None: + """The cloud tracing prompt server should forward active-node numeric progress updates.""" + observed_updates: list[dict[str, Any]] = [] + server = modal_cloud_module._TracingPromptServer( + "component-1", + {"7": {"class_type": "KSampler", "inputs": {}}}, + status_callback=observed_updates.append, + ) + + server.send_sync("executing", {"node": "7"}, None) + server.send_sync( + "progress_state", + { + "prompt_id": "component-1", + "nodes": { + "7": { + "node_id": "7", + "display_node_id": "7", + "real_node_id": "7", + "state": "running", + "value": 5, + "max": 20, + } + }, + }, + None, + ) + + assert observed_updates[0]["phase"] == "executing" + assert observed_updates[1] == { + "event_type": "node_progress", + "node_id": "7", + "display_node_id": "7", + "value": 5.0, + "max": 20.0, + } + + def test_remote_modal_consumes_streamed_tensor_result_payload( remote_modal_app_module: Any, serialization_module: Any, diff --git a/tests/test_web_modal_toggle.py b/tests/test_web_modal_toggle.py index 0996622..50c9172 100644 --- a/tests/test_web_modal_toggle.py +++ b/tests/test_web_modal_toggle.py @@ -75,6 +75,18 @@ def test_streamed_modal_progress_takes_precedence_over_proxy_events() -> None: assert "if (promptState.hasStreamedProgress && phase !== STATE_ERROR)" in source +def test_streamed_modal_node_progress_updates_active_overlay() -> None: + """The frontend should listen for numeric Modal node progress and render it on the node.""" + source = _modal_toggle_source() + + assert 'api.addEventListener("modal_progress", handleModalProgress);' in source + assert "function handleModalProgress(event)" in source + assert "function setNodeProgress(nodeIdValue, promptId, value, maxValue)" in source + assert "function clearNodeProgress(nodeIdValue, promptId)" in source + assert "state?.progress" in source + assert 'ctx.fillText(`${Math.round(progressRatio * 100)}%`' in source + + def test_prompt_cleanup_prunes_orphaned_global_status_entries() -> None: """Workflow cleanup should remove stale global badge states once prompt state is gone.""" source = _modal_toggle_source() diff --git a/web/modal_toggle.js b/web/modal_toggle.js index 384731a..c3b9dde 100644 --- a/web/modal_toggle.js +++ b/web/modal_toggle.js @@ -22,6 +22,7 @@ const STATE_ERROR = "error"; const ERROR_CLEAR_DELAY_MS = 5000; const modalNodeStates = new Map(); +const modalNodeProgress = new Map(); const modalNodeClearTimers = new Map(); const modalPromptStates = new Map(); const syntheticPromptUiStates = new Map(); @@ -431,6 +432,9 @@ function setNodesPhase(nodeIds, phase, promptId, errorMessage) { if (phase === STATE_ERROR) { scheduleNodeClear(currentNodeId, promptId, ERROR_CLEAR_DELAY_MS); } + if (phase === STATE_COMPLETE || phase === STATE_ERROR) { + clearNodeProgress(currentNodeId, promptId); + } for (const ancestorNodeId of ancestorNodeIds(currentNodeId)) { affectedAncestorNodeIds.add(ancestorNodeId); } @@ -563,6 +567,53 @@ function setPromptActiveNode(promptId, activeNodeId) { promptState.activeNodeId = activeNodeId ? String(activeNodeId) : null; } +/** + * Remove stored numeric progress for one node and its visible ancestors. + * @param {string} nodeIdValue + * @param {string | undefined} promptId + */ +function clearNodeProgress(nodeIdValue, promptId) { + const progressNodeIds = [String(nodeIdValue), ...ancestorNodeIds(nodeIdValue)]; + for (const progressNodeId of progressNodeIds) { + const progressState = modalNodeProgress.get(progressNodeId); + if (!progressState) { + continue; + } + if (promptId && progressState.promptId !== promptId) { + continue; + } + modalNodeProgress.delete(progressNodeId); + } +} + +/** + * Record numeric progress for one node and its visible ancestors. + * @param {string} nodeIdValue + * @param {string} promptId + * @param {number} value + * @param {number} maxValue + */ +function setNodeProgress(nodeIdValue, promptId, value, maxValue) { + const safeMaxValue = Math.max(1, Number(maxValue) || 1); + const safeValue = Math.max(0, Math.min(safeMaxValue, Number(value) || 0)); + const progressNodeIds = [String(nodeIdValue), ...ancestorNodeIds(nodeIdValue)]; + + for (const progressNodeId of progressNodeIds) { + if (!shouldApplyPromptState(progressNodeId, promptId)) { + continue; + } + modalNodeProgress.set(progressNodeId, { + promptId, + value: safeValue, + max: safeMaxValue, + updatedAt: nowMs(), + }); + } + + ensureAnimationLoop(); + app.graph?.setDirtyCanvas(true, true); +} + /** * Return remote component node ids for an executing proxy node event. * @param {string} promptId @@ -614,9 +665,11 @@ function getRemoteVisualState(node) { return state; } const promptState = modalPromptStates.get(state.promptId); + const progressState = modalNodeProgress.get(nodeId(node)) ?? null; return { ...state, isActiveRemoteNode: promptState?.activeNodeId === nodeId(node), + progress: progressState?.promptId === state.promptId ? progressState : null, }; } @@ -691,6 +744,26 @@ function drawRemoteNodeDecoration(node, ctx) { node.size[1] + titleHeight + borderWidth, ); ctx.restore(); + + if (!state?.progress || state.phase !== STATE_ACTIVE) { + return; + } + + const progressRatio = Math.max(0, Math.min(1, state.progress.value / state.progress.max)); + const progressWidth = Math.max(0, (node.size[0] + borderWidth * 2) * progressRatio); + const progressY = node.size[1] - Math.max(10 / scale, borderWidth * 2); + + ctx.save(); + ctx.fillStyle = "rgba(15, 23, 42, 0.72)"; + ctx.fillRect(-borderWidth, progressY, node.size[0] + borderWidth * 2, 8 / scale); + ctx.fillStyle = "rgba(216, 180, 254, 0.92)"; + ctx.fillRect(-borderWidth, progressY, progressWidth, 8 / scale); + ctx.fillStyle = "#f8fafc"; + ctx.font = `${Math.max(11 / scale, 8)}px ui-sans-serif, system-ui, sans-serif`; + ctx.textAlign = "right"; + ctx.textBaseline = "bottom"; + ctx.fillText(`${Math.round(progressRatio * 100)}%`, node.size[0], progressY - 2 / scale); + ctx.restore(); } /** @@ -758,6 +831,9 @@ function handleModalStatus(event) { setGlobalStatusPhase(promptId, STATE_ERROR, nodeIds.length); setTimeout(() => clearGlobalStatusPhase(promptId), ERROR_CLEAR_DELAY_MS); setPromptActiveNode(promptId, null); + for (const nodeIdValue of nodeIds) { + clearNodeProgress(nodeIdValue, promptId); + } setNodesPhase(nodeIds, STATE_ERROR, promptId, detail.error_message); return; } @@ -770,6 +846,7 @@ function handleModalStatus(event) { setGlobalStatusPhase(promptId, EXECUTION_PHASE, nodeIds.length); setNodesPhase(nodeIds, STATE_READY, promptId); if (previousActiveNodeId && previousActiveNodeId !== nextActiveNodeId) { + clearNodeProgress(previousActiveNodeId, promptId); setNodesPhase([previousActiveNodeId], STATE_COMPLETE, promptId); } if (nextActiveNodeId) { @@ -782,6 +859,7 @@ function handleModalStatus(event) { if (detail.phase === "execution_success") { promptState.hasStreamedProgress = true; if (promptState.activeNodeId) { + clearNodeProgress(promptState.activeNodeId, promptId); setNodesPhase([promptState.activeNodeId], STATE_COMPLETE, promptId); } setPromptActiveNode(promptId, null); @@ -790,6 +868,31 @@ function handleModalStatus(event) { } } +/** + * Apply a streamed numeric Modal node-progress event. + * @param {CustomEvent} event + */ +function handleModalProgress(event) { + const detail = eventDetail(event); + const promptId = String(detail.prompt_id ?? ""); + const progressNodeId = String(detail.display_node_id ?? detail.node_id ?? ""); + if (!promptId || !progressNodeId) { + return; + } + + const promptState = ensurePromptState(promptId); + promptState.hasStreamedProgress = true; + setPromptActiveNode(promptId, progressNodeId); + setGlobalStatusPhase(promptId, EXECUTION_PHASE, promptState.remoteNodeIds.length || 1); + setNodesPhase([progressNodeId], STATE_ACTIVE, promptId); + setNodeProgress( + progressNodeId, + promptId, + Number(detail.value ?? 0), + Number(detail.max ?? 1), + ); +} + /** * Update remote component visuals from a native ComfyUI execution event. * @param {CustomEvent} event @@ -824,11 +927,17 @@ function handleExecutionPhase(event, phase) { setGlobalStatusPhase(promptId, STATE_ERROR, componentNodeIds.length); setTimeout(() => clearGlobalStatusPhase(promptId), ERROR_CLEAR_DELAY_MS); setPromptActiveNode(promptId, null); + for (const nodeIdValue of componentNodeIds) { + clearNodeProgress(nodeIdValue, promptId); + } setNodesPhase(componentNodeIds, STATE_ERROR, promptId, detail.exception_message); return; } if (phase === STATE_COMPLETE) { setPromptActiveNode(promptId, null); + for (const nodeIdValue of componentNodeIds) { + clearNodeProgress(nodeIdValue, promptId); + } setNodesPhase(componentNodeIds, STATE_COMPLETE, promptId, detail.exception_message); } } @@ -846,6 +955,7 @@ function clearPromptRemoteStates(promptId) { } for (const remoteNodeId of promptState.remoteNodeIds) { clearNodeTimer(remoteNodeId); + clearNodeProgress(remoteNodeId, promptId); const currentState = modalNodeStates.get(remoteNodeId); if (currentState?.promptId === promptId) { modalNodeStates.delete(remoteNodeId); @@ -853,6 +963,7 @@ function clearPromptRemoteStates(promptId) { } for (const ancestorNodeId of promptState.descendantNodeIdsByAncestor.keys()) { clearNodeTimer(ancestorNodeId); + clearNodeProgress(ancestorNodeId, promptId); const currentState = modalNodeStates.get(ancestorNodeId); if (currentState?.promptId === promptId) { modalNodeStates.delete(ancestorNodeId); @@ -974,6 +1085,7 @@ function registerExecutionListeners() { } api.addEventListener("modal_status", handleModalStatus); + api.addEventListener("modal_progress", handleModalProgress); api.addEventListener("executing", (event) => handleExecutionPhase(event, EXECUTION_PHASE)); api.addEventListener("executed", (event) => { endSyntheticExecutionUi(String(eventDetail(event).prompt_id ?? ""));