feat(ui): more batch data types

invokeai/app/invocations/primitives.py

@@ -544,18 +544,86 @@ def invoke(self, context: InvocationContext) -> BoundingBoxOutput:
 @invocation(
     "image_batch",
     title="Image Batch",
-    tags=["primitives", "image", "batch", "internal"],
+    tags=["primitives", "image", "batch", "special"],
     category="primitives",
     version="1.0.0",
     classification=Classification.Special,
 )
 class ImageBatchInvocation(BaseInvocation):
     """Create a batched generation, where the workflow is executed once for each image in the batch."""
 
-    images: list[ImageField] = InputField(min_length=1, description="The images to batch over", input=Input.Direct)
+    images: list[ImageField] = InputField(
+        default=[], min_length=1, description="The images to batch over", input=Input.Direct
+    )
 
     def __init__(self):
         raise NotImplementedError("This class should never be executed or instantiated directly.")
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
         raise NotImplementedError("This class should never be executed or instantiated directly.")
+
+
+@invocation(
+    "string_batch",
+    title="String Batch",
+    tags=["primitives", "string", "batch", "special"],
+    category="primitives",
+    version="1.0.0",
+    classification=Classification.Special,
+)
+class StringBatchInvocation(BaseInvocation):
+    """Create a batched generation, where the workflow is executed once for each string in the batch."""
+
+    strings: list[str] = InputField(
+        default=[], min_length=1, description="The strings to batch over", input=Input.Direct
+    )
+
+    def __init__(self):
+        raise NotImplementedError("This class should never be executed or instantiated directly.")
+
+    def invoke(self, context: InvocationContext) -> StringOutput:
+        raise NotImplementedError("This class should never be executed or instantiated directly.")
+
+
+@invocation(
+    "integer_batch",
+    title="Integer Batch",
+    tags=["primitives", "integer", "number", "batch", "special"],
+    category="primitives",
+    version="1.0.0",
+    classification=Classification.Special,
+)
+class IntegerBatchInvocation(BaseInvocation):
+    """Create a batched generation, where the workflow is executed once for each integer in the batch."""
+
+    integers: list[int] = InputField(
+        default=[], min_length=1, description="The integers to batch over", input=Input.Direct
+    )
+
+    def __init__(self):
+        raise NotImplementedError("This class should never be executed or instantiated directly.")
+
+    def invoke(self, context: InvocationContext) -> IntegerOutput:
+        raise NotImplementedError("This class should never be executed or instantiated directly.")
+
+
+@invocation(
+    "float_batch",
+    title="Float Batch",
+    tags=["primitives", "float", "number", "batch", "special"],
+    category="primitives",
+    version="1.0.0",
+    classification=Classification.Special,
+)
+class FloatBatchInvocation(BaseInvocation):
+    """Create a batched generation, where the workflow is executed once for each float in the batch."""
+
+    floats: list[float] = InputField(
+        default=[], min_length=1, description="The floats to batch over", input=Input.Direct
+    )
+
+    def __init__(self):
+        raise NotImplementedError("This class should never be executed or instantiated directly.")
+
+    def invoke(self, context: InvocationContext) -> FloatOutput:
+        raise NotImplementedError("This class should never be executed or instantiated directly.")

invokeai/app/services/session_queue/session_queue_common.py

@@ -108,8 +108,16 @@ def validate_types(cls, v: Optional[BatchDataCollection]):
             return v
         for batch_data_list in v:
             for datum in batch_data_list:
+                if not datum.items:
+                    continue
+
+                # Special handling for numbers - they can be mixed
+                # TODO(psyche): Update BatchDatum to have a `type` field to specify the type of the items, then we can have strict float and int fields
+                if all(isinstance(item, (int, float)) for item in datum.items):
+                    continue
+
                 # Get the type of the first item in the list
-                first_item_type = type(datum.items[0]) if datum.items else None
+                first_item_type = type(datum.items[0])
                 for item in datum.items:
                     if type(item) is not first_item_type:
                         raise BatchItemsTypeError("All items in a batch must have the same type")

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

@@ -177,7 +177,11 @@
         "none": "None",
         "new": "New",
         "generating": "Generating",
