From d61b90aca8c0fd7caf419cc0309c87c81566e4d4 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sat, 10 Oct 2026 16:17:17 -0400 Subject: [PATCH 1/5] fix(graph): stop If branches fed by a collected Iterate from deadlocking An If no longer waits for an unexecuted iterator whose frame a Collect closes before the If. Branch sources behind a Collect now resolve the If gate in its enclosing frame instead of silently dropping items. Iterators on a rejected branch no longer block downstream nodes. --- .../services/shared/graph_if_activation.py | 20 +- .../app/services/shared/graph_materializer.py | 86 +++++ .../test_execution_engine_differential.py | 327 ++++++++++++++++++ 3 files changed, 432 insertions(+), 1 deletion(-) diff --git a/invokeai/app/services/shared/graph_if_activation.py b/invokeai/app/services/shared/graph_if_activation.py index 80220ee7dda..764773e5a17 100644 --- a/invokeai/app/services/shared/graph_if_activation.py +++ b/invokeai/app/services/shared/graph_if_activation.py @@ -78,11 +78,29 @@ def get_source_dependencies( ActivationDependency( owner_id=source_if_id, branch=matching_fields[0], - frame=iteration_path, + frame=self._get_owner_frame(source_if_id, iteration_path), ) ) return tuple(dependencies) + def _get_owner_frame(self, if_node_id: str, iteration_path: tuple[int, ...]) -> tuple[int, ...]: + """Return the If frame that gates a branch source at ``iteration_path``. + + A Collect between a branch source and its If closes the source's inner frames, so the source runs deeper + than the If that admits it. Its gate is the If execution whose frame is the longest prefix of the source's + path. Until such an If exists the source's own path is used, which no gate satisfies. + """ + + owner_frames = { + self._state._get_iteration_path(exec_node_id) + for exec_node_id in self._state._prepared_registry().get_prepared_ids(if_node_id) + } + return max( + (frame for frame in owner_frames if iteration_path[: len(frame)] == frame), + key=len, + default=iteration_path, + ) + def is_source_admitted(self, source_node_id: str, iteration_path: tuple[int, ...] = ()) -> bool: dependencies = self.get_source_dependencies(source_node_id, iteration_path) return bool( diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 33f5f840741..5bf0ecf1459 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -1583,6 +1583,90 @@ def _attach_pending_if_inputs(self, *, enqueue: bool = True) -> None: if enqueue: self._state._enqueue_if_ready(exec_node_id) + def _is_closed_branch_iterator(self, iterator_id: str, node_id: str) -> bool: + """Return whether an If may be materialized before ``iterator_id`` executes. + + An If needs only its condition to be prepared, but otherwise waits for upstream iterators so it inherits their + frames. An iterator inside one of the If's branches waits for the If to select that branch, so waiting for it + deadlocks. An iterator cannot change the If's frames when its own frame closes before the If, and every + enclosing frame that stays open at the If also reaches it through the condition, which the If already has. + """ + + if not isinstance(self._state.graph.get_node(node_id), IfInvocation): + return False + if self._is_iterator_frame_open_at(iterator_id, node_id): + return False + condition_source_ids = [ + edge.source.node_id for edge in self._state.graph._get_input_edges(node_id, "condition") + ] + return all( + not self._is_iterator_frame_open_at(enclosing_id, node_id) + or any( + enclosing_id == condition_source_id + or self._is_iterator_frame_open_at(enclosing_id, condition_source_id) + for condition_source_id in condition_source_ids + ) + for enclosing_id in nx.ancestors(self._state._get_source_graph_flat(), iterator_id) + if isinstance(self._state.graph.get_node(enclosing_id), (ForInvocation, IterateInvocation)) + ) + + def _is_iterator_frame_open_at(self, iterator_id: str, node_id: str) -> bool: + """Return whether some path carries ``iterator_id``'s iteration frame into ``node_id``. + + Depth counts frames opened since ``iterator_id``. Per-item outputs of an Iterate or For open one. An item input + of a Collect closes one, except a For's final output, which is already in the For's parent frame and which the + Collect groups there. A Collect's collection input keeps its source's frame. A Collect that would close + ``iterator_id``'s frame counts as leaving it open unless ``iterator_id`` is the item's innermost iterator. + """ + + pending = [(iterator_id, 0)] + seen: set[tuple[str, int]] = set() + while pending: + current_id, depth = pending.pop() + current_node = self._state.graph.get_node(current_id) + for edge in self._state.graph._get_output_edges(current_id): + if edge.type != "default": + continue + is_for_final = ( + isinstance(current_node, ForInvocation) + and get_output_field_scope(current_node, edge.source.field) == OutputScope.Final + ) + next_depth = depth + if isinstance(current_node, (ForInvocation, IterateInvocation)) and not is_for_final: + next_depth += 1 + if ( + isinstance(self._state.graph.get_node(edge.destination.node_id), CollectInvocation) + and edge.destination.field == ITEM_FIELD + and not is_for_final + ): + next_depth -= 1 + if next_depth <= 0: + if not self._is_innermost_iterator(iterator_id, current_id): + return True + continue + if edge.destination.node_id == node_id: + return True + state = (edge.destination.node_id, next_depth) + if state not in seen: + seen.add(state) + pending.append(state) + return False + + def _is_innermost_iterator(self, iterator_id: str, collected_source_id: str) -> bool: + """Return whether ``iterator_id`` owns the innermost frame of a value entering a Collect. + + A Collect closes only the innermost frame of its item. A sibling iterator combined into the same value adds a + frame the depth count cannot order, so the Collect may close the sibling's frame instead. + """ + + source_iterators = self.get_node_iterators(collected_source_id) + if isinstance(self._state.graph.get_node(collected_source_id), (ForInvocation, IterateInvocation)): + source_iterators.append(collected_source_id) + source_graph = self._state._get_source_graph_flat() + return all( + other_id == iterator_id or nx.has_path(source_graph, other_id, iterator_id) for other_id in source_iterators + ) + 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() @@ -1741,6 +1825,8 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: and not any( isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) and ancestor_id not in self._state.executed + and not self._state._is_source_inactive(ancestor_id) + and not self._is_closed_branch_iterator(ancestor_id, node_id) for ancestor_id in nx.ancestors(g, node_id) ) ), diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index ce89287c311..385a5703f35 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -3236,6 +3236,333 @@ def test_nested_iterate_keeps_values_in_outer_frame_with_unrelated_if( assert state.is_complete() +def _collected_iterate_branch_graph(condition: bool, condition_from_edge: bool = False) -> Graph: + """src -> iterate -> body -> collect -> if.true_input, the shape reported in #9484.""" + graph = Graph() + graph.add_node(FixedCollectionTestInvocation(id="src", value=True)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=False if condition_from_edge else condition)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("src", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + if condition_from_edge: + graph.add_node(BooleanInvocation(id="condition", value=condition)) + graph.add_edge(create_edge("condition", "value", "if", "condition")) + return graph + + +def _doubly_collected_nested_iterate_branch_graph(condition: bool) -> Graph: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=["x", "y"])) + graph.add_node(IterateInvocation(id="outer_iterate")) + graph.add_node(FixedCollectionTestInvocation(id="inner_source")) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(CollectInvocation(id="outer_collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("outer_source", "collection", "outer_iterate", "collection")) + graph.add_edge(create_edge("outer_iterate", "item", "inner_source", "value")) + graph.add_edge(create_edge("inner_source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "outer_collect", "item")) + graph.add_edge(create_edge("outer_collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +def _for_output_iterate_branch_graph(condition: bool) -> Graph: + graph = Graph() + graph.add_node(ForInvocation(id="for", collection=["a", "b"])) + graph.add_node(ForReturnInvocation(id="return")) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("for", "item", "return", "output")) + graph.add_edge(create_loop_linkage("for", "return")) + graph.add_edge(create_edge("for", "output_collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +@pytest.mark.parametrize("condition", [True, False]) +@pytest.mark.parametrize( + ("build_graph", "selected_value"), + [ + pytest.param(_collected_iterate_branch_graph, [0, 1], id="literal-condition"), + pytest.param( + lambda condition: _collected_iterate_branch_graph(condition, condition_from_edge=True), + [0, 1], + id="edge-condition", + ), + pytest.param(_doubly_collected_nested_iterate_branch_graph, [[0, 1], [0, 1]], id="two-collects"), + pytest.param(_for_output_iterate_branch_graph, ["a", "b"], id="for-output"), + ], +) +def test_if_branch_fed_by_closed_iterate_completes( + build_graph: Callable[[bool], Graph], + selected_value: list[Any], + condition: bool, + force_compatibility_scheduler: bool, +) -> None: + """An Iterate whose frame closes before the If must neither deadlock nor lose its items (#9484).""" + graph = build_graph(condition) + graph.validate_self() + + trace, state = _run(GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete() + sink_exec_ids = state._prepared_registry().get_prepared_ids("sink") + assert [state.results[exec_id].value for exec_id in sink_exec_ids] == [ + selected_value if condition else ["fallback"] + ] + unselected_ids = {"fallback"} if condition else {"iterate", "body", "collect"} + assert not unselected_ids & set(trace) + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +@pytest.mark.parametrize( + ("outer_items", "condition_source", "expected_sinks"), + [ + pytest.param(["x", ""], ("inner_source", "condition"), {(0,): [0, 1], (1,): ""}, id="derived-condition"), + pytest.param([True, False], ("outer_iterate", "item"), {(0,): [0, 1], (1,): False}, id="item-condition"), + ], +) +def test_per_item_if_branch_fed_by_collected_inner_iterate( + outer_items: list[Any], + condition_source: tuple[str, str], + expected_sinks: dict[tuple[int, ...], Any], + force_compatibility_scheduler: bool, +) -> None: + """A per-item condition gives the If its frames, so an inner Iterate closed before the If cannot block it.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=outer_items)) + graph.add_node(IterateInvocation(id="outer_iterate")) + graph.add_node(FixedCollectionTestInvocation(id="inner_source")) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("outer_source", "collection", "outer_iterate", "collection")) + graph.add_edge(create_edge("outer_iterate", "item", "inner_source", "value")) + graph.add_edge(create_edge(*condition_source, "if", "condition")) + graph.add_edge(create_edge("inner_source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("outer_iterate", "item", "fallback", "value")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + graph.validate_self() + + _trace, state = _run(GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete() + assert { + state._get_iteration_path(exec_id): state.results[exec_id].value + for exec_id in state._prepared_registry().get_prepared_ids("sink") + } == expected_sinks + assert {state._get_iteration_path(exec_id) for exec_id in state._prepared_registry().get_prepared_ids("body")} == { + (0, 0), + (0, 1), + } + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +@pytest.mark.parametrize("selected_branch", ["collect", "collect2"]) +def test_if_branches_fed_by_collects_of_one_shared_iterate( + selected_branch: str, force_compatibility_scheduler: bool +) -> None: + graph = Graph() + graph.add_node(FixedCollectionTestInvocation(id="src")) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(AnyTypeTestInvocation(id="body2")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(CollectInvocation(id="collect2")) + graph.add_node(IfInvocation(id="if", condition=selected_branch == "collect")) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("src", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "body", "value")) + graph.add_edge(create_edge("iterate", "item", "body2", "value")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("body2", "value", "collect2", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("collect2", "collection", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + graph.validate_self() + + trace, state = _run(GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete() + assert [state.results[exec_id].value for exec_id in state._prepared_registry().get_prepared_ids("sink")] == [[0, 1]] + unselected_branch = "collect2" if selected_branch == "collect" else "collect" + assert selected_branch in trace + assert unselected_branch not in trace + + +def _per_item_collection_input_branch_graph() -> Graph: + """A Collect's collection input keeps its per-item frame.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=["x", "y"])) + graph.add_node(IterateInvocation(id="outer_iterate")) + graph.add_node(AnyTypeTestInvocation(id="side")) + graph.add_node(FixedCollectionTestInvocation(id="items")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=True)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("outer_source", "collection", "outer_iterate", "collection")) + graph.add_edge(create_edge("outer_iterate", "item", "side", "value")) + graph.add_edge(create_edge("outer_iterate", "item", "items", "value")) + graph.add_edge(create_edge("items", "collection", "collect", "collection")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +def _per_item_for_output_collect_branch_graph() -> Graph: + """A Collect groups a For's final output in the For's parent frame, so the outer item frame stays open.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=[["a", "b"], ["c"]])) + graph.add_node(IterateInvocation(id="outer_iterate")) + graph.add_node(ForInvocation(id="for")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(ForReturnInvocation(id="return")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=True)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("outer_source", "collection", "outer_iterate", "collection")) + graph.add_edge(create_edge("outer_iterate", "item", "for", "collection")) + graph.add_edge(create_edge("for", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "return", "output")) + graph.add_edge(create_loop_linkage("for", "return")) + graph.add_edge(create_edge("for", "output_collection", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +@pytest.mark.parametrize( + ("build_graph", "expected_sinks"), + [ + pytest.param(_per_item_collection_input_branch_graph, {(0,): [0, 1], (1,): [0, 1]}, id="collection-input"), + pytest.param(_per_item_for_output_collect_branch_graph, {(0,): [["a", "b"]], (1,): [["c"]]}, id="for-output"), + ], +) +def test_if_branch_through_collect_that_keeps_item_frame_runs_per_item( + build_graph: Callable[[], Graph], + expected_sinks: dict[tuple[int, ...], Any], + force_compatibility_scheduler: bool, +) -> None: + graph = build_graph() + graph.validate_self() + + _trace, state = _run(GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete() + assert { + state._get_iteration_path(exec_id): state.results[exec_id].value + for exec_id in state._prepared_registry().get_prepared_ids("sink") + } == expected_sinks + + +def _shared_outer_iterate_branch_graph(condition: bool) -> Graph: + """The outer frame reaches the If only through the deferred branch, so the If cannot see it yet.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=["x", "y"])) + graph.add_node(IterateInvocation(id="outer_iterate")) + graph.add_node(AnyTypeTestInvocation(id="side")) + graph.add_node(FixedCollectionTestInvocation(id="inner_source")) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("outer_source", "collection", "outer_iterate", "collection")) + graph.add_edge(create_edge("outer_iterate", "item", "side", "value")) + graph.add_edge(create_edge("outer_iterate", "item", "inner_source", "value")) + graph.add_edge(create_edge("inner_source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +def _sibling_iterates_branch_graph(condition: bool) -> Graph: + """The Collect closes only one of two sibling frames, so the If inherits the other.""" + graph = Graph() + graph.add_node(FixedCollectionTestInvocation(id="left_source")) + graph.add_node(FixedCollectionTestInvocation(id="right_source")) + graph.add_node(IterateInvocation(id="left_iterate")) + graph.add_node(IterateInvocation(id="right_iterate")) + graph.add_node(AddInvocation(id="body")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("left_source", "collection", "left_iterate", "collection")) + graph.add_edge(create_edge("right_source", "collection", "right_iterate", "collection")) + graph.add_edge(create_edge("left_iterate", "item", "body", "a")) + graph.add_edge(create_edge("right_iterate", "item", "body", "b")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +@pytest.mark.parametrize("condition", [True, False]) +@pytest.mark.parametrize( + "build_graph", + [ + pytest.param(_shared_outer_iterate_branch_graph, id="outer-frame-through-branch"), + pytest.param(_sibling_iterates_branch_graph, id="sibling-frame-left-open"), + ], +) +def test_if_is_not_materialized_in_parent_frame_while_branch_frame_is_open( + build_graph: Callable[[bool], Graph], condition: bool, force_compatibility_scheduler: bool +) -> None: + """The If must not be materialized in the parent frame while an unexecuted branch can still give it item frames.""" + graph = build_graph(condition) + graph.validate_self() + + _trace, state = _run(GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler) + + assert () not in { + state._get_iteration_path(exec_id) for exec_id in state._prepared_registry().get_prepared_ids("if") + } + + @pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) def test_if_collect_collection_inner_iterate_keeps_values_in_outer_frame( force_compatibility_scheduler: bool, From e8ff5dfcbcbe39f74a8c07df85bef1ca73b804c2 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Sun, 11 Oct 2026 10:29:07 -0500 Subject: [PATCH 2/5] test(graph): cover remaining If iterator frame cases Add regression cases for open iterator branches and root-frame If inputs. --- .../test_execution_engine_differential.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index 385a5703f35..d6b10a427a7 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -3563,6 +3563,99 @@ def test_if_is_not_materialized_in_parent_frame_while_branch_frame_is_open( } +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +def test_if_branch_fed_directly_by_open_iterate_runs_for_each_item( + force_compatibility_scheduler: bool, +) -> None: + graph = Graph() + graph.add_node(FixedCollectionTestInvocation(id="source", value=True)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(IfInvocation(id="if", condition=True)) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "true_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + assert {state._get_iteration_path(exec_id) for exec_id in state._prepared_registry().get_prepared_ids("if")} == { + (0,), + (1,), + } + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == [0, 1] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +def test_per_item_if_can_select_values_from_root_frame( + force_compatibility_scheduler: bool, +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="conditions", first=[True, False])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="root_true", value="true-root")) + graph.add_node(AnyTypeTestInvocation(id="root_false", value="false-root")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("conditions", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "condition")) + graph.add_edge(create_edge("root_true", "value", "if", "true_input")) + graph.add_edge(create_edge("root_false", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + assert {state._get_iteration_path(exec_id) for exec_id in state._prepared_registry().get_prepared_ids("if")} == { + (0,), + (1,), + } + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == ["true-root", "false-root"] + + +@pytest.mark.parametrize( + ("condition", "expected_paths", "expected_values"), + [ + pytest.param(True, {(0,), (1,)}, ["item:0", "item:1"], id="selected-item-frame"), + pytest.param(False, {()}, ["root-value"], id="selected-root-frame"), + ], +) +def test_literal_if_result_uses_selected_branch_frame_cardinality( + condition: bool, + expected_paths: set[tuple[int, ...]], + expected_values: list[str], +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[0, 1])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(MarkedAnyTypeTestInvocation(id="item_branch", marker="item:")) + graph.add_node(AnyTypeTestInvocation(id="root_branch", value="root-value")) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "item_branch", "value")) + graph.add_edge(create_edge("item_branch", "value", "if", "true_input")) + graph.add_edge(create_edge("root_branch", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph(GraphExecutionState(graph=graph)) + + assert state.is_complete(), f"execution stalled after {trace}" + if_paths = {state._get_iteration_path(exec_id) for exec_id in state._prepared_registry().get_prepared_ids("if")} + assert if_paths == expected_paths + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == expected_values + + @pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) def test_if_collect_collection_inner_iterate_keeps_values_in_outer_frame( force_compatibility_scheduler: bool, From e27f6fde3592d6f8355cb4fec6ac36669042a761 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 11 Oct 2026 16:07:06 -0400 Subject: [PATCH 3/5] fix(graph): decide If branches from the condition when no If execution can gate them Branch sources of an If whose frames depend on the selected branch, or that run shallower than every If frame, are admitted from the condition value; per-item Ifs broadcast parent-frame values. Branch frames are matched to If frames by structural depth, rejected iterators no longer add frames, and empty branches or conditions complete instead of stalling. An If whose condition and branch iterate unrelated collections raises UnsupportedIfFramesError instead of failing with a KeyError. --- .../services/shared/graph_if_activation.py | 194 ++++++- .../app/services/shared/graph_if_runtime.py | 25 + .../app/services/shared/graph_materializer.py | 127 ++++- .../test_execution_engine_differential.py | 484 +++++++++++++++++- 4 files changed, 778 insertions(+), 52 deletions(-) diff --git a/invokeai/app/services/shared/graph_if_activation.py b/invokeai/app/services/shared/graph_if_activation.py index 764773e5a17..50ee383a630 100644 --- a/invokeai/app/services/shared/graph_if_activation.py +++ b/invokeai/app/services/shared/graph_if_activation.py @@ -1,11 +1,24 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional +from invokeai.app.invocations.fields import OutputScope from invokeai.app.invocations.logic import IfInvocation from invokeai.app.services.shared.execution_engine.scheduler import ActivationDependency -from invokeai.app.services.shared.graph_validation import nx +from invokeai.app.services.shared.graph_validation import ( + ITEM_FIELD, + CollectInvocation, + ForInvocation, + IterateInvocation, + get_output_field_scope, + nx, +) if TYPE_CHECKING: from invokeai.app.services.shared.graph import GraphExecutionState + from invokeai.app.services.shared.graph_models import Edge + + +class UnsupportedIfFramesError(ValueError): + """An If whose condition and branch iterate over unrelated collections has no defined set of frames.""" class _IfActivationController: @@ -19,6 +32,8 @@ class _IfActivationController: def __init__(self, state: "GraphExecutionState") -> None: self._state = state self._branch_sources_cache: dict[tuple[str, str], frozenset[str]] = {} + self._branch_framed_cache: dict[str, bool] = {} + self._frame_depths: Optional[dict[str, int]] = None def _branch_sources(self, if_node_id: str, branch_field: str, source_graph: "nx.DiGraph") -> set[str]: cache_key = (if_node_id, branch_field) @@ -78,28 +93,72 @@ def get_source_dependencies( ActivationDependency( owner_id=source_if_id, branch=matching_fields[0], - frame=self._get_owner_frame(source_if_id, iteration_path), + frame=self._get_owner_frame(source_node_id, source_if_id, matching_fields[0], iteration_path), ) ) return tuple(dependencies) - def _get_owner_frame(self, if_node_id: str, iteration_path: tuple[int, ...]) -> tuple[int, ...]: - """Return the If frame that gates a branch source at ``iteration_path``. + def _get_owner_frame( + self, source_node_id: str, if_node_id: str, branch_field: str, iteration_path: tuple[int, ...] + ) -> tuple[int, ...]: + """Return the part of a branch source's frame that the If shares. - A Collect between a branch source and its If closes the source's inner frames, so the source runs deeper - than the If that admits it. Its gate is the If execution whose frame is the longest prefix of the source's - path. Until such an If exists the source's own path is used, which no gate satisfies. + A Collect between a branch source and its If closes the source's innermost frames, so the source runs deeper + than the If that admits it. Only the leading indexes up to the branch's depth at the If are shared; the rest + belong to closed frames and must not be compared with If frames from other iterators, which reuse the same + indexes. """ - owner_frames = { - self._state._get_iteration_path(exec_node_id) - for exec_node_id in self._state._prepared_registry().get_prepared_ids(if_node_id) - } - return max( - (frame for frame in owner_frames if iteration_path[: len(frame)] == frame), - key=len, - default=iteration_path, + if not iteration_path or not any( + isinstance(self._state.graph.nodes.get(node_id), (ForInvocation, IterateInvocation)) + for node_id in (*nx.ancestors(self._state._get_source_graph_flat(), source_node_id), source_node_id) + ): + return iteration_path + branch_depth = max( + (self._get_edge_depth(edge) for edge in self._state.graph._get_input_edges(if_node_id, branch_field)), + default=0, ) + return iteration_path[:branch_depth] + + def _get_edge_depth(self, edge: "Edge", depths: Optional[dict[str, int]] = None) -> int: + """Return the frame depth of the value an edge carries; a For's final output is in the For's parent frame.""" + + source_node = self._state.graph.get_node(edge.source.node_id) + depth = (depths if depths is not None else self._get_frame_depths())[edge.source.node_id] + if isinstance(source_node, ForInvocation) and get_output_field_scope(source_node, edge.source.field) == ( + OutputScope.Final + ): + return max(depth - 1, 0) + return depth + + def _get_frame_depths(self) -> dict[str, int]: + """Return the length of the iteration path each source node executes at. + + An Iterate or For runs once per item, one level below its input. A Collect's item input closes one level. + Other nodes run at the deepest level among their inputs. + """ + + if self._frame_depths is not None: + return self._frame_depths + depths: dict[str, int] = {} + for node_id in nx.topological_sort(self._state._get_source_graph_flat()): + node = self._state.graph.get_node(node_id) + input_edges = [edge for edge in self._state.graph._get_input_edges(node_id) if edge.type == "default"] + if isinstance(node, CollectInvocation): + depth = max( + ( + self._get_edge_depth(edge, depths) - (1 if edge.destination.field == ITEM_FIELD else 0) + for edge in input_edges + ), + default=0, + ) + else: + depth = max((self._get_edge_depth(edge, depths) for edge in input_edges), default=0) + if isinstance(node, (ForInvocation, IterateInvocation)): + depth += 1 + depths[node_id] = max(depth, 0) + self._frame_depths = depths + return depths def is_source_admitted(self, source_node_id: str, iteration_path: tuple[int, ...] = ()) -> bool: dependencies = self.get_source_dependencies(source_node_id, iteration_path) @@ -114,11 +173,112 @@ def is_source_rejected(self, source_node_id: str, iteration_path: tuple[int, ... for dependency in self.get_source_dependencies(source_node_id, iteration_path) ) + def get_condition_verdict(self, dependency: ActivationDependency) -> Optional[tuple[bool, bool]]: + """Return ``(satisfied, rejected)`` for a dependency no If execution can gate, else ``None``. + + Such a dependency is decided by the If's condition value in every frame related to the dependency's frame: + a parent-frame condition decides each item below it, and a parent-frame source is selected when any item + frame selects it. + """ + + if not self._is_decided_by_condition(dependency.owner_id, dependency.frame): + return None + if_node = self._state.graph.get_node(dependency.owner_id) + assert isinstance(if_node, IfInvocation) + condition_edges = self._state.graph._get_input_edges(dependency.owner_id, "condition") + if not condition_edges: + selected = "true_input" if if_node.condition else "false_input" + return selected == dependency.branch, selected != dependency.branch + + condition_edge = condition_edges[0] + condition_source_id = condition_edge.source.node_id + if condition_source_id not in self._state.source_prepared_mapping: + return False, False + materializer = self._state._materializer() + selected_branches: set[str] = set() + is_known = True + for prepared_id in materializer._get_prepared_nodes_for_source(condition_source_id): + condition_path = materializer._get_prepared_edge_iteration_path(condition_edge, prepared_id) + shared_depth = min(len(condition_path), len(dependency.frame)) + if condition_path[:shared_depth] != dependency.frame[:shared_depth]: + continue + result = self._state.results.get(prepared_id) + if result is None: + is_known = False + continue + selected_branches.add("true_input" if getattr(result, condition_edge.source.field) else "false_input") + satisfied = dependency.branch in selected_branches + return satisfied, is_known and not satisfied + + def _is_decided_by_condition(self, if_node_id: str, frame: tuple[int, ...]) -> bool: + """Return whether no If execution can gate a branch source running at ``frame``. + + Either the If takes its frames from the branch it selects, or its executions all run in frames below + ``frame``, as when a per-item condition selects between values computed once in the parent frame. + """ + + if self._is_branch_framed(if_node_id) or self._has_empty_condition(if_node_id): + return True + owner_frames = self._state._prepared_if_exec_frames(if_node_id) + return any(owner_frame[: len(frame)] == frame for owner_frame in owner_frames) and not any( + frame[: len(owner_frame)] == owner_frame for owner_frame in owner_frames + ) + + def _has_empty_condition(self, if_node_id: str) -> bool: + """Return whether the condition ran in no frame, so the If never executes and selects neither branch.""" + + condition_edges = self._state.graph._get_input_edges(if_node_id, "condition") + if not condition_edges: + return False + condition_source_id = condition_edges[0].source.node_id + return ( + condition_source_id in self._state.executed + and not self._state._materializer()._get_prepared_nodes_for_source(condition_source_id) + ) + + def _is_branch_framed(self, if_node_id: str) -> bool: + """Return whether the If's frames depend on which branch it selects. + + That holds when an iterator's frame stays open into one branch input but not the other, and the condition + does not carry it. Such an If cannot be materialized until its selected branch runs. + """ + + cached = self._branch_framed_cache.get(if_node_id) + if cached is not None: + return cached + materializer = self._state._materializer() + source_graph = self._state._get_source_graph_flat() + condition_iterator_ids: list[str] = [] + branch_iterator_ids: list[str] = [] + for iterator_id in nx.ancestors(source_graph, if_node_id): + if not isinstance(self._state.graph.get_node(iterator_id), (ForInvocation, IterateInvocation)): + continue + if materializer._is_iterator_frame_open_at(iterator_id, if_node_id, {"condition"}): + condition_iterator_ids.append(iterator_id) + elif materializer._is_iterator_frame_open_at( + iterator_id, if_node_id, {"true_input"} + ) != materializer._is_iterator_frame_open_at(iterator_id, if_node_id, {"false_input"}): + branch_iterator_ids.append(iterator_id) + for branch_iterator_id in branch_iterator_ids: + for condition_iterator_id in condition_iterator_ids: + if not nx.has_path(source_graph, condition_iterator_id, branch_iterator_id) and not nx.has_path( + source_graph, branch_iterator_id, condition_iterator_id + ): + raise UnsupportedIfFramesError( + f"If node {if_node_id} combines a condition iterated by {condition_iterator_id} with a branch " + f"that may carry frames of the unrelated {branch_iterator_id}, which is not supported" + ) + result = bool(branch_iterator_ids) + self._branch_framed_cache[if_node_id] = result + return result + def is_source_inactive(self, source_node_id: str, iteration_path: tuple[int, ...] = ()) -> bool: """Return whether a source belongs exclusively to a rejected If branch.""" + if self.is_source_rejected(source_node_id, iteration_path): + return True if iteration_path: - return self.is_source_rejected(source_node_id, iteration_path) + return False source_graph = self._state._get_source_graph_flat() relevant_if_ids = { diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index 0496f0d4bf1..0b0bff63033 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -114,6 +114,16 @@ def _get_if_iteration_paths( ) else: branch_paths = {condition_path or ()} + if ( + not branch_paths + and condition_path is not None + and any( + edge.source.node_id not in state.source_prepared_mapping and edge.source.node_id not in state.executed + for edge in selected_edges + ) + ): + # The selected branch is not materialized yet; the If still runs in the condition's frame. + branch_paths = {condition_path} for branch_path in branch_paths: if condition_path is None: @@ -237,8 +247,20 @@ def _record_compatibility_activation_token( ) +def _get_condition_verdict( + state: "GraphExecutionState", dependency: ActivationDependency +) -> Optional[tuple[bool, bool]]: + # Fresh shapes never place an If where none of its executions can gate a branch source. + if state._can_use_fresh_flat_if_activation(): + return None + return state._if_activation_controller().get_condition_verdict(dependency) + + def _is_activation_dependency_satisfied(state: "GraphExecutionState", dependency: ActivationDependency) -> bool: """Check one opaque plan requirement against durable gate and token state.""" + verdict = _get_condition_verdict(state, dependency) + if verdict is not None: + return verdict[0] matching_gate_ids = state._prepared_if_exec_ids(dependency.owner_id, dependency.frame) if not matching_gate_ids: return False @@ -276,6 +298,9 @@ def _is_activation_dependency_satisfied(state: "GraphExecutionState", dependency def _is_activation_dependency_rejected(state: "GraphExecutionState", dependency: ActivationDependency) -> bool: """Return whether a resolved gate explicitly selected another branch.""" + verdict = _get_condition_verdict(state, dependency) + if verdict is not None: + return verdict[1] 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) diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 5bf0ecf1459..f21c93624ce 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -1065,7 +1065,7 @@ def _get_collect_iteration_group_key(self, edge: Edge, sibling_depth: Optional[i return path def _get_collect_source_iterator_ids(self, source_node_id: str) -> list[str]: - iterator_node_ids = self.get_node_iterators(source_node_id) + iterator_node_ids = self._get_active_node_iterators(source_node_id) if isinstance(self._state.graph.get_node(source_node_id), IterateInvocation): iterator_node_ids.append(source_node_id) return iterator_node_ids @@ -1319,8 +1319,8 @@ def _get_parent_iteration_mappings( input_edges = input_edges or self._state.graph._get_input_edges(next_node_id) parent_node_ids = list(dict.fromkeys(edge.source.node_id for edge in input_edges)) iterator_graph = self.iterator_graph(graph) - iterator_nodes = self.get_node_iterators(next_node_id, iterator_graph) - if not iterator_nodes: + iterator_nodes = self._get_active_node_iterators(next_node_id, iterator_graph) + if not iterator_nodes or self._has_mixed_depth_if_input(input_edges): return iter(self._get_parent_iteration_mappings_without_iterators(next_node_id, input_edges)) iterator_nodes_prepared = [ @@ -1417,15 +1417,19 @@ def _has_admitted_source_mapping(self, node_id: str, graph: "nx.DiGraph") -> boo if isinstance(self._state.graph.get_node(node_id), IfInvocation): condition_edges = self._state.graph._get_input_edges(node_id, "condition") if not condition_edges: + iteration_paths = graph_if_runtime._get_if_iteration_paths( + self._state, + node_id, + set(), + self._get_prepared_nodes_for_source, + self._get_prepared_edge_iteration_path, + ) + if not iteration_paths: + # A selected branch that has run without producing a frame leaves the If with no executions. + return self._is_selected_branch_resolved_empty(node_id) return any( self._state._is_source_activation_admitted(node_id, iteration_path) - for iteration_path in graph_if_runtime._get_if_iteration_paths( - self._state, - node_id, - set(), - self._get_prepared_nodes_for_source, - self._get_prepared_edge_iteration_path, - ) + for iteration_path in iteration_paths ) mappings = self._get_if_condition_iteration_mappings(node_id, graph) input_edges = condition_edges @@ -1454,6 +1458,35 @@ def _has_admitted_source_mapping(self, node_id: str, graph: "nx.DiGraph") -> boo for iteration_mapping in mappings ) + def _is_selected_branch_resolved_empty(self, node_id: str) -> bool: + """Return whether every branch the condition selects has run without producing a frame.""" + + node = self._state.graph.get_node(node_id) + assert isinstance(node, IfInvocation) + condition_edges = self._state.graph._get_input_edges(node_id, "condition") + if not condition_edges: + selected_fields = {"true_input" if node.condition else "false_input"} + else: + condition_edge = condition_edges[0] + condition_source_id = condition_edge.source.node_id + if condition_source_id not in self._state.executed: + return False + selected_fields = { + "true_input" + if getattr(self._state.results[prepared_id], condition_edge.source.field) + else "false_input" + for prepared_id in self._get_prepared_nodes_for_source(condition_source_id) + } + selected_edges = [ + edge + for selected_field in selected_fields + for edge in self._state.graph._get_input_edges(node_id, selected_field) + ] + return bool(selected_edges) and all( + edge.source.node_id in self._state.executed and not self._get_prepared_nodes_for_source(edge.source.node_id) + for edge in selected_edges + ) + def _get_if_condition_iteration_mappings( self, node_id: str, graph: "nx.DiGraph" ) -> Iterable[list[tuple[str, str]]]: @@ -1482,6 +1515,9 @@ def _get_if_condition_iteration_mappings( ) for branch_field in ("true_input", "false_input") ): + if condition_edges: + # The condition is outside the If's iterator frames, so each If frame reads it from a parent frame. + return iter(self._get_parent_iteration_mappings_without_iterators(node_id, condition_edges)) return iter([[]]) if any(iterator_node_id not in self._state.source_prepared_mapping for iterator_node_id in iterator_nodes): @@ -1547,11 +1583,16 @@ def _attach_pending_if_inputs(self, *, enqueue: bool = True) -> None: for source_edge in selected_edges: if source_edge.source.node_id not in self._state.source_prepared_mapping: continue - candidates = [ - prepared_id - for prepared_id in self._get_ordered_prepared_nodes_for_edge(source_edge) - if self._get_prepared_edge_iteration_path(source_edge, prepared_id) == target_path - ] + # A value prepared in a parent frame is shared by every item frame below it. + candidates = sorted( + ( + prepared_id + for prepared_id in self._get_ordered_prepared_nodes_for_edge(source_edge) + if target_path[: len(self._get_prepared_edge_iteration_path(source_edge, prepared_id))] + == self._get_prepared_edge_iteration_path(source_edge, prepared_id) + ), + key=lambda prepared_id: -len(self._get_prepared_edge_iteration_path(source_edge, prepared_id)), + ) if not candidates: continue selected_exec_ids.append(candidates[0]) @@ -1610,13 +1651,15 @@ def _is_closed_branch_iterator(self, iterator_id: str, node_id: str) -> bool: if isinstance(self._state.graph.get_node(enclosing_id), (ForInvocation, IterateInvocation)) ) - def _is_iterator_frame_open_at(self, iterator_id: str, node_id: str) -> bool: + def _is_iterator_frame_open_at(self, iterator_id: str, node_id: str, via_fields: Optional[set[str]] = None) -> bool: """Return whether some path carries ``iterator_id``'s iteration frame into ``node_id``. + ``via_fields`` limits the paths to those entering ``node_id`` through one of those input fields. + Depth counts frames opened since ``iterator_id``. Per-item outputs of an Iterate or For open one. An item input of a Collect closes one, except a For's final output, which is already in the For's parent frame and which the Collect groups there. A Collect's collection input keeps its source's frame. A Collect that would close - ``iterator_id``'s frame counts as leaving it open unless ``iterator_id`` is the item's innermost iterator. + ``iterator_id``'s frame leaves it open unless ``iterator_id`` is the item's innermost iterator. """ pending = [(iterator_id, 0)] @@ -1641,11 +1684,13 @@ def _is_iterator_frame_open_at(self, iterator_id: str, node_id: str) -> bool: ): next_depth -= 1 if next_depth <= 0: - if not self._is_innermost_iterator(iterator_id, current_id): + if self._is_innermost_iterator(iterator_id, current_id): + continue + next_depth = depth + if edge.destination.node_id == node_id: + if via_fields is None or edge.destination.field in via_fields: return True continue - if edge.destination.node_id == node_id: - return True state = (edge.destination.node_id, next_depth) if state not in seen: seen.add(state) @@ -1693,6 +1738,33 @@ def get_node_iterators(self, node_id: str, it_graph: Optional["nx.DiGraph"] = No if isinstance(self._state.graph.get_node(n), (ForInvocation, IterateInvocation)) ] + def _has_mixed_depth_if_input(self, input_edges: list[Edge]) -> bool: + """Return whether an input If runs at different depths, as when only one branch keeps an item frame open. + + Its frames then match no single product of iterator executions, so they are matched by path instead. + """ + + for edge in input_edges: + if not isinstance(self._state.graph.get_node(edge.source.node_id), IfInvocation): + continue + depths = {len(frame) for frame in self._state._prepared_if_exec_frames(edge.source.node_id)} + if len(depths) > 1: + return True + return False + + def _get_active_node_iterators(self, node_id: str, it_graph: Optional["nx.DiGraph"] = None) -> list[str]: + """Return the iterators that give ``node_id`` frames. + + An iterator on a rejected If branch never runs, so it adds no frame to values selected past that If. + """ + + return [ + iterator_node_id + for iterator_node_id in self.get_node_iterators(node_id, it_graph) + if iterator_node_id in self._state.source_prepared_mapping + or not self._state._is_source_inactive(iterator_node_id) + ] + def _get_prepared_nodes_for_source(self, source_node_id: str) -> set[str]: return { exec_node_id @@ -1817,11 +1889,11 @@ 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 ( 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)) ) + # Admission reads the prepared frames of every unblocked upstream iterator. and not any( isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) and ancestor_id not in self._state.executed @@ -1829,6 +1901,7 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: and not self._is_closed_branch_iterator(ancestor_id, node_id) for ancestor_id in nx.ancestors(g, node_id) ) + and self._has_admitted_source_mapping(node_id, g) ), None, ) @@ -1873,7 +1946,15 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: or iteration_path[: len(condition_path or ())] == (condition_path or ()) ) ] - for iteration_mapping in matching_mappings: + # Only the condition is attached, so mappings that differ in other iterators create one If. + condition_source_ids = {edge.source.node_id for edge in condition_edges} + unique_mappings = { + tuple( + entry for entry in iteration_mapping if entry[0] in condition_source_ids + ): iteration_mapping + for iteration_mapping in matching_mappings + } + for iteration_mapping in unique_mappings.values(): create_results = self.create_execution_node( next_node_id, iteration_mapping, @@ -1881,7 +1962,7 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: input_fields={"condition"}, ) new_node_ids.extend(create_results) - else: + elif not self._is_selected_branch_resolved_empty(next_node_id): for iteration_mapping in iteration_mappings: iteration_path = self._get_known_iteration_path(-1, iteration_mapping, condition_edges) create_results = self.create_execution_node( diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index d6b10a427a7..6a339056e6c 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -37,6 +37,7 @@ _ExecutionScheduler, _GenericGraphSchedulerAdapter, ) +from invokeai.app.services.shared.graph_if_activation import UnsupportedIfFramesError from invokeai.app.services.shared.graph_nested_iterate_planner import get_nested_iterate_sequence_depth from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context from tests.test_nodes import ( @@ -3541,26 +3542,51 @@ def _sibling_iterates_branch_graph(condition: bool) -> Graph: @pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) -@pytest.mark.parametrize("condition", [True, False]) @pytest.mark.parametrize( - "build_graph", + ("build_graph", "condition", "expected_sinks"), [ - pytest.param(_shared_outer_iterate_branch_graph, id="outer-frame-through-branch"), - pytest.param(_sibling_iterates_branch_graph, id="sibling-frame-left-open"), + pytest.param( + _shared_outer_iterate_branch_graph, + True, + [((0,), [0, 1]), ((1,), [0, 1])], + id="outer-frame-through-branch-selected", + ), + pytest.param( + _shared_outer_iterate_branch_graph, False, [((), ["fallback"])], id="outer-frame-through-branch-rejected" + ), + pytest.param( + _sibling_iterates_branch_graph, + True, + [((0,), [0, 1]), ((1,), [1, 2])], + id="sibling-frame-left-open-selected", + ), + pytest.param( + _sibling_iterates_branch_graph, False, [((), ["fallback"])], id="sibling-frame-left-open-rejected" + ), ], ) -def test_if_is_not_materialized_in_parent_frame_while_branch_frame_is_open( - build_graph: Callable[[bool], Graph], condition: bool, force_compatibility_scheduler: bool +def test_if_takes_frames_from_selected_branch_left_open( + build_graph: Callable[[bool], Graph], + condition: bool, + expected_sinks: list[tuple[tuple[int, ...], list[int | str]]], + force_compatibility_scheduler: bool, ) -> None: - """The If must not be materialized in the parent frame while an unexecuted branch can still give it item frames.""" + """A branch that leaves an iterator frame open gives the If one execution per item only when it is selected.""" graph = build_graph(condition) graph.validate_self() - _trace, state = _run(GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler) + trace, state = _run(GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler) - assert () not in { - state._get_iteration_path(exec_id) for exec_id in state._prepared_registry().get_prepared_ids("if") - } + assert state.is_complete(), f"execution stalled after {trace}" + assert ( + sorted( + (state._get_iteration_path(exec_id), state.results[exec_id].value) + for exec_id in state._prepared_registry().get_prepared_ids("sink") + ) + == expected_sinks + ) + if not condition: + assert "body" not in trace @pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) @@ -3621,6 +3647,7 @@ def test_per_item_if_can_select_values_from_root_frame( assert state.results[collect_id].collection == ["true-root", "false-root"] +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) @pytest.mark.parametrize( ("condition", "expected_paths", "expected_values"), [ @@ -3632,6 +3659,7 @@ def test_literal_if_result_uses_selected_branch_frame_cardinality( condition: bool, expected_paths: set[tuple[int, ...]], expected_values: list[str], + force_compatibility_scheduler: bool, ) -> None: graph = Graph() graph.add_node(CollectionConcatInvocation(id="source", first=[0, 1])) @@ -3647,7 +3675,49 @@ def test_literal_if_result_uses_selected_branch_frame_cardinality( graph.add_edge(create_edge("if", "value", "collect", "item")) graph.validate_self() - trace, state = _run_graph(GraphExecutionState(graph=graph)) + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + if_paths = {state._get_iteration_path(exec_id) for exec_id in state._prepared_registry().get_prepared_ids("if")} + assert if_paths == expected_paths + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == expected_values + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +@pytest.mark.parametrize( + ("condition", "expected_paths", "expected_values"), + [ + pytest.param(True, {(0,), (1,)}, [0, 1], id="selected-item-frame"), + pytest.param(False, {()}, ["root-value"], id="selected-root-frame"), + ], +) +def test_root_condition_if_result_uses_selected_branch_frame_cardinality( + condition: bool, + expected_paths: set[tuple[int, ...]], + expected_values: list[int | str], + force_compatibility_scheduler: bool, +) -> None: + """Each If frame reads a condition computed once in the parent frame.""" + graph = Graph() + graph.add_node(FixedCollectionTestInvocation(id="source", value=True)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(FixedCollectionTestInvocation(id="condition", value=condition)) + graph.add_node(AnyTypeTestInvocation(id="root_branch", value="root-value")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("condition", "condition", "if", "condition")) + graph.add_edge(create_edge("iterate", "item", "if", "true_input")) + graph.add_edge(create_edge("root_branch", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) assert state.is_complete(), f"execution stalled after {trace}" if_paths = {state._get_iteration_path(exec_id) for exec_id in state._prepared_registry().get_prepared_ids("if")} @@ -3656,6 +3726,396 @@ def test_literal_if_result_uses_selected_branch_frame_cardinality( assert state.results[collect_id].collection == expected_values +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +def test_literal_if_selecting_empty_iterate_branch_produces_no_values(force_compatibility_scheduler: bool) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="root_branch", value="root-value")) + graph.add_node(IfInvocation(id="if", condition=True)) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "true_input")) + graph.add_edge(create_edge("root_branch", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + assert "root_branch" not in trace + assert not state._prepared_registry().get_prepared_ids("if") + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == [] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +def test_per_item_if_mixes_item_frame_and_root_frame_branches(force_compatibility_scheduler: bool) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="conditions", first=[True, False])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(MarkedAnyTypeTestInvocation(id="item_branch", marker="item:")) + graph.add_node(AnyTypeTestInvocation(id="root_branch", value="root-value")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("conditions", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "condition")) + graph.add_edge(create_edge("iterate", "item", "item_branch", "value")) + graph.add_edge(create_edge("item_branch", "value", "if", "true_input")) + graph.add_edge(create_edge("root_branch", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == ["item:True", "root-value"] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +@pytest.mark.parametrize( + ("conditions", "expected_values"), + [ + pytest.param([True, False], ["item:True", ["mp", "mq"]], id="root-collect-selected-last"), + pytest.param([False, True], [["mp", "mq"], "item:True"], id="root-collect-selected-first"), + ], +) +def test_per_item_if_selects_root_collected_iterate_without_aliasing_its_frames( + conditions: list[bool], expected_values: list[Any], force_compatibility_scheduler: bool +) -> None: + """The root Iterate's item indexes are unrelated to the If's item frames even though they share values.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="conditions", first=conditions)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(MarkedAnyTypeTestInvocation(id="item_branch", marker="item:")) + graph.add_node(CollectionConcatInvocation(id="root_source", first=["p", "q"])) + graph.add_node(IterateInvocation(id="root_iterate")) + graph.add_node(MarkedAnyTypeTestInvocation(id="root_body", marker="m")) + graph.add_node(CollectInvocation(id="root_collect")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("conditions", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "condition")) + graph.add_edge(create_edge("iterate", "item", "item_branch", "value")) + graph.add_edge(create_edge("item_branch", "value", "if", "true_input")) + graph.add_edge(create_edge("root_source", "collection", "root_iterate", "collection")) + graph.add_edge(create_edge("root_iterate", "item", "root_body", "value")) + graph.add_edge(create_edge("root_body", "value", "root_collect", "item")) + graph.add_edge(create_edge("root_collect", "collection", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == expected_values + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +@pytest.mark.parametrize( + ("condition", "expected_values"), + [pytest.param(True, [0, 1], id="item-frames"), pytest.param(False, ["root-value"], id="root-frame")], +) +def test_node_after_branch_framed_if_takes_frames_from_selected_branch( + condition: bool, expected_values: list[int | str], force_compatibility_scheduler: bool +) -> None: + graph = Graph() + graph.add_node(FixedCollectionTestInvocation(id="source", value=True)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="root_branch", value="root-value")) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "true_input")) + graph.add_edge(create_edge("root_branch", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + graph.add_edge(create_edge("sink", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == expected_values + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +def test_computed_condition_selecting_empty_iterate_branch_produces_no_values( + force_compatibility_scheduler: bool, +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(FixedCollectionTestInvocation(id="condition", value=True)) + graph.add_node(AnyTypeTestInvocation(id="root_branch", value="root-value")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("condition", "condition", "if", "condition")) + graph.add_edge(create_edge("iterate", "item", "if", "true_input")) + graph.add_edge(create_edge("root_branch", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + assert "root_branch" not in trace + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == [] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +def test_if_with_empty_per_item_condition_runs_neither_root_branch(force_compatibility_scheduler: bool) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="conditions", first=[])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="root_true", value="true-root")) + graph.add_node(AnyTypeTestInvocation(id="root_false", value="false-root")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("conditions", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "condition")) + graph.add_edge(create_edge("root_true", "value", "if", "true_input")) + graph.add_edge(create_edge("root_false", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + assert "root_true" not in trace and "root_false" not in trace + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == [] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +@pytest.mark.parametrize( + ("outer_items", "expected_sinks", "expected_collections"), + [ + pytest.param( + ["x", ""], + [((0, 0), "b0"), ((0, 1), "b1"), ((1,), "fallback")], + [((0,), ["b0", "b1"]), ((1,), ["fallback"])], + id="first-selected", + ), + pytest.param( + ["x", "x"], + [((0, 0), "b0"), ((0, 1), "b1"), ((1, 0), "b0"), ((1, 1), "b1")], + [((0,), ["b0", "b1"]), ((1,), ["b0", "b1"])], + id="all-selected", + ), + pytest.param( + ["", ""], + [((0,), "fallback"), ((1,), "fallback")], + [((), ["fallback", "fallback"])], + id="none-selected", + ), + ], +) +def test_per_outer_item_if_keeps_nested_iterate_frame_open_on_selected_branch( + outer_items: list[str], + expected_sinks: list[tuple[tuple[int, ...], str]], + expected_collections: list[tuple[tuple[int, ...], list[str]]], + force_compatibility_scheduler: bool, +) -> None: + """Each outer item's If keeps the inner item frames only when it selects the branch that carries them.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=outer_items)) + graph.add_node(IterateInvocation(id="outer_iterate")) + graph.add_node(FixedCollectionTestInvocation(id="inner_source")) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(MarkedAnyTypeTestInvocation(id="body", marker="b")) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("outer_source", "collection", "outer_iterate", "collection")) + graph.add_edge(create_edge("outer_iterate", "item", "inner_source", "value")) + graph.add_edge(create_edge("outer_iterate", "item", "condition", "value")) + graph.add_edge(create_edge("inner_source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "if", "true_input")) + graph.add_edge(create_edge("condition", "condition", "if", "condition")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + assert ( + sorted( + (state._get_iteration_path(exec_id), state.results[exec_id].value) + for exec_id in state._prepared_registry().get_prepared_ids("sink") + ) + == expected_sinks + ) + assert ( + sorted( + (state._get_iteration_path(exec_id), state.results[exec_id].collection) + for exec_id in state._prepared_registry().get_prepared_ids("collect") + ) + == expected_collections + ) + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +@pytest.mark.parametrize( + ("condition", "expected_values"), + [ + pytest.param(True, [[11, 21], [12, 22]], id="true"), + pytest.param(False, [["f11", "f21"], ["f12", "f22"]], id="false"), + ], +) +def test_if_between_collects_of_sibling_iterates_keeps_the_open_frame( + condition: bool, expected_values: list[list[int | str]], force_compatibility_scheduler: bool +) -> None: + """Each branch collects one of two zipped frames, so both branches stay in the other frame.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="left_source", first=[1, 2])) + graph.add_node(CollectionConcatInvocation(id="right_source", first=[10, 20])) + graph.add_node(IterateInvocation(id="left")) + graph.add_node(IterateInvocation(id="right")) + graph.add_node(AddInvocation(id="true_body")) + graph.add_node(CollectInvocation(id="true_collect")) + graph.add_node(AddInvocation(id="false_body")) + graph.add_node(MarkedAnyTypeTestInvocation(id="false_mark", marker="f")) + graph.add_node(CollectInvocation(id="false_collect")) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("left_source", "collection", "left", "collection")) + graph.add_edge(create_edge("right_source", "collection", "right", "collection")) + for body in ("true_body", "false_body"): + graph.add_edge(create_edge("left", "item", body, "a")) + graph.add_edge(create_edge("right", "item", body, "b")) + graph.add_edge(create_edge("true_body", "value", "true_collect", "item")) + graph.add_edge(create_edge("false_body", "value", "false_mark", "value")) + graph.add_edge(create_edge("false_mark", "value", "false_collect", "item")) + graph.add_edge(create_edge("true_collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("false_collect", "collection", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == expected_values + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) +@pytest.mark.parametrize( + ("conditions", "expected_values"), + [ + pytest.param([True, False], ["inner:True", "outer-false"], id="first-nested"), + pytest.param([False, True], ["outer-false", "inner:True"], id="second-nested"), + ], +) +def test_per_item_if_selecting_per_item_nested_if_keeps_every_item( + conditions: list[bool], expected_values: list[str], force_compatibility_scheduler: bool +) -> None: + """The outer If runs in every item frame even when an item selects a nested If that is not materialized yet.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="conditions", first=conditions)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(FixedCollectionTestInvocation(id="inner_condition")) + graph.add_node(MarkedAnyTypeTestInvocation(id="inner_true", marker="inner:")) + graph.add_node(AnyTypeTestInvocation(id="inner_false", value="inner-false")) + graph.add_node(IfInvocation(id="inner_if")) + graph.add_node(AnyTypeTestInvocation(id="outer_false", value="outer-false")) + graph.add_node(IfInvocation(id="outer_if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("conditions", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "outer_if", "condition")) + graph.add_edge(create_edge("iterate", "item", "inner_condition", "value")) + graph.add_edge(create_edge("inner_condition", "condition", "inner_if", "condition")) + graph.add_edge(create_edge("iterate", "item", "inner_true", "value")) + graph.add_edge(create_edge("inner_true", "value", "inner_if", "true_input")) + graph.add_edge(create_edge("inner_false", "value", "inner_if", "false_input")) + graph.add_edge(create_edge("inner_if", "value", "outer_if", "true_input")) + graph.add_edge(create_edge("outer_false", "value", "outer_if", "false_input")) + graph.add_edge(create_edge("outer_if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"execution stalled after {trace}" + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == expected_values + + +def test_long_chain_into_if_branch_completes() -> None: + """Frame depths are derived without recursing once per upstream node, which overflowed near 330 nodes.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="conditions", first=[True, False])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(MarkedAnyTypeTestInvocation(id="false_branch", marker="f")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("conditions", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "condition")) + graph.add_edge(create_edge("iterate", "item", "false_branch", "value")) + graph.add_edge(create_edge("false_branch", "value", "if", "false_input")) + previous = ("iterate", "item") + for index in range(400): + graph.add_node(AnyTypeTestInvocation(id=f"chain_{index}")) + graph.add_edge(create_edge(previous[0], previous[1], f"chain_{index}", "value")) + previous = (f"chain_{index}", "value") + graph.add_edge(create_edge(previous[0], previous[1], "if", "true_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + + trace, state = _run_graph(GraphExecutionState(graph=graph)) + + assert state.is_complete(), f"execution stalled after {trace[-5:]}" + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == [True, "fFalse"] + + +def test_if_with_condition_and_branch_from_unrelated_iterates_is_rejected() -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="conditions", first=[True, False])) + graph.add_node(IterateInvocation(id="condition_iterate")) + graph.add_node(FixedCollectionTestInvocation(id="branch_source", value=True)) + graph.add_node(IterateInvocation(id="branch_iterate")) + graph.add_node(AnyTypeTestInvocation(id="root_branch", value="root-value")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("conditions", "collection", "condition_iterate", "collection")) + graph.add_edge(create_edge("condition_iterate", "item", "if", "condition")) + graph.add_edge(create_edge("branch_source", "collection", "branch_iterate", "collection")) + graph.add_edge(create_edge("branch_iterate", "item", "if", "true_input")) + graph.add_edge(create_edge("root_branch", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + with pytest.raises(UnsupportedIfFramesError, match="unrelated branch_iterate"): + _run_graph(GraphExecutionState(graph=graph)) + + @pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) def test_if_collect_collection_inner_iterate_keeps_values_in_outer_frame( force_compatibility_scheduler: bool, From 86a3660fef1fdc523a3264979d4d028ff8cd920a Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 18 Oct 2026 14:11:35 -0400 Subject: [PATCH 4/5] fix(graph): keep literal If over unselected For axes and reject unsafe reopens A literal If again runs in its selected branch's frames when an unrelated axis reaches it only through the unselected branch's For output. A literal If pairing an axis reopened after a Collect now requires a direct Collect-to-Iterate reopen over one nested chain of axes, and rejects reopens whose length can differ. Co-Authored-By: Claude Opus 5.5 --- invokeai/app/services/shared/README.md | 10 +- .../app/services/shared/graph_if_runtime.py | 167 +++++- .../test_execution_engine_differential.py | 479 +++++++++++++++++- 3 files changed, 628 insertions(+), 28 deletions(-) diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index 2363fb29907..2be4ec6ae64 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -821,8 +821,14 @@ Run `C` -> `D:0` -> enqueue `D`. Run `D` -> done. `Collect.item` input closes an axis), unless the input that decides the `If`'s frames carries both: the condition, or the known selected branch of a literal condition. A cross product computed by one selected node therefore keeps its frames, while one `Iterate` per branch, or a root value selected over a cross product, is rejected. A literal - condition selecting an axis reopened downstream of a `Collect` that closed the other branch's axis pairs the two - axes by index. + condition whose selected branch carries one of the two also runs when the other comes only from the unselected + branch and either the unselected one reaches the `If` only through a `For` final output, which adds no frames, or + the selected axis, or an iterator enclosing it, reopens the unselected one. An `Iterate` reopens every axis of its + `Collect`'s item (the one the `Collect` closes and the ones it keeps), and transitively what those reopen, when + its `collection` comes directly from a `Collect` with a single `item` input and nothing appended through its + `collection`: each item becomes one frame, so the selected frames pair with those axes by index. Any other reopen + (a node between the `Collect` and the `Iterate`, or extra collected items) can change the length, and selecting + the upstream axis over its reopening pairs it with per-group frames; both would drop or invent frames. - A branch whose open axes, or the axes a branch `Collect` keeps, are deeper than a per-item condition's, unless both branches carry the same deeper axes. - A root-frame condition choosing between branches with different deeper axes. When both branches share nested diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index 39e2bf625cd..67f333a03f2 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -15,7 +15,6 @@ ForInvocation, IterateInvocation, get_output_field_scope, - nx, ) if TYPE_CHECKING: @@ -29,12 +28,13 @@ class UnsupportedIfIterationAxesError(ValueError): def validate_if_iteration_axes(state: "GraphExecutionState") -> None: """Reject unrelated condition and branch axes using graph topology alone.""" iterator_ancestry = _IteratorAncestry(state) + collect_axes_cache: dict[str, Optional[tuple[str, ...]]] = {} for node_id, node in state.graph.nodes.items(): if not isinstance(node, IfInvocation): continue condition_edges = state.graph._get_input_edges(node_id, "condition") if not condition_edges: - _validate_frame_axes(state, iterator_ancestry, node_id, condition_edges) + _validate_frame_axes(state, iterator_ancestry, node_id, condition_edges, collect_axes_cache) continue condition_iterators = iterator_ancestry.get_iterator_ancestors_for_edge(condition_edges[0]) branch_edges = { @@ -72,19 +72,29 @@ def validate_if_iteration_axes(state: "GraphExecutionState") -> None: f"unrelated open iterator axes: condition {sorted(condition_iterators)}, " f"branch {sorted(branch_iterators)}" ) - _validate_frame_axes(state, iterator_ancestry, node_id, condition_edges) + _validate_frame_axes(state, iterator_ancestry, node_id, condition_edges, collect_axes_cache) def _validate_frame_axes( - state: "GraphExecutionState", iterator_ancestry: _IteratorAncestry, node_id: str, condition_edges: list[Edge] + state: "GraphExecutionState", + iterator_ancestry: _IteratorAncestry, + node_id: str, + condition_edges: list[Edge], + collect_axes_cache: dict[str, Optional[tuple[str, ...]]], ) -> None: """Reject frame sets an If cannot run in. The If's frames come from the iterators open at it, including one reached through a For's final output. Two of them that are unrelated (neither stays open inside the other) form a cross product, which only the input deciding - the If's frames can carry: the condition, or the known selected branch of a literal condition. One exception keeps - the existing pairing for a literal condition whose selected branch reopens an axis downstream of the other one, - after a Collect closed it. A root-frame condition cannot select between branches with different deeper axes. + the If's frames can carry: the condition, or the known selected branch of a literal condition. A literal condition + whose selected branch carries one of the two axes also runs when the other comes only from the unselected branch + and either: + - the unselected one reaches the If only through a For's final output, so it adds no frames and the If runs in + the selected branch's frames; or + - the selected axis, or an iterator enclosing it, reopens the unselected one one-to-one (see + `_get_reopened_source_axes`), so the selected frames pair with it by index. Selecting the upstream axis over its + reopening is not supported. + A root-frame condition cannot select between branches with different deeper axes. """ if condition_edges: _validate_collect_branch_axes(state, iterator_ancestry, node_id, condition_edges[0]) @@ -96,7 +106,6 @@ def _validate_frame_axes( assert isinstance(node, IfInvocation) deciding_edges = state.graph._get_input_edges(node_id, "true_input" if node.condition else "false_input") deciding_axes = {axis_id for edge in deciding_edges for axis_id in iterator_ancestry.get_frame_axes_for_edge(edge)} - source_graph = state._get_source_graph_flat() for index, axis_id in enumerate(frame_iterators): for other_axis_id in frame_iterators[index + 1 :]: if axis_id in iterator_ancestry.get_frame_iterators( @@ -105,13 +114,16 @@ def _validate_frame_axes( continue if axis_id in deciding_axes and other_axis_id in deciding_axes: continue - if not condition_edges and any( - selected_id in deciding_axes - and unselected_id not in deciding_axes - and nx.has_path(source_graph, unselected_id, selected_id) - for selected_id, unselected_id in ((axis_id, other_axis_id), (other_axis_id, axis_id)) - ): - continue + if not condition_edges and (axis_id in deciding_axes) != (other_axis_id in deciding_axes): + selected_id, unselected_id = ( + (axis_id, other_axis_id) if axis_id in deciding_axes else (other_axis_id, axis_id) + ) + if unselected_id not in iterator_ancestry.get_ordered_iterator_ancestors(node_id): + continue + if unselected_id in _get_reopened_source_axes( + state, iterator_ancestry, selected_id, collect_axes_cache + ): + continue raise UnsupportedIfIterationAxesError( f"If node '{node_id}' cannot run in the frames of unrelated open iterator axes: " f"{sorted((axis_id, other_axis_id))}" @@ -134,6 +146,49 @@ def _validate_frame_axes( ) +def _get_reopened_source_axes( + state: "GraphExecutionState", + iterator_ancestry: _IteratorAncestry, + axis_id: str, + collect_axes_cache: dict[str, Optional[tuple[str, ...]]], +) -> set[str]: + """Return the axes whose frames an axis, or an iterator enclosing it, reopens one-to-one. + + An Iterate reopens a Collect's item axes when its collection comes straight from that Collect, the Collect + gathers exactly one item per frame of them (a single item input and nothing appended through `collection`), and + those axes form one nested chain. Each item then becomes one frame of the Iterate, in the item frames' order, so a + selected axis can pair by index with those axes, and with the axes they in turn reopen. A node between the Collect + and the Iterate, extra collected items, or an item crossing unrelated axes can break that pairing, and the If + would then pair frames of unequal axes by index. + """ + sources: set[str] = set() + visited: set[str] = set() + pending = [axis_id] + while pending: + current_id = pending.pop() + for reopened_id in (current_id, *iterator_ancestry.get_frame_iterators(current_id)): + if reopened_id in visited: + continue + visited.add(reopened_id) + if not isinstance(state.graph.get_node(reopened_id), IterateInvocation): + continue + # Graph validation gives every Iterate exactly one collection input. + [collection_edge] = state.graph._get_input_edges(reopened_id, COLLECTION_FIELD) + collect_id = collection_edge.source.node_id + collect = state.graph.get_node(collect_id) + if not isinstance(collect, CollectInvocation) or collect.collection: + continue + collect_inputs = state.graph._get_input_edges(collect_id) + if len(collect_inputs) != 1 or collect_inputs[0].destination.field != ITEM_FIELD: + continue + item_axes = _get_edge_frame_axes(state, iterator_ancestry, collect_inputs[0], collect_axes_cache) + if item_axes is None or not _is_nested_chain(iterator_ancestry, item_axes): + continue + sources.update(item_axes) + pending.extend(item_axes) + return sources + + def _validate_collect_branch_axes( state: "GraphExecutionState", iterator_ancestry: _IteratorAncestry, node_id: str, condition_edge: Edge ) -> None: @@ -165,17 +220,83 @@ def _get_collect_kept_axes( """Return the axes a Collect's output stays in: its item's axes except the innermost one, which it closes.""" kept: set[str] = set() for edge in state.graph._get_input_edges(collect_id): - axes = iterator_ancestry.get_frame_axes_for_edge(edge) - source = state.graph.get_node(edge.source.node_id) - is_for_final = isinstance(source, ForInvocation) and ( - get_output_field_scope(source, edge.source.field) == OutputScope.Final - ) - if edge.destination.field == ITEM_FIELD and not is_for_final: - axes = axes[:-1] - kept.update(axes) + kept.update(_without_closed_axis(state, edge, iterator_ancestry.get_frame_axes_for_edge(edge))) return kept +def _without_closed_axis(state: "GraphExecutionState", edge: Edge, axes: tuple[str, ...]) -> tuple[str, ...]: + """Return the axes a Collect input keeps: an item closes its innermost axis unless it is a For's final output.""" + source = state.graph.get_node(edge.source.node_id) + is_for_final = isinstance(source, ForInvocation) and ( + get_output_field_scope(source, edge.source.field) == OutputScope.Final + ) + if edge.destination.field == ITEM_FIELD and not is_for_final: + return axes[:-1] + return axes + + +def _is_nested_chain(iterator_ancestry: _IteratorAncestry, axes: tuple[str, ...]) -> bool: + """Return whether each axis encloses the next one.""" + return all( + outer_id in iterator_ancestry.get_frame_iterators(inner_id) + for outer_id, inner_id in zip(axes, axes[1:], strict=False) + ) + + +def _get_edge_frame_axes( + state: "GraphExecutionState", + iterator_ancestry: _IteratorAncestry, + edge: Edge, + collect_axes_cache: dict[str, Optional[tuple[str, ...]]], +) -> Optional[tuple[str, ...]]: + """Return the frames open on an edge, outer to inner, following a Collect source into the axes it keeps. + + Iterator ancestry stops at a Collect, so a Collect's output would otherwise look like it runs in no frames. A + Collect whose kept axes are not one nested chain, or that collects such a Collect, has no frame order a reopen could + pair with, so it yields None. Results are cached per Collect for one validation run and resolved with an explicit + stack, so a lattice or a long chain of Collects costs one pass per node. + """ + if not isinstance(state.graph.get_node(edge.source.node_id), CollectInvocation): + return iterator_ancestry.get_frame_axes_for_edge(edge) + pending = [edge.source.node_id] + while pending: + collect_id = pending[-1] + if collect_id in collect_axes_cache: + pending.pop() + continue + input_edges = state.graph._get_input_edges(collect_id) + unresolved = [ + input_edge.source.node_id + for input_edge in input_edges + if isinstance(state.graph.get_node(input_edge.source.node_id), CollectInvocation) + and input_edge.source.node_id not in collect_axes_cache + ] + if unresolved: + pending.extend(unresolved) + continue + pending.pop() + kept: Optional[set[str]] = set() + for input_edge in input_edges: + source_id = input_edge.source.node_id + if isinstance(state.graph.get_node(source_id), CollectInvocation): + input_axes = collect_axes_cache[source_id] + else: + input_axes = iterator_ancestry.get_frame_axes_for_edge(input_edge) + if input_axes is None: + kept = None + break + kept.update(_without_closed_axis(state, input_edge, input_axes)) + ordered = None + if kept is not None: + ordered = tuple( + sorted(kept, key=lambda axis_id: (len(iterator_ancestry.get_frame_iterators(axis_id)), axis_id)) + ) + if not _is_nested_chain(iterator_ancestry, ordered): + ordered = None + collect_axes_cache[collect_id] = ordered + return collect_axes_cache[edge.source.node_id] + + def _get_effective_iteration_paths_for_edge( state: "GraphExecutionState", edge: Edge, diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index 1f492b81f92..17c7aa74a89 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -3721,7 +3721,9 @@ def _cross_product_branch_graph(condition: bool, condition_from_branch: bool = F return _graph_from(nodes, edges) -def _collect_then_reopen_graph(condition: bool, extra_items: tuple[str, ...] = ()) -> Graph: +def _collect_then_reopen_graph( + condition: bool, extra_items: tuple[str, ...] = (), appended_items: tuple[str, ...] = () +) -> Graph: nodes = [ CollectionConcatInvocation(id="src", first=["a1", "a2"]), IterateInvocation(id="first"), @@ -3738,6 +3740,9 @@ def _collect_then_reopen_graph(condition: bool, extra_items: tuple[str, ...] = ( ("first", "item", "body", "value"), ("body", "value", "closing_collect", "item"), ] + if appended_items: + nodes.append(CollectionConcatInvocation(id="appended", first=list(appended_items))) + edges.append(("appended", "collection", "closing_collect", "collection")) if extra_items: nodes.append(CollectionConcatInvocation(id="extend", second=list(extra_items))) edges += [ @@ -3756,6 +3761,223 @@ def _collect_then_reopen_graph(condition: bool, extra_items: tuple[str, ...] = ( return _graph_from(nodes, edges) +def _collect_with_literal_items_then_reopen_graph() -> Graph: + graph = _collect_then_reopen_graph(True) + graph.get_node("closing_collect").collection = ["x"] + return graph + + +def _collect_of_collect_reopen_graph( + groups: list[list[str]], unselected_axis: str, extra_collect: bool = False, append_unrelated_axis: bool = False +) -> Graph: + """An Iterate reopens a Collect of per-group Collects, so its frames are the groups of the outer axis.""" + nodes: list[Any] = [ + CollectionConcatInvocation(id="groups_src", first=groups), + IterateInvocation(id="group"), + CollectionConcatInvocation(id="members_src"), + IterateInvocation(id="member"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + CollectInvocation(id="group_collect"), + CollectInvocation(id="closing_collect"), + IterateInvocation(id="reopened"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ] + edges: list[Any] = [ + ("groups_src", "collection", "group", "collection"), + ("group", "item", "members_src", "first"), + ("members_src", "collection", "member", "collection"), + ("member", "item", "body", "value"), + ("body", "value", "group_collect", "item"), + ("group_collect", "collection", "closing_collect", "item"), + ] + if append_unrelated_axis: + nodes += [ + CollectionConcatInvocation(id="appended_src", first=[["v1"], ["v2"], ["v3"]]), + IterateInvocation(id="appended"), + CollectionConcatInvocation(id="appended_items"), + ] + edges += [ + ("appended_src", "collection", "appended", "collection"), + ("appended", "item", "appended_items", "first"), + ("appended_items", "collection", "group_collect", "collection"), + ] + if extra_collect: + nodes.append(CollectInvocation(id="outer_collect")) + edges += [ + ("closing_collect", "collection", "outer_collect", "item"), + ("outer_collect", "collection", "reopened", "collection"), + ] + else: + edges.append(("closing_collect", "collection", "reopened", "collection")) + edges += [ + ("reopened", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + (unselected_axis, "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ] + return _graph_from(nodes, edges) + + +def _per_item_condition_with_collect_of_collects_branch_graph() -> Graph: + """A per-item condition on the outer axis chooses between a Collect over a deeper axis and a Collect of one.""" + nodes: list[Any] = [ + CollectionConcatInvocation(id="src", first=[[[["a"], ["b"]], [["c"]]], [[["x"]], []], [[["d", "e"]]]]), + ] + edges: list[Any] = [] + previous = ("src", "collection") + for axis_id in ("outer", "middle", "deep", "inner"): + nodes.append(IterateInvocation(id=axis_id)) + edges.append((*previous, axis_id, "collection")) + if axis_id != "inner": + nodes.append(CollectionConcatInvocation(id=f"{axis_id}_items")) + edges.append((axis_id, "item", f"{axis_id}_items", "first")) + previous = (f"{axis_id}_items", "collection") + nodes += [ + MarkedAnyTypeTestInvocation(id="inner_body", marker="b:"), + CollectInvocation(id="inner_collect"), + CollectInvocation(id="deep_collect"), + MarkedAnyTypeTestInvocation(id="deep_body", marker="n:"), + CollectInvocation(id="deep_item_collect"), + FixedCollectionTestInvocation(id="condition"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ] + edges += [ + ("inner", "item", "inner_body", "value"), + ("inner_body", "value", "inner_collect", "item"), + ("inner_collect", "collection", "deep_collect", "item"), + ("deep", "item", "deep_body", "value"), + ("deep_body", "value", "deep_item_collect", "item"), + ("outer", "index", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("deep_item_collect", "collection", "if", "true_input"), + ("deep_collect", "collection", "if", "false_input"), + ("if", "value", "collect", "item"), + ] + return _graph_from(nodes, edges) + + +def _reopen_of_a_cross_product_graph(swap_add_inputs: bool, condition: bool = False) -> Graph: + """A literal If chooses between an Iterate reopening a Collect over two crossed axes and the crossed inner axis.""" + first_axis, second_axis = ("right", "left") if swap_add_inputs else ("left", "right") + return _graph_from( + [ + CollectionConcatInvocation(id="left_src", first=[1, 2]), + IterateInvocation(id="left"), + CollectionConcatInvocation(id="right_src", first=[10, 20, 30]), + IterateInvocation(id="right"), + AddInvocation(id="add"), + CollectInvocation(id="closing_collect"), + IterateInvocation(id="reopened"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=condition), + CollectInvocation(id="collect"), + ], + [ + ("left_src", "collection", "left", "collection"), + ("right_src", "collection", "right", "collection"), + (first_axis, "item", "add", "a"), + (second_axis, "item", "add", "b"), + ("add", "value", "closing_collect", "item"), + ("closing_collect", "collection", "reopened", "collection"), + ("reopened", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + (second_axis, "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + + +def _reopen_of_a_different_closed_axis_graph() -> Graph: + return _graph_from( + [ + CollectionConcatInvocation(id="closed_src", first=["z1", "z2", "z3"]), + IterateInvocation(id="closed"), + CollectInvocation(id="closing_collect"), + IterateInvocation(id="reopened"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + CollectionConcatInvocation(id="other_src", first=["u1", "u2"]), + IterateInvocation(id="other"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("closed_src", "collection", "closed", "collection"), + ("closed", "item", "closing_collect", "item"), + ("closing_collect", "collection", "reopened", "collection"), + ("reopened", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("other_src", "collection", "other", "collection"), + ("other", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + + +def _for_output_connected_to_selected_axis_graph(selected_is_downstream: bool) -> Graph: + """A For's final output on the unselected branch, under an axis connected to the selected one by a reopen.""" + upstream_ids = ("upstream_src", "upstream") + downstream_ids = ("downstream_src", "downstream") + for_parent = upstream_ids[1] if selected_is_downstream else downstream_ids[1] + selected_axis = downstream_ids[1] if selected_is_downstream else upstream_ids[1] + return _graph_from( + [ + CollectionConcatInvocation(id="upstream_src", first=[["a0", "a1"], ["a2"]]), + IterateInvocation(id="upstream"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + CollectInvocation(id="closing_collect"), + CollectionConcatInvocation(id="downstream_src", second=[["x"]]), + IterateInvocation(id="downstream"), + FixedCollectionTestInvocation(id="for_items"), + ForInvocation(id="for"), + ForReturnInvocation(id="return"), + MarkedAnyTypeTestInvocation(id="selected_value", marker="s:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("upstream_src", "collection", "upstream", "collection"), + ("upstream", "item", "body", "value"), + ("body", "value", "closing_collect", "item"), + ("closing_collect", "collection", "downstream_src", "first"), + ("downstream_src", "collection", "downstream", "collection"), + (for_parent, "item", "for_items", "value"), + ("for_items", "collection", "for", "collection"), + ("for", "item", "return", "output"), + create_loop_linkage("for", "return"), + ("for", "output_collection", "if", "false_input"), + (selected_axis, "item", "selected_value", "value"), + ("selected_value", "value", "if", "true_input"), + ("if", "value", "collect", "item"), + ], + ) + + +@pytest.mark.parametrize( + ("selected_is_downstream", "expected"), + [ + (True, ["s:b:['a0', 'a1']", "s:b:['a2']", "s:['x']"]), + (False, ["s:['a0', 'a1']", "s:['a2']"]), + ], + ids=["selects-reopened-axis", "selects-closed-axis"], +) +def test_literal_if_ignores_unselected_for_output_under_an_axis_connected_to_the_selected_one( + selected_is_downstream: bool, expected: list[str] +) -> None: + """An unselected For final output adds no frames even when its axis feeds, or is fed by, the selected axis.""" + _trace, state = _run_to_completion(_for_output_connected_to_selected_axis_graph(selected_is_downstream)) + + assert _source_values(state, "collect") == {(): expected} + + @pytest.mark.parametrize("condition_from_branch", [False, True], ids=["literal", "condition-from-branch"]) def test_if_over_cross_product_carried_by_its_deciding_input_completes(condition_from_branch: bool) -> None: """Unrelated axes combined by one input are that input's own frames, so the If keeps the cross product.""" @@ -3764,13 +3986,199 @@ def test_if_over_cross_product_carried_by_its_deciding_input_completes(condition assert _source_values(state, "collect") == {(0,): [11, 21, 31], (1,): [12, 22, 32]} -def test_literal_if_selecting_axis_reopened_after_a_collect_keeps_its_frames() -> None: - """A selected axis reopened downstream of a Collect that closed the other branch's axis keeps its pairing.""" +def test_literal_if_selecting_axis_reopened_directly_from_a_collect_pairs_frames() -> None: + """An Iterate fed straight from the Collect that closed the other axis has its length, so frames pair by index.""" _trace, state = _run_to_completion(_collect_then_reopen_graph(True)) assert _source_values(state, "collect") == {(0,): ["t:b:a1"], (1,): ["t:b:a2"]} +def test_literal_if_selecting_axis_reopened_through_two_collects_pairs_frames() -> None: + """A chain of direct reopens keeps the first axis's length, so the last reopened axis pairs with it.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="src", first=["a1", "a2", "a3"]), + IterateInvocation(id="first"), + MarkedAnyTypeTestInvocation(id="first_body", marker="b:"), + CollectInvocation(id="first_collect"), + IterateInvocation(id="middle"), + MarkedAnyTypeTestInvocation(id="middle_body", marker="c:"), + CollectInvocation(id="middle_collect"), + IterateInvocation(id="last"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("src", "collection", "first", "collection"), + ("first", "item", "first_body", "value"), + ("first_body", "value", "first_collect", "item"), + ("first_collect", "collection", "middle", "collection"), + ("middle", "item", "middle_body", "value"), + ("middle_body", "value", "middle_collect", "item"), + ("middle_collect", "collection", "last", "collection"), + ("last", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("first", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + _trace, state = _run_to_completion(graph) + + assert _source_values(state, "collect") == {(0,): ["t:c:b:a1"], (1,): ["t:c:b:a2"], (2,): ["t:c:b:a3"]} + + +def test_literal_if_selecting_axis_reopened_from_nested_collects_pairs_with_the_outer_axis() -> None: + """Each Collect in a chain closes one more axis, so the last one gathers one item per outermost frame.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="outer_src", first=[[["a"], ["b"]], [["c", "d"]]]), + IterateInvocation(id="outer"), + CollectionConcatInvocation(id="middle_src"), + IterateInvocation(id="middle"), + CollectionConcatInvocation(id="inner_src"), + IterateInvocation(id="inner"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + CollectInvocation(id="inner_collect"), + CollectInvocation(id="middle_collect"), + CollectInvocation(id="outer_collect"), + IterateInvocation(id="reopened"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("outer_src", "collection", "outer", "collection"), + ("outer", "item", "middle_src", "first"), + ("middle_src", "collection", "middle", "collection"), + ("middle", "item", "inner_src", "first"), + ("inner_src", "collection", "inner", "collection"), + ("inner", "item", "body", "value"), + ("body", "value", "inner_collect", "item"), + ("inner_collect", "collection", "middle_collect", "item"), + ("middle_collect", "collection", "outer_collect", "item"), + ("outer_collect", "collection", "reopened", "collection"), + ("reopened", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("outer", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + _trace, state = _run_to_completion(graph) + + assert _source_values(state, "collect") == { + (0,): ["t:[['b:a'], ['b:b']]"], + (1,): ["t:[['b:c', 'b:d']]"], + } + + +def test_literal_if_selecting_axis_nested_under_a_reopened_axis_pairs_frames() -> None: + """An Iterate nested under a direct reopen runs in the reopened frames, which pair with the closed axis.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="src", first=[["x1", "x2"], ["y1"]]), + IterateInvocation(id="first"), + AnyTypeTestInvocation(id="body"), + CollectInvocation(id="closing_collect"), + IterateInvocation(id="reopened"), + CollectionConcatInvocation(id="nested_src"), + IterateInvocation(id="nested"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("src", "collection", "first", "collection"), + ("first", "item", "body", "value"), + ("body", "value", "closing_collect", "item"), + ("closing_collect", "collection", "reopened", "collection"), + ("reopened", "item", "nested_src", "first"), + ("nested_src", "collection", "nested", "collection"), + ("nested", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("first", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + _trace, state = _run_to_completion(graph) + + assert _source_values(state, "collect") == {(0, 0): ["t:x1"], (0, 1): ["t:x2"], (1, 0): ["t:y1"]} + + +def test_literal_if_selecting_axis_reopened_within_kept_outer_frames_pairs_both_axes() -> None: + """A Collect closing an inner axis keeps the outer one; its reopening pairs with the inner axis per outer frame.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="outer_src", first=[["a1", "a2"], ["b1"]]), + IterateInvocation(id="outer"), + CollectionConcatInvocation(id="inner_src"), + IterateInvocation(id="inner"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + CollectInvocation(id="closing_collect"), + IterateInvocation(id="reopened"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("outer_src", "collection", "outer", "collection"), + ("outer", "item", "inner_src", "first"), + ("inner_src", "collection", "inner", "collection"), + ("inner", "item", "body", "value"), + ("body", "value", "closing_collect", "item"), + ("closing_collect", "collection", "reopened", "collection"), + ("reopened", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("inner", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + _trace, state = _run_to_completion(graph) + + assert _source_values(state, "collect") == {(0, 0): ["t:b:a1"], (0, 1): ["t:b:a2"], (1, 0): ["t:b:b1"]} + + +@pytest.mark.parametrize( + ("condition", "expected"), + [(False, [["a0", "a1"], ["a2"]]), (True, [["b0"], ["b1"], ["b2"]])], + ids=["selects-a", "selects-b"], +) +def test_literal_if_ignores_independent_axis_reaching_it_through_unselected_for_output( + condition: bool, expected: list[list[str]] +) -> None: + """An unselected For final output adds no frames, so the If runs once per item of the selected branch's axis.""" + nodes: list[Any] = [IfInvocation(id="if", condition=condition), CollectInvocation(id="collect")] + edges: list[Any] = [("if", "value", "collect", "item")] + for prefix, values, branch in ( + ("a", [["a0", "a1"], ["a2"]], "false_input"), + ("b", [["b0"], ["b1"], ["b2"]], "true_input"), + ): + nodes += [ + CollectionConcatInvocation(id=f"{prefix}_src", first=values), + IterateInvocation(id=f"{prefix}_iter"), + ForInvocation(id=f"{prefix}_for"), + ForReturnInvocation(id=f"{prefix}_return"), + ] + edges += [ + (f"{prefix}_src", "collection", f"{prefix}_iter", "collection"), + (f"{prefix}_iter", "item", f"{prefix}_for", "collection"), + (f"{prefix}_for", "item", f"{prefix}_return", "output"), + create_loop_linkage(f"{prefix}_for", f"{prefix}_return"), + (f"{prefix}_for", "output_collection", "if", branch), + ] + _trace, state = _run_to_completion(_graph_from(nodes, edges)) + + assert _source_values(state, "collect") == {(): expected} + + @pytest.mark.parametrize("condition", [True, False]) def test_root_condition_if_broadcasts_over_nested_axes_both_branches_share(condition: bool) -> None: """A root-frame condition selects between branches that run in the same nested frames, once per inner item.""" @@ -4003,11 +4411,76 @@ def test_per_item_if_gates_branch_for_together_with_the_for_nested_in_its_body(f "unrelated open iterator axes", id="literal-root-value-over-cross-product", ), + pytest.param( + _per_item_condition_with_collect_of_collects_branch_graph, + "deeper open iterator axes kept by a Collect", + id="per-item-condition-deeper-collect-beside-collect-of-collects", + ), pytest.param( lambda: _collect_then_reopen_graph(condition=False, extra_items=("x",)), "unrelated open iterator axes", id="literal-selects-axis-a-collect-closed-before-the-other", ), + pytest.param( + lambda: _collect_then_reopen_graph(condition=True, extra_items=("x",)), + "unrelated open iterator axes", + id="literal-selects-axis-reopened-through-a-node-after-the-collect", + ), + pytest.param( + lambda: _collect_then_reopen_graph(condition=True, appended_items=("x",)), + "unrelated open iterator axes", + id="literal-selects-axis-reopened-from-a-collect-that-appends-items", + ), + pytest.param( + _collect_with_literal_items_then_reopen_graph, + "unrelated open iterator axes", + id="literal-selects-axis-reopened-from-a-collect-with-literal-items", + ), + pytest.param( + lambda: _collect_then_reopen_graph(condition=False), + "unrelated open iterator axes", + id="literal-selects-axis-an-iterate-reopens", + ), + pytest.param( + lambda: _reopen_of_a_cross_product_graph(swap_add_inputs=False), + "unrelated open iterator axes", + id="literal-selects-closed-axis-of-a-reopened-cross-product", + ), + pytest.param( + lambda: _reopen_of_a_cross_product_graph(swap_add_inputs=True), + "unrelated open iterator axes", + id="literal-selects-closed-axis-of-a-reopened-swapped-cross-product", + ), + pytest.param( + lambda: _reopen_of_a_cross_product_graph(swap_add_inputs=True, condition=True), + "unrelated open iterator axes", + id="literal-selects-axis-reopening-a-cross-product", + ), + pytest.param( + lambda: _collect_of_collect_reopen_graph([["a"], [], ["b", "c"]], "member"), + "unrelated open iterator axes", + id="literal-selects-collect-of-collects-reopening-over-the-inner-axis", + ), + pytest.param( + lambda: _collect_of_collect_reopen_graph([["a"], ["b", "c"]], "group", extra_collect=True), + "unrelated open iterator axes", + id="literal-selects-reopening-of-more-collects-than-axes", + ), + *( + pytest.param( + lambda unselected_axis=unselected_axis: _collect_of_collect_reopen_graph( + [["a"], ["b", "c"]], unselected_axis, extra_collect=True, append_unrelated_axis=True + ), + "unrelated open iterator axes", + id=f"literal-selects-reopening-of-collects-keeping-unrelated-axes-over-{unselected_axis}", + ) + for unselected_axis in ("group", "appended") + ), + pytest.param( + _reopen_of_a_different_closed_axis_graph, + "unrelated open iterator axes", + id="literal-selects-axis-reopening-a-different-closed-axis", + ), ], ) def test_if_without_one_nested_chain_of_frames_is_rejected_before_any_node_runs( From a219e95e8d837833da6000c6be4f2ea7e6b36075 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Sun, 18 Oct 2026 18:27:29 -0500 Subject: [PATCH 5/5] fix(graph): preserve index-only If frames --- .../app/services/shared/graph_if_runtime.py | 8 ++++-- .../test_execution_engine_differential.py | 27 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index 67f333a03f2..c9a75e9438b 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -480,9 +480,13 @@ def _get_if_iteration_paths_with_ancestry( get_prepared_nodes_for_source(innermost_id) if innermost_id in state.source_prepared_mapping else set() ) iterator_edges = state.graph._get_output_edges(innermost_id, ITEM_FIELD) - if not prepared_ids or not iterator_edges: + if not prepared_ids: return [] - frame_paths = {get_prepared_edge_iteration_path(iterator_edges[0], prepared_id) for prepared_id in prepared_ids} + frame_paths = ( + {get_prepared_edge_iteration_path(iterator_edges[0], prepared_id) for prepared_id in prepared_ids} + if iterator_edges + else {state._get_iteration_path(prepared_id) for prepared_id in prepared_ids} + ) paths = { frame_path if frame_path[: len(branch_path)] == branch_path else branch_path for branch_path in paths diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index 17c7aa74a89..38a42f86ed2 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -11241,6 +11241,33 @@ def test_literal_if_branch_feeds_per_item_if_without_stranding_unselected_root( assert "per_item_unselected" not in trace +def test_per_item_if_condition_from_iterate_index_selects_literal_values() -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[10, 20])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="false_value", value="no")) + graph.add_node(AnyTypeTestInvocation(id="true_value", value="yes")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="output")) + for source, source_field, destination, destination_field in ( + ("source", "collection", "iterate", "collection"), + ("iterate", "index", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("false_value", "value", "if", "false_input"), + ("true_value", "value", "if", "true_input"), + ("if", "value", "output", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + trace, state = _run_graph(GraphExecutionState(graph=graph)) + + assert state.is_complete(), f"graph stalled after {trace}" + output_exec_id = next(iter(state.source_prepared_mapping["output"])) + assert state.results[output_exec_id].collection == ["no", "yes"] + assert trace.count("if") == 2 + + @pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) def test_deep_branch_root_false_with_two_outer_frames_and_nested_collect( force_compatibility_scheduler: bool,