Removed nodes_v1_test.py, removed both v1/v3 test nodes from being attempted to be loaded in nodes.py

`ComfyExtension` Example (PR to v3-definition)

comfy_api/latest/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
 from typing import Type, TYPE_CHECKING
 from comfy_api.internal import ComfyAPIBase
 from comfy_api.internal.singleton import ProxiedSingleton
@@ -75,6 +76,19 @@ class ComfyAPI_latest(ComfyAPIBase):
 
     execution: Execution
 
+class ComfyExtension(ABC):
+    async def on_load(self) -> None:
+        """
+        Called when an extension is loaded.
+        This should be used to initialize any global resources neeeded by the extension.
+        """
+
+    @abstractmethod
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        """
+        Returns a list of nodes that this extension provides.
+        """
+
 class Input:
     Image = ImageInput
     Audio = AudioInput
@@ -106,4 +120,5 @@ __all__ = [
     "Input",
     "InputImpl",
     "Types",
+    "ComfyExtension",
 ]

comfy_api/latest/_io.py

@@ -6,26 +6,27 @@ from abc import ABC, abstractmethod
 from collections import Counter
 from dataclasses import asdict, dataclass
 from enum import Enum
-from typing import Any, Callable, Literal, TypedDict, TypeVar
+from typing import Any, Callable, Literal, TypedDict, TypeVar, TYPE_CHECKING
+from typing_extensions import NotRequired, final
 
 # used for type hinting
 import torch
-from spandrel import ImageModelDescriptor
-from typing_extensions import NotRequired, final
 
-from comfy.clip_vision import ClipVisionModel
+if TYPE_CHECKING:
-from comfy.clip_vision import Output as ClipVisionOutput_
+    from spandrel import ImageModelDescriptor
-from comfy.controlnet import ControlNet
+    from comfy.clip_vision import ClipVisionModel
-from comfy.hooks import HookGroup, HookKeyframeGroup
+    from comfy.clip_vision import Output as ClipVisionOutput_
-from comfy.model_patcher import ModelPatcher
+    from comfy.controlnet import ControlNet
-from comfy.samplers import CFGGuider, Sampler
+    from comfy.hooks import HookGroup, HookKeyframeGroup
-from comfy.sd import CLIP, VAE
+    from comfy.model_patcher import ModelPatcher
-from comfy.sd import StyleModel as StyleModel_
+    from comfy.samplers import CFGGuider, Sampler
-from comfy_api.input import VideoInput
+    from comfy.sd import CLIP, VAE
+    from comfy.sd import StyleModel as StyleModel_
+    from comfy_api.input import VideoInput
 from comfy_api.internal import (_ComfyNodeInternal, _NodeOutputInternal, classproperty, copy_class, first_real_override, is_class,
     prune_dict, shallow_clone_class)
 from comfy_api.latest._resources import Resources, ResourcesLocal
-from comfy_execution.graph import ExecutionBlocker
+from comfy_execution.graph_utils import ExecutionBlocker
 
 # from comfy_extras.nodes_images import SVG as SVG_ # NOTE: needs to be moved before can be imported due to circular reference
 
@@ -543,7 +544,8 @@ class Conditioning(ComfyTypeIO):
 
 @comfytype(io_type="SAMPLER")
 class Sampler(ComfyTypeIO):
-    Type = Sampler
+    if TYPE_CHECKING:
+        Type = Sampler
 
 @comfytype(io_type="SIGMAS")
 class Sigmas(ComfyTypeIO):
@@ -555,44 +557,54 @@ class Noise(ComfyTypeIO):
 
 @comfytype(io_type="GUIDER")
 class Guider(ComfyTypeIO):
-    Type = CFGGuider
+    if TYPE_CHECKING:
+        Type = CFGGuider
 
 @comfytype(io_type="CLIP")
 class Clip(ComfyTypeIO):
-    Type = CLIP
+    if TYPE_CHECKING:
+        Type = CLIP
 
 @comfytype(io_type="CONTROL_NET")
 class ControlNet(ComfyTypeIO):
-    Type = ControlNet
+    if TYPE_CHECKING:
+        Type = ControlNet
 
 @comfytype(io_type="VAE")
 class Vae(ComfyTypeIO):
-    Type = VAE
+    if TYPE_CHECKING:
+        Type = VAE
 
 @comfytype(io_type="MODEL")
 class Model(ComfyTypeIO):
-    Type = ModelPatcher
+    if TYPE_CHECKING:
+        Type = ModelPatcher
 
 @comfytype(io_type="CLIP_VISION")
 class ClipVision(ComfyTypeIO):
