Defer networkx import for schema generation

invokeai/app/services/shared/graph.py

@@ -7,6 +7,7 @@
 from dataclasses import dataclass
 from functools import wraps
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Concatenate,
@@ -22,7 +23,6 @@
     get_origin,
 )
 
-import networkx as nx
 from pydantic import (
     BaseModel,
     ConfigDict,
@@ -54,6 +54,26 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.misc import uuid_string
 
+if TYPE_CHECKING:
+    import networkx as nx
+else:
+
+    class _LazyNetworkX:
+        _module: Any | None = None
+
+        def _load(self) -> Any:
+            if self._module is None:
+                import networkx
+
+                self._module = networkx
+            return self._module
+
+        def __getattr__(self, name: str) -> Any:
+            return getattr(self._load(), name)
+
+    nx = _LazyNetworkX()
+
+
 # in 3.10 this would be "from types import NoneType"
 NoneType = type(None)
 
@@ -641,7 +661,7 @@ def _index_prepared_nodes_by_iteration_path(self, prepared_nodes: set[str]) -> d
         return prepared_nodes_by_iteration_path
 
     def _get_target_iteration_path(
-        self, source_node_id: str, graph: nx.DiGraph, prepared_iterator_nodes: tuple[str, ...]
+        self, source_node_id: str, graph: "nx.DiGraph", prepared_iterator_nodes: tuple[str, ...]
     ) -> Optional[tuple[int, ...]]:
         parent_iterators = self._get_parent_iterator_exec_nodes(source_node_id, graph, list(prepared_iterator_nodes))
         parent_paths = [self._state._get_iteration_path(prepared_id) for prepared_id, _ in parent_iterators]
@@ -656,7 +676,7 @@ def _get_target_iteration_path(
     def _get_indexed_iteration_node(
         self,
         source_node_id: str,
-        graph: nx.DiGraph,
+        graph: "nx.DiGraph",
         prepared_iterator_nodes: tuple[str, ...],
         prepared_nodes_by_iteration_path: dict[tuple[int, ...], list[str]],
     ) -> Optional[str]:
@@ -672,7 +692,7 @@ def _get_indexed_iteration_node(
                 return None
         return None
 
-    def _get_parent_iteration_mappings(self, next_node_id: str, graph: nx.DiGraph) -> Iterable[list[tuple[str, str]]]:
+    def _get_parent_iteration_mappings(self, next_node_id: str, graph: "nx.DiGraph") -> Iterable[list[tuple[str, str]]]:
         parent_node_ids = [source_id for source_id, _ in graph.in_edges(next_node_id)]
         iterator_graph = self.iterator_graph(graph)
         iterator_nodes = self.get_node_iterators(next_node_id, iterator_graph)
@@ -692,7 +712,7 @@ def _get_parent_iteration_mappings(self, next_node_id: str, graph: nx.DiGraph) -
         }
 
         def iter_mappings() -> Iterable[list[tuple[str, str]]]:
-            execution_graph: Optional[nx.DiGraph] = None
+            execution_graph: Optional["nx.DiGraph"] = None
             for prepared_iterators in itertools.product(*iterator_nodes_prepared):
                 mapping: list[tuple[str, str]] = []
                 for node_id in parent_node_ids:
@@ -750,7 +770,7 @@ def create_execution_node(
 
         return new_nodes
 
-    def iterator_graph(self, base: Optional[nx.DiGraph] = None) -> nx.DiGraph:
+    def iterator_graph(self, base: Optional["nx.DiGraph"] = None) -> "nx.DiGraph":
         """Gets a DiGraph with edges to collectors removed so an ancestor search produces all active iterators for any node"""
         g = base.copy() if base is not None else self._state.graph.nx_graph_flat()
         collectors = (
@@ -760,7 +780,7 @@ def iterator_graph(self, base: Optional[nx.DiGraph] = None) -> nx.DiGraph:
             g.remove_edges_from(list(g.in_edges(c)))
         return g
 
-    def get_node_iterators(self, node_id: str, it_graph: Optional[nx.DiGraph] = None) -> list[str]:
+    def get_node_iterators(self, node_id: str, it_graph: Optional["nx.DiGraph"] = None) -> list[str]:
         g = it_graph or self.iterator_graph()
         return [n for n in nx.ancestors(g, node_id) if isinstance(self._state.graph.get_node(n), IterateInvocation)]
 
@@ -772,7 +792,7 @@ def _get_prepared_nodes_for_source(self, source_node_id: str) -> set[str]:
         }
 
     def _get_parent_iterator_exec_nodes(
-        self, source_node_id: str, graph: nx.DiGraph, prepared_iterator_nodes: list[str]
+        self, source_node_id: str, graph: "nx.DiGraph", prepared_iterator_nodes: list[str]
     ) -> list[tuple[str, str]]:
         iterator_source_node_mapping = [
             (prepared_exec_node_id, self._state.prepared_source_mapping[prepared_exec_node_id])
@@ -785,7 +805,7 @@ def _get_parent_iterator_exec_nodes(
         ]
 
     def _matches_parent_iterators(
-        self, candidate_exec_node_id: str, parent_iterators: list[tuple[str, str]], execution_graph: nx.DiGraph
+        self, candidate_exec_node_id: str, parent_iterators: list[tuple[str, str]], execution_graph: "nx.DiGraph"
     ) -> bool:
         return all(
             nx.has_path(execution_graph, parent_iterator_exec_id, candidate_exec_node_id)
@@ -797,7 +817,7 @@ def _get_direct_prepared_iterator_match(
         prepared_nodes: set[str],
         prepared_iterator_nodes: list[str],
         parent_iterators: list[tuple[str, str]],
-        execution_graph: nx.DiGraph,
+        execution_graph: "nx.DiGraph",
     ) -> Optional[str]:
         prepared_iterator = next((node_id for node_id in prepared_iterator_nodes if node_id in prepared_nodes), None)
         if prepared_iterator is None:
@@ -807,7 +827,7 @@ def _get_direct_prepared_iterator_match(
         return None
 
     def _find_prepared_node_matching_iterators(
-        self, prepared_nodes: set[str], parent_iterators: list[tuple[str, str]], execution_graph: nx.DiGraph
+        self, prepared_nodes: set[str], parent_iterators: list[tuple[str, str]], execution_graph: "nx.DiGraph"
     ) -> Optional[str]:
         return next(
             (
@@ -821,8 +841,8 @@ def _find_prepared_node_matching_iterators(
     def get_iteration_node(
         self,
         source_node_id: str,
-        graph: nx.DiGraph,
-        execution_graph: nx.DiGraph,
+        graph: "nx.DiGraph",
+        execution_graph: "nx.DiGraph",
         prepared_iterator_nodes: list[str],
         prepared_nodes: Optional[set[str]] = None,
     ) -> Optional[str]:
@@ -846,7 +866,7 @@ def get_iteration_node(
 
         return self._find_prepared_node_matching_iterators(prepared_nodes, parent_iterators, execution_graph)
 
-    def prepare(self, base_g: Optional[nx.DiGraph] = None) -> Optional[str]:
+    def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]:
         g = base_g or self._state.graph.nx_graph_flat()
         next_node_id = next(
             (
@@ -1065,7 +1085,7 @@ def _get_ordered_iterator_sources(self, source_node_id: str) -> list[str]:
         return iterator_sources
 
     def _get_iterator_exec_id(
-        self, iterator_source_id: str, exec_node_id: str, execution_graph: nx.DiGraph
+        self, iterator_source_id: str, exec_node_id: str, execution_graph: "nx.DiGraph"
     ) -> Optional[str]:
         prepared = self._state.source_prepared_mapping.get(iterator_source_id)
         if not prepared:
@@ -2286,15 +2306,15 @@ def _is_collector_connection_valid(
 
         return None
 
-    def nx_graph(self) -> nx.DiGraph:
+    def nx_graph(self) -> "nx.DiGraph":
         """Returns a NetworkX DiGraph representing the layout of this graph"""
         # TODO: Cache this?
         g = nx.DiGraph()
         g.add_nodes_from(list(self.nodes.keys()))
         g.add_edges_from({(e.source.node_id, e.destination.node_id) for e in self.edges})
         return g
 
-    def nx_graph_flat(self, nx_graph: Optional[nx.DiGraph] = None) -> nx.DiGraph:
+    def nx_graph_flat(self, nx_graph: Optional["nx.DiGraph"] = None) -> "nx.DiGraph":
         """Returns a flattened NetworkX DiGraph, including all subgraphs (but not with iterations expanded)"""
         g = nx_graph or nx.DiGraph()
 
@@ -2756,20 +2776,20 @@ def create_child_workflow_execution_state(self, graph: Graph, frame: WorkflowCal
     def _create_execution_node(self, node_id: str, iteration_node_map: list[tuple[str, str]]) -> list[str]:
         return self._materializer().create_execution_node(node_id, iteration_node_map)
 
-    def _iterator_graph(self, base: Optional[nx.DiGraph] = None) -> nx.DiGraph:
+    def _iterator_graph(self, base: Optional["nx.DiGraph"] = None) -> "nx.DiGraph":
         return self._materializer().iterator_graph(base)
 
-    def _get_node_iterators(self, node_id: str, it_graph: Optional[nx.DiGraph] = None) -> list[str]:
+    def _get_node_iterators(self, node_id: str, it_graph: Optional["nx.DiGraph"] = None) -> list[str]:
         return self._materializer().get_node_iterators(node_id, it_graph)
 
-    def _prepare(self, base_g: Optional[nx.DiGraph] = None) -> Optional[str]:
+    def _prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]:
         return self._materializer().prepare(base_g)
 
     def _get_iteration_node(
         self,
         source_node_id: str,
-        graph: nx.DiGraph,
-        execution_graph: nx.DiGraph,
+        graph: "nx.DiGraph",
+        execution_graph: "nx.DiGraph",
         prepared_iterator_nodes: list[str],
     ) -> Optional[str]:
         return self._materializer().get_iteration_node(source_node_id, graph, execution_graph, prepared_iterator_nodes)

tests/test_imports.py

@@ -1,5 +1,8 @@
 import importlib
 import pkgutil
+import subprocess
+import sys
+import textwrap
 
 import invokeai
 
@@ -33,3 +36,29 @@ def test_invokeai_imports():
                 failed_to_import.add(mod)
 
     assert not failed_to_import, f"Modules failed to import: {failed_to_import}"
+
+
+def test_graph_module_import_does_not_require_networkx():
+    script = """
+    import builtins
+
+    real_import = builtins.__import__
+
+    def blocked_import(name, globals=None, locals=None, fromlist=(), level=0):
+        if name == "networkx" or name.startswith("networkx."):
+            raise ModuleNotFoundError("No module named 'networkx'")
+        return real_import(name, globals, locals, fromlist, level)
+
+    builtins.__import__ = blocked_import
+    import invokeai.app.services.shared.graph
+    """
+
+    result = subprocess.run(
+        [sys.executable, "-c", textwrap.dedent(script)],
+        capture_output=True,
+        text=True,
+        timeout=60,
+        check=False,
+    )
+
+    assert result.returncode == 0, result.stderr