-        "warnings": "Warnings"
+        "warnings": "Warnings",
+        "start": "Start",
+        "count": "Count",
+        "step": "Step",
+        "values": "Values"
     },
     "hrf": {
         "hrf": "High Resolution Fix",
@@ -989,7 +993,11 @@
         "imageAccessError": "Unable to find image {{image_name}}, resetting to default",
         "boardAccessError": "Unable to find board {{board_id}}, resetting to default",
         "modelAccessError": "Unable to find model {{key}}, resetting to default",
-        "saveToGallery": "Save To Gallery"
+        "saveToGallery": "Save To Gallery",
+        "addItem": "Add Item",
+        "generateValues": "Generate Values",
+        "floatRangeGenerator": "Float Range Generator",
+        "integerRangeGenerator": "Integer Range Generator"
     },
     "parameters": {
         "aspect": "Aspect",
@@ -1024,11 +1032,18 @@
             "addingImagesTo": "Adding images to",
             "invoke": "Invoke",
             "missingFieldTemplate": "Missing field template",
-            "missingInputForField": "{{nodeLabel}} -> {{fieldLabel}}: missing input",
+            "missingInputForField": "missing input",
             "missingNodeTemplate": "Missing node template",
-            "collectionEmpty": "{{nodeLabel}} -> {{fieldLabel}} empty collection",
-            "collectionTooFewItems": "{{nodeLabel}} -> {{fieldLabel}}: too few items, minimum {{minItems}}",
-            "collectionTooManyItems": "{{nodeLabel}} -> {{fieldLabel}}: too many items, maximum {{maxItems}}",
+            "collectionEmpty": "empty collection",
+            "collectionTooFewItems": "too few items, minimum {{minItems}}",
+            "collectionTooManyItems": "too many items, maximum {{maxItems}}",
+            "collectionStringTooLong": "too long, max {{maxLength}}",
+            "collectionStringTooShort": "too short, min {{minLength}}",
+            "collectionNumberGTMax": "{{value}} > {{maximum}} (inc max)",
+            "collectionNumberLTMin": "{{value}} < {{minimum}} (inc min)",
+            "collectionNumberGTExclusiveMax": "{{value}} >= {{exclusiveMaximum}} (exc max)",
+            "collectionNumberLTExclusiveMin": "{{value}} <= {{exclusiveMinimum}} (exc min)",
+            "collectionNumberNotMultipleOf": "{{value}} not multiple of {{multipleOf}}",
             "noModelSelected": "No model selected",
             "noT5EncoderModelSelected": "No T5 Encoder model selected for FLUX generation",
             "noFLUXVAEModelSelected": "No VAE model selected for FLUX generation",

invokeai/frontend/web/src/app/components/App.tsx

@@ -22,6 +22,8 @@ import { DynamicPromptsModal } from 'features/dynamicPrompts/components/DynamicP
 import DeleteBoardModal from 'features/gallery/components/Boards/DeleteBoardModal';
 import { ImageContextMenu } from 'features/gallery/components/ImageContextMenu/ImageContextMenu';
 import { useStarterModelsToast } from 'features/modelManagerV2/hooks/useStarterModelsToast';
+import { FloatRangeGeneratorModal } from 'features/nodes/components/FloatRangeGeneratorModal';
+import { IntegerRangeGeneratorModal } from 'features/nodes/components/IntegerRangeGeneratorModal';
 import { ShareWorkflowModal } from 'features/nodes/components/sidePanel/WorkflowListMenu/ShareWorkflowModal';
 import { ClearQueueConfirmationsAlertDialog } from 'features/queue/components/ClearQueueConfirmationAlertDialog';
 import { DeleteStylePresetDialog } from 'features/stylePresets/components/DeleteStylePresetDialog';
@@ -110,6 +112,8 @@ const App = ({ config = DEFAULT_CONFIG, studioInitAction }: Props) => {
       <ImageContextMenu />
       <FullscreenDropzone />
       <VideosModal />
+      <FloatRangeGeneratorModal />
+      <IntegerRangeGeneratorModal />
     </ErrorBoundary>
   );
 };

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedNodes.ts

@@ -2,7 +2,14 @@ import { logger } from 'app/logging/logger';
 import { enqueueRequested } from 'app/store/actions';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { selectNodesSlice } from 'features/nodes/store/selectors';
