Merge branch 'invoke-ai:main' into dype-paper

Author: JPPhoto

invokeai/app/services/shared/graph.py

@@ -123,6 +123,23 @@ def is_any(t: Any) -> bool:
     return t == Any or Any in get_args(t)
 
 
+def extract_collection_item_types(t: Any) -> set[Any]:
+    """Extracts list item types from a collection annotation, including unions containing list branches."""
+    if is_any(t):
+        return {Any}
+
+    if get_origin(t) is list:
+        return {arg for arg in get_args(t) if arg != NoneType}
+
+    item_types: set[Any] = set()
+    for arg in get_args(t):
+        if is_any(arg):
+            item_types.add(Any)
+        elif get_origin(arg) is list:
+            item_types.update(item_arg for item_arg in get_args(arg) if item_arg != NoneType)
+    return item_types
+
+
 def are_connection_types_compatible(from_type: Any, to_type: Any) -> bool:
     if not from_type or not to_type:
         return False
@@ -280,7 +297,7 @@ class CollectInvocationOutput(BaseInvocationOutput):
     )
 
 
-@invocation("collect", version="1.0.0")
+@invocation("collect", version="1.1.0")
 class CollectInvocation(BaseInvocation):
     """Collects values into a collection"""
 
@@ -292,7 +309,10 @@ class CollectInvocation(BaseInvocation):
         input=Input.Connection,
     )
     collection: list[Any] = InputField(
-        description="The collection, will be provided on execution", default=[], ui_hidden=True
+        description="An optional collection to append to",
+        default=[],
+        ui_type=UIType._Collection,
+        input=Input.Connection,
     )
 
     def invoke(self, context: InvocationContext) -> CollectInvocationOutput:
@@ -520,7 +540,9 @@ def _validate_edge(self, edge: Edge):
 
         # Validate that an edge to this node+field doesn't already exist
         input_edges = self._get_input_edges(edge.destination.node_id, edge.destination.field)
-        if len(input_edges) > 0 and not isinstance(to_node, CollectInvocation):
+        if len(input_edges) > 0 and (
+            not isinstance(to_node, CollectInvocation) or edge.destination.field != ITEM_FIELD
+        ):
             raise InvalidEdgeError(f"Edge already exists ({edge})")
 
         # Validate that no cycles would be created
@@ -546,8 +568,10 @@ def _validate_edge(self, edge: Edge):
                 raise InvalidEdgeError(f"Iterator output type does not match iterator input type ({edge}): {err}")
 
         # Validate if collector input type matches output type (if this edge results in both being set)
-        if isinstance(to_node, CollectInvocation) and edge.destination.field == ITEM_FIELD:
-            err = self._is_collector_connection_valid(edge.destination.node_id, new_input=edge.source)
+        if isinstance(to_node, CollectInvocation) and edge.destination.field in (ITEM_FIELD, COLLECTION_FIELD):
+            err = self._is_collector_connection_valid(
+                edge.destination.node_id, new_input=edge.source, new_input_field=edge.destination.field
+            )
             if err is not None:
                 raise InvalidEdgeError(f"Collector output type does not match collector input type ({edge}): {err}")
 
