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 }
   }
 

src/services/litegraphService.ts

@@ -10,15 +10,16 @@ import { Vector2 } from '@comfyorg/litegraph'
 import { IBaseWidget, IWidget } from '@comfyorg/litegraph/dist/types/widgets'
 
 import { st } from '@/i18n'
-import { api } from '@/scripts/api'
 import { ANIM_PREVIEW_WIDGET, ComfyApp, app } from '@/scripts/app'
 import { $el } from '@/scripts/ui'
 import { calculateImageGrid, createImageHost } from '@/scripts/ui/imagePreview'
 import { useCanvasStore } from '@/stores/graphStore'
+import { useNodeOutputStore } from '@/stores/imagePreviewStore'
 import { useToastStore } from '@/stores/toastStore'
-import { ComfyNodeDef, ExecutedWsMessage } from '@/types/apiTypes'
+import { ComfyNodeDef } from '@/types/apiTypes'
 import type { NodeId } from '@/types/comfyWorkflow'
 import { normalizeI18nKey } from '@/utils/formatUtil'
+import { is_all_same_aspect_ratio } from '@/utils/imageUtil'
 import { isImageNode } from '@/utils/litegraphUtil'
 
 import { useExtensionService } from './extensionService'
@@ -404,30 +405,22 @@ export const useLitegraphService = () => {
     ) {
       if (this.flags.collapsed) return
 
-      const imgURLs: (string[] | string)[] = []
-      let imagesChanged = false
-
-      const output: ExecutedWsMessage['output'] = app.nodeOutputs[this.id + '']
-      if (output?.images && this.images !== output.images) {
-        this.animatedImages = output?.animated?.find(Boolean)
+      let imgURLs: string[] = []
+      const nodeOutputStore = useNodeOutputStore()
+      const output = nodeOutputStore.getNodeOutputs(this)
+      let imagesChanged = nodeOutputStore.isImagesChanged(this)
+      if (output && imagesChanged) {
+        this.animatedImages = output.animated?.find(Boolean)
         this.images = output.images
-        imagesChanged = true
-        const preview = this.animatedImages ? '' : app.getPreviewFormatParam()
-
-        for (const params of output.images) {
-          const imgUrlPart = new URLSearchParams(params).toString()
-          const rand = app.getRandParam()
-          const imgUrl = api.apiURL(`/view?${imgUrlPart}${preview}${rand}`)
-          imgURLs.push(imgUrl)
-        }
+        imgURLs = nodeOutputStore.getNodeImageUrls(this)
       }
 
-      const preview = app.nodePreviewImages[this.id + '']
+      const preview = nodeOutputStore.getNodePreviews(this)
       if (this.preview !== preview) {
         this.preview = preview
         imagesChanged = true
         if (preview != null) {
-          imgURLs.push(preview)
+          imgURLs.push(...preview)
         }
       }
 
@@ -454,22 +447,10 @@ export const useLitegraphService = () => {
             }
           })
         } else {
-          this.imgs = null
+          this.imgs = undefined
         }
       }
 
-      const is_all_same_aspect_ratio = (imgs: HTMLImageElement[]) => {
-        // assume: imgs.length >= 2
-        const ratio = imgs[0].naturalWidth / imgs[0].naturalHeight
-
-        for (let i = 1; i < imgs.length; i++) {
-          const this_ratio = imgs[i].naturalWidth / imgs[i].naturalHeight
-          if (ratio != this_ratio) return false
-        }
-
-        return true
-      }
-
       // Nothing to do
       if (!this.imgs?.length) return
 

src/stores/imagePreviewStore.ts

