Support batch image upload (#2597)

src/composables/widgets/useImageUploadWidget.ts

@@ -3,25 +3,23 @@ import { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets'
 
 import { useNodeImage } from '@/composables/useNodeImage'
 import { useNodeImageUpload } from '@/composables/useNodeImageUpload'
+import { useValueTransform } from '@/composables/useValueTransform'
 import type { ComfyWidgetConstructor } from '@/scripts/widgets'
 import type { ComfyApp } from '@/types'
 import type { InputSpec, ResultItem } from '@/types/apiTypes'
 import { createAnnotatedPath } from '@/utils/formatUtil'
+import { addToComboValues } from '@/utils/litegraphUtil'
+
+type InternalFile = string | ResultItem
+type InternalValue = InternalFile | InternalFile[]
+type ExposedValue = string | string[]
 
 const isImageFile = (file: File) => file.type.startsWith('image/')
 
-const findFileComboWidget = (node: LGraphNode, inputData: InputSpec) =>
-  node.widgets?.find(
-    (w) => w.name === (inputData[1]?.widget ?? 'image') && w.type === 'combo'
-  ) as IComboWidget & { value: string }
-
-const addToComboValues = (widget: IComboWidget, path: string) => {
-  if (!widget.options) widget.options = { values: [] }
-  if (!widget.options.values) widget.options.values = []
-  if (!widget.options.values.includes(path)) {
-    widget.options.values.push(path)
+const findFileComboWidget = (node: LGraphNode, inputName: string) =>
+  node.widgets!.find((w) => w.name === inputName) as IComboWidget & {
+    value: ExposedValue
   }
-}
 
 export const useImageUploadWidget = () => {
   const widgetConstructor: ComfyWidgetConstructor = (
@@ -30,43 +28,46 @@ export const useImageUploadWidget = () => {
     inputData: InputSpec,
     app: ComfyApp
   ) => {
-    // TODO: specify upload widget via input spec rather than input name
-    const fileComboWidget = findFileComboWidget(node, inputData)
-    const { allow_batch, image_folder = 'input' } = inputData[1] ?? {}
+    const inputOptions = inputData[1] ?? {}
+    const { imageInputName, allow_batch, image_folder = 'input' } = inputOptions
+
+    const { showImage } = useNodeImage(node)
+
+    const fileComboWidget = findFileComboWidget(node, imageInputName)
     const initialFile = `${fileComboWidget.value}`
-    const { showImage } = useNodeImage(node, { allowBatch: allow_batch })
+    const formatPath = (value: InternalFile) =>
+      createAnnotatedPath(value, { rootFolder: image_folder })
 
-    let internalValue: string | ResultItem = initialFile
+    const transform = (internalValue: InternalValue): ExposedValue => {
+      if (!internalValue) return initialFile
+      if (Array.isArray(internalValue))
+        return allow_batch
+          ? internalValue.map(formatPath)
+          : formatPath(internalValue[0])
+      return formatPath(internalValue)
+    }
 
-    // Setup getter/setter that transforms from `ResultItem` to string and formats paths
-    Object.defineProperty(fileComboWidget, 'value', {
-      set: function (value: string | ResultItem) {
-        internalValue = value
-      },
-      get: function () {
-        if (!internalValue) return initialFile
-        if (typeof internalValue === 'string')
-          return createAnnotatedPath(internalValue, {
-            rootFolder: image_folder
-          })
-        if (!internalValue.filename) return initialFile
-        return createAnnotatedPath(internalValue)
-      }
-    })
+    Object.defineProperty(
+      fileComboWidget,
+      'value',
+      useValueTransform(transform, initialFile)
+    )
 
     // Setup file upload handling
-    const { fileInput } = useNodeImageUpload(node, {
+    const { openFileSelection } = useNodeImageUpload(node, {
+      allow_batch,
       fileFilter: isImageFile,
       onUploadComplete: (output) => {
         output.forEach((path) => addToComboValues(fileComboWidget, path))
-        fileComboWidget.value = output[0]
+        // @ts-expect-error litegraph combo value type does not support arrays yet
+        fileComboWidget.value = output
         fileComboWidget.callback?.(output)
       }
     })
 
     // Create the button widget for selecting the files
     const uploadWidget = node.addWidget('button', inputName, 'image', () =>
-      fileInput.click()
+      openFileSelection()
     )
     uploadWidget.label = 'choose file to upload'
     // @ts-expect-error serialize is not typed