@@ -676,76 +700,152 @@ def _is_iterator_connection_valid(
 
         # Collector input type must match all iterator output types
         if isinstance(input_node, CollectInvocation):
-            collector_inputs = self._get_input_edges(input_node.id, ITEM_FIELD)
-            if len(collector_inputs) == 0:
-                return "Iterator input collector must have at least one item input edge"
-
-            # Traverse the graph to find the first collector input edge. Collectors validate that their collection
-            # inputs are all of the same type, so we can use the first input edge to determine the collector's type
-            first_collector_input_edge = collector_inputs[0]
-            first_collector_input_type = get_output_field_type(
-                self.get_node(first_collector_input_edge.source.node_id), first_collector_input_edge.source.field
-            )
-            resolved_collector_type = (
-                first_collector_input_type
-                if get_origin(first_collector_input_type) is None
-                else get_args(first_collector_input_type)
-            )
-            if not all((are_connection_types_compatible(resolved_collector_type, t) for t in output_field_types)):
+            input_root_type = self._get_collector_input_root_type(input_node.id)
+            if input_root_type is None:
+                return "Iterator input collector must have at least one item or collection input edge"
+            if not all((are_connection_types_compatible(input_root_type, t) for t in output_field_types)):
                 return "Iterator collection type must match all iterator output types"
 
         return None
 
+    def _resolve_collector_input_types(self, node_id: str, visited: Optional[set[str]] = None) -> set[Any]:
+        """Resolves possible item types for a collector's inputs, recursively following chained collectors."""
+        visited = visited or set()
+        if node_id in visited:
+            return set()
+        visited.add(node_id)
+
+        input_types: set[Any] = set()
+
+        for edge in self._get_input_edges(node_id, ITEM_FIELD):
+            input_field_type = get_output_field_type(self.get_node(edge.source.node_id), edge.source.field)
+            resolved_types = [input_field_type] if get_origin(input_field_type) is None else get_args(input_field_type)
+            input_types.update(t for t in resolved_types if t != NoneType)
+
+        for edge in self._get_input_edges(node_id, COLLECTION_FIELD):
+            source_node = self.get_node(edge.source.node_id)
+            if isinstance(source_node, CollectInvocation) and edge.source.field == COLLECTION_FIELD:
+                input_types.update(self._resolve_collector_input_types(source_node.id, visited.copy()))
+                continue
+
+            input_field_type = get_output_field_type(source_node, edge.source.field)
+            input_types.update(extract_collection_item_types(input_field_type))
+
+        return input_types
+
+    def _get_collector_input_root_type(self, node_id: str) -> Any | None:
+        input_types = self._resolve_collector_input_types(node_id)
+        non_any_input_types = {t for t in input_types if t != Any}
+        if len(non_any_input_types) == 0 and Any in input_types:
+            return Any
+        if len(non_any_input_types) == 0:
+            return None
+
+        type_tree = nx.DiGraph()
+        type_tree.add_nodes_from(non_any_input_types)
+        type_tree.add_edges_from([e for e in itertools.permutations(non_any_input_types, 2) if issubclass(e[1], e[0])])
+        type_degrees = type_tree.in_degree(type_tree.nodes)
+        root_types = [t[0] for t in type_degrees if t[1] == 0]  # type: ignore
+        if len(root_types) != 1:
+            return Any
+        return root_types[0]
+
     def _is_collector_connection_valid(
         self,
         node_id: str,
         new_input: Optional[EdgeConnection] = None,
+        new_input_field: Optional[str] = None,
         new_output: Optional[EdgeConnection] = None,
     ) -> str | None:
-        inputs = [e.source for e in self._get_input_edges(node_id, ITEM_FIELD)]
+        item_inputs = [e.source for e in self._get_input_edges(node_id, ITEM_FIELD)]
+        collection_inputs = [e.source for e in self._get_input_edges(node_id, COLLECTION_FIELD)]
         outputs = [e.destination for e in self._get_output_edges(node_id, COLLECTION_FIELD)]
 
         if new_input is not None:
-            inputs.append(new_input)
+            field = new_input_field or ITEM_FIELD
+            if field == ITEM_FIELD:
+                item_inputs.append(new_input)
+            elif field == COLLECTION_FIELD:
+                collection_inputs.append(new_input)
         if new_output is not None:
             outputs.append(new_output)
 
-        # Get input and output fields (the fields linked to the iterator's input/output)
-        input_field_types = [get_output_field_type(self.get_node(e.node_id), e.field) for e in inputs]
+        if len(item_inputs) == 0 and len(collection_inputs) == 0:
+            return "Collector must have at least one item or collection input edge"
+
+        # Get input and output fields (the fields linked to the collector's input/output)
+        item_input_field_types = [get_output_field_type(self.get_node(e.node_id), e.field) for e in item_inputs]
+        collection_input_field_types = [
+            get_output_field_type(self.get_node(e.node_id), e.field) for e in collection_inputs
+        ]
         output_field_types = [get_input_field_type(self.get_node(e.node_id), e.field) for e in outputs]
 
+        if not all((is_list_or_contains_list(t) or is_any(t) for t in collection_input_field_types)):
+            return "Collector collection input must be a collection"
+
         # Validate that all inputs are derived from or match a single type
         input_field_types = {
             resolved_type
-            for input_field_type in input_field_types
+            for input_field_type in item_input_field_types
             for resolved_type in (
                 [input_field_type] if get_origin(input_field_type) is None else get_args(input_field_type)
             )
             if resolved_type != NoneType
         }  # Get unique types
+
+        for input_conn, input_field_type in zip(collection_inputs, collection_input_field_types, strict=False):
+            source_node = self.get_node(input_conn.node_id)
+            if isinstance(source_node, CollectInvocation) and input_conn.field == COLLECTION_FIELD:
+                input_field_types.update(self._resolve_collector_input_types(source_node.id))
+                continue
+            input_field_types.update(extract_collection_item_types(input_field_type))
+
+        non_any_input_field_types = {t for t in input_field_types if t != Any}
         type_tree = nx.DiGraph()
-        type_tree.add_nodes_from(input_field_types)
-        type_tree.add_edges_from([e for e in itertools.permutations(input_field_types, 2) if issubclass(e[1], e[0])])
+        type_tree.add_nodes_from(non_any_input_field_types)
+        type_tree.add_edges_from(
+            [e for e in itertools.permutations(non_any_input_field_types, 2) if issubclass(e[1], e[0])]
+        )
         type_degrees = type_tree.in_degree(type_tree.nodes)
-        if sum((t[1] == 0 for t in type_degrees)) != 1:  # type: ignore
+        root_types = [t[0] for t in type_degrees if t[1] == 0]  # type: ignore
+        if len(root_types) > 1:
             return "Collector input collection items must be of a single type"
 
-        # Get the input root type
-        input_root_type = next(t[0] for t in type_degrees if t[1] == 0)  # type: ignore
+        # Get the input root type (if known)
+        input_root_type = root_types[0] if len(root_types) == 1 else None
 
         # Verify that all outputs are lists
         if not all(is_list_or_contains_list(t) or is_any(t) for t in output_field_types):
             return "Collector output must connect to a collection input"
 
         # Verify that all outputs match the input type (are a base class or the same class)
-        if not all(
-            is_any(t)
-            or is_union_subtype(input_root_type, get_args(t)[0])
-            or issubclass(input_root_type, get_args(t)[0])
-            for t in output_field_types
-        ):
+        if input_root_type is not None:
+            if not all(
+                is_any(t)
+                or is_union_subtype(input_root_type, get_args(t)[0])
+                or issubclass(input_root_type, get_args(t)[0])
+                for t in output_field_types
+            ):
+                return "Collector outputs must connect to a collection input with a matching type"
+        elif any(not is_any(t) and get_args(t)[0] != Any for t in output_field_types):
             return "Collector outputs must connect to a collection input with a matching type"
 
+        # If this collector outputs to another collector's collection input, validate against the downstream
+        # collector's resolved input type (if available).
+        for output in outputs:
+            output_node = self.get_node(output.node_id)
+            if not isinstance(output_node, CollectInvocation) or output.field != COLLECTION_FIELD:
+                continue
+            output_root_type = self._get_collector_input_root_type(output_node.id)
+            if output_root_type is None:
+                continue
+            if input_root_type is None:
+                if output_root_type != Any:
+                    return "Collector outputs must connect to a collection input with a matching type"
+                continue
+            if not are_connection_types_compatible(input_root_type, output_root_type):
+                return "Collector outputs must connect to a collection input with a matching type"
+
         return None
 
     def nx_graph(self) -> nx.DiGraph:
@@ -1211,8 +1311,19 @@ def _prepare_inputs(self, node: BaseInvocation):
         if isinstance(node, CollectInvocation):
             item_edges = [e for e in input_edges if e.destination.field == ITEM_FIELD]
             item_edges.sort(key=lambda e: (self._get_iteration_path(e.source.node_id), e.source.node_id))
-
-            output_collection = [copydeep(getattr(self.results[e.source.node_id], e.source.field)) for e in item_edges]
+            collection_edges = [e for e in input_edges if e.destination.field == COLLECTION_FIELD]
+            collection_edges.sort(key=lambda e: (self._get_iteration_path(e.source.node_id), e.source.node_id))
+
+            output_collection = []
+            for edge in collection_edges:
+                source_value = copydeep(getattr(self.results[edge.source.node_id], edge.source.field))
+                if isinstance(source_value, list):
+                    output_collection.extend(source_value)
+                else:
+                    output_collection.append(source_value)
+            output_collection.extend(
+                copydeep(getattr(self.results[e.source.node_id], e.source.field)) for e in item_edges
+            )
             node.collection = output_collection
         else:
             for edge in input_edges:

invokeai/frontend/web/public/locales/en.json

@@ -3107,9 +3107,9 @@
     "whatsNew": {
         "whatsNewInInvoke": "What's New in Invoke",
         "items": [
-            "FLUX.2 Klein Support: InvokeAI now supports the new FLUX.2 Klein models (4B and 9B variants) with GGUF, FP8, and Diffusers formats. Features include txt2img, img2img, inpainting, and outpainting. See 'Starter Models' to get started.",
-            "DyPE support for FLUX models improves high-resolution  (>1536 px up to 4K) images. Go to the 'Advanced Options' section to activate.",
-            "Z-Image Turbo diversity: Active 'Seed Variance Enhancer' under 'Advanced Options' to add diversitiy to your ZiT gens."
+            "Multi-user mode supports multiple isolated users on the same server.",
+            "Enhanced support for Z-Image and FLUX.2 Models.",
+            "Multiple user interface enhancements and new canvas features."
         ],
         "takeUserSurvey": "📣 Let us know how you like InvokeAI. Take our User Experience Survey!",
         "readReleaseNotes": "Read Release Notes",

invokeai/frontend/web/src/features/nodes/store/util/getCollectItemType.test.ts

@@ -41,4 +41,14 @@ describe(getCollectItemType.name, () => {
     const result = getCollectItemType({ add: addWithoutOutputValue, collect }, [n2, n1], [e1], n1.id);
     expect(result).toBeNull();
   });
+
+  it('should return the upstream collect item type for chained collects', () => {
+    const n1 = buildNode(collect);
+    const n2 = buildNode(collect);
+    const n3 = buildNode(add);
+    const e1 = buildEdge(n3.id, 'value', n1.id, 'item');
+    const e2 = buildEdge(n1.id, 'collection', n2.id, 'collection');
+    const result = getCollectItemType(templates, [n1, n2, n3], [e1, e2], n2.id);
+    expect(result).toEqual<FieldType>({ name: 'IntegerField', cardinality: 'SINGLE', batch: false });
+  });
 });

