From 1939ccde1e4a21a0e974ebd261dc3d580ee9feca Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Fri, 7 Mar 2025 10:40:41 +1000 Subject: [PATCH] feat(app): less janky custom node loading - We don't need to copy the init file. Just crawl the custom nodes dir for modules and import them all. Dunno why I didn't do this initially. - Pass the logger in as an arg. There was a race condition where if we got the logger directly in the load_custom_nodes function, the config would not have been loaded fully yet and we'd end up with the wrong custom nodes path! - Remove permissions-setting logic, I do not believe it is relevant for custom nodes - Minor cleanup of the utility --- invokeai/app/invocations/custom_nodes/init.py | 64 -------------- invokeai/app/invocations/load_custom_nodes.py | 87 ++++++++++++++----- invokeai/app/run_app.py | 2 +- 3 files changed, 66 insertions(+), 87 deletions(-) delete mode 100644 invokeai/app/invocations/custom_nodes/init.py diff --git a/invokeai/app/invocations/custom_nodes/init.py b/invokeai/app/invocations/custom_nodes/init.py deleted file mode 100644 index 171a3786904..00000000000 --- a/invokeai/app/invocations/custom_nodes/init.py +++ /dev/null @@ -1,64 +0,0 @@ -""" -Invoke-managed custom node loader. See README.md for more information. -""" - -import sys -import traceback -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() -loaded_packs: list[str] = [] -failed_packs: list[str] = [] - -custom_nodes_dir = Path(__file__).parent - -for d in custom_nodes_dir.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(module_name, 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 {module_name}") - - try: - module = module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - - loaded_packs.append(module_name) - except Exception: - failed_packs.append(module_name) - full_error = traceback.format_exc() - logger.error(f"Failed to load node pack {module_name} (may have partially loaded):\n{full_error}") - - del init, module_name - -loaded_count = len(loaded_packs) -if loaded_count > 0: - logger.info( - f"Loaded {loaded_count} node pack{'s' if loaded_count != 1 else ''} from {custom_nodes_dir}: {', '.join(loaded_packs)}" - ) diff --git a/invokeai/app/invocations/load_custom_nodes.py b/invokeai/app/invocations/load_custom_nodes.py index 993237478ea..a3a8194a3b9 100644 --- a/invokeai/app/invocations/load_custom_nodes.py +++ b/invokeai/app/invocations/load_custom_nodes.py @@ -1,40 +1,83 @@ +import logging import shutil import sys +import traceback from importlib.util import module_from_spec, spec_from_file_location from pathlib import Path -def load_custom_nodes(custom_nodes_path: Path): +def load_custom_nodes(custom_nodes_path: Path, logger: logging.Logger): """ Loads all custom nodes from the custom_nodes_path directory. - This function copies a custom __init__.py file to the custom_nodes_path directory, effectively turning it into a - python module. + If custom_nodes_path does not exist, it creates it. - The custom __init__.py file itself imports all the custom node packs as python modules from the custom_nodes_path - directory. + It also copies the custom_nodes/README.md file to the custom_nodes_path directory. Because this file may change, + it is _always_ copied to the custom_nodes_path directory. - Then,the custom __init__.py file is programmatically imported using importlib. As it executes, it imports all the - custom node packs as python modules. + Then, it crawls the custom_nodes_path directory and imports all top-level directories as python modules. + + If the directory does not contain an __init__.py file or starts with an `_` or `.`, it is skipped. """ + + # create the custom nodes directory if it does not exist custom_nodes_path.mkdir(parents=True, exist_ok=True) - custom_nodes_init_path = str(custom_nodes_path / "__init__.py") - custom_nodes_readme_path = str(custom_nodes_path / "README.md") + # Copy the README file to the custom nodes directory + source_custom_nodes_readme_path = Path(__file__).parent / "custom_nodes/README.md" + target_custom_nodes_readme_path = Path(custom_nodes_path) / "README.md" - # copy our custom nodes __init__.py to the custom nodes directory - shutil.copy(Path(__file__).parent / "custom_nodes/init.py", custom_nodes_init_path) - shutil.copy(Path(__file__).parent / "custom_nodes/README.md", custom_nodes_readme_path) + # copy our custom nodes README to the custom nodes directory + shutil.copy(source_custom_nodes_readme_path, target_custom_nodes_readme_path) - # set the same permissions as the destination directory, in case our source is read-only, - # so that the files are user-writable - for p in custom_nodes_path.glob("**/*"): - p.chmod(custom_nodes_path.stat().st_mode) + loaded_packs: list[str] = [] + failed_packs: list[str] = [] # 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) + for d in custom_nodes_path.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 + spec = spec_from_file_location(module_name, init.absolute()) + + if spec is None or spec.loader is None: + logger.warning(f"Could not load {init}") + continue + + logger.info(f"Loading node pack {module_name}") + + try: + module = module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + + loaded_packs.append(module_name) + except Exception: + failed_packs.append(module_name) + full_error = traceback.format_exc() + logger.error(f"Failed to load node pack {module_name} (may have partially loaded):\n{full_error}") + + del init, module_name + + loaded_count = len(loaded_packs) + if loaded_count > 0: + logger.info( + f"Loaded {loaded_count} node pack{'s' if loaded_count != 1 else ''} from {custom_nodes_path}: {', '.join(loaded_packs)}" + ) diff --git a/invokeai/app/run_app.py b/invokeai/app/run_app.py index f8972cf4f11..3cc8f9d787a 100644 --- a/invokeai/app/run_app.py +++ b/invokeai/app/run_app.py @@ -59,7 +59,7 @@ def run_app() -> None: # Load custom nodes. This must be done after importing the Graph class, which itself imports all modules from the # invocations module. The ordering here is implicit, but important - we want to load custom nodes after all the # core nodes have been imported so that we can catch when a custom node clobbers a core node. - load_custom_nodes(custom_nodes_path=app_config.custom_nodes_path) + load_custom_nodes(custom_nodes_path=app_config.custom_nodes_path, logger=logger) # Start the server. config = uvicorn.Config(