Skip to content

Commit d40f2fa

Browse files
feat(app): improved custom load loading ordering
Previously, custom node loading occurred _during module imports_. A consequence of this is that when a custom node import fails (e.g. its type clobbers an existing node), the app fails to start up. In fact, any time we import basically anything from the app, we trigger custom node imports! Not good. This logic is now in its own function, called as the API app starts up. If a custom node load fails for any reason, it no longer prevents the app from starting up. One other bonus we get from this is that we can now ensure custom nodes are loaded _after_ core nodes. Any clobbering that may occur while loading custom nodes is now guaranteed to be a custom node clobbering a core node's type - and not the other way round.
1 parent 933f4f6 commit d40f2fa

File tree

3 files changed

+46
-28
lines changed

3 files changed

+46
-28
lines changed

invokeai/app/api_app.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
workflows,
3737
)
3838
from invokeai.app.api.sockets import SocketIO
39+
from invokeai.app.invocations.load_custom_nodes import load_custom_nodes
3940
from invokeai.app.services.config.config_default import get_config
4041
from invokeai.app.util.custom_openapi import get_openapi_func
4142
from invokeai.backend.util.devices import TorchDevice
@@ -63,6 +64,11 @@
6364
# the correct port when the server starts in the lifespan handler.
6465
port = app_config.port
6566

67+
# Load custom nodes. This must be done after importing the Graph class, which itself imports all modules from the
68+
# invocations module. The ordering here is implicit, but important - we want to load custom nodes after all the
69+
# core nodes have been imported so that we can catch when a custom node clobbers a core node.
70+
load_custom_nodes(custom_nodes_path=app_config.custom_nodes_path)
71+
6672

6773
@asynccontextmanager
6874
async def lifespan(app: FastAPI):
Lines changed: 0 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,5 @@
1-
import shutil
2-
import sys
3-
from importlib.util import module_from_spec, spec_from_file_location
41
from pathlib import Path
52

6-
from invokeai.app.services.config.config_default import get_config
7-
8-
custom_nodes_path = Path(get_config().custom_nodes_path)
9-
custom_nodes_path.mkdir(parents=True, exist_ok=True)
10-
11-
custom_nodes_init_path = str(custom_nodes_path / "__init__.py")
12-
custom_nodes_readme_path = str(custom_nodes_path / "README.md")
13-
14-
# copy our custom nodes __init__.py to the custom nodes directory
15-
shutil.copy(Path(__file__).parent / "custom_nodes/init.py", custom_nodes_init_path)
16-
shutil.copy(Path(__file__).parent / "custom_nodes/README.md", custom_nodes_readme_path)
17-
18-
# set the same permissions as the destination directory, in case our source is read-only,
19-
# so that the files are user-writable
20-
for p in custom_nodes_path.glob("**/*"):
21-
p.chmod(custom_nodes_path.stat().st_mode)
22-
23-
# Import custom nodes, see https://docs.python.org/3/library/importlib.html#importing-programmatically
24-
spec = spec_from_file_location("custom_nodes", custom_nodes_init_path)
25-
if spec is None or spec.loader is None:
26-
raise RuntimeError(f"Could not load custom nodes from {custom_nodes_init_path}")
27-
module = module_from_spec(spec)
28-
sys.modules[spec.name] = module
29-
spec.loader.exec_module(module)
30-
313
# add core nodes to __all__
324
python_files = filter(lambda f: not f.name.startswith("_"), Path(__file__).parent.glob("*.py"))
335
__all__ = [f.stem for f in python_files] # type: ignore
Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import shutil
2+
import sys
3+
from importlib.util import module_from_spec, spec_from_file_location
4+
from pathlib import Path
5+
6+
7+
def load_custom_nodes(custom_nodes_path: Path):
8+
"""
9+
Loads all custom nodes from the custom_nodes_path directory.
10+
11+
This function copies a custom __init__.py file to the custom_nodes_path directory, effectively turning it into a
12+
python module.
13+
14+
The custom __init__.py file itself imports all the custom node packs as python modules from the custom_nodes_path
15+
directory.
16+
17+
Then,the custom __init__.py file is programmatically imported using importlib. As it executes, it imports all the
18+
custom node packs as python modules.
19+
"""
20+
custom_nodes_path.mkdir(parents=True, exist_ok=True)
21+
22+
custom_nodes_init_path = str(custom_nodes_path / "__init__.py")
23+
custom_nodes_readme_path = str(custom_nodes_path / "README.md")
24+
25+
# copy our custom nodes __init__.py to the custom nodes directory
26+
shutil.copy(Path(__file__).parent / "custom_nodes/init.py", custom_nodes_init_path)
27+
shutil.copy(Path(__file__).parent / "custom_nodes/README.md", custom_nodes_readme_path)
28+
29+
# set the same permissions as the destination directory, in case our source is read-only,
30+
# so that the files are user-writable
31+
for p in custom_nodes_path.glob("**/*"):
32+
p.chmod(custom_nodes_path.stat().st_mode)
33+
34+
# Import custom nodes, see https://docs.python.org/3/library/importlib.html#importing-programmatically
35+
spec = spec_from_file_location("custom_nodes", custom_nodes_init_path)
36+
if spec is None or spec.loader is None:
37+
raise RuntimeError(f"Could not load custom nodes from {custom_nodes_init_path}")
38+
module = module_from_spec(spec)
39+
sys.modules[spec.name] = module
40+
spec.loader.exec_module(module)

0 commit comments

Comments
 (0)