mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-10 06:10:50 +08:00
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:
parent
c176b214cc
commit
59c565adf3
34
nodes.py
34
nodes.py
@ -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)
|
||||
|
||||
|
||||
126
tests-unit/nodes_test/test_duplicate_import.py
Normal file
126
tests-unit/nodes_test/test_duplicate_import.py
Normal 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"])
|
||||
|
||||
Loading…
Reference in New Issue
Block a user