-    Type = ClipVisionModel
+    if TYPE_CHECKING:
+        Type = ClipVisionModel
 
 @comfytype(io_type="CLIP_VISION_OUTPUT")
 class ClipVisionOutput(ComfyTypeIO):
-    Type = ClipVisionOutput_
+    if TYPE_CHECKING:
+        Type = ClipVisionOutput_
 
 @comfytype(io_type="STYLE_MODEL")
 class StyleModel(ComfyTypeIO):
-    Type = StyleModel_
+    if TYPE_CHECKING:
+        Type = StyleModel_
 
 @comfytype(io_type="GLIGEN")
 class Gligen(ComfyTypeIO):
     '''ModelPatcher that wraps around a 'Gligen' model.'''
-    Type = ModelPatcher
+    if TYPE_CHECKING:
+        Type = ModelPatcher
 
 @comfytype(io_type="UPSCALE_MODEL")
 class UpscaleModel(ComfyTypeIO):
-    Type = ImageModelDescriptor
+    if TYPE_CHECKING:
+        Type = ImageModelDescriptor
 
 @comfytype(io_type="AUDIO")
 class Audio(ComfyTypeIO):
@@ -603,7 +615,8 @@ class Audio(ComfyTypeIO):
 
 @comfytype(io_type="VIDEO")
 class Video(ComfyTypeIO):
-    Type = VideoInput
+    if TYPE_CHECKING:
+        Type = VideoInput
 
 @comfytype(io_type="SVG")
 class SVG(ComfyTypeIO):
@@ -629,11 +642,13 @@ class Mesh(ComfyTypeIO):
 
 @comfytype(io_type="HOOKS")
 class Hooks(ComfyTypeIO):
-    Type = HookGroup
+    if TYPE_CHECKING:
+        Type = HookGroup
 
 @comfytype(io_type="HOOK_KEYFRAMES")
 class HookKeyframes(ComfyTypeIO):
-    Type = HookKeyframeGroup
+    if TYPE_CHECKING:
+        Type = HookKeyframeGroup
 
 @comfytype(io_type="TIMESTEPS_RANGE")
 class TimestepsRange(ComfyTypeIO):

comfy_api/v0_0_2/__init__.py

@@ -6,7 +6,7 @@ from comfy_api.latest import (
 )
 from typing import Type, TYPE_CHECKING
 from comfy_api.internal.async_to_sync import create_sync_class
-from comfy_api.latest import io, ui  #noqa: F401
+from comfy_api.latest import io, ui, ComfyExtension  #noqa: F401
 
 
 class ComfyAPIAdapter_v0_0_2(ComfyAPI_latest):
@@ -41,4 +41,5 @@ __all__ = [
     "Input",
     "InputImpl",
     "Types",
+    "ComfyExtension",
 ]

comfy_execution/graph.py

@@ -4,9 +4,12 @@ from typing import Type, Literal
 import nodes
 import asyncio
 import inspect
-from comfy_execution.graph_utils import is_link
+from comfy_execution.graph_utils import is_link, ExecutionBlocker
 from comfy.comfy_types.node_typing import ComfyNodeABC, InputTypeDict, InputTypeOptions
 
+# NOTE: ExecutionBlocker code got moved to graph_utils.py to prevent torch being imported too soon during unit tests
+ExecutionBlocker = ExecutionBlocker
+
 class DependencyCycleError(Exception):
     pass
 
@@ -294,21 +297,3 @@ class ExecutionList(TopologicalSort):
                 del blocked_by[node_id]
             to_remove = [node_id for node_id in blocked_by if len(blocked_by[node_id]) == 0]
         return list(blocked_by.keys())
-
-class ExecutionBlocker:
-    """
-    Return this from a node and any users will be blocked with the given error message.
-    If the message is None, execution will be blocked silently instead.
-    Generally, you should avoid using this functionality unless absolutely necessary. Whenever it's
-    possible, a lazy input will be more efficient and have a better user experience.
-    This functionality is useful in two cases:
-    1. You want to conditionally prevent an output node from executing. (Particularly a built-in node
-       like SaveImage. For your own output nodes, I would recommend just adding a BOOL input and using
-       lazy evaluation to let it conditionally disable itself.)
-    2. You have a node with multiple possible outputs, some of which are invalid and should not be used.
-       (I would recommend not making nodes like this in the future -- instead, make multiple nodes with
-       different outputs. Unfortunately, there are several popular existing nodes using this pattern.)
-    """
-    def __init__(self, message):
-        self.message = message
-

comfy_execution/graph_utils.py

@@ -137,3 +137,19 @@ def add_graph_prefix(graph, outputs, prefix):
 
     return new_graph, tuple(new_outputs)
 
