mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-03-12 22:02:14 +00:00
Small refactor of custom node loading code.
This commit is contained in:
parent
1f717903bc
commit
8c4ccb55d1
43
nodes.py
43
nodes.py
@ -948,8 +948,29 @@ NODE_CLASS_MAPPINGS = {
|
|||||||
"VAEDecodeTiled": VAEDecodeTiled,
|
"VAEDecodeTiled": VAEDecodeTiled,
|
||||||
}
|
}
|
||||||
|
|
||||||
CUSTOM_NODE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "custom_nodes")
|
def load_custom_node(module_path):
|
||||||
|
module_name = os.path.basename(module_path)
|
||||||
|
if os.path.isfile(module_path):
|
||||||
|
sp = os.path.splitext(module_path)
|
||||||
|
module_name = sp[0]
|
||||||
|
try:
|
||||||
|
if os.path.isfile(module_path):
|
||||||
|
module_spec = importlib.util.spec_from_file_location(module_name, module_path)
|
||||||
|
else:
|
||||||
|
module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py"))
|
||||||
|
module = importlib.util.module_from_spec(module_spec)
|
||||||
|
sys.modules[module_name] = module
|
||||||
|
module_spec.loader.exec_module(module)
|
||||||
|
if hasattr(module, "NODE_CLASS_MAPPINGS") and getattr(module, "NODE_CLASS_MAPPINGS") is not None:
|
||||||
|
NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
|
||||||
|
else:
|
||||||
|
print(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.")
|
||||||
|
except Exception as e:
|
||||||
|
print(traceback.format_exc())
|
||||||
|
print(f"Cannot import {module_path} module for custom nodes:", e)
|
||||||
|
|
||||||
def load_custom_nodes():
|
def load_custom_nodes():
|
||||||
|
CUSTOM_NODE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "custom_nodes")
|
||||||
possible_modules = os.listdir(CUSTOM_NODE_PATH)
|
possible_modules = os.listdir(CUSTOM_NODE_PATH)
|
||||||
if "__pycache__" in possible_modules:
|
if "__pycache__" in possible_modules:
|
||||||
possible_modules.remove("__pycache__")
|
possible_modules.remove("__pycache__")
|
||||||
@ -957,22 +978,6 @@ def load_custom_nodes():
|
|||||||
for possible_module in possible_modules:
|
for possible_module in possible_modules:
|
||||||
module_path = os.path.join(CUSTOM_NODE_PATH, possible_module)
|
module_path = os.path.join(CUSTOM_NODE_PATH, possible_module)
|
||||||
if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue
|
if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue
|
||||||
|
load_custom_node(module_path)
|
||||||
|
|
||||||
module_name = possible_module
|
load_custom_nodes()
|
||||||
try:
|
|
||||||
if os.path.isfile(module_path):
|
|
||||||
module_spec = importlib.util.spec_from_file_location(module_name, module_path)
|
|
||||||
else:
|
|
||||||
module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py"))
|
|
||||||
module = importlib.util.module_from_spec(module_spec)
|
|
||||||
sys.modules[module_name] = module
|
|
||||||
module_spec.loader.exec_module(module)
|
|
||||||
if hasattr(module, "NODE_CLASS_MAPPINGS") and getattr(module, "NODE_CLASS_MAPPINGS") is not None:
|
|
||||||
NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
|
|
||||||
else:
|
|
||||||
print(f"Skip {possible_module} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.")
|
|
||||||
except Exception as e:
|
|
||||||
print(traceback.format_exc())
|
|
||||||
print(f"Cannot import {possible_module} module for custom nodes:", e)
|
|
||||||
|
|
||||||
load_custom_nodes()
|
|
||||||
|
Loading…
Reference in New Issue
Block a user