Fix nested collector iteration scope

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

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)