File size: 2,113 Bytes
c6535db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | from pathlib import Path
import sys
import os
import importlib.util
__version__ = "3.0.0"
# Locate current and node directories
current_dir = Path(__file__).parent
nodes_dir = current_dir / "py"
# Add both current and nodes directories to sys.path
for path in [current_dir, nodes_dir]:
if str(path) not in sys.path:
sys.path.insert(0, str(path))
# Initialize mappings
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
WEB_DIRECTORY = "./web"
def load_nodes():
"""Automatically discover and load node definitions"""
for base in (current_dir, nodes_dir):
if not base.exists():
continue
for file in base.rglob("*.py"):
if file.stem == "__init__":
continue
models_dir = current_dir / "models"
try:
if file.is_relative_to(models_dir):
continue
except Exception:
if str(models_dir) in str(file):
continue
try:
spec = importlib.util.spec_from_file_location(file.stem, file)
if spec and spec.loader:
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
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)
if hasattr(module, "Paths") and hasattr(module.Paths, "LLM_DIR"):
os.makedirs(module.Paths.LLM_DIR, exist_ok=True)
except Exception as e:
print(f"Error loading {file}: {e}")
# Load all nodes
load_nodes()
__all__ = ["NODE_CLASS_MAPPINGS","NODE_DISPLAY_NAME_MAPPINGS"]
print(f'\033[34m[ComfyUI-RMBG]\033[0m v\033[93m{__version__}\033[0m | '
f'\033[93m{len(NODE_CLASS_MAPPINGS)} nodes\033[0m \033[92mLoaded\033[0m')
|