+class ExecutionBlocker:
+    """
+    Return this from a node and any users will be blocked with the given error message.
+    If the message is None, execution will be blocked silently instead.
+    Generally, you should avoid using this functionality unless absolutely necessary. Whenever it's
+    possible, a lazy input will be more efficient and have a better user experience.
+    This functionality is useful in two cases:
+    1. You want to conditionally prevent an output node from executing. (Particularly a built-in node
+       like SaveImage. For your own output nodes, I would recommend just adding a BOOL input and using
+       lazy evaluation to let it conditionally disable itself.)
+    2. You have a node with multiple possible outputs, some of which are invalid and should not be used.
+       (I would recommend not making nodes like this in the future -- instead, make multiple nodes with
+       different outputs. Unfortunately, there are several popular existing nodes using this pattern.)
+    """
+    def __init__(self, message):
+        self.message = message

comfy_extras/nodes_v1_test.py

@@ -1,77 +0,0 @@
-import torch
-from comfy.comfy_types.node_typing import ComfyNodeABC, IO
-import asyncio
-from comfy.utils import ProgressBar
-import time
-
-
-class TestNode(ComfyNodeABC):
-
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "image": (IO.IMAGE,),
-                "some_int": (IO.INT, {"display_name": "new_name",
-                                      "min": 0, "max": 127, "default": 42,
-                                      "tooltip": "My tooltip 😎", "display": "slider"}),
-                "combo": (IO.COMBO, {"options": ["a", "b", "c"], "tooltip": "This is a combo input"}),
-                "combo2": (IO.COMBO, {"options": ["a", "b", "c"], "multi_select": True, "tooltip": "This is a combo input"}),
-            },
-            "optional": {
-                "xyz": ("XYZ",),
-                "mask": (IO.MASK,),
-            }
-        }
-
-    RETURN_TYPES = (IO.INT, IO.IMAGE)
-    RETURN_NAMES = ("INT", "img🖼️")
-    OUTPUT_TOOLTIPS = (None, "This is an image")
-    FUNCTION = "do_thing"
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "v3 nodes"
-
-    def do_thing(self, image: torch.Tensor, some_int: int, combo: str, combo2: list[str], xyz=None, mask: torch.Tensor=None):
-        return (some_int, image)
-
-
-class TestSleep(ComfyNodeABC):
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "value": (IO.ANY, {}),
-                "seconds": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 9999.0, "step": 0.01, "tooltip": "The amount of seconds to sleep."}),
-            },
-            "hidden": {
-                "unique_id": "UNIQUE_ID",
-            },
-        }
-    RETURN_TYPES = (IO.ANY,)
-    FUNCTION = "sleep"
-
-    CATEGORY = "_for_testing"
-
-    async def sleep(self, value, seconds, unique_id):
-        pbar = ProgressBar(seconds, node_id=unique_id)
-        start = time.time()
-        expiration = start + seconds
-        now = start
-        while now < expiration:
-            now = time.time()
-            pbar.update_absolute(now - start)
-            await asyncio.sleep(0.02)
-        return (value,)
-
-
-NODE_CLASS_MAPPINGS = {
-    "V1TestNode1": TestNode,
-    "V1TestSleep": TestSleep,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "V1TestNode1": "V1 Test Node",
-    "V1TestSleep": "V1 Test Sleep",
-}

comfy_extras/nodes_v3_test.py

@@ -1,9 +1,11 @@
 import torch
 import time
 from comfy_api.latest import io, ui, _io
+from comfy_api.latest import ComfyExtension
 import logging  # noqa
 import comfy.utils
 import asyncio
+from typing_extensions import override
 
 @io.comfytype(io_type="XYZ")
 class XYZ(io.ComfyTypeIO):
@@ -34,22 +36,6 @@ class V3TestNode(io.ComfyNode):
                 io.MultiCombo.Input("combo2", options=["a","b","c"]),
                 io.MultiType.Input(io.Int.Input("int_multitype", display_name="haha"), types=[io.Float]),
                 io.MultiType.Input("multitype", types=[io.Mask, io.Float, io.Int], optional=True),
