Honor supplied graph in If path fallback

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)

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:

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,
+    ) == [()]