[refactor] Refactor and type image upload options (#4185)

Author: christian-byrne

src/composables/node/useNodeImageUpload.ts

@@ -3,15 +3,29 @@ import type { LGraphNode } from '@comfyorg/litegraph'
 import { useNodeDragAndDrop } from '@/composables/node/useNodeDragAndDrop'
 import { useNodeFileInput } from '@/composables/node/useNodeFileInput'
 import { useNodePaste } from '@/composables/node/useNodePaste'
+import type { ResultItemType } from '@/schemas/apiSchema'
 import { api } from '@/scripts/api'
 import { useToastStore } from '@/stores/toastStore'
 
 const PASTED_IMAGE_EXPIRY_MS = 2000
 
-const uploadFile = async (file: File, isPasted: boolean) => {
+interface ImageUploadFormFields {
+  /**
+   * The folder to upload the file to.
+   * @example 'input', 'output', 'temp'
+   */
+  type: ResultItemType
+}
+
+const uploadFile = async (
+  file: File,
+  isPasted: boolean,
+  formFields: Partial<ImageUploadFormFields> = {}
+) => {
   const body = new FormData()
   body.append('image', file)
   if (isPasted) body.append('subfolder', 'pasted')
+  if (formFields.type) body.append('type', formFields.type)
 
   const resp = await api.fetchApi('/upload/image', {
     method: 'POST',
@@ -53,7 +67,7 @@ export const useNodeImageUpload = (
 
   const handleUpload = async (file: File) => {
     try {
-      const path = await uploadFile(file, isPastedFile(file))
+      const path = await uploadFile(file, isPastedFile(file), {})
       if (!path) return
       return path
     } catch (error) {

src/composables/widgets/useImageUploadWidget.ts

@@ -6,7 +6,7 @@ import { useNodeImageUpload } from '@/composables/node/useNodeImageUpload'
 import { useValueTransform } from '@/composables/useValueTransform'
 import { t } from '@/i18n'
 import type { ResultItem } from '@/schemas/apiSchema'
-import type { InputSpec } from '@/schemas/nodeDefSchema'
+import type { ComboInputOptions, InputSpec } from '@/schemas/nodeDefSchema'
 import type { ComfyWidgetConstructor } from '@/scripts/widgets'
 import { useNodeOutputStore } from '@/stores/imagePreviewStore'
 import { createAnnotatedPath } from '@/utils/formatUtil'
@@ -33,7 +33,7 @@ export const useImageUploadWidget = () => {
     inputName: string,
     inputData: InputSpec
   ) => {
-    const inputOptions = inputData[1] ?? {}
+    const inputOptions = (inputData[1] ?? {}) as ComboInputOptions
     const { imageInputName, allow_batch, image_folder = 'input' } = inputOptions
     const nodeOutputStore = useNodeOutputStore()
 
@@ -43,11 +43,9 @@ export const useImageUploadWidget = () => {
     const { showPreview } = isVideo ? useNodeVideo(node) : useNodeImage(node)
 
     const fileFilter = isVideo ? isVideoFile : isImageFile
-    // @ts-expect-error InputSpec is not typed correctly
-    const fileComboWidget = findFileComboWidget(node, imageInputName)
+    const fileComboWidget = findFileComboWidget(node, imageInputName ?? '')
     const initialFile = `${fileComboWidget.value}`
     const formatPath = (value: InternalFile) =>
-      // @ts-expect-error InputSpec is not typed correctly
       createAnnotatedPath(value, { rootFolder: image_folder })
 
     const transform = (internalValue: InternalValue): ExposedValue => {
@@ -67,7 +65,6 @@ export const useImageUploadWidget = () => {
 
     // Setup file upload handling
     const { openFileSelection } = useNodeImageUpload(node, {
-      // @ts-expect-error InputSpec is not typed correctly
       allow_batch,
       fileFilter,
       accept,

src/extensions/core/uploadAudio.ts

@@ -5,15 +5,14 @@ import { useNodeDragAndDrop } from '@/composables/node/useNodeDragAndDrop'
 import { useNodeFileInput } from '@/composables/node/useNodeFileInput'
 import { useNodePaste } from '@/composables/node/useNodePaste'
 import { t } from '@/i18n'
+import type { ResultItemType } from '@/schemas/apiSchema'
 import type { ComfyNodeDef } from '@/schemas/nodeDefSchema'
 import type { DOMWidget } from '@/scripts/domWidget'
 import { useToastStore } from '@/stores/toastStore'
 
 import { api } from '../../scripts/api'
 import { app } from '../../scripts/app'
 
-type FolderType = 'input' | 'output' | 'temp'
-
 function splitFilePath(path: string): [string, string] {
   const folder_separator = path.lastIndexOf('/')
   if (folder_separator === -1) {
@@ -28,7 +27,7 @@ function splitFilePath(path: string): [string, string] {
 function getResourceURL(
   subfolder: string,
   filename: string,
-  type: FolderType = 'input'
+  type: ResultItemType = 'input'
 ): string {
   const params = [
     'filename=' + encodeURIComponent(filename),

src/schemas/apiSchema.ts

@@ -8,14 +8,16 @@ import { zKeybinding } from '@/schemas/keyBindingSchema'
 import { NodeBadgeMode } from '@/types/nodeSource'
 import { LinkReleaseTriggerAction } from '@/types/searchBoxTypes'
 
+export const resultItemType = z.enum(['input', 'output', 'temp'])
 const zNodeType = z.string()
 const zQueueIndex = z.number()
 const zPromptId = z.string()
 const zResultItem = z.object({
   filename: z.string().optional(),
   subfolder: z.string().optional(),
-  type: z.string().optional()
+  type: resultItemType.optional()
 })
+export type ResultItemType = z.infer<typeof resultItemType>
 export type ResultItem = z.infer<typeof zResultItem>
 const zOutputs = z
   .object({

src/schemas/nodeDefSchema.ts

@@ -1,6 +1,8 @@
 import { z } from 'zod'
 import { fromZodError } from 'zod-validation-error'
 
+import { resultItemType } from './apiSchema'
+
 const zComboOption = z.union([z.string(), z.number()])
 const zRemoteWidgetConfig = z.object({
   route: z.string().url().or(z.string().startsWith('/')),
@@ -72,14 +74,19 @@ export const zStringInputOptions = zBaseInputOptions.extend({
 export const zComboInputOptions = zBaseInputOptions.extend({
   control_after_generate: z.boolean().optional(),
   image_upload: z.boolean().optional(),
-  image_folder: z.enum(['input', 'output', 'temp']).optional(),
+  image_folder: resultItemType.optional(),
   allow_batch: z.boolean().optional(),
   video_upload: z.boolean().optional(),
   animated_image_upload: z.boolean().optional(),
   options: z.array(zComboOption).optional(),
   remote: zRemoteWidgetConfig.optional(),
   /** Whether the widget is a multi-select widget. */
-  multi_select: zMultiSelectOption.optional()
+  multi_select: zMultiSelectOption.optional(),
+  /**
+   * The name of the image upload input (filename combo) if it exists
+   * @example 'image'
+   */
+  imageInputName: z.string().optional()
 })
 
 const zIntInputSpec = z.tuple([z.literal('INT'), zIntInputOptions.optional()])

src/stores/imagePreviewStore.ts

@@ -1,15 +1,19 @@
 import { LGraphNode } from '@comfyorg/litegraph'
 import { defineStore } from 'pinia'
 
-import { ExecutedWsMessage, ResultItem } from '@/schemas/apiSchema'
+import {
+  ExecutedWsMessage,
+  ResultItem,
+  ResultItemType
+} from '@/schemas/apiSchema'
 import { api } from '@/scripts/api'
 import { app } from '@/scripts/app'
 import { parseFilePath } from '@/utils/formatUtil'
 import { isVideoNode } from '@/utils/litegraphUtil'
 
 const createOutputs = (
   filenames: string[],
-  type: string,
+  type: ResultItemType,
   isAnimated: boolean
 ): ExecutedWsMessage['output'] => {
   return {
@@ -88,7 +92,7 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
     {
       folder = 'input',
       isAnimated = false
-    }: { folder?: string; isAnimated?: boolean } = {}
+    }: { folder?: ResultItemType; isAnimated?: boolean } = {}
   ) {
     if (!filenames || !node) return