invokeai/frontend/web/src/features/nodes/store/util/getCollectItemType.ts

@@ -2,6 +2,16 @@ import type { Templates } from 'features/nodes/store/types';
 import type { FieldType } from 'features/nodes/types/field';
 import type { AnyEdge, AnyNode } from 'features/nodes/types/invocation';
 
+const toItemType = (fieldType: FieldType): FieldType | null => {
+  if (fieldType.name === 'CollectionField') {
+    return null;
+  }
+  if (fieldType.cardinality === 'COLLECTION' || fieldType.cardinality === 'SINGLE_OR_COLLECTION') {
+    return { ...fieldType, cardinality: 'SINGLE' };
+  }
+  return fieldType;
+};
+
 /**
  * Given a collect node, return the type of the items it collects. The graph is traversed to find the first node and
  * field connected to the collector's `item` input. The field type of that field is returned, else null if there is no
@@ -18,21 +28,56 @@ export const getCollectItemType = (
   edges: AnyEdge[],
   nodeId: string
 ): FieldType | null => {
-  const firstEdgeToCollect = edges.find((edge) => edge.target === nodeId && edge.targetHandle === 'item');
-  if (!firstEdgeToCollect?.sourceHandle) {
-    return null;
-  }
-  const node = nodes.find((n) => n.id === firstEdgeToCollect.source);
-  if (!node) {
-    return null;
-  }
-  const template = templates[node.data.type];
-  if (!template) {
-    return null;
-  }
-  const fieldTemplate = template.outputs[firstEdgeToCollect.sourceHandle];
-  if (!fieldTemplate) {
+  const getCollectItemTypeInternal = (currentNodeId: string, visited: Set<string>): FieldType | null => {
+    if (visited.has(currentNodeId)) {
+      return null;
+    }
+    visited.add(currentNodeId);
+
+    const firstItemEdgeToCollect = edges.find((edge) => edge.target === currentNodeId && edge.targetHandle === 'item');
+    if (firstItemEdgeToCollect?.sourceHandle) {
+      const node = nodes.find((n) => n.id === firstItemEdgeToCollect.source);
+      if (!node) {
+        return null;
+      }
+      const template = templates[node.data.type];
+      if (!template) {
+        return null;
+      }
+      const fieldTemplate = template.outputs[firstItemEdgeToCollect.sourceHandle];
+      if (!fieldTemplate) {
+        return null;
+      }
+      return toItemType(fieldTemplate.type);
+    }
+
+    const firstCollectionEdgeToCollect = edges.find(
+      (edge) => edge.target === currentNodeId && edge.targetHandle === 'collection'
+    );
+    if (!firstCollectionEdgeToCollect?.sourceHandle) {
+      return null;
+    }
+    const sourceNode = nodes.find((n) => n.id === firstCollectionEdgeToCollect.source);
+    if (!sourceNode) {
+      return null;
+    }
+    if (sourceNode.data.type === 'collect' && firstCollectionEdgeToCollect.sourceHandle === 'collection') {
+      return getCollectItemTypeInternal(sourceNode.id, visited);
+    }
+    const sourceTemplate = templates[sourceNode.data.type];
+    if (!sourceTemplate) {
+      return null;
+    }
+    const sourceFieldTemplate = sourceTemplate.outputs[firstCollectionEdgeToCollect.sourceHandle];
+    if (!sourceFieldTemplate) {
+      return null;
+    }
+    return toItemType(sourceFieldTemplate.type);
+  };
+
+  const itemType = getCollectItemTypeInternal(nodeId, new Set());
+  if (!itemType) {
     return null;
   }
-  return fieldTemplate.type;
+  return itemType;
 };