From 9f7958f26d5266ca09b020da61255389e0ed15f9 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 7 Jul 2026 18:47:42 -0500 Subject: [PATCH 1/2] Fix nested collector iteration scope --- invokeai/app/services/shared/README.md | 42 ++++++--- invokeai/app/services/shared/graph.py | 122 ++++++++++++++++++++++--- tests/test_graph_execution_state.py | 82 ++++++++++++++++- 3 files changed, 218 insertions(+), 28 deletions(-) diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index c6603f2c5d2..317a4e16b0c 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -131,7 +131,8 @@ mutation helpers. Those helpers reject changes once the affected nodes have alre - iteration path - runtime state such as pending, ready, executed, or skipped - **Ready queues grouped by class** (private attrs): `_ready_queues: dict[class_name, deque[str]]`, - `_active_class: Optional[str]`. Optional `ready_order: list[str]` to prioritize classes. + `_active_class: Optional[str]`. Optional `ready_order: list[str]` to prioritize classes. Queues are rebuilt from + persisted execution state when a session is deserialized. ### 4.2 Core methods @@ -179,13 +180,19 @@ Workflow-call note: - `_PreparedExecRegistry` Owns the relationship between source graph nodes and prepared execution graph nodes, plus cached metadata such as iteration path and runtime state. - `_ExecutionMaterializer` Expands source graph nodes into concrete execution graph nodes when the scheduler runs out of - ready work. When matching prepared parents for a downstream exec node, skipped prepared exec nodes are ignored and - cannot be selected as live inputs. + ready work. It owns iterator expansion, collector grouping, prepared-parent selection, and creation of execution-graph + edges. When matching prepared parents for a downstream exec node, skipped prepared exec nodes are ignored and cannot + be selected as live inputs. - `_ExecutionScheduler` Owns indegree transitions, ready queues, class batching, and downstream release on completion. -- `_ExecutionRuntime` Owns iteration-path lookup and input hydration for prepared exec nodes. +- `_ExecutionRuntime` Owns iteration-path lookup, collect input ordering, and input hydration for prepared exec nodes. - `_IfBranchScheduler` Applies lazy `If` semantics by deferring branch-local work until the condition is known, then releasing the selected branch and skipping the unselected branch. +`GraphExecutionState.model_post_init()` rehydrates private runtime helpers and caches after normal construction or a +JSON/model round trip. Rehydration reconstructs prepared exec metadata, cached iteration paths, resolved `If` branch +state when the condition is already available, and ready queues from `execution_graph`, `indegree`, `executed`, and +`results`. This keeps persisted sessions resumable without persisting private helper objects. + ### 4.4 Preparation (`_prepare()`) - Build a flat DAG from the **source** graph. @@ -196,14 +203,21 @@ Workflow-call note: 1. if it is an iterator, *its inputs are already executed*, 1. it has *no unexecuted iterator ancestors*. -- If the node is a **CollectInvocation**: collapse all prepared parents into one mapping and create **one** exec node. +- If the node is a **CollectInvocation**: group prepared parent exec nodes by iteration path and create one collector + exec node per group. A collector collapses the immediate iterator that feeds its `item` input, but preserves enclosing + iterator paths. This lets a shape such as `outer_iter -> inner_collection -> inner_iter -> collect -> consumer` + produce one collected result per outer iteration instead of mixing all inner items into one global collection. + Incoming `collection` inputs are treated as ancestor groups and are copied into each matching descendant item group. - Otherwise: compute all combinations of prepared iterator ancestors. For each combination, choose the prepared parent - for each upstream by matching iterator ancestry, then create **one** exec node. + for each upstream by matching iterator ancestry, then create **one** exec node. If a node no longer has visible + iterator ancestors because the source path crosses a collector, prepared parent iteration paths are still used to + materialize one downstream exec node for each preserved collector path. - For each new exec node: - Deep-copy the source node; assign a fresh ID (and `index` for iterators). + - Cache the preserved iteration path when the materializer has one, such as for grouped collectors. - Wire edges from chosen prepared parents. - Set `indegree = number of unmet inputs` (i.e., parents not yet executed). - Try to resolve any `If`-specific scheduling state. @@ -213,9 +227,10 @@ Workflow-call note: - `_enqueue_if_ready(nid)` enqueues by class name only when `indegree == 0`, the node has not already executed, and the node is not deferred by an unresolved `If`. -- `_get_next_node()` drains the `_active_class` queue FIFO; when empty, selects the next nonempty class queue (by - `ready_order` if set, else alphabetical), and continues. Optional fairness knobs can limit batch size per class; - default is drain fully. +- `_get_next_node()` drains the `_active_class` queue; when empty, selects the next nonempty class queue (by + `ready_order` if set, else alphabetical), and continues. Within each class queue, ready exec nodes are ordered by + iteration path so expanded iterator work runs in a stable outer-to-inner order. Optional fairness knobs can limit + batch size per class; default is drain fully. #### 4.5.1 Indegree (what it is and how it's used) @@ -232,9 +247,10 @@ Run `C` -> `D:0` -> enqueue `D`. Run `D` -> done. ### 4.6 Input hydration (`_prepare_inputs()`) -- For **CollectInvocation**: gather all incoming `item` values into `collection`, sorting inputs by iteration path so - collected results are stable across expanded iterations. Incoming `collection` values are merged first, then incoming - `item` values are appended. +- For **CollectInvocation**: gather the materialized incoming `item` values into `collection`, sorting inputs by + iteration path so collected results are stable across expanded iterations. Incoming `collection` values are merged + first, then incoming `item` values are appended. By the time hydration runs, the materializer has already selected the + iteration group for this collector exec node, so hydration only sees inputs that belong to that group. - For **IfInvocation**: hydrate only `condition` and the selected branch input. As a defensive guard against inconsistent runtime or deserialized session state, the runtime raises if the selected input edge points at an exec node with no stored runtime output. In normal scheduling this path should be unreachable. @@ -282,6 +298,8 @@ In normal execution, all runtime expansion occurs in `execution_graph` with trac - Nodes are enqueued only when `indegree == 0` and they are not deferred by an unresolved `If`. - `results` and `errors` are keyed by **exec node id**. - Collectors aggregate `item` inputs and may also merge incoming `collection` inputs during runtime hydration. + Collectors nested under iterators preserve enclosing iteration paths, so downstream consumers materialize per enclosing + iteration instead of receiving a mixed collection from unrelated outer iterations. - Branch-exclusive nodes behind an unselected `If` branch are skipped, not failed. ## 7) Extensibility diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index bebdde4a7cb..542d91aa078 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -411,17 +411,106 @@ def _initialize_execution_node(self, exec_node_id: str) -> None: self._state._try_resolve_if_node(exec_node_id) self._state._enqueue_if_ready(exec_node_id) - def _get_collect_iteration_mappings(self, parent_node_ids: list[str]) -> list[tuple[str, str]]: - all_iteration_mappings: list[tuple[str, str]] = [] - for source_node_id in parent_node_ids: - prepared_nodes = self._get_prepared_nodes_for_source(source_node_id) - all_iteration_mappings.extend((source_node_id, prepared_id) for prepared_id in prepared_nodes) - return all_iteration_mappings + def _get_collect_iteration_group_key(self, edge: Edge) -> tuple[int, ...]: + path = self._state._get_iteration_path(edge.source.node_id) + if edge.destination.field == ITEM_FIELD: + return path[:-1] + return path + + def _get_ordered_prepared_nodes_for_source(self, source_node_id: str) -> list[str]: + return sorted( + self._get_prepared_nodes_for_source(source_node_id), + key=lambda exec_node_id: (self._state._get_iteration_path(exec_node_id), exec_node_id), + ) + + def _get_collect_iteration_mapping_groups( + self, input_edges: list[Edge] + ) -> list[tuple[tuple[int, ...], list[tuple[str, str]]]]: + grouped_mappings: dict[tuple[int, ...], list[tuple[str, str]]] = {} + for edge in input_edges: + prepared_nodes = self._get_ordered_prepared_nodes_for_source(edge.source.node_id) + for prepared_id in prepared_nodes: + prepared_edge = Edge( + source=EdgeConnection(node_id=prepared_id, field=edge.source.field), + destination=edge.destination, + ) + group_key = self._get_collect_iteration_group_key(prepared_edge) + grouped_mappings.setdefault(group_key, []).append((edge.source.node_id, prepared_id)) + + final_group_keys = sorted( + group_key + for group_key in grouped_mappings + if not any( + group_key != other_group_key and other_group_key[: len(group_key)] == group_key + for other_group_key in grouped_mappings + ) + ) + + return [ + ( + group_key, + [ + mapping + for source_group_key, mappings in sorted(grouped_mappings.items()) + if group_key[: len(source_group_key)] == source_group_key + for mapping in mappings + ], + ) + for group_key in final_group_keys + ] + + def _get_parent_iteration_mappings_without_iterators( + self, parent_node_ids: list[str] + ) -> list[list[tuple[str, str]]]: + parent_prepared_nodes = { + node_id: self._get_ordered_prepared_nodes_for_source(node_id) for node_id in parent_node_ids + } + iteration_paths = sorted( + { + self._state._get_iteration_path(prepared_id) + for prepared_nodes in parent_prepared_nodes.values() + for prepared_id in prepared_nodes + if self._state._get_iteration_path(prepared_id) != () + } + ) + if not iteration_paths: + iteration_paths = [()] + + mappings: list[list[tuple[str, str]]] = [] + for iteration_path in iteration_paths: + mapping: list[tuple[str, str]] = [] + for node_id, prepared_nodes in parent_prepared_nodes.items(): + matching_prepared_node = next( + ( + prepared_id + for prepared_id in prepared_nodes + if self._state._get_iteration_path(prepared_id) == iteration_path + ), + None, + ) + if matching_prepared_node is None: + matching_prepared_node = next( + ( + prepared_id + for prepared_id in prepared_nodes + if self._state._get_iteration_path(prepared_id) == () + ), + None, + ) + if matching_prepared_node is None: + break + mapping.append((node_id, matching_prepared_node)) + if len(mapping) == len(parent_node_ids): + mappings.append(mapping) + return mappings 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) iterator_nodes = self.get_node_iterators(next_node_id, iterator_graph) + if not iterator_nodes: + return self._get_parent_iteration_mappings_without_iterators(parent_node_ids) + iterator_nodes_prepared = [list(self._state.source_prepared_mapping[node_id]) for node_id in iterator_nodes] iterator_node_prepared_combinations = list(itertools.product(*iterator_nodes_prepared)) @@ -439,7 +528,12 @@ def _get_parent_iteration_mappings(self, next_node_id: str, graph: nx.DiGraph) - if all(prepared_id is not None for _, prepared_id in mapping) ] - def create_execution_node(self, node_id: str, iteration_node_map: list[tuple[str, str]]) -> list[str]: + def create_execution_node( + self, + node_id: str, + iteration_node_map: list[tuple[str, str]], + iteration_path: Optional[tuple[int, ...]] = None, + ) -> list[str]: """Prepares an iteration node and connects all edges, returning the new node id""" node = self._state.graph.get_node(node_id) @@ -451,6 +545,8 @@ def create_execution_node(self, node_id: str, iteration_node_map: list[tuple[str new_nodes: list[str] = [] 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) self._attach_execution_edges(new_node.id, new_edges) self._initialize_execution_node(new_node.id) new_nodes.append(new_node.id) @@ -578,12 +674,12 @@ def prepare(self, base_g: Optional[nx.DiGraph] = None) -> Optional[str]: new_node_ids: list[str] = [] if isinstance(next_node, CollectInvocation): - next_node_parents = [source_id for source_id, _ in g.in_edges(next_node_id)] - create_results = self.create_execution_node( - next_node_id, self._get_collect_iteration_mappings(next_node_parents) - ) - if create_results is not None: - new_node_ids.extend(create_results) + for iteration_path, iteration_mappings in self._get_collect_iteration_mapping_groups( + self._state.graph._get_input_edges(next_node_id) + ): + 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) else: for iteration_mappings in self._get_parent_iteration_mappings(next_node_id, g): create_results = self.create_execution_node(next_node_id, iteration_mappings) diff --git a/tests/test_graph_execution_state.py b/tests/test_graph_execution_state.py index 3ef77b90a1a..9fa6dcc5d7f 100644 --- a/tests/test_graph_execution_state.py +++ b/tests/test_graph_execution_state.py @@ -6,6 +6,7 @@ from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext from invokeai.app.invocations.collections import RangeInvocation +from invokeai.app.invocations.fields import InputField, OutputField from invokeai.app.invocations.logic import IfInvocation, IfInvocationOutput from invokeai.app.invocations.math import AddInvocation, MultiplyInvocation from invokeai.app.invocations.primitives import ( @@ -33,6 +34,25 @@ ) +class IntegerCollectionTestInvocationOutput(BaseInvocationOutput): + collection: list[int] = OutputField(default=[]) + + +class IntegerCollectionFromItemTestInvocation(BaseInvocation): + value: int = InputField(default=0) + + def invoke(self, context: InvocationContext) -> IntegerCollectionTestInvocationOutput: + base = self.value * 10 + return IntegerCollectionTestInvocationOutput(collection=[base, base + 1]) + + +class IntegerCollectionPassthroughTestInvocation(BaseInvocation): + collection: list[int] = InputField(default=[]) + + def invoke(self, context: InvocationContext) -> IntegerCollectionTestInvocationOutput: + return IntegerCollectionTestInvocationOutput(collection=self.collection.copy()) + + @pytest.fixture def simple_graph() -> Graph: g = Graph() @@ -702,6 +722,60 @@ def test_graph_nested_iterate_execution_order(execution_number: int): assert sum_values == [0, 1, 10, 11] +def test_graph_collector_nested_under_outer_iterator_collects_only_current_outer_iteration_items(): + 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(IntegerCollectionFromItemTestInvocation(id="inner_collection")) + graph.add_node(IterateInvocation(id="inner_iter")) + graph.add_node(AddInvocation(id="inner_item", b=0)) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(IntegerCollectionPassthroughTestInvocation(id="per_outer_consumer")) + + 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", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "per_outer_consumer", "collection")) + + g = GraphExecutionState(graph=graph) + execute_all_nodes(g) + + prepared_consumer_ids = g.source_prepared_mapping["per_outer_consumer"] + consumer_collections = sorted(g.results[node_id].collection for node_id in prepared_consumer_ids) + + assert consumer_collections == [[0, 1], [10, 11]] + + +def test_graph_collector_reuses_outer_collection_input_for_each_nested_iterator_group(): + graph = Graph() + + graph.add_node(RangeInvocation(id="base_collection", start=100, stop=101, step=1)) + graph.add_node(RangeInvocation(id="outer_range", start=0, stop=2, step=1)) + graph.add_node(IterateInvocation(id="outer_iter")) + graph.add_node(IntegerCollectionFromItemTestInvocation(id="inner_collection")) + graph.add_node(IterateInvocation(id="inner_iter")) + graph.add_node(AddInvocation(id="inner_item", b=0)) + graph.add_node(CollectInvocation(id="collect")) + + graph.add_edge(create_edge("base_collection", "collection", "collect", "collection")) + 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", "collect", "item")) + + g = GraphExecutionState(graph=graph) + execute_all_nodes(g) + + prepared_collect_ids = g.source_prepared_mapping["collect"] + collect_results = sorted(g.results[node_id].collection for node_id in prepared_collect_ids) + + assert collect_results == [[100, 0, 1], [100, 10, 11]] + + def test_graph_validate_self_iterator_without_collection_input_raises_invalid_edge_error(): """Iterator nodes with no collection input should fail validation cleanly. @@ -1212,9 +1286,11 @@ def test_prepare_if_inputs_raises_when_selected_branch_source_has_no_result(): assert "iteration_path=()" in message -def test_get_collect_iteration_mappings_ignores_skipped_prepared_exec_nodes(): +def test_get_collect_iteration_mapping_groups_ignores_skipped_prepared_exec_nodes(): graph = Graph() graph.add_node(AnyTypeTestInvocation(id="parent", value="value")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("parent", "value", "collect", "item")) g = GraphExecutionState(graph=graph) @@ -1222,9 +1298,9 @@ def test_get_collect_iteration_mappings_ignores_skipped_prepared_exec_nodes(): active_exec_id = g._create_execution_node("parent", [])[0] g._set_prepared_exec_state(skipped_exec_id, "skipped") - mappings = g._materializer()._get_collect_iteration_mappings(["parent"]) + mappings = g._materializer()._get_collect_iteration_mapping_groups(graph._get_input_edges("collect")) - assert mappings == [("parent", active_exec_id)] + assert mappings == [((), [("parent", active_exec_id)])] def test_get_iteration_node_ignores_skipped_prepared_exec_nodes(): From 89fff2b0f076fbb63378f3ce679f8a338b6cb40a Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 7 Jul 2026 19:05:55 -0500 Subject: [PATCH 2/2] Test deeper nested collector iteration paths --- tests/test_graph_execution_state.py | 31 +++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/test_graph_execution_state.py b/tests/test_graph_execution_state.py index 9fa6dcc5d7f..4b35a2fe968 100644 --- a/tests/test_graph_execution_state.py +++ b/tests/test_graph_execution_state.py @@ -776,6 +776,37 @@ def test_graph_collector_reuses_outer_collection_input_for_each_nested_iterator_ assert collect_results == [[100, 0, 1], [100, 10, 11]] +def test_graph_collector_nested_under_three_iterators_preserves_outer_iteration_paths(): + 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(IntegerCollectionFromItemTestInvocation(id="middle_collection")) + graph.add_node(IterateInvocation(id="middle_iter")) + graph.add_node(IntegerCollectionFromItemTestInvocation(id="inner_collection")) + graph.add_node(IterateInvocation(id="inner_iter")) + graph.add_node(AddInvocation(id="inner_item", b=0)) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(IntegerCollectionPassthroughTestInvocation(id="per_middle_consumer")) + + graph.add_edge(create_edge("outer_range", "collection", "outer_iter", "collection")) + graph.add_edge(create_edge("outer_iter", "item", "middle_collection", "value")) + graph.add_edge(create_edge("middle_collection", "collection", "middle_iter", "collection")) + graph.add_edge(create_edge("middle_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", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "per_middle_consumer", "collection")) + + g = GraphExecutionState(graph=graph) + execute_all_nodes(g) + + prepared_consumer_ids = g.source_prepared_mapping["per_middle_consumer"] + consumer_collections = sorted(g.results[node_id].collection for node_id in prepared_consumer_ids) + + assert consumer_collections == [[0, 1], [10, 11], [100, 101], [110, 111]] + + def test_graph_validate_self_iterator_without_collection_input_raises_invalid_edge_error(): """Iterator nodes with no collection input should fail validation cleanly.