Refactor node image upload and preview (#2580)

Co-authored-by: huchenlei <huchenlei@proton.me>

src/composables/useNodeDragAndDrop.ts

@@ -0,0 +1,49 @@
+import type { LGraphNode } from '@comfyorg/litegraph'
+
+type DragHandler = (e: DragEvent) => boolean
+type DropHandler<T> = (files: File[]) => Promise<T[]>
+
+interface DragAndDropOptions<T> {
+  onDragOver?: DragHandler
+  onDrop: DropHandler<T>
+  fileFilter?: (file: File) => boolean
+}
+
+/**
+ * Adds drag and drop file handling to a node
+ */
+export const useNodeDragAndDrop = <T>(
+  node: LGraphNode,
+  options: DragAndDropOptions<T>
+) => {
+  const { onDragOver, onDrop, fileFilter = () => true } = options
+
+  const hasFiles = (items: DataTransferItemList) =>
+    !!Array.from(items).find((f) => f.kind === 'file')
+
+  const filterFiles = (files: FileList) => Array.from(files).filter(fileFilter)
+
+  const hasValidFiles = (files: FileList) => filterFiles(files).length > 0
+
+  const isDraggingFiles = (e: DragEvent | undefined) => {
+    if (!e?.dataTransfer?.items) return false
+    return onDragOver?.(e) ?? hasFiles(e.dataTransfer.items)
+  }
+
+  const isDraggingValidFiles = (e: DragEvent | undefined) => {
+    if (!e?.dataTransfer?.files) return false
+    return hasValidFiles(e.dataTransfer.files)
+  }
+
+  node.onDragOver = isDraggingFiles
+
+  node.onDragDrop = function (e: DragEvent) {
+    if (!isDraggingValidFiles(e)) return false
+
+    const files = filterFiles(e.dataTransfer!.files)
+    onDrop(files).then((results) => {
+      if (!results?.length) return
+    })
+    return true
+  }
+}

src/composables/useNodeImage.ts

@@ -0,0 +1,31 @@
+import type { LGraphNode } from '@comfyorg/litegraph'
+
+import { useNodeOutputStore } from '@/stores/imagePreviewStore'
+
+interface NodeImageOptions {
+  allowBatch?: boolean
+}
+
+/**
+ * Attaches a preview image to a node.
+ */
+export const useNodeImage = (node: LGraphNode, options: NodeImageOptions) => {
+  const { allowBatch = false } = options
+  const nodeOutputStore = useNodeOutputStore()
+
+  /** Displays output image(s) on the node. */
+  function showImage(output: string | string[]) {
+    if (!output) return
+    if (allowBatch || typeof output === 'string') {
+      nodeOutputStore.setNodeOutputs(node, output)
+    } else {
+      nodeOutputStore.setNodeOutputs(node, output[0])
+    }
+    node.setSizeForImage?.()
+    node.graph?.setDirtyCanvas(true)
+  }
+
+  return {
+    showImage
+  }
+}

src/composables/useNodeImageUpload.ts

@@ -0,0 +1,107 @@
+import type { LGraphNode } from '@comfyorg/litegraph'
+
+import { api } from '@/scripts/api'
+import { useToastStore } from '@/stores/toastStore'
+
+import { useNodeDragAndDrop } from './useNodeDragAndDrop'
+import { useNodePaste } from './useNodePaste'
+
+const ACCEPTED_IMAGE_TYPES = 'image/jpeg,image/png,image/webp'
+const PASTED_IMAGE_EXPIRY_MS = 2000
+
+const createFileInput = () => {
+  const fileInput = document.createElement('input')
+  fileInput.type = 'file'
+  fileInput.accept = ACCEPTED_IMAGE_TYPES
+  return fileInput
+}
+
+const uploadFile = async (file: File, isPasted: boolean) => {
+  const body = new FormData()
+  body.append('image', file)
+  if (isPasted) body.append('subfolder', 'pasted')
+
+  const resp = await api.fetchApi('/upload/image', {
+    method: 'POST',
+    body
+  })
+
+  if (resp.status !== 200) {
+    useToastStore().addAlert(resp.status + ' - ' + resp.statusText)
+    return
+  }
+
+  const data = await resp.json()
+  return data.subfolder ? `${data.subfolder}/${data.name}` : data.name
+}
+
+interface ImageUploadOptions {
+  fileFilter?: (file: File) => boolean
+  onUploadComplete: (paths: string[]) => void
+}
+
+export const useNodeImageUpload = (
+  node: LGraphNode,
+  options: ImageUploadOptions
+) => {
+  const { fileFilter = () => true, onUploadComplete } = options
+
+  const isPastedFile = (file: File): boolean =>
+    file.name === 'image.png' &&
+    file.lastModified - Date.now() < PASTED_IMAGE_EXPIRY_MS
+
+  const handleUpload = async (file: File) => {
+    try {
+      const path = await uploadFile(file, isPastedFile(file))
+      if (!path) return
+      return path
+    } catch (error) {
+      useToastStore().addAlert(String(error))
+    }
+  }
+
+  // Handle drag & drop
+  useNodeDragAndDrop(node, {
+    fileFilter,
+    onDrop: async (files) => {
+      const paths = await Promise.all(files.map(handleUpload))
+      const validPaths = paths.filter((p): p is string => !!p)
+      if (validPaths.length) {
+        onUploadComplete(validPaths)
+      }
+      return validPaths
+    }
+  })
+
+  // Handle paste
+  useNodePaste(node, {
+    fileFilter,
+    onPaste: async (file) => {
+      const path = await handleUpload(file)
+      if (path) {
+        onUploadComplete([path])
+      }
+      return path
+    }
+  })
+
+  // Handle file input
+  const fileInput = createFileInput()
+  fileInput.onchange = async () => {
+    if (fileInput.files?.length) {
+      const paths = await Promise.all(
+        Array.from(fileInput.files).filter(fileFilter).map(handleUpload)
+      )
+      const validPaths = paths.filter((p): p is string => !!p)
+      if (validPaths.length) {
+        onUploadComplete(validPaths)
+      }
+    }
+  }
+  document.body.append(fileInput)
+
+  return {
+    fileInput,
+    handleUpload
+  }
+}

src/composables/useNodePaste.ts

@@ -0,0 +1,27 @@
+import type { LGraphNode } from '@comfyorg/litegraph'
+
+type PasteHandler<T> = (file: File) => Promise<T>
+
+interface NodePasteOptions<T> {
+  onPaste: PasteHandler<T>
+  fileFilter?: (file: File) => boolean
+}
+
+/**
+ * Adds paste handling to a node
+ */
+export const useNodePaste = <T>(
+  node: LGraphNode,
+  options: NodePasteOptions<T>
+) => {
+  const { onPaste, fileFilter = () => true } = options
+
+  node.pasteFile = function (file: File) {
+    if (!fileFilter(file)) return false
+
+    onPaste(file).then((result) => {
+      if (!result) return
+    })
+    return true
+  }
+}

src/composables/widgets/useImageUploadWidget.ts

@@ -1,11 +1,27 @@
 import type { LGraphNode } from '@comfyorg/litegraph'
-import type { IStringWidget } from '@comfyorg/litegraph/dist/types/widgets'
+import { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets'
 
-import { api } from '@/scripts/api'
+import { useNodeImage } from '@/composables/useNodeImage'
+import { useNodeImageUpload } from '@/composables/useNodeImageUpload'
 import type { ComfyWidgetConstructor } from '@/scripts/widgets'
-import { useToastStore } from '@/stores/toastStore'
 import type { ComfyApp } from '@/types'
-import type { InputSpec } from '@/types/apiTypes'
+import type { InputSpec, ResultItem } from '@/types/apiTypes'
+import { createAnnotatedPath } from '@/utils/formatUtil'
+
+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)
+  }
+}
 
 export const useImageUploadWidget = () => {
   const widgetConstructor: ComfyWidgetConstructor = (
@@ -14,174 +30,64 @@ export const useImageUploadWidget = () => {
     inputData: InputSpec,
     app: ComfyApp
   ) => {
-    const imageWidget = node.widgets?.find(
-      (w) => w.name === (inputData[1]?.widget ?? 'image')
-    ) as IStringWidget
-    const { image_folder = 'input' } = inputData[1] ?? {}
+    // 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 initialFile = `${fileComboWidget.value}`
+    const { showImage } = useNodeImage(node, { allowBatch: allow_batch })
 
-    function showImage(name: string) {
-      const img = new Image()
-      img.onload = () => {
-        node.imgs = [img]
-        app.graph.setDirtyCanvas(true)
-      }
-      const folder_separator = name.lastIndexOf('/')
-      let subfolder = ''
-      if (folder_separator > -1) {
-        subfolder = name.substring(0, folder_separator)
-        name = name.substring(folder_separator + 1)
-      }
-      img.src = api.apiURL(
-        `/view?filename=${encodeURIComponent(name)}&type=${image_folder}&subfolder=${subfolder}${app.getPreviewFormatParam()}${app.getRandParam()}`
-      )
-      node.setSizeForImage?.()
-    }
+    let internalValue: string | ResultItem = initialFile
 
-    const default_value = imageWidget.value
-    Object.defineProperty(imageWidget, 'value', {
-      set: function (value) {
-        this._real_value = value
+    // 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 (!this._real_value) {
-          return default_value
-        }
-
-        let value = this._real_value
-        if (value.filename) {
-          const real_value = value
-          value = ''
-          if (real_value.subfolder) {
-            value = real_value.subfolder + '/'
-          }
-
-          value += real_value.filename
-
-          if (real_value.type && real_value.type !== 'input')
-            value += ` [${real_value.type}]`
-        }
-        return value
+        if (!internalValue) return initialFile
+        if (typeof internalValue === 'string')
+          return createAnnotatedPath(internalValue, {
+            rootFolder: image_folder
+          })
+        if (!internalValue.filename) return initialFile
+        return createAnnotatedPath(internalValue)
       }
     })
 
-    // Add our own callback to the combo widget to render an image when it changes
+    // Setup file upload handling
+    const { fileInput } = useNodeImageUpload(node, {
+      fileFilter: isImageFile,
+      onUploadComplete: (output) => {
+        output.forEach((path) => addToComboValues(fileComboWidget, path))
+        fileComboWidget.value = output[0]
+        fileComboWidget.callback?.(output)
+      }
+    })
+
+    // Create the button widget for selecting the files
+    const uploadWidget = node.addWidget('button', inputName, 'image', () =>
+      fileInput.click()
+    )
+    uploadWidget.label = 'choose file to upload'
+    // @ts-expect-error serialize is not typed
+    uploadWidget.serialize = false
+
     // TODO: Explain this?
     // @ts-expect-error LGraphNode.callback is not typed
+    // Add our own callback to the combo widget to render an image when it changes
     const cb = node.callback
-    imageWidget.callback = function (...args) {
-      showImage(imageWidget.value)
-      if (cb) {
-        return cb.apply(this, args)
-      }
+    fileComboWidget.callback = function (...args) {
+      showImage(fileComboWidget.value)
+      if (cb) return cb.apply(this, args)
     }
 
     // On load if we have a value then render the image
     // The value isnt set immediately so we need to wait a moment
     // No change callbacks seem to be fired on initial setting of the value
     requestAnimationFrame(() => {
-      if (imageWidget.value) {
-        showImage(imageWidget.value)
-      }
+      showImage(fileComboWidget.value)
     })
 
-    // Add types for upload parameters
-    async function uploadFile(file: File, updateNode: boolean, pasted = false) {
-      try {
-        // Wrap file in formdata so it includes filename
-        const body = new FormData()
-        body.append('image', file)
-        if (pasted) body.append('subfolder', 'pasted')
-        const resp = await api.fetchApi('/upload/image', {
-          method: 'POST',
-          body
-        })
-
-        if (resp.status === 200) {
-          const data = await resp.json()
-          // Add the file to the dropdown list and update the widget value
-          let path = data.name
-          if (data.subfolder) path = data.subfolder + '/' + path
-
-          if (!imageWidget.options) {
-            imageWidget.options = { values: [] }
-          }
-          if (!imageWidget.options.values) {
-            imageWidget.options.values = []
-          }
-          if (!imageWidget.options.values.includes(path)) {
-            imageWidget.options.values.push(path)
-          }
-
-          if (updateNode) {
-            showImage(path)
-            imageWidget.value = path
-          }
-        } else {
-          useToastStore().addAlert(resp.status + ' - ' + resp.statusText)
-        }
-      } catch (error) {
-        useToastStore().addAlert(String(error))
-      }
-    }
-
-    const fileInput = document.createElement('input')
-    Object.assign(fileInput, {
-      type: 'file',
-      accept: 'image/jpeg,image/png,image/webp',
-      style: 'display: none',
-      onchange: async () => {
-        // Add null check for files
-        if (fileInput.files && fileInput.files.length) {
-          await uploadFile(fileInput.files[0], true)
-        }
-      }
-    })
-    document.body.append(fileInput)
-
-    // Create the button widget for selecting the files
-    const uploadWidget = node.addWidget('button', inputName, 'image', () => {
-      fileInput.click()
-    })
-    uploadWidget.label = 'choose file to upload'
-    // @ts-expect-error IWidget.serialize is not typed
-    uploadWidget.serialize = false
-
-    // Add handler to check if an image is being dragged over our node
-    node.onDragOver = function (e: DragEvent) {
-      if (e.dataTransfer && e.dataTransfer.items) {
-        const image = [...e.dataTransfer.items].find((f) => f.kind === 'file')
-        return !!image
-      }
-
-      return false
-    }
-
-    // On drop upload files
-    node.onDragDrop = function (e: DragEvent) {
-      console.log('onDragDrop called')
-      let handled = false
-      if (e.dataTransfer?.files) {
-        for (const file of e.dataTransfer.files) {
-          if (file.type.startsWith('image/')) {
-            uploadFile(file, !handled)
-            handled = true
-          }
-        }
-      }
-      return handled
-    }
-
-    node.pasteFile = function (file: File) {
-      if (file.type.startsWith('image/')) {
-        const is_pasted =
-          file.name === 'image.png' && file.lastModified - Date.now() < 2000
-        uploadFile(file, true, is_pasted)
-        return true
-      }
-      return false
-    }
-
     return { widget: uploadWidget }
   }