Fix graph stall for collectors downstream of empty iterators

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:

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))

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()