Merge pull request #12959 from pollockjj/wheel-support-pr

feat(isolation): wheel support for isolated custom nodes

.pyisolate_venvs/ComfyUI-DepthAnythingV2/cache/cache_key

@@ -1 +0,0 @@
-f03e4c88e21504c3

.pyisolate_venvs/ComfyUI-DepthAnythingV2/cache/node_info.json

@@ -1,81 +0,0 @@
-{
-  "DepthAnything_V2": {
-    "input_types": {
-      "required": {
-        "da_model": {
-          "__pyisolate_tuple__": [
-            "DAMODEL"
-          ]
-        },
-        "images": {
-          "__pyisolate_tuple__": [
-            "IMAGE"
-          ]
-        }
-      }
-    },
-    "return_types": [
-      "IMAGE"
-    ],
-    "return_names": [
-      "image"
-    ],
-    "function": "process",
-    "category": "DepthAnythingV2",
-    "output_node": false,
-    "output_is_list": null,
-    "is_v3": false,
-    "display_name": "Depth Anything V2"
-  },
-  "DownloadAndLoadDepthAnythingV2Model": {
-    "input_types": {
-      "required": {
-        "model": {
-          "__pyisolate_tuple__": [
-            [
-              "depth_anything_v2_vits_fp16.safetensors",
-              "depth_anything_v2_vits_fp32.safetensors",
-              "depth_anything_v2_vitb_fp16.safetensors",
-              "depth_anything_v2_vitb_fp32.safetensors",
-              "depth_anything_v2_vitl_fp16.safetensors",
-              "depth_anything_v2_vitl_fp32.safetensors",
-              "depth_anything_v2_vitg_fp32.safetensors",
-              "depth_anything_v2_metric_hypersim_vitl_fp32.safetensors",
-              "depth_anything_v2_metric_vkitti_vitl_fp32.safetensors"
-            ],
-            {
-              "default": "depth_anything_v2_vitl_fp32.safetensors"
-            }
-          ]
-        }
-      },
-      "optional": {
-        "precision": {
-          "__pyisolate_tuple__": [
-            [
-              "auto",
-              "bf16",
-              "fp16",
-              "fp32"
-            ],
-            {
-              "default": "auto"
-            }
-          ]
-        }
-      }
-    },
-    "return_types": [
-      "DAMODEL"
-    ],
-    "return_names": [
-      "da_v2_model"
-    ],
-    "function": "loadmodel",
-    "category": "DepthAnythingV2",
-    "output_node": false,
-    "output_is_list": null,
-    "is_v3": false,
-    "display_name": "DownloadAndLoadDepthAnythingV2Model"
-  }
-}

.pyisolate_venvs/ComfyUI-IsolationToolkit/cache/cache_key

@@ -1 +0,0 @@
-4b90e6876f4c0b8c

comfy/isolation/extension_loader.py

@@ -27,6 +27,47 @@ except ImportError:
 logger = logging.getLogger(__name__)
 
 
+def _register_web_directory(extension_name: str, node_dir: Path) -> None:
+    """Register an isolated extension's web directory on the host side."""
+    import nodes
+
+    # Method 1: pyproject.toml [tool.comfy] web field
+    pyproject = node_dir / "pyproject.toml"
+    if pyproject.exists():
+        try:
+            with pyproject.open("rb") as f:
+                data = tomllib.load(f)
+            web_dir_name = data.get("tool", {}).get("comfy", {}).get("web")
+            if web_dir_name:
+                web_dir_path = str(node_dir / web_dir_name)
+                if os.path.isdir(web_dir_path):
+                    nodes.EXTENSION_WEB_DIRS[extension_name] = web_dir_path
+                    logger.debug("][ Registered web dir for isolated %s: %s", extension_name, web_dir_path)
+                    return
+        except Exception:
+            pass
+
+    # Method 2: __init__.py WEB_DIRECTORY constant (parse without importing)
+    init_file = node_dir / "__init__.py"
+    if init_file.exists():
+        try:
+            source = init_file.read_text()
+            for line in source.splitlines():
+                stripped = line.strip()
+                if stripped.startswith("WEB_DIRECTORY"):
+                    # Parse: WEB_DIRECTORY = "./web" or WEB_DIRECTORY = "web"
+                    _, _, value = stripped.partition("=")
+                    value = value.strip().strip("\"'")
+                    if value:
+                        web_dir_path = str((node_dir / value).resolve())
+                        if os.path.isdir(web_dir_path):
+                            nodes.EXTENSION_WEB_DIRS[extension_name] = web_dir_path
+                            logger.debug("][ Registered web dir for isolated %s: %s", extension_name, web_dir_path)
+                            return
+        except Exception:
+            pass
+
+
 async def _stop_extension_safe(
     extension: ComfyNodeExtension, extension_name: str
 ) -> None:
