From f90beb73f8b574f7542e44a22f48099b1998df47 Mon Sep 17 00:00:00 2001
From: Jedrzej Kosinski <kosinkadink1@gmail.com>
Date: Wed, 30 Jul 2025 19:42:44 -0700
Subject: [PATCH] Fix ComfyExtension registration not working

---
 comfy_api/v0_0_2/__init__.py  |  3 ++-
 comfy_extras/nodes_v3_test.py |  3 ++-
 nodes.py                      | 16 +++-------------
 3 files changed, 7 insertions(+), 15 deletions(-)

diff --git a/comfy_api/v0_0_2/__init__.py b/comfy_api/v0_0_2/__init__.py
index 8ddd772a5..de0f95001 100644
--- a/comfy_api/v0_0_2/__init__.py
+++ b/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",
 ]
diff --git a/comfy_extras/nodes_v3_test.py b/comfy_extras/nodes_v3_test.py
index c7bdf1da1..978c0575c 100644
--- a/comfy_extras/nodes_v3_test.py
+++ b/comfy_extras/nodes_v3_test.py
@@ -1,6 +1,7 @@
 import torch
 import time
-from comfy_api.latest import io, ui, _io, ComfyExtension
+from comfy_api.latest import io, ui, _io
+from comfy_api.latest import ComfyExtension
 import logging  # noqa
 import comfy.utils
 import asyncio
diff --git a/nodes.py b/nodes.py
index 262d695ba..c8d99ddd8 100644
--- a/nodes.py
+++ b/nodes.py
@@ -30,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
 
@@ -2174,7 +2174,7 @@ async def load_custom_node(module_path: str, ignore=set(), module_parent="custom
                     extension = await entrypoint()
                 else:
                     extension = entrypoint()
-                if not isinstance(extension, io.ComfyExtension):
+                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()
@@ -2189,20 +2189,10 @@ async def load_custom_node(module_path: str, ignore=set(), module_parent="custom
                         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
-        # V3 node definition
-        elif getattr(module, "NODES_LIST", None) is not None:
-            for node_cls in module.NODES_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
         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