feat(nodes): simple custom nodes

Custom nodes may be places in `$INVOKEAI_ROOT/nodes/` (configurable with `custom_nodes_dir` option).

On app startup, an `__init__.py` is copied into the custom nodes dir, which recursively loads all python files in the directory as modules (files starting with `_` are ignored). The custom nodes dir is now a python module itself.

When we `from invocations import *` to load init all invocations, we load the custom nodes dir, registering all custom nodes.
This commit is contained in:
psychedelicious 2023-10-19 17:51:55 +11:00
parent b7f63a4065
commit 8604943e89
3 changed files with 68 additions and 6 deletions

View File

@ -1,8 +1,25 @@
import os
import shutil
import sys
from importlib.util import module_from_spec, spec_from_file_location
from pathlib import Path
__all__ = []
from invokeai.app.services.config.config_default import InvokeAIAppConfig
dirname = os.path.dirname(os.path.abspath(__file__))
for f in os.listdir(dirname):
if f != "__init__.py" and os.path.isfile("%s/%s" % (dirname, f)) and f[-3:] == ".py":
__all__.append(f[:-3])
custom_nodes_path = Path(InvokeAIAppConfig.get_config().custom_nodes_path.absolute())
custom_nodes_path.mkdir(parents=True, exist_ok=True)
custom_nodes_init_path = str(custom_nodes_path / "__init__.py")
# copy our custom nodes __init__.py to the custom nodes directory
shutil.copy(Path(__file__).parent / "_custom_nodes_init.py", custom_nodes_init_path)
# Import custom nodes, see https://docs.python.org/3/library/importlib.html#importing-programmatically
spec = spec_from_file_location("custom_nodes", custom_nodes_init_path)
if spec is None or spec.loader is None:
raise RuntimeError(f"Could not load custom nodes from {custom_nodes_init_path}")
module = module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
# add core nodes to __all__
python_files = filter(lambda f: not f.name.startswith("_"), Path(__file__).parent.rglob("*.py"))
__all__ = list(f.stem for f in python_files) # type: ignore

View File

@ -0,0 +1,37 @@
"""
InvokeAI custom nodes initialization
This file is responsible for loading all custom nodes from this directory.
All python files are loaded on app startup. Custom nodes will be initialized and available for use
in workflows.
The app must be restarted for changes to be picked up.
This file is overwritten on launch. Do not edit this file directly.
"""
import sys
from importlib import import_module
from importlib.util import module_from_spec, spec_from_file_location
from pathlib import Path
from invokeai.backend.util.logging import InvokeAILogger
logger = InvokeAILogger.get_logger()
count = 0
for f in Path(__file__).parent.rglob("*.py"):
module_name = f.stem
if (not module_name.startswith("_")) and (module_name not in globals()):
spec = spec_from_file_location(module_name, f.absolute())
if spec is None or spec.loader is None:
logger.warn(f"Could not load {f}")
continue
module = module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
count += 1
del f, module_name
logger.info(f"Loaded {count} modules from {Path(__file__).parent}")
del import_module, Path

View File

@ -243,6 +243,7 @@ class InvokeAIAppConfig(InvokeAISettings):
db_dir : Optional[Path] = Field(default=Path('databases'), description='Path to InvokeAI databases directory', json_schema_extra=Categories.Paths)
outdir : Optional[Path] = Field(default=Path('outputs'), description='Default folder for output images', json_schema_extra=Categories.Paths)
use_memory_db : bool = Field(default=False, description='Use in-memory database for storing image metadata', json_schema_extra=Categories.Paths)
custom_nodes_dir : Path = Field(default=Path('nodes'), description='Path to directory for custom nodes', json_schema_extra=Categories.Paths)
from_file : Optional[Path] = Field(default=None, description='Take command input from the indicated file (command-line client only)', json_schema_extra=Categories.Paths)
# LOGGING
@ -410,6 +411,13 @@ class InvokeAIAppConfig(InvokeAISettings):
"""
return self._resolve(self.models_dir)
@property
def custom_nodes_path(self) -> Path:
"""
Path to the custom nodes directory
"""
return self._resolve(self.custom_nodes_dir)
# the following methods support legacy calls leftover from the Globals era
@property
def full_precision(self) -> bool: