diff --git a/README.md b/README.md index dfa825a..08b1fd2 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ ## What is implemented - A frontend extension in [web/modal_toggle.js](/home/user/git/ComfyUI-Modal/web/modal_toggle.js) that injects a `Run on Modal` toggle into every non-internal node, draws live remote-state overlays on marked nodes, and submits queued prompts to `/modal/queue_prompt`. -- A backend route in [api_intercept.py](/home/user/git/ComfyUI-Modal/api_intercept.py) that reads the workflow snapshot from `extra_pnginfo.workflow`, identifies connected remote-marked node components, syncs referenced assets, and rewrites each component into one signature-preserving Modal proxy node. +- A backend route in [api_intercept.py](/home/user/git/ComfyUI-Modal/api_intercept.py) that reads the workflow snapshot from `extra_pnginfo.workflow`, recursively discovers Modal markers even inside nested subgraph workflow metadata, identifies connected remote-marked node components, syncs referenced assets, and rewrites each component into one signature-preserving Modal proxy node. - A dynamic proxy registry in [modal_executor_node.py](/home/user/git/ComfyUI-Modal/modal_executor_node.py) that mirrors the exported output count and output types of each remote component so ComfyUI validation still succeeds after rewrite. - A content-addressable sync engine in [sync_engine.py](/home/user/git/ComfyUI-Modal/sync_engine.py) that mirrors model files and a zipped `custom_nodes/` bundle into storage keyed by SHA256. - A remote runtime skeleton in [remote/modal_app.py](/home/user/git/ComfyUI-Modal/remote/modal_app.py) that supports local fallback execution for tests and can call into Modal when the SDK is available. @@ -35,7 +35,7 @@ ## Runtime behavior 1. Queueing a prompt posts to `/modal/queue_prompt`. -2. The route loads the workflow metadata from `extra_data.extra_pnginfo.workflow`. +2. The route loads the workflow metadata from `extra_data.extra_pnginfo.workflow`, including nested saved subgraph fragments. 3. Connected regions of nodes with `properties.is_modal_remote = true` are replaced with generated `ModalUniversalExecutor_` proxy nodes. 4. File-like widget inputs ending in `.safetensors`, `.ckpt`, `.pt`, or `.vae` are resolved to local files and mirrored into storage. 5. The executor serializes the component boundary inputs with `safetensors`, forwards the remote component payload to the runtime, and deserializes the exported boundary outputs back into ComfyUI values. @@ -171,6 +171,7 @@ Use the normal queue button in ComfyUI. The frontend extension intercepts queue submission and sends the prompt to `/modal/queue_prompt` instead of `/prompt`. On the backend: 1. The request handler inspects the workflow snapshot in `extra_pnginfo.workflow`. + Modal marker discovery now walks nested saved subgraph fragments too, so nodes marked inside ComfyUI subgraphs still participate in remote rewrite. 2. Connected remote-marked nodes are partitioned into remote components and each component is replaced with an internal `ModalUniversalExecutor_` proxy node. 3. Referenced model assets are mirrored into storage. 4. If custom-node syncing is enabled, the local `custom_nodes/` directory is zipped and mirrored if its content hash changed. diff --git a/api_intercept.py b/api_intercept.py index 51059c5..7e2cc8e 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -8,7 +8,7 @@ import time import uuid from collections import defaultdict from dataclasses import dataclass, field -from typing import Any +from typing import Any, Iterator from aiohttp import web @@ -173,6 +173,42 @@ def _is_link(value: Any) -> bool: ) +def _looks_like_workflow_node(fragment: dict[str, Any]) -> bool: + """Return whether a JSON fragment resembles a saved ComfyUI workflow node.""" + return "id" in fragment and "properties" in fragment + + +def _iter_workflow_nodes( + workflow_fragment: Any, + visited_object_ids: set[int] | None = None, +) -> Iterator[dict[str, Any]]: + """Yield workflow node dictionaries from a nested saved workflow fragment.""" + if visited_object_ids is None: + visited_object_ids = set() + + if isinstance(workflow_fragment, dict): + object_id = id(workflow_fragment) + if object_id in visited_object_ids: + return + visited_object_ids.add(object_id) + + if _looks_like_workflow_node(workflow_fragment): + yield workflow_fragment + + for value in workflow_fragment.values(): + yield from _iter_workflow_nodes(value, visited_object_ids) + return + + if isinstance(workflow_fragment, list): + object_id = id(workflow_fragment) + if object_id in visited_object_ids: + return + visited_object_ids.add(object_id) + + for item in workflow_fragment: + yield from _iter_workflow_nodes(item, visited_object_ids) + + def extract_remote_node_ids( workflow: dict[str, Any] | None, settings: ModalSyncSettings | None = None, @@ -183,7 +219,7 @@ def extract_remote_node_ids( marker = (settings or get_settings()).marker_property remote_node_ids: set[str] = set() - for node in workflow.get("nodes", []): + for node in _iter_workflow_nodes(workflow): properties = node.get("properties") or {} if properties.get(marker): remote_node_ids.add(str(node.get("id"))) diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index c89cb4f..17350dd 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -3,6 +3,7 @@ ## 2026-04-17 - The workflow metadata already rides along in `extra_data.extra_pnginfo.workflow`, so the backend can discover `node.properties.is_modal_remote` without any extra frontend metadata channel. +- That workflow metadata is not always flat. ComfyUI can embed saved node records inside nested subgraph fragments, so Modal marker discovery has to recurse through workflow JSON instead of only scanning the top-level `workflow["nodes"]` list. - A single wildcard-output proxy node is not enough for server-side graph rewriting. ComfyUI validates downstream link indices and output arity against the replacement class, so the backend has to generate proxy node classes that mirror the original node’s output signature before validation runs. - Keeping the sync engine content-addressable makes prompt rewrites cheap after the first upload. The route can block on sync safely as long as repeated requests collapse to a hash marker check. - Hashing and zipping the entire `custom_nodes/` tree inside the queue route is too expensive for `COMFY_MODAL_EXECUTION_MODE=local`. Local mode should skip custom-node bundle sync by default, with an explicit override when that path needs to be exercised. diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index 9a2fb17..948f38f 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -153,6 +153,47 @@ def test_rewrite_groups_connected_remote_nodes_into_single_proxy( assert summary.synced_assets[0].uploaded is True +def test_extract_remote_node_ids_recurses_into_nested_subgraph_workflows( + api_intercept_module: Any, + settings_module: Any, +) -> None: + """Modal marker extraction should find nodes nested inside saved subgraph metadata.""" + 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=Path("/tmp/storage"), + remote_storage_root="/storage", + custom_nodes_archive_name="custom_nodes_bundle.zip", + comfyui_root=None, + custom_nodes_dir=Path("/tmp/custom_nodes"), + ) + + workflow = { + "nodes": [ + { + "id": 100, + "properties": {"is_modal_remote": False}, + "subgraph": { + "nodes": [ + {"id": 11, "properties": {"is_modal_remote": True}}, + {"id": 12, "properties": {"is_modal_remote": False}}, + ] + }, + } + ] + } + + assert api_intercept_module.extract_remote_node_ids(workflow, settings) == {"11"} + + def test_rewrite_rejects_non_transportable_remote_inputs( api_intercept_module: Any, settings_module: Any, @@ -226,6 +267,95 @@ def test_rewrite_rejects_non_transportable_remote_inputs( assert summary.remote_component_ids == ["1"] +def test_rewrite_detects_remote_marker_inside_nested_subgraph_workflow( + api_intercept_module: Any, + settings_module: Any, + sync_engine_module: Any, + tmp_path: Path, +) -> None: + """Prompt rewrite should honor Modal markers found inside nested subgraph metadata.""" + custom_nodes_dir = tmp_path / "custom_nodes" + custom_nodes_dir.mkdir() + (custom_nodes_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") + + 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=custom_nodes_dir, + ) + sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) + fake_nodes_module = type( + "FakeNodesModule", + (), + { + "NODE_CLASS_MAPPINGS": { + "ModelSource": _FakeRemoteModelNode, + "RemoteConsumer": _FakeRemoteSamplerNode, + "LocalConsumer": _FakeLocalSinkNode, + }, + "NODE_DISPLAY_NAME_MAPPINGS": {}, + }, + )() + + workflow = { + "nodes": [ + { + "id": 99, + "properties": {"is_modal_remote": False}, + "subgraph": { + "nodes": [ + {"id": 1, "properties": {"is_modal_remote": False}}, + {"id": 2, "properties": {"is_modal_remote": True}}, + ] + }, + }, + {"id": 3, "properties": {"is_modal_remote": False}}, + ] + } + prompt = { + "1": { + "class_type": "ModelSource", + "inputs": {}, + "_meta": {"title": "Model Source"}, + }, + "2": { + "class_type": "RemoteConsumer", + "inputs": {"model": ["1", 0]}, + "_meta": {"title": "Remote Consumer"}, + }, + "3": { + "class_type": "LocalConsumer", + "inputs": {"latent": ["2", 0]}, + "_meta": {"title": "Local Consumer"}, + }, + } + + 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, + ) + + assert list(rewritten_prompt) == ["1", "3"] + assert rewritten_prompt["3"]["inputs"]["latent"] == ["1", 0] + assert summary.remote_node_ids == ["1", "2"] + assert summary.remote_component_ids == ["1"] + + def test_rewrite_auto_expands_upstream_non_transportable_dependencies( api_intercept_module: Any, settings_module: Any,