diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index 80c6cae951a..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 @@ -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.""" diff --git a/tests/test_graph_execution_state.py b/tests/test_graph_execution_state.py index 0156e0319ad..95f1e41e0e0 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, @@ -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() diff --git a/tests/test_node_graph.py b/tests/test_node_graph.py index 4f3b262204a..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 @@ -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()