-                # ComboInput("combo", image_upload=True, image_folder=FolderType.output,
-                #             remote=RemoteOptions(
-                #                 route="/internal/files/output",
-                #                 refresh_button=True,
-                #             ),
-                #             tooltip="This is a combo input"),
-                # IntegerInput("some_int", display_name="new_name", min=0, tooltip="My tooltip 😎", display=NumberDisplay.slider, ),
-                # ComboDynamicInput("mask", behavior=InputBehavior.optional),
-                # IntegerInput("some_int", display_name="new_name", min=0, tooltip="My tooltip 😎", display=NumberDisplay.slider,
-                #              dependent_inputs=[ComboDynamicInput("mask", behavior=InputBehavior.optional)],
-                #              dependent_values=[lambda my_value: IO.STRING if my_value < 5 else IO.NUMBER],
-                #              ),
-                # ["option1", "option2". "option3"]
-                # ComboDynamicInput["sdfgjhl", [ComboDynamicOptions("option1", [IntegerInput("some_int", display_name="new_name", min=0, tooltip="My tooltip 😎", display=NumberDisplay.slider, ImageInput(), MaskInput(), String()]),
-                #                              CombyDynamicOptons("option2", [])
-                #                                                   ]]
             ],
             outputs=[
                 io.Int.Output(),
@@ -271,7 +257,6 @@ class V3DummyEndInherit(V3DummyEnd):
         logging.info(f"V3DummyEndInherit: {cls.COOL_VALUE}")
         return super().execute(xyz)
 
-
 NODES_LIST: list[type[io.ComfyNode]] = [
     V3TestNode,
     # V3LoraLoader,
@@ -281,3 +266,11 @@ NODES_LIST: list[type[io.ComfyNode]] = [
     V3DummyEnd,
     V3DummyEndInherit,
 ]
+
+class v3TestExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return NODES_LIST
+
+async def comfy_entrypoint() -> v3TestExtension:
+    return v3TestExtension()

nodes.py

@@ -6,6 +6,7 @@ import os
 import sys
 import json
 import hashlib
+import inspect
 import traceback
 import math
 import time
@@ -29,7 +30,7 @@ import comfy.controlnet
 from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict, FileLocator
 from comfy_api.internal import register_versions, ComfyAPIWithVersion
 from comfy_api.version_list import supported_versions
-from comfy_api.latest import io
+from comfy_api.latest import io, ComfyExtension
 
 import comfy.clip_vision
 
@@ -2162,17 +2163,36 @@ async def load_custom_node(module_path: str, ignore=set(), module_parent="custom
             if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and getattr(module, "NODE_DISPLAY_NAME_MAPPINGS") is not None:
                 NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
             return True
-        # V3 node definition
+        # V3 Extension Definition
-        elif getattr(module, "NODES_LIST", None) is not None:
+        elif hasattr(module, "comfy_entrypoint"):
-            for node_cls in module.NODES_LIST:
+            entrypoint = getattr(module, "comfy_entrypoint")
-                node_cls: io.ComfyNode
+            if not callable(entrypoint):
-                schema = node_cls.GET_SCHEMA()
+                logging.warning(f"comfy_entrypoint in {module_path} is not callable, skipping.")
-                if schema.node_id not in ignore:
+                return False
-                    NODE_CLASS_MAPPINGS[schema.node_id] = node_cls
+            try:
-                    node_cls.RELATIVE_PYTHON_MODULE = "{}.{}".format(module_parent, get_module_name(module_path))
+                if inspect.iscoroutinefunction(entrypoint):
-                if schema.display_name is not None:
+                    extension = await entrypoint()
-                    NODE_DISPLAY_NAME_MAPPINGS[schema.node_id] = schema.display_name
+                else:
-            return True
+                    extension = entrypoint()
+                if not isinstance(extension, ComfyExtension):
+                    logging.warning(f"comfy_entrypoint in {module_path} did not return a ComfyExtension, skipping.")
+                    return False
+                node_list = await extension.get_node_list()
+                if not isinstance(node_list, list):
+                    logging.warning(f"comfy_entrypoint in {module_path} did not return a list of nodes, skipping.")
+                    return False
+                for node_cls in node_list:
+                    node_cls: io.ComfyNode
+                    schema = node_cls.GET_SCHEMA()
+                    if schema.node_id not in ignore:
+                        NODE_CLASS_MAPPINGS[schema.node_id] = node_cls
+                        node_cls.RELATIVE_PYTHON_MODULE = "{}.{}".format(module_parent, get_module_name(module_path))
+                    if schema.display_name is not None:
+                        NODE_DISPLAY_NAME_MAPPINGS[schema.node_id] = schema.display_name
+                return True
+            except Exception as e:
+                logging.warning(f"Error while calling comfy_entrypoint in {module_path}: {e}")
+                return False
         else:
             logging.warning(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS or NODES_LIST (need one).")
             return False
@@ -2300,8 +2320,6 @@ async def init_builtin_extra_nodes():
         "nodes_camera_trajectory.py",
         "nodes_edit_model.py",
         "nodes_tcfg.py",
-        "nodes_v3_test.py", # TODO: remove
-        "nodes_v1_test.py", # TODO: remove
     ]
 
     import_failed = []