41 lines
1.6 KiB
Python
41 lines
1.6 KiB
Python
import os
|
|
import importlib.util
|
|
import sys
|
|
|
|
# 初始化 Mappings 字典
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
# 告知 ComfyUI 前端文件位置
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
# 获取当前文件 (__init__.py) 所在的目录
|
|
current_dir = os.path.dirname(__file__)
|
|
|
|
# 遍历当前目录下的所有文件
|
|
for filename in os.listdir(current_dir):
|
|
# 筛选出 .py 文件，排除 __init__.py 和以 _ 开头的系统文件
|
|
if filename.endswith(".py") and filename != "__init__.py" and not filename.startswith("_"):
|
|
module_name = filename[:-3] # 去掉 .py 后缀
|
|
file_path = os.path.join(current_dir, filename)
|
|
|
|
try:
|
|
# 动态加载模块
|
|
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
|
|
# 检查模块中是否有 NODE_CLASS_MAPPINGS
|
|
if hasattr(module, "NODE_CLASS_MAPPINGS") and isinstance(module.NODE_CLASS_MAPPINGS, dict):
|
|
NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
|
|
print(f"Loaded nodes from {filename}")
|
|
|
|
# 检查模块中是否有 NODE_DISPLAY_NAME_MAPPINGS (可选)
|
|
if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and isinstance(module.NODE_DISPLAY_NAME_MAPPINGS, dict):
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
except Exception as e:
|
|
print(f"Error loading module {filename}: {e}")
|
|
|
|
# 导出供 ComfyUI 使用
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] |