188 lines
7.5 KiB
Python
188 lines
7.5 KiB
Python
import os
|
|
from .nodes.nodes import NODE_CLASS_MAPPINGS as MAIN_NODE_CLASS_MAPPINGS
|
|
from .nodes.nodes import NODE_DISPLAY_NAME_MAPPINGS as MAIN_NODE_DISPLAY_NAME_MAPPINGS
|
|
from .nodes.helper_nodes import HELPER_NODE_CLASS_MAPPINGS
|
|
from .nodes.helper_nodes import HELPER_NODE_DISPLAY_NAME_MAPPINGS
|
|
|
|
# Import control panel node
|
|
try:
|
|
from .nodes.utils.control_panel import CONTROL_PANEL_NODE_CLASS_MAPPINGS
|
|
from .nodes.utils.control_panel import CONTROL_PANEL_NODE_DISPLAY_NAME_MAPPINGS
|
|
except ImportError as e:
|
|
print(f"Warning: Could not import control panel node: {e}")
|
|
CONTROL_PANEL_NODE_CLASS_MAPPINGS = {}
|
|
CONTROL_PANEL_NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
# Import civit metadata helper node
|
|
try:
|
|
from .nodes.utils.civit_metadata_helper import CIVIT_METADATA_HELPER_NODE_CLASS_MAPPINGS
|
|
from .nodes.utils.civit_metadata_helper import CIVIT_METADATA_HELPER_NODE_DISPLAY_NAME_MAPPINGS
|
|
except ImportError as e:
|
|
print(f"Warning: Could not import civit metadata helper node: {e}")
|
|
CIVIT_METADATA_HELPER_NODE_CLASS_MAPPINGS = {}
|
|
CIVIT_METADATA_HELPER_NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
# Import Azure Storage Upload node
|
|
try:
|
|
from .nodes.utils.azure_storage_upload import NODE_CLASS_MAPPINGS as AZURE_STORAGE_NODE_CLASS_MAPPINGS
|
|
from .nodes.utils.azure_storage_upload import NODE_DISPLAY_NAME_MAPPINGS as AZURE_STORAGE_NODE_DISPLAY_NAME_MAPPINGS
|
|
except ImportError as e:
|
|
print(f"Warning: Could not import Azure Storage Upload node: {e}")
|
|
AZURE_STORAGE_NODE_CLASS_MAPPINGS = {}
|
|
AZURE_STORAGE_NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
# Import Prompt Builder node
|
|
try:
|
|
from .nodes.utils.prompt_builder import NODE_CLASS_MAPPINGS as PROMPT_BUILDER_NODE_CLASS_MAPPINGS
|
|
from .nodes.utils.prompt_builder import NODE_DISPLAY_NAME_MAPPINGS as PROMPT_BUILDER_NODE_DISPLAY_NAME_MAPPINGS
|
|
except ImportError as e:
|
|
print(f"Warning: Could not import Prompt Builder node: {e}")
|
|
PROMPT_BUILDER_NODE_CLASS_MAPPINGS = {}
|
|
PROMPT_BUILDER_NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
# Import LoRA Info Extractor node
|
|
try:
|
|
from .nodes.utils.lora_info_extractor import NODE_CLASS_MAPPINGS as LORA_INFO_NODE_CLASS_MAPPINGS
|
|
from .nodes.utils.lora_info_extractor import NODE_DISPLAY_NAME_MAPPINGS as LORA_INFO_NODE_DISPLAY_NAME_MAPPINGS
|
|
except ImportError as e:
|
|
print(f"Warning: Could not import LoRA Info Extractor node: {e}")
|
|
LORA_INFO_NODE_CLASS_MAPPINGS = {}
|
|
LORA_INFO_NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
|
|
# Import media_selection nodes
|
|
try:
|
|
from .nodes.media_selection.media_selection import MediaSelection
|
|
|
|
MEDIA_SELECTION_NODE_CLASS_MAPPINGS = {
|
|
"MediaSelection": MediaSelection,
|
|
}
|
|
MEDIA_SELECTION_NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"MediaSelection": "Media Selection",
|
|
}
|
|
except ImportError as e:
|
|
print(f"Warning: Could not import media_selection nodes: {e}")
|
|
MEDIA_SELECTION_NODE_CLASS_MAPPINGS = {}
|
|
MEDIA_SELECTION_NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
# Import vace_annotators nodes
|
|
try:
|
|
from .nodes.vace_annotators import NODE_CLASS_MAPPINGS as VACE_ANNOTATORS_NODE_CLASS_MAPPINGS
|
|
from .nodes.vace_annotators import NODE_DISPLAY_NAME_MAPPINGS as VACE_ANNOTATORS_NODE_DISPLAY_NAME_MAPPINGS
|
|
except ImportError as e:
|
|
print(f"Warning: Could not import vace_annotators nodes: {e}")
|
|
VACE_ANNOTATORS_NODE_CLASS_MAPPINGS = {}
|
|
VACE_ANNOTATORS_NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
# Register config API routes
|
|
try:
|
|
from .nodes.config_api import register_config_routes
|
|
from .nodes.restart_api import register_restart_routes
|
|
from server import PromptServer
|
|
app = getattr(PromptServer.instance, "app", None) or PromptServer.instance
|
|
if app:
|
|
register_config_routes(app)
|
|
register_restart_routes(app)
|
|
print("Swiss Army Knife: Registered config API routes")
|
|
except Exception as e:
|
|
print(f"Swiss Army Knife: Could not register config API routes: {e}")
|
|
|
|
|
|
# Get version from pyproject.toml for cache busting
|
|
def get_version():
|
|
try:
|
|
import tomllib
|
|
pyproject_path = os.path.join(os.path.dirname(__file__), "pyproject.toml")
|
|
with open(pyproject_path, "rb") as f:
|
|
data = tomllib.load(f)
|
|
return data["project"]["version"]
|
|
except Exception:
|
|
# Fallback to timestamp if version reading fails
|
|
import time
|
|
return str(int(time.time()))
|
|
|
|
# Combine main nodes, helper nodes, control panel, civit metadata helper, azure storage, lora_manager, media_selection, and vace_annotators nodes
|
|
NODE_CLASS_MAPPINGS = {
|
|
**MAIN_NODE_CLASS_MAPPINGS,
|
|
**HELPER_NODE_CLASS_MAPPINGS,
|
|
**CONTROL_PANEL_NODE_CLASS_MAPPINGS,
|
|
**CIVIT_METADATA_HELPER_NODE_CLASS_MAPPINGS,
|
|
**AZURE_STORAGE_NODE_CLASS_MAPPINGS,
|
|
**PROMPT_BUILDER_NODE_CLASS_MAPPINGS,
|
|
**LORA_INFO_NODE_CLASS_MAPPINGS,
|
|
|
|
**MEDIA_SELECTION_NODE_CLASS_MAPPINGS,
|
|
**VACE_ANNOTATORS_NODE_CLASS_MAPPINGS
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
**MAIN_NODE_DISPLAY_NAME_MAPPINGS,
|
|
**HELPER_NODE_DISPLAY_NAME_MAPPINGS,
|
|
**CONTROL_PANEL_NODE_DISPLAY_NAME_MAPPINGS,
|
|
**CIVIT_METADATA_HELPER_NODE_DISPLAY_NAME_MAPPINGS,
|
|
**AZURE_STORAGE_NODE_DISPLAY_NAME_MAPPINGS,
|
|
**PROMPT_BUILDER_NODE_DISPLAY_NAME_MAPPINGS,
|
|
**LORA_INFO_NODE_DISPLAY_NAME_MAPPINGS,
|
|
|
|
**MEDIA_SELECTION_NODE_DISPLAY_NAME_MAPPINGS,
|
|
**VACE_ANNOTATORS_NODE_DISPLAY_NAME_MAPPINGS
|
|
}
|
|
|
|
WEB_DIRECTORY = "./web"
|
|
VERSION = get_version()
|
|
|
|
# Load DEBUG setting from environment
|
|
DEBUG = os.environ.get("DEBUG", "false").lower() in ("true", "1", "yes")
|
|
|
|
# Initialize profiler (must happen before execution starts)
|
|
try:
|
|
from .nodes.utils.resource_monitor.prestartup import inject_profiling
|
|
from .nodes.utils.resource_monitor.profiler_core import ProfilerManager
|
|
from .nodes.utils.resource_monitor.profiler_api import register_profiler_routes, set_profiler_manager
|
|
|
|
# Inject profiling hooks into ComfyUI execution
|
|
if inject_profiling():
|
|
# Get profiler instance
|
|
profiler_manager = ProfilerManager.get_instance()
|
|
set_profiler_manager(profiler_manager)
|
|
|
|
# Register profiler API routes
|
|
try:
|
|
from server import PromptServer
|
|
server = PromptServer.instance
|
|
register_profiler_routes(server)
|
|
print("[SwissArmyKnife] Profiler initialized successfully")
|
|
except Exception as e:
|
|
print(f"[SwissArmyKnife] Warning: Could not register profiler API routes: {e}")
|
|
else:
|
|
print("[SwissArmyKnife] Warning: Profiler initialization failed")
|
|
except Exception as e:
|
|
print(f"[SwissArmyKnife] Warning: Could not initialize profiler: {e}")
|
|
|
|
# Initialize resource monitor service
|
|
try:
|
|
from .nodes.monitor_api import register_monitor_routes, set_monitor_service
|
|
from .nodes.utils.resource_monitor import MonitorService
|
|
from .nodes.restart_api import register_restart_routes
|
|
|
|
# Import PromptServer
|
|
try:
|
|
from server import PromptServer
|
|
|
|
# Register API routes
|
|
server = PromptServer.instance
|
|
register_monitor_routes(server)
|
|
register_restart_routes(server)
|
|
|
|
# Create and start monitor service
|
|
monitor_service = MonitorService(server, interval=2.0)
|
|
set_monitor_service(monitor_service)
|
|
monitor_service.start()
|
|
|
|
print("[SwissArmyKnife] Resource monitor initialized successfully")
|
|
except Exception as e:
|
|
print(f"[SwissArmyKnife] Warning: Could not initialize resource monitor: {e}")
|
|
except ImportError as e:
|
|
print(f"[SwissArmyKnife] Warning: Resource monitor dependencies not available: {e}")
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY", "VERSION", "DEBUG"]
|