-import { isImageFieldCollectionInputInstance } from 'features/nodes/types/field';
+import type { ImageField } from 'features/nodes/types/common';
+import {
+  isFloatFieldCollectionInputInstance,
+  isImageFieldCollectionInputInstance,
+  isIntegerFieldCollectionInputInstance,
+  isStringFieldCollectionInputInstance,
+} from 'features/nodes/types/field';
+import type { InvocationNodeEdge } from 'features/nodes/types/invocation';
 import { isInvocationNode } from 'features/nodes/types/invocation';
 import { buildNodesGraph } from 'features/nodes/util/graph/buildNodesGraph';
 import { buildWorkflowWithValidation } from 'features/nodes/util/workflow/buildWorkflow';
@@ -33,29 +40,92 @@ export const addEnqueueRequestedNodes = (startAppListening: AppStartListening) =
 
       const data: Batch['data'] = [];
 
-      // Skip edges from batch nodes - these should not be in the graph, they exist only in the UI
-      const imageBatchNodes = nodes.nodes.filter(isInvocationNode).filter((node) => node.data.type === 'image_batch');
-      for (const node of imageBatchNodes) {
-        const images = node.data.inputs['images'];
-        if (!isImageFieldCollectionInputInstance(images)) {
-          log.warn({ nodeId: node.id }, 'Image batch images field is not an image collection');
-          break;
-        }
-        const edgesFromImageBatch = nodes.edges.filter((e) => e.source === node.id && e.sourceHandle === 'image');
+      const addBatchDataCollectionItem = (edges: InvocationNodeEdge[], items?: ImageField[] | string[] | number[]) => {
         const batchDataCollectionItem: NonNullable<Batch['data']>[number] = [];
-        for (const edge of edgesFromImageBatch) {
+        for (const edge of edges) {
           if (!edge.targetHandle) {
             break;
           }
           batchDataCollectionItem.push({
             node_path: edge.target,
             field_name: edge.targetHandle,
-            items: images.value,
+            items,
           });
         }
         if (batchDataCollectionItem.length > 0) {
           data.push(batchDataCollectionItem);
         }
+      };
+
+      // Grab image batch nodes for special handling
+      const imageBatchNodes = nodes.nodes.filter(isInvocationNode).filter((node) => node.data.type === 'image_batch');
+
+      for (const node of imageBatchNodes) {
+        // Satisfy TS
+        const images = node.data.inputs['images'];
+        if (!isImageFieldCollectionInputInstance(images)) {
+          log.warn({ nodeId: node.id }, 'Image batch images field is not an image collection');
+          break;
+        }
+
+        // Find outgoing edges from the batch node, we will remove these from the graph and create batch data collection items from them instead
+        const edgesFromImageBatch = nodes.edges.filter((e) => e.source === node.id && e.sourceHandle === 'image');
+        addBatchDataCollectionItem(edgesFromImageBatch, images.value);
+      }
+
+      // Grab string batch nodes for special handling
+      const stringBatchNodes = nodes.nodes.filter(isInvocationNode).filter((node) => node.data.type === 'string_batch');
+      for (const node of stringBatchNodes) {
+        // Satisfy TS
+        const strings = node.data.inputs['strings'];
+        if (!isStringFieldCollectionInputInstance(strings)) {
+          log.warn({ nodeId: node.id }, 'String batch strings field is not a string collection');
+          break;
+        }
+
+        // Find outgoing edges from the batch node, we will remove these from the graph and create batch data collection items from them instead
+        const edgesFromStringBatch = nodes.edges.filter((e) => e.source === node.id && e.sourceHandle === 'value');
+        addBatchDataCollectionItem(edgesFromStringBatch, strings.value);
+      }
+
+      // Grab integer batch nodes for special handling
+      const integerBatchNodes = nodes.nodes
+        .filter(isInvocationNode)
+        .filter((node) => node.data.type === 'integer_batch');
+      for (const node of integerBatchNodes) {
+        // Satisfy TS
+        const integers = node.data.inputs['integers'];
+        if (!isIntegerFieldCollectionInputInstance(integers)) {
+          log.warn({ nodeId: node.id }, 'Integer batch integers field is not an integer collection');
+          break;
+        }
+        if (!integers.value) {
+          log.warn({ nodeId: node.id }, 'Integer batch integers field is empty');
+          break;
+        }
+
+        // Find outgoing edges from the batch node, we will remove these from the graph and create batch data collection items from them instead
+        const edgesFromStringBatch = nodes.edges.filter((e) => e.source === node.id && e.sourceHandle === 'value');
+        addBatchDataCollectionItem(edgesFromStringBatch, integers.value);
+      }
+
+      // Grab float batch nodes for special handling
+      const floatBatchNodes = nodes.nodes.filter(isInvocationNode).filter((node) => node.data.type === 'float_batch');
+      for (const node of floatBatchNodes) {
+        // Satisfy TS
+        const floats = node.data.inputs['floats'];
+        if (!isFloatFieldCollectionInputInstance(floats)) {
+          log.warn({ nodeId: node.id }, 'Float batch floats field is not a float collection');
+          break;
+        }
+        if (!floats.value) {
+          log.warn({ nodeId: node.id }, 'Float batch floats field is empty');
+          break;
+        }
+
+        // Find outgoing edges from the batch node, we will remove these from the graph and create batch data collection items from them instead
+        const edgesFromStringBatch = nodes.edges.filter((e) => e.source === node.id && e.sourceHandle === 'value');
+        addBatchDataCollectionItem(edgesFromStringBatch, floats.value);
       }
 
       const batchConfig: BatchConfig = {

invokeai/frontend/web/src/features/dnd/dnd.ts

@@ -1,3 +1,4 @@
+import { logger } from 'app/logging/logger';
 import type { AppDispatch, RootState } from 'app/store/store';
 import { getPrefixedId } from 'features/controlLayers/konva/util';
 import type {
@@ -9,7 +10,6 @@ import { selectComparisonImages } from 'features/gallery/components/ImageViewer/
 import type { BoardId } from 'features/gallery/store/types';
 import {
   addImagesToBoard,
-  addImagesToNodeImageFieldCollectionAction,
   createNewCanvasEntityFromImage,
   removeImagesFromBoard,
   replaceCanvasEntityObjectsWithImage,
@@ -19,10 +19,14 @@ import {
   setRegionalGuidanceReferenceImage,
   setUpscaleInitialImage,
 } from 'features/imageActions/actions';
-import type { FieldIdentifier } from 'features/nodes/types/field';
+import { fieldImageCollectionValueChanged } from 'features/nodes/store/nodesSlice';
+import { selectFieldInputInstance, selectNodesSlice } from 'features/nodes/store/selectors';
+import { type FieldIdentifier, isImageFieldCollectionInputInstance } from 'features/nodes/types/field';
 import type { ImageDTO } from 'services/api/types';
 import type { JsonObject } from 'type-fest';
 
+const log = logger('dnd');
+
 type RecordUnknown = Record<string | symbol, unknown>;
 
 type DndData<
@@ -268,15 +272,27 @@ export const addImagesToNodeImageFieldCollectionDndTarget: DndTarget<
     }
 
     const { fieldIdentifier } = targetData.payload;
-    const imageDTOs: ImageDTO[] = [];
+
+    const fieldInputInstance = selectFieldInputInstance(
+      selectNodesSlice(getState()),
+      fieldIdentifier.nodeId,
+      fieldIdentifier.fieldName
+    );
+
+    if (!isImageFieldCollectionInputInstance(fieldInputInstance)) {
+      log.warn({ fieldIdentifier }, 'Attempted to add images to a non-image field collection');
+      return;
+    }
+
+    const newValue = fieldInputInstance.value ? [...fieldInputInstance.value] : [];
 
     if (singleImageDndSource.typeGuard(sourceData)) {
-      imageDTOs.push(sourceData.payload.imageDTO);
+      newValue.push({ image_name: sourceData.payload.imageDTO.image_name });
     } else {
-      imageDTOs.push(...sourceData.payload.imageDTOs);
+      newValue.push(...sourceData.payload.imageDTOs.map(({ image_name }) => ({ image_name })));
     }
 
-    addImagesToNodeImageFieldCollectionAction({ fieldIdentifier, imageDTOs, dispatch, getState });
+    dispatch(fieldImageCollectionValueChanged({ ...fieldIdentifier, value: newValue }));
   },
 };
 //#endregion