From 7120d980469fd653ed33d5d4da96848dda80ced9 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Fri, 10 Jul 2026 19:41:39 -0500 Subject: [PATCH 1/3] fix empty collector graph stalls --- invokeai/app/services/shared/graph.py | 12 ++++++++++++ tests/test_graph_execution_state.py | 23 +++++++++++++++++++++++ tests/test_node_graph.py | 14 +++++++------- 3 files changed, 42 insertions(+), 7 deletions(-) diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index 80c6cae951a..f7c97104585 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -523,6 +523,11 @@ def _get_parent_iteration_mappings_without_iterators( mappings.append(mapping) return mappings + def _mark_source_node_empty(self, source_node_id: str) -> None: + self._state.source_prepared_mapping[source_node_id] = set() + self._state.executed.add(source_node_id) + self._state.executed_history.append(source_node_id) + def _get_parent_iteration_mappings(self, next_node_id: str, graph: nx.DiGraph) -> list[list[tuple[str, str]]]: parent_node_ids = [source_id for source_id, _ in graph.in_edges(next_node_id)] iterator_graph = self.iterator_graph(graph) @@ -705,6 +710,10 @@ def prepare(self, base_g: Optional[nx.DiGraph] = None) -> Optional[str]: if create_results is not None: new_node_ids.extend(create_results) + if not new_node_ids: + self._mark_source_node_empty(next_node_id) + return next_node_id + return next(iter(new_node_ids), None) @@ -1522,6 +1531,9 @@ def _validate_collector_edge_rules( and edge.source.field == COLLECTION_FIELD and not self._is_destination_field_list_of_Any(edge) and not self._is_destination_field_Any(edge) + # A materialized collector may legitimately have no input edges when its upstream iterator had zero + # iterations. The source graph has already validated the collector's input and output types. + and self._get_input_edges(source_node.id) ): err = self._is_collector_connection_valid(edge.source.node_id, new_output=edge.destination) if err is not None: diff --git a/tests/test_graph_execution_state.py b/tests/test_graph_execution_state.py index 0156e0319ad..52dae2613e5 100644 --- a/tests/test_graph_execution_state.py +++ b/tests/test_graph_execution_state.py @@ -14,6 +14,7 @@ BooleanCollectionOutput, BooleanInvocation, BooleanOutput, + IntegerCollectionInvocation, ) from invokeai.app.services.shared.graph import ( CollectInvocation, @@ -459,6 +460,28 @@ def test_graph_state_collects(): assert sorted(g.results[n6[0].id].collection) == sorted(test_prompts) +def test_graph_state_empty_iterator_collects_and_completes(): + graph = Graph() + graph.add_node(IntegerCollectionInvocation(id="collection", collection=[])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AddInvocation(id="add", b=1)) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(IntegerCollectionInvocation(id="consumer")) + graph.add_edge(create_edge("collection", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "add", "a")) + graph.add_edge(create_edge("add", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "consumer", "collection")) + + state = GraphExecutionState(graph=graph) + execute_all_nodes(state) + + assert state.is_complete() + prepared_collect_id = next(iter(state.source_prepared_mapping["collect"])) + assert state.results[prepared_collect_id].collection == [] + prepared_consumer_id = next(iter(state.source_prepared_mapping["consumer"])) + assert state.results[prepared_consumer_id].collection == [] + + def test_graph_state_resumes_partially_executed_session_after_json_round_trip(): graph = Graph() graph.add_node(RangeInvocation(id="c", start=1, stop=5, step=1)) diff --git a/tests/test_node_graph.py b/tests/test_node_graph.py index 4f3b262204a..c153c44404b 100644 --- a/tests/test_node_graph.py +++ b/tests/test_node_graph.py @@ -1070,10 +1070,10 @@ def test_iterator_collector_iterator_chain_with_empty_collection(): session = GraphExecutionState(graph=g) run_session_with_mock_context(session) - # With empty collection, iterators don't create execution nodes, so collectors don't execute - # Verify that the final collector was never prepared (which is correct behavior) - assert n7.id not in session.source_prepared_mapping - - # Verify only the source collection node executed - assert n1.id in session.source_prepared_mapping - assert len(session.source_prepared_mapping[n1.id]) == 1 + first_output = get_single_output_from_session(session, n4.id) + final_output = get_single_output_from_session(session, n7.id) + assert isinstance(first_output, CollectInvocationOutput) + assert isinstance(final_output, CollectInvocationOutput) + assert first_output.collection == [] + assert final_output.collection == [] + assert session.is_complete() From 16b417b81cbbc2d19d7fd27d2f0b788cda5e4a72 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Sat, 11 Jul 2026 17:27:57 -0500 Subject: [PATCH 2/3] fix empty collector after grouped materialization --- invokeai/app/services/shared/graph.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index f7c97104585..bb31cb231f3 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -698,9 +698,12 @@ def prepare(self, base_g: Optional[nx.DiGraph] = None) -> Optional[str]: new_node_ids: list[str] = [] if isinstance(next_node, CollectInvocation): - for iteration_path, iteration_mappings in self._get_collect_iteration_mapping_groups( + iteration_mapping_groups = self._get_collect_iteration_mapping_groups( self._state.graph._get_input_edges(next_node_id) - ): + ) + if not iteration_mapping_groups: + iteration_mapping_groups = [((), [])] + for iteration_path, iteration_mappings in iteration_mapping_groups: create_results = self.create_execution_node(next_node_id, iteration_mappings, iteration_path) if create_results is not None: new_node_ids.extend(create_results) From 433397a5072e19090153130e0e85cfe9c6df0636 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Mon, 20 Jul 2026 22:20:23 -0500 Subject: [PATCH 3/3] address empty collector review feedback --- invokeai/app/services/shared/graph.py | 75 +++++++++++++++++++++------ tests/test_graph_execution_state.py | 69 ++++++++++++++++++++++++ tests/test_node_graph.py | 9 ++++ 3 files changed, 136 insertions(+), 17 deletions(-) diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index bb31cb231f3..d65e5213cc2 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -397,7 +397,7 @@ def _create_execution_node_copy(self, node: BaseInvocation, node_id: str, iterat def _attach_execution_edges(self, exec_node_id: str, new_edges: list[Edge]) -> None: for edge in new_edges: - self._state.execution_graph.add_edge( + self._state.execution_graph._add_execution_edge( Edge( source=edge.source, destination=EdgeConnection(node_id=exec_node_id, field=edge.destination.field), @@ -423,12 +423,32 @@ def _get_ordered_prepared_nodes_for_source(self, source_node_id: str) -> list[st key=lambda exec_node_id: (self._state._get_iteration_path(exec_node_id), exec_node_id), ) + def _get_collect_candidate_group_keys(self, edge: Edge) -> set[tuple[int, ...]]: + source_node_id = edge.source.node_id + iterator_node_ids = self.get_node_iterators(source_node_id) + if isinstance(self._state.graph.get_node(source_node_id), IterateInvocation): + iterator_node_ids.append(source_node_id) + + group_depth = len(iterator_node_ids) + if edge.destination.field == ITEM_FIELD: + group_depth = max(group_depth - 1, 0) + if group_depth == 0: + return {()} + + return { + iteration_path[:group_depth] + for iterator_node_id in iterator_node_ids + for prepared_id in self._get_ordered_prepared_nodes_for_source(iterator_node_id) + if len(iteration_path := self._state._get_iteration_path(prepared_id)) >= group_depth + } + def _get_collect_iteration_mapping_groups( self, input_edges: list[Edge] ) -> list[tuple[tuple[int, ...], list[tuple[str, str]]]]: prepared_inputs: list[tuple[Edge, str, str, tuple[int, ...]]] = [] group_keys: set[tuple[int, ...]] = set() for edge in input_edges: + group_keys.update(self._get_collect_candidate_group_keys(edge)) prepared_nodes = self._get_ordered_prepared_nodes_for_source(edge.source.node_id) for prepared_id in prepared_nodes: prepared_edge = Edge( @@ -441,6 +461,9 @@ def _get_collect_iteration_mapping_groups( (prepared_edge, edge.source.node_id, prepared_id, self._state._get_iteration_path(prepared_id)) ) + if not group_keys: + group_keys.add(()) + final_group_keys = sorted( group_key for group_key in group_keys @@ -570,7 +593,10 @@ def create_execution_node( for iteration_index in iteration_indexes: new_node = self._create_execution_node_copy(node, node_id, iteration_index) if iteration_path is not None: - self._state._prepared_registry().set_iteration_path(new_node.id, iteration_path) + prepared_iteration_path = iteration_path + if isinstance(node, IterateInvocation): + prepared_iteration_path += (iteration_index,) + self._state._prepared_registry().set_iteration_path(new_node.id, prepared_iteration_path) self._attach_execution_edges(new_node.id, new_edges) self._initialize_execution_node(new_node.id) new_nodes.append(new_node.id) @@ -701,23 +727,27 @@ def prepare(self, base_g: Optional[nx.DiGraph] = None) -> Optional[str]: iteration_mapping_groups = self._get_collect_iteration_mapping_groups( self._state.graph._get_input_edges(next_node_id) ) - if not iteration_mapping_groups: - iteration_mapping_groups = [((), [])] for iteration_path, iteration_mappings in iteration_mapping_groups: create_results = self.create_execution_node(next_node_id, iteration_mappings, iteration_path) - if create_results is not None: - new_node_ids.extend(create_results) + new_node_ids.extend(create_results) else: + parent_iterator_nodes = self.get_node_iterators(next_node_id) for iteration_mappings in self._get_parent_iteration_mappings(next_node_id, g): - create_results = self.create_execution_node(next_node_id, iteration_mappings) - if create_results is not None: - new_node_ids.extend(create_results) + iteration_path = None + if not parent_iterator_nodes: + iteration_path = max( + (self._state._get_iteration_path(prepared_id) for _, prepared_id in iteration_mappings), + key=lambda path: (len(path), path), + default=(), + ) + create_results = self.create_execution_node(next_node_id, iteration_mappings, iteration_path) + new_node_ids.extend(create_results) if not new_node_ids: self._mark_source_node_empty(next_node_id) return next_node_id - return next(iter(new_node_ids), None) + return new_node_ids[0] class _ExecutionScheduler: @@ -1344,7 +1374,13 @@ def add_edge(self, edge: Edge) -> None: :raises InvalidEdgeError: the provided edge is invalid. """ - self._validate_edge(edge) + self._add_edge(edge, allow_inputless_source_collector=False) + + def _add_execution_edge(self, edge: Edge) -> None: + self._add_edge(edge, allow_inputless_source_collector=True) + + def _add_edge(self, edge: Edge, allow_inputless_source_collector: bool) -> None: + self._validate_edge(edge, allow_inputless_source_collector) if edge not in self.edges: self.edges.append(edge) else: @@ -1520,7 +1556,11 @@ def _validate_iterator_edge_rules( raise InvalidEdgeError(f"Iterator output type does not match iterator input type ({edge}): {err}") def _validate_collector_edge_rules( - self, edge: Edge, source_node: BaseInvocation, destination_node: BaseInvocation + self, + edge: Edge, + source_node: BaseInvocation, + destination_node: BaseInvocation, + allow_inputless_source_collector: bool, ) -> None: if isinstance(destination_node, CollectInvocation) and edge.destination.field in (ITEM_FIELD, COLLECTION_FIELD): err = self._is_collector_connection_valid( @@ -1534,22 +1574,23 @@ def _validate_collector_edge_rules( and edge.source.field == COLLECTION_FIELD and not self._is_destination_field_list_of_Any(edge) and not self._is_destination_field_Any(edge) - # A materialized collector may legitimately have no input edges when its upstream iterator had zero - # iterations. The source graph has already validated the collector's input and output types. - and self._get_input_edges(source_node.id) ): + if allow_inputless_source_collector and not any( + edge.destination.node_id == source_node.id for edge in self.edges + ): + return err = self._is_collector_connection_valid(edge.source.node_id, new_output=edge.destination) if err is not None: raise InvalidEdgeError(f"Collector input type does not match collector output type ({edge}): {err}") - def _validate_edge(self, edge: Edge): + def _validate_edge(self, edge: Edge, allow_inputless_source_collector: bool = False): """Validates that a new edge doesn't create a cycle in the graph""" source_node, destination_node = self._get_edge_nodes(edge) self._validate_edge_destination_uniqueness(edge, destination_node) self._validate_edge_would_not_create_cycle(edge) self._validate_edge_field_compatibility(edge, source_node, destination_node) self._validate_iterator_edge_rules(edge, source_node, destination_node) - self._validate_collector_edge_rules(edge, source_node, destination_node) + self._validate_collector_edge_rules(edge, source_node, destination_node, allow_inputless_source_collector) def has_node(self, node_id: str) -> bool: """Determines whether or not a node exists in the graph.""" diff --git a/tests/test_graph_execution_state.py b/tests/test_graph_execution_state.py index 52dae2613e5..95f1e41e0e0 100644 --- a/tests/test_graph_execution_state.py +++ b/tests/test_graph_execution_state.py @@ -47,6 +47,16 @@ def invoke(self, context: InvocationContext) -> IntegerCollectionTestInvocationO return IntegerCollectionTestInvocationOutput(collection=[base, base + 1]) +class MaybeEmptyIntegerCollectionTestInvocation(BaseInvocation): + value: int = InputField(default=0) + always_empty: bool = InputField(default=False) + + def invoke(self, context: InvocationContext) -> IntegerCollectionTestInvocationOutput: + if self.always_empty or self.value == 0: + return IntegerCollectionTestInvocationOutput(collection=[]) + return IntegerCollectionTestInvocationOutput(collection=[self.value]) + + class IntegerCollectionPassthroughTestInvocation(BaseInvocation): collection: list[int] = InputField(default=[]) @@ -482,6 +492,30 @@ def test_graph_state_empty_iterator_collects_and_completes(): assert state.results[prepared_consumer_id].collection == [] +def test_graph_state_multiple_empty_iterator_branches_complete(): + graph = Graph() + for branch in ("first", "second"): + graph.add_node(IntegerCollectionInvocation(id=f"{branch}_collection", collection=[])) + graph.add_node(IterateInvocation(id=f"{branch}_iterate")) + graph.add_node(AddInvocation(id=f"{branch}_add", b=1)) + graph.add_node(CollectInvocation(id=f"{branch}_collect")) + graph.add_node(IntegerCollectionPassthroughTestInvocation(id=f"{branch}_consumer")) + graph.add_edge(create_edge(f"{branch}_collection", "collection", f"{branch}_iterate", "collection")) + graph.add_edge(create_edge(f"{branch}_iterate", "item", f"{branch}_add", "a")) + graph.add_edge(create_edge(f"{branch}_add", "value", f"{branch}_collect", "item")) + graph.add_edge(create_edge(f"{branch}_collect", "collection", f"{branch}_consumer", "collection")) + + state = GraphExecutionState(graph=graph) + execute_all_nodes(state) + + assert state.is_complete() + assert state.next() is None + assert state.next() is None + for branch in ("first", "second"): + prepared_consumer_id = next(iter(state.source_prepared_mapping[f"{branch}_consumer"])) + assert state.results[prepared_consumer_id].collection == [] + + def test_graph_state_resumes_partially_executed_session_after_json_round_trip(): graph = Graph() graph.add_node(RangeInvocation(id="c", start=1, stop=5, step=1)) @@ -780,6 +814,41 @@ def test_graph_collector_nested_under_outer_iterator_collects_only_current_outer assert consumer_collections == [[0, 1], [10, 11]] +@pytest.mark.parametrize( + ("always_empty", "expected_collection"), + [(True, [[], []]), (False, [[], [1]])], +) +def test_graph_collector_nested_under_outer_iterator_preserves_empty_groups( + always_empty: bool, expected_collection: list[list[int]] +): + graph = Graph() + graph.add_node(RangeInvocation(id="outer_range", start=0, stop=2, step=1)) + graph.add_node(IterateInvocation(id="outer_iter")) + graph.add_node(MaybeEmptyIntegerCollectionTestInvocation(id="inner_collection", always_empty=always_empty)) + graph.add_node(IterateInvocation(id="inner_iter")) + graph.add_node(AddInvocation(id="inner_item", b=0)) + graph.add_node(CollectInvocation(id="inner_collect")) + graph.add_node(IntegerCollectionPassthroughTestInvocation(id="per_outer_consumer")) + graph.add_node(CollectInvocation(id="outer_collect")) + + graph.add_edge(create_edge("outer_range", "collection", "outer_iter", "collection")) + graph.add_edge(create_edge("outer_iter", "item", "inner_collection", "value")) + graph.add_edge(create_edge("inner_collection", "collection", "inner_iter", "collection")) + graph.add_edge(create_edge("inner_iter", "item", "inner_item", "a")) + graph.add_edge(create_edge("inner_item", "value", "inner_collect", "item")) + graph.add_edge(create_edge("inner_collect", "collection", "per_outer_consumer", "collection")) + graph.add_edge(create_edge("per_outer_consumer", "collection", "outer_collect", "item")) + + state = GraphExecutionState(graph=graph) + execute_all_nodes(state) + + prepared_inner_collect_ids = state.source_prepared_mapping["inner_collect"] + assert sorted(state._get_iteration_path(node_id) for node_id in prepared_inner_collect_ids) == [(0,), (1,)] + prepared_outer_collect_id = next(iter(state.source_prepared_mapping["outer_collect"])) + assert state.results[prepared_outer_collect_id].collection == expected_collection + assert state.is_complete() + + def test_graph_collector_reuses_outer_collection_input_for_each_nested_iterator_group(): graph = Graph() diff --git a/tests/test_node_graph.py b/tests/test_node_graph.py index c153c44404b..f8a826b4e4e 100644 --- a/tests/test_node_graph.py +++ b/tests/test_node_graph.py @@ -276,6 +276,15 @@ def test_graph_connects_collector(): g.add_edge(e3) +def test_graph_rejects_collector_output_edge_before_input_edge(): + graph = Graph() + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(ListPassThroughInvocation(id="consumer")) + + with pytest.raises(InvalidEdgeError, match="Collector must have at least one item or collection input edge"): + graph.add_edge(create_edge("collect", "collection", "consumer", "collection")) + + # TODO: test that derived types mixed with base types are compatible