@@ -255,6 +296,7 @@ async def load_isolated_node(
         "dependencies": dependencies,
         "share_torch": share_torch,
         "share_cuda_ipc": share_cuda_ipc,
+        "sandbox_mode": host_policy["sandbox_mode"],
         "sandbox": sandbox_config,
     }
     if cuda_wheels is not None:
@@ -263,6 +305,11 @@ async def load_isolated_node(
     extension = manager.load_extension(extension_config)
     register_dummy_module(extension_name, node_dir)
 
+    # Register web directory on the host — only when sandbox is disabled.
+    # In sandbox mode, serving untrusted JS to the browser is not safe.
+    if host_policy["sandbox_mode"] == "disabled":
+        _register_web_directory(extension_name, node_dir)
+
     # Try cache first (lazy spawn)
     if is_cache_valid(node_dir, manifest_path, venv_root):
         cached_data = load_from_cache(node_dir, venv_root)

comfy/isolation/extension_wrapper.py

@@ -237,6 +237,9 @@ class ComfyNodeExtension(ExtensionBase):
                         "not_idempotent": bool(
                             getattr(node_cls, "NOT_IDEMPOTENT", False)
                         ),
+                        "accept_all_inputs": bool(
+                            getattr(node_cls, "ACCEPT_ALL_INPUTS", False)
+                        ),
                     }
                 )
             except Exception as exc:
@@ -369,14 +372,20 @@ class ComfyNodeExtension(ExtensionBase):
         handler = getattr(instance, function_name)
 
         try:
+            import torch
             if asyncio.iscoroutinefunction(handler):
-                result = await handler(**resolved_inputs)
+                with torch.inference_mode():
+                    result = await handler(**resolved_inputs)
             else:
                 import functools
 
