Skip to content

Commit 59c565a

Browse files
committed
Fix duplicate custom node import when cross-referenced
When one custom node imports another custom node (e.g., via 'import custom_nodes.other_node'), the ComfyUI custom node loader would previously load the imported node twice: 1. First time: via the explicit import statement (module name: custom_nodes.other_node) 2. Second time: via the automatic custom node scanner (module name: custom_nodes/other_node) This caused issues including: - Module initialization code being executed twice - Route handlers being registered multiple times - Potential state inconsistencies and data loss This fix adds a check before loading a custom node to detect if it's already loaded in sys.modules under either the standard import name or the path-based name. If already loaded, it reuses the existing module instead of re-executing it. Benefits: - Prevents duplicate initialization of custom nodes - Allows custom nodes to safely import other custom nodes - Improves performance by avoiding redundant module loading - Maintains backward compatibility with existing custom nodes Also includes linting fix: use lazy % formatting in logging functions.
1 parent c176b21 commit 59c565a

File tree

2 files changed

+155
-5
lines changed

2 files changed

+155
-5
lines changed

nodes.py

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2135,18 +2135,42 @@ async def load_custom_node(module_path: str, ignore=set(), module_parent="custom
21352135
elif os.path.isdir(module_path):
21362136
sys_module_name = module_path.replace(".", "_x_")
21372137

2138+
# Check if module is already loaded to prevent duplicate imports
2139+
# This can happen when one custom node imports another custom node
2140+
standard_module_name = f"{module_parent}.{get_module_name(module_path)}"
2141+
already_loaded = False
2142+
module = None
2143+
2144+
if sys_module_name in sys.modules:
2145+
logging.debug("Custom node %s already loaded, reusing existing module", sys_module_name)
2146+
module = sys.modules[sys_module_name]
2147+
already_loaded = True
2148+
elif standard_module_name in sys.modules:
2149+
logging.debug("Custom node %s already loaded via standard import, reusing existing module", standard_module_name)
2150+
module = sys.modules[standard_module_name]
2151+
# Register the module under sys_module_name as well to avoid future conflicts
2152+
sys.modules[sys_module_name] = module
2153+
already_loaded = True
2154+
21382155
try:
21392156
logging.debug("Trying to load custom node {}".format(module_path))
2157+
2158+
# Determine module_dir regardless of whether module is already loaded
21402159
if os.path.isfile(module_path):
2141-
module_spec = importlib.util.spec_from_file_location(sys_module_name, module_path)
21422160
module_dir = os.path.split(module_path)[0]
21432161
else:
2144-
module_spec = importlib.util.spec_from_file_location(sys_module_name, os.path.join(module_path, "__init__.py"))
21452162
module_dir = module_path
2163+
2164+
# Only execute module loading if not already loaded
2165+
if not already_loaded:
2166+
if os.path.isfile(module_path):
2167+
module_spec = importlib.util.spec_from_file_location(sys_module_name, module_path)
2168+
else:
2169+
module_spec = importlib.util.spec_from_file_location(sys_module_name, os.path.join(module_path, "__init__.py"))
21462170

2147-
module = importlib.util.module_from_spec(module_spec)
2148-
sys.modules[sys_module_name] = module
2149-
module_spec.loader.exec_module(module)
2171+
module = importlib.util.module_from_spec(module_spec)
2172+
sys.modules[sys_module_name] = module
2173+
module_spec.loader.exec_module(module)
21502174

21512175
LOADED_MODULE_DIRS[module_name] = os.path.abspath(module_dir)
21522176

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
"""
2+
Test for preventing duplicate custom node imports.
3+
4+
This test verifies that when a custom node is imported by another custom node,
5+
the module loading mechanism correctly detects and reuses the already-loaded module
6+
instead of loading it again.
7+
"""
8+
import sys
9+
import os
10+
import pytest
11+
from unittest.mock import MagicMock, patch
12+
13+
# Mock the required modules before importing nodes
14+
mock_comfy_api = MagicMock()
15+
mock_comfy_api.latest.io.ComfyNode = MagicMock
16+
mock_comfy_api.latest.ComfyExtension = MagicMock
17+
18+
sys.modules['comfy_api'] = mock_comfy_api
19+
sys.modules['comfy_api.latest'] = mock_comfy_api.latest
20+
sys.modules['comfy_api.latest.io'] = mock_comfy_api.latest.io
21+
22+
# Mock folder_paths
23+
mock_folder_paths = MagicMock()
24+
sys.modules['folder_paths'] = mock_folder_paths
25+
26+
# Mock comfy modules
27+
sys.modules['comfy'] = MagicMock()
28+
sys.modules['comfy.model_management'] = MagicMock()
29+
30+
# Now we can import nodes
31+
import nodes
32+
33+
34+
@pytest.mark.asyncio
35+
async def test_no_duplicate_import_when_already_loaded():
36+
"""
37+
Test that load_custom_node detects and reuses already-loaded modules.
38+
39+
Scenario:
40+
1. Custom node A is loaded by another custom node (e.g., via direct import)
41+
2. ComfyUI's custom node scanner encounters custom node A again
42+
3. The scanner should detect that A is already loaded and reuse it
43+
"""
44+
# Create a mock module
45+
mock_module = MagicMock()
46+
mock_module.NODE_CLASS_MAPPINGS = {}
47+
mock_module.WEB_DIRECTORY = None
48+
49+
# Simulate that the module was already imported with standard naming
50+
module_name = "custom_nodes.test_node"
51+
sys.modules[module_name] = mock_module
52+
53+
# Track if exec_module is called (should not be called for already-loaded modules)
54+
exec_called = False
55+
56+
def mock_exec_module(module):
57+
nonlocal exec_called
58+
exec_called = True
59+
60+
# Patch the importlib methods
61+
with patch('importlib.util.spec_from_file_location') as mock_spec_func, \
62+
patch('importlib.util.module_from_spec') as mock_module_func:
63+
64+
mock_spec = MagicMock()
65+
mock_spec.loader.exec_module = mock_exec_module
66+
mock_spec_func.return_value = mock_spec
67+
mock_module_func.return_value = MagicMock()
68+
69+
# Create a temporary test directory to simulate the custom node path
70+
import tempfile
71+
with tempfile.TemporaryDirectory() as tmpdir:
72+
test_node_dir = os.path.join(tmpdir, "test_node")
73+
os.makedirs(test_node_dir)
74+
75+
# Create an __init__.py file
76+
init_file = os.path.join(test_node_dir, "__init__.py")
77+
with open(init_file, 'w') as f:
78+
f.write("NODE_CLASS_MAPPINGS = {}\n")
79+
80+
# Attempt to load the custom node
81+
# Since we mocked sys.modules with 'custom_nodes.test_node',
82+
# the function should detect it and not execute the module again
83+
result = await nodes.load_custom_node(test_node_dir)
84+
85+
# The function should return True (successful load)
86+
assert result == True
87+
88+
# exec_module should NOT have been called because module was already loaded
89+
assert exec_called == False, "exec_module should not be called for already-loaded modules"
90+
91+
92+
@pytest.mark.asyncio
93+
async def test_load_new_module_when_not_loaded():
94+
"""
95+
Test that load_custom_node properly loads new modules that haven't been imported yet.
96+
"""
97+
import tempfile
98+
99+
# Create a temporary test directory
100+
with tempfile.TemporaryDirectory() as tmpdir:
101+
test_node_dir = os.path.join(tmpdir, "new_test_node")
102+
os.makedirs(test_node_dir)
103+
104+
# Create an __init__.py file with required attributes
105+
init_file = os.path.join(test_node_dir, "__init__.py")
106+
with open(init_file, 'w') as f:
107+
f.write("NODE_CLASS_MAPPINGS = {}\n")
108+
109+
# Clear any existing module with this name
110+
sys_module_name = test_node_dir.replace(".", "_x_")
111+
if sys_module_name in sys.modules:
112+
del sys.modules[sys_module_name]
113+
114+
# Load the custom node
115+
result = await nodes.load_custom_node(test_node_dir)
116+
117+
# Should return True for successful load
118+
assert result == True
119+
120+
# Module should now be in sys.modules
121+
assert sys_module_name in sys.modules
122+
123+
124+
if __name__ == "__main__":
125+
pytest.main([__file__, "-v"])
126+

0 commit comments

Comments
 (0)