Fix API V2 runtime invariants in cloud mode

Author: benjcooley

comfy_api/latest/_sdk.py

@@ -53,8 +53,6 @@
 if TYPE_CHECKING:  # keep this module import-safe / torch-free at import time
     import torch
 
-    from ._io import NodeOutput
-
 logger = logging.getLogger(__name__)
 
 # Env var an operator points at a directory/module implementing ``register``.
@@ -16088,13 +16086,13 @@ async def _model_ground_image(
                 if not isinstance(raw_box, dict):
                     raise RuntimeError("SAM3 grounding returned an invalid box")
                 box: dict[str, float] = {}
-                for field in ("x", "y", "width", "height", "score"):
-                    value = raw_box.get(field)
+                for box_field in ("x", "y", "width", "height", "score"):
+                    value = raw_box.get(box_field)
                     if (type(value) not in (int, float)
                             or not math.isfinite(float(value))):
                         raise RuntimeError(
-                            f"SAM3 grounding box has invalid {field}")
-                    box[field] = float(value)
+                            f"SAM3 grounding box has invalid {box_field}")
+                    box[box_field] = float(value)
                 if (box["width"] < 0.0 or box["height"] < 0.0
                         or not 0.0 <= box["score"] <= 1.0
                         or abs(box["x"]) > width * 4
@@ -16579,6 +16577,7 @@ def __init__(self) -> None:
         self.ops_provider: OpsProvider = InProcessOps()
         self.ref_resolver_factory: Callable[[], RefResolver] = InProcessRefResolver
         self._overlay_name: Optional[str] = None
+        self._extension_host_module_url: Optional[str] = None
 
     # Overlay entry points -------------------------------------------------- #
     def register_execution_backend(self, impl: ExecutionBackend) -> None:
@@ -16597,6 +16596,21 @@ def register_ref_resolver_factory(self, factory: Callable[[], RefResolver]) -> N
         logger.info("SDK: ref resolver -> %s", getattr(factory, "__name__", factory))
         self.ref_resolver_factory = factory
 
+    def register_extension_host(self, module_url: str) -> None:
+        if not module_url:
+            raise ValueError("extension host module URL must not be empty")
+        self._extension_host_module_url = module_url
+
+    @property
+    def frontend_runtime_config(self) -> dict[str, Any]:
+        if self._extension_host_module_url is None:
+            return {}
+        return {
+            "extension_host": {
+                "module_url": self._extension_host_module_url,
+            }
+        }
+
     @property
     def overlay_active(self) -> bool:
         return self._overlay_name is not None
@@ -16643,15 +16657,21 @@ def load_overlay(spec: Optional[str] = None) -> bool:
         module = importlib.import_module(spec)  # importable module name
 
     if module is None:
-        logger.error("SDK overlay %r could not be loaded", spec)
-        return False
+        raise RuntimeError(f"SDK overlay {spec!r} could not be loaded")
 
     register = getattr(module, "register", None)
     if not callable(register):
-        logger.error("SDK overlay %r has no register(providers) entrypoint", spec)
-        return False
+        raise RuntimeError(
+            f"SDK overlay {spec!r} has no register(providers) entrypoint"
+        )
 
     register(providers)
     providers._overlay_name = getattr(module, "__name__", spec)
     logger.info("SDK overlay loaded: %s", providers._overlay_name)
     return True
+
+
+def should_load_legacy_custom_nodes(
+    *, secure_mode: bool, disabled: bool, has_whitelist: bool
+) -> bool:
+    return not secure_mode and (not disabled or has_whitelist)

execution.py

@@ -821,6 +821,7 @@ def __init__(self, server, cache_type=False, cache_args=None):
         self.cache_args = cache_args
         self.cache_type = cache_type
         self.server = server
+        self.prompt_model_tracker = comfy.model_patcher.PromptModelTracker()
         self._runner = asyncio.Runner()
         self._active_loop = None
         self._active_task = None

main.py

@@ -558,16 +558,13 @@ def start_comfyui(asyncio_loop=None):
     if args.enable_manager and not args.disable_manager_ui:
         comfyui_manager.start()
 
-    # Custom-node SDK overlay: if COMFY_OVERLAY_MODULE points at a module, it
-    # registers isolated execution/ctx/ref providers before nodes load. Unset =
-    # pure in-process open-source behavior.
-    try:
-        from comfy_api.latest import sdk
-        sdk.providers  # ensure module imported
-        from comfy_api.latest._sdk import load_overlay
-        load_overlay()
-    except Exception as e:
-        logging.error(f"Custom-node SDK overlay failed to load: {e}")
+    from comfy_api.latest import sdk
+    sdk.providers
+    from comfy_api.latest._sdk import (
+        load_overlay,
+        should_load_legacy_custom_nodes,
+    )
+    secure_custom_nodes = load_overlay()
 
     hook_breaker_ac10a0.save_functions()
     asyncio_loop.run_until_complete(nodes.init_extra_nodes(