diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index 44b876d1651..99c96cc0c50 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -669,6 +669,14 @@ materialization. - Try to resolve any `If`-specific scheduling state. - If the node is ready and not deferred by an unresolved `If`, enqueue it into its class queue. +Literal-condition `If` fallback checks use the private `_IteratorAncestry` context. One context belongs to one +compatibility `prepare()` call and its source DAG, shared by candidate admission and selected-node materialization; +a standalone iteration-path lookup creates its own context and passes it through both `If` and ordinary-node +recursion. The filtered iterator graph is built only when an unresolved selected branch needs the ancestry fallback. +The context computes boolean proper-ancestor reachability with collector and final-`For` boundaries preserved, and +reuses that analysis within the call. No context is retained on execution state or serialized, so later calls observe +new supplied graphs and source mutations without cache invalidation. + ### 4.5 Readiness and class ordering - `_enqueue_if_ready(nid)` applies generic readiness: `indegree == 0`, not executed or claimed, and, for an `If` diff --git a/invokeai/app/services/shared/execution_engine/primitives.py b/invokeai/app/services/shared/execution_engine/primitives.py index dc1bec185a9..78ca5049c67 100644 --- a/invokeai/app/services/shared/execution_engine/primitives.py +++ b/invokeai/app/services/shared/execution_engine/primitives.py @@ -51,6 +51,8 @@ class _InternalModel(BaseModel): def __eq__(self, other: object) -> bool: """Compare durable fields, excluding process-local synchronization locks.""" + if self is other: + return True if not isinstance(other, type(self)): return NotImplemented return self.model_dump(mode="python") == other.model_dump(mode="python") diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index c7a88f5e7d7..23f612d6ba3 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -1,7 +1,7 @@ import copy # noqa: F401 import json import weakref # noqa: F401 -from functools import wraps # noqa: F401 +from functools import lru_cache, wraps # noqa: F401 from typing import ( TYPE_CHECKING, # noqa: F401 Any, @@ -296,6 +296,17 @@ def __init__(self) -> None: self.fresh_flat_if_activation: Optional[bool] = None +@lru_cache(maxsize=4096) +def _engine_frame_for(state_id: str, workflow_call_depth: int, iteration_path: tuple[int, ...]) -> EngineExecutionFrame: + """Share one frozen frame per identity so gate scope checks compare by identity instead of by dump.""" + return EngineExecutionFrame( + state_id=state_id, + frame_id=f"{state_id}:{workflow_call_depth}:{','.join(str(i) for i in iteration_path)}", + iteration_path=iteration_path, + workflow_call_depth=workflow_call_depth, + ) + + class GraphExecutionState(BaseModel): """Tracks source-graph expansion, execution progress, and runtime results.""" @@ -1459,13 +1470,7 @@ def _clear_transient_runtime(self) -> None: self._generic_graph_scheduler = None def _engine_frame(self, iteration_path: tuple[int, ...]) -> EngineExecutionFrame: - frame_id = f"{self.id}:{len(self.workflow_call_stack)}:{','.join(str(i) for i in iteration_path)}" - return EngineExecutionFrame( - state_id=self.id, - frame_id=frame_id, - iteration_path=iteration_path, - workflow_call_depth=len(self.workflow_call_stack), - ) + return _engine_frame_for(self.id, len(self.workflow_call_stack), iteration_path) def _activation_gate(self, exec_node_id: str) -> ActivationGate: return graph_if_runtime._activation_gate(self, exec_node_id) diff --git a/invokeai/app/services/shared/graph_execution_runtime.py b/invokeai/app/services/shared/graph_execution_runtime.py index ec7a6fdbca2..c4ef3b8cdc0 100644 --- a/invokeai/app/services/shared/graph_execution_runtime.py +++ b/invokeai/app/services/shared/graph_execution_runtime.py @@ -51,7 +51,7 @@ def _get_iteration_source_node_id(self, exec_node_id: str) -> Optional[str]: return self._state._prepared_registry().get_source_node_id(exec_node_id) def _get_ordered_iterator_sources(self, source_node_id: str) -> list[str]: - iterator_graph = self._state._iterator_graph(self._state.graph.nx_graph()) + iterator_graph = self._state._iterator_graph(self._state._get_source_graph_flat()) iterator_sources = [ node_id for node_id in nx.ancestors(iterator_graph, source_node_id) diff --git a/invokeai/app/services/shared/graph_if_activation.py b/invokeai/app/services/shared/graph_if_activation.py index 80220ee7dda..7a97d3b3d08 100644 --- a/invokeai/app/services/shared/graph_if_activation.py +++ b/invokeai/app/services/shared/graph_if_activation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional from invokeai.app.invocations.logic import IfInvocation from invokeai.app.services.shared.execution_engine.scheduler import ActivationDependency @@ -19,6 +19,8 @@ class _IfActivationController: def __init__(self, state: "GraphExecutionState") -> None: self._state = state self._branch_sources_cache: dict[tuple[str, str], frozenset[str]] = {} + self._if_node_ids: Optional[tuple[str, ...]] = None + self._branch_memberships_by_source: dict[str, tuple[tuple[str, tuple[str, ...]], ...]] = {} def _branch_sources(self, if_node_id: str, branch_field: str, source_graph: "nx.DiGraph") -> set[str]: cache_key = (if_node_id, branch_field) @@ -46,11 +48,39 @@ def _branch_sources(self, if_node_id: str, branch_field: str, source_graph: "nx. self._branch_sources_cache[cache_key] = frozenset(branch_sources) return branch_sources + def _get_if_node_ids(self) -> tuple[str, ...]: + if self._if_node_ids is None: + self._if_node_ids = tuple( + node_id for node_id, node in self._state.graph.nodes.items() if isinstance(node, IfInvocation) + ) + return self._if_node_ids + + def _branch_memberships(self, source_node_id: str) -> tuple[tuple[str, tuple[str, ...]], ...]: + """Return each If with the branch fields whose exclusive sources include this source.""" + cached = self._branch_memberships_by_source.get(source_node_id) + if cached is not None: + return cached + + source_graph = self._state._get_source_graph_flat() + memberships = tuple( + (if_node_id, fields) + for if_node_id in self._get_if_node_ids() + if ( + fields := tuple( + branch_field + for branch_field in ("true_input", "false_input") + if source_node_id in self._branch_sources(if_node_id, branch_field, source_graph) + ) + ) + ) + self._branch_memberships_by_source[source_node_id] = memberships + return memberships + def get_dependencies(self, exec_node_id: str) -> tuple[ActivationDependency, ...]: """Return activation requirements for one concrete prepared node.""" source_node_id = self._state._prepared_registry().get_source_node_id(exec_node_id) - if not any(isinstance(node, IfInvocation) for node in self._state.graph.nodes.values()): + if not self._get_if_node_ids(): return () return self.get_source_dependencies(source_node_id, self._state._get_iteration_path(exec_node_id)) @@ -59,29 +89,11 @@ def get_source_dependencies( ) -> tuple[ActivationDependency, ...]: """Return branch requirements before a source node is materialized.""" - if not any(isinstance(node, IfInvocation) for node in self._state.graph.nodes.values()): - return () - source_graph = self._state._get_source_graph_flat() - dependencies: list[ActivationDependency] = [] - for source_if_id, source_if_node in self._state.graph.nodes.items(): - if not isinstance(source_if_node, IfInvocation): - continue - - matching_fields = tuple( - branch_field - for branch_field in ("true_input", "false_input") - if source_node_id in self._branch_sources(source_if_id, branch_field, source_graph) - ) - if len(matching_fields) != 1: - continue - dependencies.append( - ActivationDependency( - owner_id=source_if_id, - branch=matching_fields[0], - frame=iteration_path, - ) - ) - return tuple(dependencies) + return tuple( + ActivationDependency(owner_id=if_node_id, branch=fields[0], frame=iteration_path) + for if_node_id, fields in self._branch_memberships(source_node_id) + if len(fields) == 1 + ) def is_source_admitted(self, source_node_id: str, iteration_path: tuple[int, ...] = ()) -> bool: dependencies = self.get_source_dependencies(source_node_id, iteration_path) @@ -102,16 +114,7 @@ def is_source_inactive(self, source_node_id: str, iteration_path: tuple[int, ... if iteration_path: return self.is_source_rejected(source_node_id, iteration_path) - source_graph = self._state._get_source_graph_flat() - relevant_if_ids = { - if_node_id - for if_node_id, if_node in self._state.graph.nodes.items() - if isinstance(if_node, IfInvocation) - and any( - source_node_id in self._branch_sources(if_node_id, branch_field, source_graph) - for branch_field in ("true_input", "false_input") - ) - } + relevant_if_ids = {if_node_id for if_node_id, _ in self._branch_memberships(source_node_id)} if not relevant_if_ids: return False diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index c2ea5d9fa9a..780712d314b 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -5,8 +5,9 @@ from invokeai.app.invocations.logic import IfInvocation from invokeai.app.services.shared.execution_engine.primitives import ActivationGate from invokeai.app.services.shared.execution_engine.scheduler import ActivationDependency +from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry from invokeai.app.services.shared.graph_models import Edge, ExecutionToken -from invokeai.app.services.shared.graph_validation import ForInvocation, IterateInvocation, nx +from invokeai.app.services.shared.graph_validation import IterateInvocation, nx if TYPE_CHECKING: from invokeai.app.services.shared.graph import GraphExecutionState @@ -19,6 +20,7 @@ def _get_effective_iteration_paths_for_edge( get_prepared_nodes_for_source: Callable[[str], set[str]], get_prepared_edge_iteration_path: Callable[[Edge, str], tuple[int, ...]], source_graph: Optional["nx.DiGraph"] = None, + iterator_ancestry: Optional[_IteratorAncestry] = None, ) -> set[tuple[int, ...]]: source_node_id = edge.source.node_id prepared_nodes: set[str] = set() @@ -41,6 +43,7 @@ def _get_effective_iteration_paths_for_edge( get_prepared_nodes_for_source, get_prepared_edge_iteration_path, source_graph, + iterator_ancestry, ) ) @@ -57,6 +60,7 @@ def _get_effective_iteration_paths_for_edge( get_prepared_nodes_for_source, get_prepared_edge_iteration_path, source_graph, + iterator_ancestry, ) for input_edge in input_edges ) @@ -70,8 +74,11 @@ def _get_if_iteration_paths( get_prepared_nodes_for_source: Callable[[str], set[str]], get_prepared_edge_iteration_path: Callable[[Edge, str], tuple[int, ...]], source_graph: Optional["nx.DiGraph"] = None, + iterator_ancestry: Optional[_IteratorAncestry] = None, ) -> list[tuple[int, ...]]: """Infer the selected branch's effective iteration paths for an If output.""" + if iterator_ancestry is None: + iterator_ancestry = _IteratorAncestry(state, source_graph) node = state.graph.get_node(node_id) assert isinstance(node, IfInvocation) next_visited_source_ids = {*visited_source_ids, node_id} @@ -113,6 +120,7 @@ def _get_if_iteration_paths( get_prepared_nodes_for_source, get_prepared_edge_iteration_path, source_graph, + iterator_ancestry, ) for edge in selected_edges ) @@ -132,14 +140,7 @@ def _get_if_iteration_paths( if not has_selected_edges: return [()] if not condition_edges and has_unresolved_selected_source: - iterator_graph = state._iterator_graph( - source_graph if source_graph is not None else state._get_source_graph_flat() - ) - has_iterator_ancestor = any( - isinstance(state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) - for ancestor_id in nx.ancestors(iterator_graph, node_id) - ) - if not has_iterator_ancestor: + if not iterator_ancestry.has_iterator_ancestor(node_id): return [()] return [] return sorted(path for path in paths if not any(path != other and other[: len(path)] == path for other in paths)) @@ -286,19 +287,31 @@ def _is_activation_dependency_rejected(state: "GraphExecutionState", dependency: matching_gate_ids = state._prepared_if_exec_ids(dependency.owner_id, dependency.frame) for gate_id in matching_gate_ids: gate = state._activation_gate(gate_id) - expected_frame = state._expected_execution_ref(gate_id).frame - if ( - gate.frame.state_id != expected_frame.state_id - or gate.frame.frame_id != expected_frame.frame_id - or gate.frame.iteration_path != expected_frame.iteration_path - or gate.frame.workflow_call_depth != expected_frame.workflow_call_depth - ): + if not _gate_matches_expected_frame(state, gate_id, gate): return False if gate.resolved and gate.selected_branch != dependency.branch: return True return False +def _gate_matches_expected_frame(state: "GraphExecutionState", gate_id: str, gate: ActivationGate) -> bool: + # Gates registered in this process share the frame instance built for the prepared path, which makes the + # field-by-field scope check below trivially true; rehydrated or mismatched gates take the full check. + if ( + gate_id in state.execution_graph.nodes + and state.prepared_source_mapping.get(gate_id) is not None + and gate.frame is state._engine_frame(state.prepared_iteration_paths.get(gate_id, ())) + ): + return True + expected_frame = state._expected_execution_ref(gate_id).frame + return ( + gate.frame.state_id == expected_frame.state_id + and gate.frame.frame_id == expected_frame.frame_id + and gate.frame.iteration_path == expected_frame.iteration_path + and gate.frame.workflow_call_depth == expected_frame.workflow_call_depth + ) + + def _resolve_activation_gate(state: "GraphExecutionState", exec_node_id: str, branch: str) -> bool: runtime = state._generic_runtime() gate = state._activation_gate(exec_node_id) diff --git a/invokeai/app/services/shared/graph_iterator_ancestry.py b/invokeai/app/services/shared/graph_iterator_ancestry.py new file mode 100644 index 00000000000..70ab4d7ffa2 --- /dev/null +++ b/invokeai/app/services/shared/graph_iterator_ancestry.py @@ -0,0 +1,35 @@ +from typing import TYPE_CHECKING, Optional + +from invokeai.app.invocations.loops import ForInvocation +from invokeai.app.services.shared.graph_validation import IterateInvocation, nx + +if TYPE_CHECKING: + from invokeai.app.services.shared.graph import GraphExecutionState + + +class _IteratorAncestry: + """Lazily index strict iterator ancestry for one source-graph traversal.""" + + def __init__(self, state: "GraphExecutionState", source_graph: Optional["nx.DiGraph"] = None) -> None: + self._state = state + self._source_graph = source_graph + self._has_iterator_ancestor: dict[str, bool] | None = None + + def has_iterator_ancestor(self, node_id: str) -> bool: + if self._has_iterator_ancestor is None: + source_graph = self._source_graph + if source_graph is None: + source_graph = self._state._get_source_graph_flat() + iterator_graph = self._state._iterator_graph(source_graph) + has_iterator_ancestor = dict.fromkeys(iterator_graph, False) + for source_node_id in nx.topological_sort(iterator_graph): + if has_iterator_ancestor[source_node_id] or isinstance( + self._state.graph.get_node(source_node_id), (ForInvocation, IterateInvocation) + ): + for successor_id in iterator_graph.successors(source_node_id): + has_iterator_ancestor[successor_id] = True + self._has_iterator_ancestor = has_iterator_ancestor + + if node_id not in self._has_iterator_ancestor: + raise nx.NetworkXError(f"The node {node_id} is not in the digraph.") + return self._has_iterator_ancestor[node_id] diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 26a6c34c91c..7aa255deb5a 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -14,6 +14,7 @@ LoopState, ) from invokeai.app.services.shared import graph_if_runtime +from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry from invokeai.app.services.shared.graph_models import Edge, EdgeConnection from invokeai.app.services.shared.graph_validation import ( COLLECTION_FIELD, @@ -43,6 +44,9 @@ class _ExecutionNodeBuilder: def __init__(self, state: "GraphExecutionState") -> None: self._state = state self._iteration_axes_by_source: dict[str, tuple[str, ...]] = {} + # Keyed by the identity of the cached flat source graph, which every graph edit replaces. + self._flat_iterator_graph: Optional[tuple["nx.DiGraph", "nx.DiGraph"]] = None + self._node_iterators: dict[str, tuple[str, ...]] = {} def _get_iteration_axes(self, source_node_id: str) -> tuple[str, ...]: cached = self._iteration_axes_by_source.get(source_node_id) @@ -1411,7 +1415,9 @@ def create_execution_node( return new_nodes - def _has_admitted_source_mapping(self, node_id: str, graph: "nx.DiGraph") -> bool: + def _has_admitted_source_mapping( + self, node_id: str, graph: "nx.DiGraph", iterator_ancestry: Optional[_IteratorAncestry] = None + ) -> bool: if not any(isinstance(node, IfInvocation) for node in self._state.graph.nodes.values()): return True if isinstance(self._state.graph.get_node(node_id), IfInvocation): @@ -1426,6 +1432,7 @@ def _has_admitted_source_mapping(self, node_id: str, graph: "nx.DiGraph") -> boo self._get_prepared_nodes_for_source, self._get_prepared_edge_iteration_path, graph, + iterator_ancestry, ) ) mappings = self._get_if_condition_iteration_mappings(node_id, graph) @@ -1585,8 +1592,20 @@ def _attach_pending_if_inputs(self, *, enqueue: bool = True) -> None: self._state._enqueue_if_ready(exec_node_id) 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._get_source_graph_flat().copy() + """Gets a DiGraph with edges to collectors removed so an ancestor search produces all active iterators for any node + + The result for the flat source graph is shared; callers must not mutate it. + """ + source_graph = self._state._get_source_graph_flat() + if base is not None and base is not source_graph: + return self._build_iterator_graph(base) + if self._flat_iterator_graph is None or self._flat_iterator_graph[0] is not source_graph: + self._flat_iterator_graph = (source_graph, self._build_iterator_graph(source_graph)) + self._node_iterators = {} + return self._flat_iterator_graph[1] + + def _build_iterator_graph(self, base: "nx.DiGraph") -> "nx.DiGraph": + g = base.copy() collectors = ( n for n in self._state.graph.nodes if isinstance(self._state.graph.get_node(n), CollectInvocation) ) @@ -1603,10 +1622,19 @@ def iterator_graph(self, base: Optional["nx.DiGraph"] = None) -> "nx.DiGraph": return g def get_node_iterators(self, node_id: str, it_graph: Optional["nx.DiGraph"] = None) -> list[str]: - g = it_graph or self.iterator_graph() + flat_iterator_graph = self.iterator_graph() + if it_graph is not None and it_graph is not flat_iterator_graph: + return self._find_node_iterators(node_id, it_graph) + cached = self._node_iterators.get(node_id) + if cached is None: + cached = tuple(self._find_node_iterators(node_id, flat_iterator_graph)) + self._node_iterators[node_id] = cached + return list(cached) + + def _find_node_iterators(self, node_id: str, it_graph: "nx.DiGraph") -> list[str]: return [ n - for n in nx.ancestors(g, node_id) + for n in nx.ancestors(it_graph, node_id) if isinstance(self._state.graph.get_node(n), (ForInvocation, IterateInvocation)) ] @@ -1705,6 +1733,7 @@ def get_iteration_node( def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: g = base_g if base_g is not None else self._state._get_source_graph_flat() + iterator_ancestry = _IteratorAncestry(self._state, g) self._attach_pending_if_inputs() next_node_id = next( ( @@ -1734,7 +1763,7 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: and self._is_if_condition_ready(node_id) or not isinstance(self._state.graph.get_node(node_id), IfInvocation) ) - and self._has_admitted_source_mapping(node_id, g) + and self._has_admitted_source_mapping(node_id, g, iterator_ancestry) and ( not isinstance(self._state.graph.get_node(node_id), (ForInvocation, IterateInvocation)) or all(source_id in self._state.executed for source_id, _ in g.in_edges(node_id)) @@ -1769,6 +1798,7 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: self._get_prepared_nodes_for_source, self._get_prepared_edge_iteration_path, g, + iterator_ancestry, ) condition_edges = self._state.graph._get_input_edges(next_node_id, "condition") if condition_edges: diff --git a/tests/app/services/shared/test_graph_if_runtime.py b/tests/app/services/shared/test_graph_if_runtime.py index ad596f2812a..fac41fccf75 100644 --- a/tests/app/services/shared/test_graph_if_runtime.py +++ b/tests/app/services/shared/test_graph_if_runtime.py @@ -172,3 +172,71 @@ def test_nested_if_iteration_paths_use_supplied_source_graph_recursively() -> No lambda _edge, _prepared_id: (), constrained_graph, ) == [()] + + +def test_unresolved_nested_literal_ifs_share_one_iterator_graph(monkeypatch: pytest.MonkeyPatch) -> None: + nodes = {"iterator": IterateInvocation(id="iterator"), "guarded": IfInvocation(id="guarded")} + nodes.update({f"if-{i}": IfInvocation(id=f"if-{i}", condition=True) for i in range(24)}) + edges = [ + Edge( + source=EdgeConnection(node_id="iterator", field="item"), + destination=EdgeConnection(node_id="guarded", field="condition"), + ) + ] + previous = "guarded" + for i in range(24): + edges.append( + Edge( + source=EdgeConnection(node_id=previous, field="value"), + destination=EdgeConnection(node_id=f"if-{i}", field="true_input"), + ) + ) + previous = f"if-{i}" + state = GraphExecutionState(graph=Graph()) + state.graph = Graph(nodes=nodes, edges=edges) + original = GraphExecutionState._iterator_graph + builds = [] + + def count_build(self, graph=None): + builds.append(graph) + return original(self, graph) + + monkeypatch.setattr(GraphExecutionState, "_iterator_graph", count_build) + assert graph_if_runtime._get_if_iteration_paths(state, previous, set(), lambda _: set(), lambda *_: ()) == [] + assert len(builds) == 1 + + +def test_resolved_literal_if_does_not_build_iterator_graph(monkeypatch: pytest.MonkeyPatch) -> None: + state = GraphExecutionState(graph=Graph(nodes={"if": IfInvocation(id="if", condition=True, true_input=42)})) + + def unexpected_build(*_args, **_kwargs): + pytest.fail("resolved paths must not build the iterator graph") + + monkeypatch.setattr(GraphExecutionState, "_iterator_graph", unexpected_build) + assert graph_if_runtime._get_if_iteration_paths(state, "if", set(), lambda _: set(), lambda *_: ()) == [()] + + +@pytest.mark.parametrize("executed", [False, True]) +def test_prepared_selected_source_preserves_paths_without_ancestry( + executed: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + state = GraphExecutionState( + graph=Graph(nodes={"if": IfInvocation(id="if", condition=True), "source": AddInvocation(id="source")}) + ) + state.graph.add_edge( + Edge( + source=EdgeConnection(node_id="source", field="value"), + destination=EdgeConnection(node_id="if", field="true_input"), + ) + ) + state.source_prepared_mapping["source"] = set() if executed else {"source-exec"} + if executed: + state.executed.add("source") + + def unexpected_build(*_args, **_kwargs): + pytest.fail("prepared selected sources must not build the iterator graph") + + monkeypatch.setattr(GraphExecutionState, "_iterator_graph", unexpected_build) + assert graph_if_runtime._get_if_iteration_paths( + state, "if", set(), lambda source_id: state.source_prepared_mapping[source_id], lambda *_: (2,) + ) == ([] if executed else [(2,)]) diff --git a/tests/app/services/shared/test_graph_iterator_ancestry.py b/tests/app/services/shared/test_graph_iterator_ancestry.py new file mode 100644 index 00000000000..931a3c69496 --- /dev/null +++ b/tests/app/services/shared/test_graph_iterator_ancestry.py @@ -0,0 +1,89 @@ +import random + +import networkx as nx +import pytest + +from invokeai.app.invocations.loops import ForInvocation +from invokeai.app.invocations.math import AddInvocation +from invokeai.app.services.shared.graph import Edge, EdgeConnection, Graph, GraphExecutionState +from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry +from invokeai.app.services.shared.graph_validation import CollectInvocation, IterateInvocation + + +@pytest.mark.parametrize("seed", range(5)) +def test_iterator_ancestry_matches_filtered_networkx_on_generated_dags(seed: int) -> None: + randomizer = random.Random(seed) + classes = (AddInvocation, IterateInvocation, ForInvocation, CollectInvocation) + nodes = {str(i): randomizer.choice(classes)(id=str(i)) for i in range(40)} + edges = [] + for i in range(40): + for j in range(i + 1, 40): + if randomizer.random() < 0.12: + source = nodes[str(i)] + field = randomizer.choice(("item", "final_state")) if isinstance(source, ForInvocation) else "value" + edges.append( + Edge( + source=EdgeConnection(node_id=str(i), field=field), + destination=EdgeConnection(node_id=str(j), field="a"), + ) + ) + # These are topology fixtures, not executable invocation workflows. + state = GraphExecutionState(graph=Graph()) + state.graph = Graph(nodes=nodes, edges=edges) + source_graph = state._get_source_graph_flat() + for graph in (source_graph, nx.subgraph_view(source_graph, filter_edge=lambda u, v: (int(u) + int(v)) % 3 != 0)): + filtered = state._iterator_graph(graph) + ancestry = _IteratorAncestry(state, graph) + for node_id in reversed(list(graph)): + expected = any( + isinstance(nodes[ancestor], (ForInvocation, IterateInvocation)) + for ancestor in nx.ancestors(filtered, node_id) + ) + assert ancestry.has_iterator_ancestor(node_id) == expected + assert ancestry.has_iterator_ancestor(node_id) == expected + + +def test_iterator_self_collectors_and_final_for_outputs_are_not_ancestors() -> None: + nodes = { + "iterate": IterateInvocation(id="iterate"), + "for": ForInvocation(id="for"), + "collect": CollectInvocation(id="collect"), + "after_collect": AddInvocation(id="after_collect"), + "final": AddInvocation(id="final"), + "body": AddInvocation(id="body"), + } + edges = [ + Edge( + source=EdgeConnection(node_id=source, field=field), + destination=EdgeConnection(node_id=destination, field="a"), + ) + for source, field, destination in ( + ("iterate", "item", "collect"), + ("collect", "collection", "after_collect"), + ("for", "final_state", "final"), + ("for", "item", "body"), + ) + ] + state = GraphExecutionState(graph=Graph()) + state.graph = Graph(nodes=nodes, edges=edges) + ancestry = _IteratorAncestry(state) + assert {node_id: ancestry.has_iterator_ancestor(node_id) for node_id in nodes} == { + "iterate": False, + "for": False, + "collect": False, + "after_collect": False, + "final": False, + "body": True, + } + + +def test_new_context_observes_supplied_graph_mutations() -> None: + state = GraphExecutionState(graph=Graph()) + state.graph = Graph(nodes={"iterator": IterateInvocation(id="iterator"), "sink": AddInvocation(id="sink")}) + graph = nx.DiGraph() + graph.add_nodes_from(state.graph.nodes) + assert not _IteratorAncestry(state, graph).has_iterator_ancestor("sink") + graph.add_edge("iterator", "sink") + assert _IteratorAncestry(state, graph).has_iterator_ancestor("sink") + graph.remove_edge("iterator", "sink") + assert not _IteratorAncestry(state, graph).has_iterator_ancestor("sink") diff --git a/tests/app/services/shared/test_graph_materializer.py b/tests/app/services/shared/test_graph_materializer.py index e407af81998..941aaaa293c 100644 --- a/tests/app/services/shared/test_graph_materializer.py +++ b/tests/app/services/shared/test_graph_materializer.py @@ -68,3 +68,55 @@ def test_materializer_preserves_execution_edges_and_source_nodes() -> None: assert sink.a == 0 state.complete(prepared_sink.id, IntegerOutput(value=12)) assert state.is_complete() + + +def test_prepare_shares_iterator_ancestry_across_admission_and_materialization(monkeypatch: pytest.MonkeyPatch) -> None: + from invokeai.app.invocations.logic import IfInvocation + from invokeai.app.services.shared.graph_validation import IterateInvocation, nx + + nodes = { + "iterator": IterateInvocation(id="iterator"), + "condition-a": AddInvocation(id="condition-a"), + "condition-b": AddInvocation(id="condition-b"), + "inner-a": IfInvocation(id="inner-a"), + "inner-b": IfInvocation(id="inner-b"), + "blocked-a": IfInvocation(id="blocked-a", condition=True), + "blocked-b": IfInvocation(id="blocked-b", condition=True), + "selected": IfInvocation(id="selected", condition=True), + } + edges = [ + graph_facade.Edge( + source=graph_facade.EdgeConnection(node_id=source, field=field), + destination=graph_facade.EdgeConnection(node_id=destination, field=input_field), + ) + for source, field, destination, input_field in ( + ("iterator", "item", "inner-a", "true_input"), + ("condition-a", "value", "inner-a", "condition"), + ("condition-b", "value", "inner-b", "condition"), + ("inner-a", "value", "blocked-a", "true_input"), + ("inner-a", "value", "blocked-b", "true_input"), + ("inner-b", "value", "selected", "true_input"), + ) + ] + state = graph_facade.GraphExecutionState(graph=graph_facade.Graph()) + state.graph = graph_facade.Graph(nodes=nodes, edges=edges) + state.executed.add("iterator") + graph = state._get_source_graph_flat() + # Keep condition sources outside this preparation view; their If nodes stay unresolved. + preparation_graph = nx.subgraph_view(graph, filter_node=lambda node_id: not node_id.startswith("condition-")) + original = graph_facade.GraphExecutionState._iterator_graph + builds = [] + + def count_build(self, source_graph=None): + builds.append(source_graph) + return original(self, source_graph) + + monkeypatch.setattr(graph_facade.GraphExecutionState, "_iterator_graph", count_build) + prepared_id = state._materializer().prepare(preparation_graph) + assert prepared_id in state.source_prepared_mapping["selected"] + assert len(state.source_prepared_mapping["selected"]) == 1 + assert "blocked-a" not in state.source_prepared_mapping + assert "blocked-b" not in state.source_prepared_mapping + assert "inner-a" not in state.source_prepared_mapping + assert "inner-b" not in state.source_prepared_mapping + assert len(builds) == 1