[Refactor] useImageUploadWidget composable (#2515)

src/composables/widgets/useImageUploadWidget.ts

@@ -0,0 +1,188 @@
+import type { LGraphNode } from '@comfyorg/litegraph'
+import type { IStringWidget } from '@comfyorg/litegraph/dist/types/widgets'
+
+import { api } from '@/scripts/api'
+import { useToastStore } from '@/stores/toastStore'
+import type { ComfyApp } from '@/types'
+import type { InputSpec } from '@/types/apiTypes'
+
+export const useImageUploadWidget = () => {
+  const widgetConstructor = (
+    node: LGraphNode,
+    inputName: string,
+    inputData: InputSpec,
+    app: ComfyApp
+  ) => {
+    const imageWidget = node.widgets?.find(
+      (w) => w.name === (inputData[1]?.widget ?? 'image')
+    ) as IStringWidget
+    const { image_folder = 'input' } = inputData[1] ?? {}
+
+    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?.()
+    }
+
+    const default_value = imageWidget.value
+    Object.defineProperty(imageWidget, 'value', {
+      set: function (value) {
+        this._real_value = 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
+      }
+    })
+
+    // Add our own callback to the combo widget to render an image when it changes
+    // TODO: Explain this?
+    // @ts-expect-error LGraphNode.callback is not typed
+    const cb = node.callback
+    imageWidget.callback = function (...args) {
+      showImage(imageWidget.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)
+      }
+    })
+
+    // 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 }
+  }
+
+  return widgetConstructor
+}

src/scripts/widgets.ts

@@ -1,23 +1,19 @@
 // @ts-strict-ignore
 import type { LGraphNode } from '@comfyorg/litegraph'
 import type { IWidget } from '@comfyorg/litegraph'
-import type {
+import type { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets'
-  IComboWidget,
-  IStringWidget
-} from '@comfyorg/litegraph/dist/types/widgets'
 
 import { useFloatWidget } from '@/composables/widgets/useFloatWidget'
+import { useImageUploadWidget } from '@/composables/widgets/useImageUploadWidget'
 import { useIntWidget } from '@/composables/widgets/useIntWidget'
 import { useMarkdownWidget } from '@/composables/widgets/useMarkdownWidget'
 import { useRemoteWidget } from '@/composables/widgets/useRemoteWidget'
 import { useSeedWidget } from '@/composables/widgets/useSeedWidget'
 import { useStringWidget } from '@/composables/widgets/useStringWidget'
 import { useSettingStore } from '@/stores/settingStore'
-import { useToastStore } from '@/stores/toastStore'
 import { useWidgetStore } from '@/stores/widgetStore'
-import { InputSpec } from '@/types/apiTypes'
+import type { InputSpec } from '@/types/apiTypes'
 
-import { api } from './api'
 import type { ComfyApp } from './app'
 import './domWidget'
 
@@ -313,167 +309,5 @@ export const ComfyWidgets: Record<string, ComfyWidgetConstructor> = {
     }
     return res
   },
-  IMAGEUPLOAD(node: LGraphNode, inputName: string, inputData: InputSpec, app) {
+  IMAGEUPLOAD: useImageUploadWidget()
-    // TODO make image upload handle a custom node type?
-    const imageWidget = node.widgets.find(
-      (w) => w.name === (inputData[1]?.widget ?? 'image')
-    ) as IStringWidget
-    let uploadWidget
-    const { image_folder = 'input' } = inputData[1] ?? {}
-
-    function showImage(name) {
-      const img = new Image()
-      img.onload = () => {
-        node.imgs = [img]
-        app.graph.setDirtyCanvas(true)
-      }
-      let 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?.()
-    }
-
-    var default_value = imageWidget.value
-    Object.defineProperty(imageWidget, 'value', {
-      set: function (value) {
-        this._real_value = value
-      },
-
-      get: function () {
-        if (!this._real_value) {
-          return default_value
-        }
-
-        let value = this._real_value
-        if (value.filename) {
-          let 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
-      }
-    })
-
-    // Add our own callback to the combo widget to render an image when it changes
-    // TODO: Explain this?
-    // @ts-expect-error
-    const cb = node.callback
-    imageWidget.callback = function () {
-      showImage(imageWidget.value)
-      if (cb) {
-        return cb.apply(this, arguments)
-      }
-    }
-
-    // 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)
-      }
-    })
-
-    async function uploadFile(file, updateNode, 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.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(error)
-      }
-    }
-
-    const fileInput = document.createElement('input')
-    Object.assign(fileInput, {
-      type: 'file',
-      accept: 'image/jpeg,image/png,image/webp',
-      style: 'display: none',
-      onchange: async () => {
-        if (fileInput.files.length) {
-          await uploadFile(fileInput.files[0], true)
-        }
-      }
-    })
-    document.body.append(fileInput)
-
-    // Create the button widget for selecting the files
-    uploadWidget = node.addWidget('button', inputName, 'image', () => {
-      fileInput.click()
-    })
-    uploadWidget.label = 'choose file to upload'
-    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
-      for (const file of e.dataTransfer.files) {
-        if (file.type.startsWith('image/')) {
-          uploadFile(file, !handled) // Dont await these, any order is fine, only update on first one
-          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 }
-  }
 }