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.
This commit is contained in:
haoshu.gu 2025-12-24 05:51:37 +08:00
parent c176b214cc
commit 59c565adf3
2 changed files with 155 additions and 5 deletions

View File

@ -2135,18 +2135,42 @@ async def load_custom_node(module_path: str, ignore=set(), module_parent="custom
elif os.path.isdir(module_path):
sys_module_name = module_path.replace(".", "_x_")
# Check if module is already loaded to prevent duplicate imports
# This can happen when one custom node imports another custom node
standard_module_name = f"{module_parent}.{get_module_name(module_path)}"
already_loaded = False
module = None
if sys_module_name in sys.modules:
logging.debug("Custom node %s already loaded, reusing existing module", sys_module_name)
module = sys.modules[sys_module_name]
already_loaded = True
elif standard_module_name in sys.modules:
logging.debug("Custom node %s already loaded via standard import, reusing existing module", standard_module_name)
module = sys.modules[standard_module_name]
# Register the module under sys_module_name as well to avoid future conflicts
sys.modules[sys_module_name] = module
already_loaded = True
try:
logging.debug("Trying to load custom node {}".format(module_path))
# Determine module_dir regardless of whether module is already loaded
if os.path.isfile(module_path):
module_spec = importlib.util.spec_from_file_location(sys_module_name, module_path)
module_dir = os.path.split(module_path)[0]
else:
module_spec = importlib.util.spec_from_file_location(sys_module_name, os.path.join(module_path, "__init__.py"))
module_dir = module_path
# Only execute module loading if not already loaded
if not already_loaded:
if os.path.isfile(module_path):
module_spec = importlib.util.spec_from_file_location(sys_module_name, module_path)
else:
module_spec = importlib.util.spec_from_file_location(sys_module_name, os.path.join(module_path, "__init__.py"))
module = importlib.util.module_from_spec(module_spec)
sys.modules[sys_module_name] = module
module_spec.loader.exec_module(module)
module = importlib.util.module_from_spec(module_spec)
sys.modules[sys_module_name] = module
module_spec.loader.exec_module(module)
LOADED_MODULE_DIRS[module_name] = os.path.abspath(module_dir)

View File

@ -0,0 +1,126 @@
"""
Test for preventing duplicate custom node imports.
This test verifies that when a custom node is imported by another custom node,
the module loading mechanism correctly detects and reuses the already-loaded module
instead of loading it again.
"""
import sys
import os
import pytest
from unittest.mock import MagicMock, patch
# Mock the required modules before importing nodes
mock_comfy_api = MagicMock()
mock_comfy_api.latest.io.ComfyNode = MagicMock
mock_comfy_api.latest.ComfyExtension = MagicMock
sys.modules['comfy_api'] = mock_comfy_api
sys.modules['comfy_api.latest'] = mock_comfy_api.latest
sys.modules['comfy_api.latest.io'] = mock_comfy_api.latest.io
# Mock folder_paths
mock_folder_paths = MagicMock()
sys.modules['folder_paths'] = mock_folder_paths
# Mock comfy modules
sys.modules['comfy'] = MagicMock()
sys.modules['comfy.model_management'] = MagicMock()
# Now we can import nodes
import nodes
@pytest.mark.asyncio
async def test_no_duplicate_import_when_already_loaded():
"""
Test that load_custom_node detects and reuses already-loaded modules.
Scenario:
1. Custom node A is loaded by another custom node (e.g., via direct import)
2. ComfyUI's custom node scanner encounters custom node A again
3. The scanner should detect that A is already loaded and reuse it
"""
# Create a mock module
mock_module = MagicMock()
mock_module.NODE_CLASS_MAPPINGS = {}
mock_module.WEB_DIRECTORY = None
# Simulate that the module was already imported with standard naming
module_name = "custom_nodes.test_node"
sys.modules[module_name] = mock_module
# Track if exec_module is called (should not be called for already-loaded modules)
exec_called = False
def mock_exec_module(module):
nonlocal exec_called
exec_called = True
# Patch the importlib methods
with patch('importlib.util.spec_from_file_location') as mock_spec_func, \
patch('importlib.util.module_from_spec') as mock_module_func:
mock_spec = MagicMock()
mock_spec.loader.exec_module = mock_exec_module
mock_spec_func.return_value = mock_spec
mock_module_func.return_value = MagicMock()
# Create a temporary test directory to simulate the custom node path
import tempfile
with tempfile.TemporaryDirectory() as tmpdir:
test_node_dir = os.path.join(tmpdir, "test_node")
os.makedirs(test_node_dir)
# Create an __init__.py file
init_file = os.path.join(test_node_dir, "__init__.py")
with open(init_file, 'w') as f:
f.write("NODE_CLASS_MAPPINGS = {}\n")
# Attempt to load the custom node
# Since we mocked sys.modules with 'custom_nodes.test_node',
# the function should detect it and not execute the module again
result = await nodes.load_custom_node(test_node_dir)
# The function should return True (successful load)
assert result == True
# exec_module should NOT have been called because module was already loaded
assert exec_called == False, "exec_module should not be called for already-loaded modules"
@pytest.mark.asyncio
async def test_load_new_module_when_not_loaded():
"""
Test that load_custom_node properly loads new modules that haven't been imported yet.
"""
import tempfile
# Create a temporary test directory
with tempfile.TemporaryDirectory() as tmpdir:
test_node_dir = os.path.join(tmpdir, "new_test_node")
os.makedirs(test_node_dir)
# Create an __init__.py file with required attributes
init_file = os.path.join(test_node_dir, "__init__.py")
with open(init_file, 'w') as f:
f.write("NODE_CLASS_MAPPINGS = {}\n")
# Clear any existing module with this name
sys_module_name = test_node_dir.replace(".", "_x_")
if sys_module_name in sys.modules:
del sys.modules[sys_module_name]
# Load the custom node
result = await nodes.load_custom_node(test_node_dir)
# Should return True for successful load
assert result == True
# Module should now be in sys.modules
assert sys_module_name in sys.modules
if __name__ == "__main__":
pytest.main([__file__, "-v"])