+                def _run_with_inference_mode(**kwargs):
+                    with torch.inference_mode():
+                        return handler(**kwargs)
+
                 loop = asyncio.get_running_loop()
                 result = await loop.run_in_executor(
-                    None, functools.partial(handler, **resolved_inputs)
+                    None, functools.partial(_run_with_inference_mode, **resolved_inputs)
                 )
         except Exception:
             logger.exception(
@@ -388,7 +397,17 @@ class ComfyNodeExtension(ExtensionBase):
             raise
 
         if type(result).__name__ == "NodeOutput":
-            result = result.args
+            node_output_dict = {
+                "__node_output__": True,
+                "args": self._wrap_unpicklable_objects(result.args),
+            }
+            if result.ui is not None:
+                node_output_dict["ui"] = result.ui
+            if getattr(result, "expand", None) is not None:
+                node_output_dict["expand"] = result.expand
+            if getattr(result, "block_execution", None) is not None:
+                node_output_dict["block_execution"] = result.block_execution
+            return node_output_dict
         if self._is_comfy_protocol_return(result):
             wrapped = self._wrap_unpicklable_objects(result)
             return wrapped

comfy/isolation/host_policy.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import logging
+import os
 from pathlib import Path
 from typing import Dict, List, TypedDict
 
@@ -12,8 +13,12 @@ except ImportError:
 
 logger = logging.getLogger(__name__)
 
+HOST_POLICY_PATH_ENV = "COMFY_HOST_POLICY_PATH"
+VALID_SANDBOX_MODES = frozenset({"required", "disabled"})
+
 
 class HostSecurityPolicy(TypedDict):
+    sandbox_mode: str
     allow_network: bool
     writable_paths: List[str]
     readonly_paths: List[str]
@@ -21,6 +26,7 @@ class HostSecurityPolicy(TypedDict):
 
 
 DEFAULT_POLICY: HostSecurityPolicy = {
+    "sandbox_mode": "required",
     "allow_network": False,
     "writable_paths": ["/dev/shm", "/tmp"],
     "readonly_paths": [],
@@ -30,6 +36,7 @@ DEFAULT_POLICY: HostSecurityPolicy = {
 
 def _default_policy() -> HostSecurityPolicy:
     return {
+        "sandbox_mode": DEFAULT_POLICY["sandbox_mode"],
         "allow_network": DEFAULT_POLICY["allow_network"],
         "writable_paths": list(DEFAULT_POLICY["writable_paths"]),
         "readonly_paths": list(DEFAULT_POLICY["readonly_paths"]),
@@ -38,7 +45,8 @@ def _default_policy() -> HostSecurityPolicy:
 
 
 def load_host_policy(comfy_root: Path) -> HostSecurityPolicy:
-    config_path = comfy_root / "pyproject.toml"
+    config_override = os.environ.get(HOST_POLICY_PATH_ENV)
+    config_path = Path(config_override) if config_override else comfy_root / "pyproject.toml"
     policy = _default_policy()
 
     if not config_path.exists():
@@ -61,6 +69,19 @@ def load_host_policy(comfy_root: Path) -> HostSecurityPolicy:
         logger.debug("No [tool.comfy.host] section found, using defaults.")
         return policy
 
+    sandbox_mode = tool_config.get("sandbox_mode")
+    if isinstance(sandbox_mode, str):
+        normalized_sandbox_mode = sandbox_mode.strip().lower()
+        if normalized_sandbox_mode in VALID_SANDBOX_MODES:
+            policy["sandbox_mode"] = normalized_sandbox_mode
+        else:
+            logger.warning(
+                "Invalid host sandbox_mode %r in %s, using default %r.",
+                sandbox_mode,
+                config_path,
+                DEFAULT_POLICY["sandbox_mode"],
+            )
+
     if "allow_network" in tool_config:
         policy["allow_network"] = bool(tool_config["allow_network"])
 
@@ -75,7 +96,10 @@ def load_host_policy(comfy_root: Path) -> HostSecurityPolicy:
         policy["whitelist"] = {str(k): str(v) for k, v in whitelist_raw.items()}
 
     logger.debug(
-        f"Loaded Host Policy: {len(policy['whitelist'])} whitelisted nodes, Network={policy['allow_network']}"
+        "Loaded Host Policy: %d whitelisted nodes, Sandbox=%s, Network=%s",
+        len(policy["whitelist"]),
+        policy["sandbox_mode"],
+        policy["allow_network"],
     )
     return policy
 

comfy/isolation/runtime_helpers.py

@@ -215,6 +215,19 @@ def build_stub_class(
                 node_name,
                 node_unique_id or "-",
             )
+            # Reconstruct NodeOutput if the child serialized one
+            if isinstance(result, dict) and result.get("__node_output__"):
+                from comfy_api.latest import io as latest_io
+                args_raw = result.get("args", ())
+                deserialized_args = await deserialize_from_isolation(args_raw, extension)
+                deserialized_args = _detach_shared_cpu_tensors(deserialized_args)
+                scan_shm_forensics("RUNTIME:post_execute", refresh_model_context=True)
+                return latest_io.NodeOutput(
+                    *deserialized_args,
+                    ui=result.get("ui"),
+                    expand=result.get("expand"),
+                    block_execution=result.get("block_execution"),
+                )
             deserialized = await deserialize_from_isolation(result, extension)
             scan_shm_forensics("RUNTIME:post_execute", refresh_model_context=True)
             return _detach_shared_cpu_tensors(deserialized)
@@ -278,7 +291,12 @@ def build_stub_class(
         return True
 
     def _get_node_info_v1(cls):
-        return info.get("schema_v1", {})
+        node_info = copy.deepcopy(info.get("schema_v1", {}))
+        relative_python_module = node_info.get("python_module")
+        if not isinstance(relative_python_module, str) or not relative_python_module:
+            relative_python_module = f"custom_nodes.{extension.name}"
+        node_info["python_module"] = relative_python_module
+        return node_info
 
     def _get_base_class(cls):
         return latest_io.ComfyNode
@@ -308,6 +326,8 @@ def build_stub_class(
         attributes["DEPRECATED"] = info.get("deprecated", False)
         attributes["API_NODE"] = info.get("api_node", False)
         attributes["NOT_IDEMPOTENT"] = info.get("not_idempotent", False)
+        attributes["ACCEPT_ALL_INPUTS"] = info.get("accept_all_inputs", False)
+        attributes["_ACCEPT_ALL_INPUTS"] = info.get("accept_all_inputs", False)
         attributes["INPUT_IS_LIST"] = info.get("input_is_list", False)
 
     class_name = f"PyIsolate_{node_name}".replace(" ", "_")

pyproject.toml

@@ -11,6 +11,7 @@ repository = "https://github.com/comfyanonymous/ComfyUI"
 documentation = "https://docs.comfy.org/"
 
 [tool.comfy.host]
+sandbox_mode = "required"
 allow_network = false
 writable_paths = ["/dev/shm", "/tmp"]
 

tests/isolation/test_cuda_wheels_and_env_flags.py

@@ -69,6 +69,7 @@ flash_attn = "flash-attn-special"
     monkeypatch.setattr(
         "comfy.isolation.extension_loader.load_host_policy",
         lambda base_path: {
+            "sandbox_mode": "required",
             "allow_network": False,
             "writable_paths": [],
             "readonly_paths": [],
@@ -95,6 +96,7 @@ flash_attn = "flash-attn-special"
     )
 
     assert len(specs) == 1
+    assert captured["sandbox_mode"] == "required"
     assert captured["cuda_wheels"] == {
         "index_url": "https://example.invalid/cuda-wheels/",
         "packages": ["flash-attn", "sageattention"],
@@ -171,6 +173,7 @@ can_isolate = true
     monkeypatch.setattr(
         "comfy.isolation.extension_loader.load_host_policy",
         lambda base_path: {
+            "sandbox_mode": "disabled",
             "allow_network": False,
             "writable_paths": [],
             "readonly_paths": [],
@@ -196,6 +199,7 @@ can_isolate = true
         )
     )
 
+    assert captured["sandbox_mode"] == "disabled"
     assert "cuda_wheels" not in captured
 
 

tests/isolation/test_host_policy.py

@@ -10,6 +10,7 @@ def test_load_host_policy_defaults_when_pyproject_missing(tmp_path):
 
     policy = load_host_policy(tmp_path)
 
+    assert policy["sandbox_mode"] == DEFAULT_POLICY["sandbox_mode"]
     assert policy["allow_network"] == DEFAULT_POLICY["allow_network"]
     assert policy["writable_paths"] == DEFAULT_POLICY["writable_paths"]
     assert policy["readonly_paths"] == DEFAULT_POLICY["readonly_paths"]
@@ -28,6 +29,7 @@ name = "ComfyUI"
     )
 
     policy = load_host_policy(tmp_path)
+    assert policy["sandbox_mode"] == DEFAULT_POLICY["sandbox_mode"]
     assert policy["allow_network"] == DEFAULT_POLICY["allow_network"]
     assert policy["whitelist"] == {}
 
@@ -39,6 +41,7 @@ def test_load_host_policy_reads_values(tmp_path):
         tmp_path / "pyproject.toml",
         """
 [tool.comfy.host]
+sandbox_mode = "disabled"
 allow_network = true
 writable_paths = ["/tmp/a", "/tmp/b"]
 readonly_paths = ["/opt/readonly"]
@@ -49,6 +52,7 @@ ExampleNode = "*"
     )
 
     policy = load_host_policy(tmp_path)
+    assert policy["sandbox_mode"] == "disabled"
     assert policy["allow_network"] is True
     assert policy["writable_paths"] == ["/tmp/a", "/tmp/b"]
     assert policy["readonly_paths"] == ["/opt/readonly"]
@@ -70,3 +74,40 @@ whitelist = ["bad"]
     policy = load_host_policy(tmp_path)
     assert policy["allow_network"] is True
     assert policy["whitelist"] == DEFAULT_POLICY["whitelist"]
+
+
+def test_load_host_policy_ignores_invalid_sandbox_mode(tmp_path):
+    from comfy.isolation.host_policy import DEFAULT_POLICY, load_host_policy
+
+    _write_pyproject(
+        tmp_path / "pyproject.toml",
+        """
+[tool.comfy.host]
+sandbox_mode = "surprise"
+""".strip(),
+    )
+
+    policy = load_host_policy(tmp_path)
+
+    assert policy["sandbox_mode"] == DEFAULT_POLICY["sandbox_mode"]
+
+
+def test_load_host_policy_uses_env_override_path(tmp_path, monkeypatch):
+    from comfy.isolation.host_policy import load_host_policy
+
+    override_path = tmp_path / "host_policy_override.toml"
+    _write_pyproject(
+        override_path,
+        """
+[tool.comfy.host]
+sandbox_mode = "disabled"
+allow_network = true
+""".strip(),
+    )
+
+    monkeypatch.setenv("COMFY_HOST_POLICY_PATH", str(override_path))
+
+    policy = load_host_policy(tmp_path / "missing-root")
+
+    assert policy["sandbox_mode"] == "disabled"
+    assert policy["allow_network"] is True