diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index 0496f0d4bf1..c2ea5d9fa9a 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -18,6 +18,7 @@ def _get_effective_iteration_paths_for_edge( visited_source_ids: set[str], 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, ) -> set[tuple[int, ...]]: source_node_id = edge.source.node_id prepared_nodes: set[str] = set() @@ -39,6 +40,7 @@ def _get_effective_iteration_paths_for_edge( visited_source_ids, get_prepared_nodes_for_source, get_prepared_edge_iteration_path, + source_graph, ) ) @@ -54,6 +56,7 @@ def _get_effective_iteration_paths_for_edge( next_visited_source_ids, get_prepared_nodes_for_source, get_prepared_edge_iteration_path, + source_graph, ) for input_edge in input_edges ) @@ -66,6 +69,7 @@ def _get_if_iteration_paths( visited_source_ids: set[str], 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, ) -> list[tuple[int, ...]]: """Infer the selected branch's effective iteration paths for an If output.""" node = state.graph.get_node(node_id) @@ -108,6 +112,7 @@ def _get_if_iteration_paths( next_visited_source_ids, get_prepared_nodes_for_source, get_prepared_edge_iteration_path, + source_graph, ) for edge in selected_edges ) @@ -127,7 +132,9 @@ 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(state._get_source_graph_flat()) + 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) diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 33f5f840741..26a6c34c91c 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -1425,6 +1425,7 @@ def _has_admitted_source_mapping(self, node_id: str, graph: "nx.DiGraph") -> boo set(), self._get_prepared_nodes_for_source, self._get_prepared_edge_iteration_path, + graph, ) ) mappings = self._get_if_condition_iteration_mappings(node_id, graph) @@ -1767,6 +1768,7 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: set(), self._get_prepared_nodes_for_source, self._get_prepared_edge_iteration_path, + g, ) 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 2ff72717bf0..ad596f2812a 100644 --- a/tests/app/services/shared/test_graph_if_runtime.py +++ b/tests/app/services/shared/test_graph_if_runtime.py @@ -1,10 +1,13 @@ import pytest -from invokeai.app.invocations.logic import IfInvocation +from invokeai.app.invocations.logic import BooleanLogicInvocation, IfInvocation from invokeai.app.invocations.math import AddInvocation +from invokeai.app.invocations.primitives import BooleanOutput +from invokeai.app.services.shared import graph_if_runtime from invokeai.app.services.shared.execution_engine.scheduler import ActivationDependency from invokeai.app.services.shared.graph import Edge, EdgeConnection, Graph, GraphExecutionState from invokeai.app.services.shared.graph_runtime_records import _ApplyTransaction +from invokeai.app.services.shared.graph_validation import IterateInvocation from tests.app.services.shared.import_test_utils import assert_module_imports_without_graph @@ -103,3 +106,69 @@ def test_apply_rollback_reconstructs_pending_if_nodes(state: GraphExecutionState state._reset_apply_derived_caches() assert state._pending_if_exec_nodes == {"if-exec"} + + +def test_nested_if_iteration_paths_use_supplied_source_graph_recursively() -> None: + state = GraphExecutionState( + graph=Graph( + nodes={ + "if": IfInvocation(id="if", condition=True), + "inner": IfInvocation(id="inner", condition=True), + "guarded": IfInvocation(id="guarded"), + "condition": BooleanLogicInvocation(id="condition", a=True, b=True), + "bridge": AddInvocation(id="bridge"), + } + ) + ) + state.execution_graph.add_node(BooleanLogicInvocation(id="condition-exec", a=True, b=True)) + state._prepared_registry().register("condition-exec", "condition") + state.results["condition-exec"] = BooleanOutput(value=True) + state.graph.add_node(IterateInvocation(id="iterator")) + state.graph.edges.extend( + [ + Edge( + source=EdgeConnection(node_id="iterator", field="item"), + destination=EdgeConnection(node_id="guarded", field="condition"), + ), + Edge( + source=EdgeConnection(node_id="guarded", field="value"), + destination=EdgeConnection(node_id="inner", field="true_input"), + ), + Edge( + source=EdgeConnection(node_id="inner", field="value"), + destination=EdgeConnection(node_id="bridge", field="a"), + ), + Edge( + source=EdgeConnection(node_id="bridge", field="value"), + destination=EdgeConnection(node_id="if", field="true_input"), + ), + Edge( + source=EdgeConnection(node_id="condition", field="value"), + destination=EdgeConnection(node_id="if", field="condition"), + ), + ] + ) + state.graph._invalidate_edge_indexes() + source_graph = state._get_source_graph_flat() + constrained_graph = source_graph.copy() + constrained_graph.remove_edge("iterator", "guarded") + + assert ( + graph_if_runtime._get_if_iteration_paths( + state, + "if", + set(), + lambda source_id: {"condition-exec"} if source_id == "condition" else set(), + lambda _edge, _prepared_id: (), + source_graph, + ) + == [] + ) + assert graph_if_runtime._get_if_iteration_paths( + state, + "if", + set(), + lambda source_id: {"condition-exec"} if source_id == "condition" else set(), + lambda _edge, _prepared_id: (), + constrained_graph, + ) == [()]