InvokeAI/invokeai/app/invocations/custom_nodes/init.py
psychedelicious 4af4486dd9 feat(nodes,ui): add detection of custom nodes
Custom nodes have a new attribute `node_pack` indicating the node pack they came from.

- This is displayed in the UI in the icon icon tooltip.
- If a workflow is loaded and a node is unavailable, its node pack will be displayed (if it is known).
- If a workflow is migrated from v1 to v2, and the node is unknown, it falls back to "Unknown". If the missing node pack is installed and the node is updated, the node pack will be updated as expected.
2023-11-29 10:49:31 +11:00

55 lines
1.4 KiB
Python

"""
Invoke-managed custom node loader. See README.md for more information.
"""
import sys
from importlib.util import module_from_spec, spec_from_file_location
from pathlib import Path
from invokeai.app.invocations.baseinvocation import CUSTOM_NODE_PACK_SUFFIX
from invokeai.backend.util.logging import InvokeAILogger
logger = InvokeAILogger.get_logger()
loaded_count = 0
for d in Path(__file__).parent.iterdir():
# skip files
if not d.is_dir():
continue
# skip hidden directories
if d.name.startswith("_") or d.name.startswith("."):
continue
# skip directories without an `__init__.py`
init = d / "__init__.py"
if not init.exists():
continue
module_name = init.parent.stem
# skip if already imported
if module_name in globals():
continue
# load the module, appending adding a suffix to identify it as a custom node pack
spec = spec_from_file_location(f"{module_name}{CUSTOM_NODE_PACK_SUFFIX}", init.absolute())
if spec is None or spec.loader is None:
logger.warn(f"Could not load {init}")
continue
logger.info(f"Loading node pack {spec.name}")
module = module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
loaded_count += 1
del init, module_name
if loaded_count > 0:
logger.info(f"Loaded {loaded_count} node packs from {Path(__file__).parent}")