Fix graph stall for collectors downstream of empty iterators

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
@@ -523,6 +546,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)
@@ -565,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)
@@ -693,19 +724,30 @@ 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)
-            ):
+            )
+            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:
@@ -1332,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:
@@ -1508,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(
@@ -1523,18 +1575,22 @@ def _validate_collector_edge_rules(
             and not self._is_destination_field_list_of_Any(edge)
             and not self._is_destination_field_Any(edge)
         ):
+            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."""

tests/test_graph_execution_state.py

@@ -14,6 +14,7 @@
     BooleanCollectionOutput,
     BooleanInvocation,
     BooleanOutput,
+    IntegerCollectionInvocation,
 )
 from invokeai.app.services.shared.graph import (
     CollectInvocation,
@@ -46,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=[])
 
@@ -459,6 +470,52 @@ 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_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))
@@ -757,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()
 

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
 
 
@@ -1070,10 +1079,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()