diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index 80c6cae951a..bb31cb231f3 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) @@ -693,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) @@ -705,6 +713,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 +1534,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()