feat(app): less janky custom node loading

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)}"
+        )

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(