129 lines
4.9 KiB
Python
129 lines
4.9 KiB
Python
import os
|
|
import importlib
|
|
import importlib.util
|
|
import sys
|
|
import logging
|
|
import atexit
|
|
import signal
|
|
|
|
# 初始化映射字典
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
# 获取当前目录
|
|
current_dir = os.path.dirname(os.path.realpath(__file__))
|
|
|
|
# 自动发现和导入 nodes 目录下的所有节点模块
|
|
nodes_dir = os.path.join(current_dir, "nodes")
|
|
if os.path.exists(nodes_dir) and os.path.isdir(nodes_dir):
|
|
# 将 nodes 目录添加到 Python 路径
|
|
if nodes_dir not in sys.path:
|
|
sys.path.append(nodes_dir)
|
|
|
|
# 遍历 nodes 目录下的所有 Python 文件
|
|
for file in os.listdir(nodes_dir):
|
|
if file.endswith(".py") and file != "__init__.py":
|
|
module_name = file[:-3] # 去掉 .py 后缀
|
|
try:
|
|
# 导入模块
|
|
module = importlib.import_module(f".{module_name}", package=f"{__name__}.nodes")
|
|
|
|
# 如果模块有节点映射，则添加到全局映射中
|
|
if hasattr(module, "NODE_CLASS_MAPPINGS"):
|
|
NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
|
|
if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"):
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
|
|
except Exception as e:
|
|
print(f"导入节点模块 {module_name} 时出错: {e}")
|
|
|
|
# 添加这行来支持文档
|
|
WEB_DIRECTORY = os.path.join(current_dir, "web")
|
|
|
|
# 注册前端扩展目录（JS），用于动态输入端口等功能
|
|
try:
|
|
import nodes
|
|
js_dir = os.path.join(current_dir, "js")
|
|
if os.path.isdir(js_dir):
|
|
# 通过 EXTENSION_WEB_DIRS 显式注册，确保浏览器侧加载到扩展脚本
|
|
nodes.EXTENSION_WEB_DIRS["ComfyUI-1hewNodes"] = js_dir
|
|
except Exception:
|
|
# 在服务端启动早期阶段或旧版本 ComfyUI 上，nodes 可能尚未可用
|
|
pass
|
|
|
|
# 全局变量保存监控模块引用
|
|
_workflow_watcher = None
|
|
_prev_sigint_handler = None
|
|
_prev_sigterm_handler = None
|
|
|
|
# 自动启动工作流监控功能
|
|
def _start_workflow_monitor():
|
|
"""启动工作流监控功能"""
|
|
global _workflow_watcher
|
|
try:
|
|
# 导入工作流监控模块（从utils/workflow文件夹）
|
|
from .utils.workflow import monitor
|
|
_workflow_watcher = monitor
|
|
|
|
# 启动监控
|
|
success = monitor.start_workflow_monitor()
|
|
if success:
|
|
# 注册关闭钩子
|
|
_register_shutdown_hooks()
|
|
else:
|
|
logging.warning("[ComfyUI-1hewNodes] 工作流监控启动失败")
|
|
except Exception as e:
|
|
logging.error(f"[ComfyUI-1hewNodes] 启动工作流监控时出错：{str(e)}")
|
|
|
|
def _stop_workflow_monitor():
|
|
"""停止工作流监控功能"""
|
|
global _workflow_watcher
|
|
try:
|
|
if _workflow_watcher and _workflow_watcher.is_monitoring():
|
|
_workflow_watcher.stop_workflow_monitor()
|
|
except Exception as e:
|
|
logging.error(f"[ComfyUI-1hewNodes] 停止工作流监控时出错：{str(e)}")
|
|
|
|
def _register_shutdown_hooks():
|
|
"""注册关闭钩子和信号处理"""
|
|
# 注册atexit钩子
|
|
atexit.register(_stop_workflow_monitor)
|
|
|
|
# 注册信号处理（仅在主线程中有效）
|
|
try:
|
|
global _prev_sigint_handler, _prev_sigterm_handler
|
|
if hasattr(signal, 'SIGINT'):
|
|
_prev_sigint_handler = signal.getsignal(signal.SIGINT)
|
|
signal.signal(signal.SIGINT, _signal_handler)
|
|
if hasattr(signal, 'SIGTERM'):
|
|
_prev_sigterm_handler = signal.getsignal(signal.SIGTERM)
|
|
signal.signal(signal.SIGTERM, _signal_handler)
|
|
except ValueError:
|
|
# 在非主线程中注册信号处理会失败，这是正常的
|
|
pass
|
|
|
|
def _signal_handler(signum, frame):
|
|
"""信号处理函数"""
|
|
logging.info(f"[ComfyUI-1hewNodes] 接收到信号 {signum}，正在停止工作流监控...")
|
|
_stop_workflow_monitor()
|
|
# 保持默认终止行为或链式调用原处理器
|
|
try:
|
|
if hasattr(signal, 'SIGINT') and signum == signal.SIGINT:
|
|
# 如果原处理器是默认或未设置，则触发默认 KeyboardInterrupt
|
|
if _prev_sigint_handler in (signal.SIG_DFL, None, signal.default_int_handler):
|
|
raise KeyboardInterrupt
|
|
# 如果原处理器不是忽略，则链式调用原处理器
|
|
if _prev_sigint_handler is not None and _prev_sigint_handler != signal.SIG_IGN:
|
|
_prev_sigint_handler(signum, frame)
|
|
elif hasattr(signal, 'SIGTERM') and signum == signal.SIGTERM:
|
|
if _prev_sigterm_handler not in (None, signal.SIG_IGN):
|
|
_prev_sigterm_handler(signum, frame)
|
|
except KeyboardInterrupt:
|
|
# 让上层按默认流程处理退出
|
|
raise
|
|
|
|
# 自动启动工作流监控
|
|
_start_workflow_monitor()
|
|
|
|
# 导出映射
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|