@@ -0,0 +1,85 @@
+import { LGraphNode } from '@comfyorg/litegraph'
+import { defineStore } from 'pinia'
+
+import { api } from '@/scripts/api'
+import { ExecutedWsMessage, ResultItem } from '@/types/apiTypes'
+
+const toOutputs = (
+  filenames: string[],
+  type: string
+): ExecutedWsMessage['output'] => {
+  return {
+    images: filenames.map((image) => {
+      return { filename: image, subfolder: '', type }
+    })
+  }
+}
+
+const getPreviewParam = (node: LGraphNode) => {
+  if (node.animatedImages) return ''
+  return app.getPreviewFormatParam()
+}
+
+export const useNodeOutputStore = defineStore('nodeOutput', () => {
+  function getNodeOutputs(node: LGraphNode): ExecutedWsMessage['output'] {
+    return app.nodeOutputs[node.id + '']
+  }
+
+  function getNodePreviews(node: LGraphNode): string[] {
+    return app.nodePreviewImages[node.id + '']
+  }
+
+  function getNodeImageUrls(node: LGraphNode): string[] {
+    const outputs = getNodeOutputs(node)
+    if (!outputs?.images?.length) return []
+
+    return outputs.images.map((image) => {
+      const imgUrlPart = new URLSearchParams(image)
+      const rand = app.getRandParam()
+      const previewParam = getPreviewParam(node)
+      return api.apiURL(`/view?${imgUrlPart}${previewParam}${rand}`)
+    })
+  }
+
+  /**
+   * Checks if the node's images have changed from what's stored
+   * @returns true if images have changed, false otherwise
+   */
+  function isImagesChanged(node: LGraphNode): boolean {
+    const currentImages = node.images || []
+    const { images: newImages } = getNodeOutputs(node) ?? {}
+    if (!newImages?.length) return false
+
+    return currentImages !== newImages
+  }
+
+  function setNodeOutputs(
+    node: LGraphNode,
+    filenames: string | string[] | ResultItem,
+    options: { folder?: string } = {}
+  ) {
+    if (!filenames) return
+
+    const { folder = 'input' } = options
+    const nodeId = node.id + ''
+
+    if (typeof filenames === 'string') {
+      app.nodeOutputs[nodeId] = toOutputs([filenames], folder)
+    } else if (!Array.isArray(filenames)) {
+      app.nodeOutputs[nodeId] = filenames
+    } else {
+      const resultItems = toOutputs(filenames, folder)
+      if (!resultItems?.images?.length) return
+
+      app.nodeOutputs[nodeId] = resultItems
+    }
+  }
+
+  return {
+    getNodeOutputs,
+    getNodeImageUrls,
+    getNodePreviews,
+    setNodeOutputs,
+    isImagesChanged
+  }
+})

src/utils/formatUtil.ts

@@ -1,3 +1,5 @@
+import { ResultItem } from '@/types/apiTypes'
+
 export function formatCamelCase(str: string): string {
   // Check if the string is camel case
   const isCamelCase = /^([A-Z][a-z]*)+$/.test(str)
@@ -213,3 +215,20 @@ export function isValidUrl(url: string): boolean {
     return false
   }
 }
+
+const createAnnotation = (rootFolder = 'input'): string =>
+  rootFolder !== 'input' ? ` [${rootFolder}]` : ''
+
+const createPath = (filename: string, subfolder = ''): string =>
+  subfolder ? `${subfolder}/${filename}` : filename
+
+/** Creates annotated filepath in format used by folder_paths.py */
+export function createAnnotatedPath(
+  item: string | ResultItem,
+  options: { rootFolder?: string; subfolder?: string } = {}
+): string {
+  const { rootFolder = 'input', subfolder } = options
+  if (typeof item === 'string')
+    return `${createPath(item, subfolder)}${createAnnotation(rootFolder)}`
+  return `${createPath(item.filename ?? '', item.subfolder)}${createAnnotation(item.type)}`
+}

src/utils/imageUtil.ts

@@ -0,0 +1,12 @@
+export const is_all_same_aspect_ratio = (imgs: HTMLImageElement[]): boolean => {
+  if (!imgs.length || imgs.length === 1) return true
+
+  const ratio = imgs[0].naturalWidth / imgs[0].naturalHeight
+
+  for (let i = 1; i < imgs.length; i++) {
+    const this_ratio = imgs[i].naturalWidth / imgs[i].naturalHeight
+    if (ratio != this_ratio) return false
+  }
+
+  return true
+}