Use v2 input spec for combo widget (#2878)

Author: huchenlei

src/composables/widgets/useComboWidget.ts

@@ -1,70 +1,76 @@
 import type { LGraphNode } from '@comfyorg/litegraph'
 import type { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets'
 
+import { transformInputSpecV2ToV1 } from '@/schemas/nodeDef/migration'
 import {
+  ComboInputSpec,
   type InputSpec,
-  getComboSpecComboOptions,
   isComboInputSpec
-} from '@/schemas/nodeDefSchema'
-import { addValueControlWidgets } from '@/scripts/widgets'
-import type { ComfyWidgetConstructor } from '@/scripts/widgets'
-import { useWidgetStore } from '@/stores/widgetStore'
+} from '@/schemas/nodeDef/nodeDefSchemaV2'
+import {
+  type ComfyWidgetConstructorV2,
+  addValueControlWidgets
+} from '@/scripts/widgets'
 
 import { useRemoteWidget } from './useRemoteWidget'
 
+const getDefaultValue = (inputSpec: ComboInputSpec) => {
+  if (inputSpec.default) return inputSpec.default
+  if (inputSpec.options?.length) return inputSpec.options[0]
+  if (inputSpec.remote) return 'Loading...'
+  return undefined
+}
+
 export const useComboWidget = () => {
-  const widgetConstructor: ComfyWidgetConstructor = (
+  const widgetConstructor: ComfyWidgetConstructorV2 = (
     node: LGraphNode,
-    inputName: string,
-    inputData: InputSpec
+    inputSpec: InputSpec
   ) => {
-    if (!isComboInputSpec(inputData)) {
-      throw new Error(`Invalid input data: ${inputData}`)
+    if (!isComboInputSpec(inputSpec)) {
+      throw new Error(`Invalid input data: ${inputSpec}`)
     }
 
-    const widgetStore = useWidgetStore()
-    const inputOptions = inputData[1] ?? {}
-    const comboOptions = getComboSpecComboOptions(inputData)
-
-    const defaultValue = widgetStore.getDefaultValue(inputData)
+    const comboOptions = inputSpec.options ?? []
+    const defaultValue = getDefaultValue(inputSpec)
 
-    const res = {
-      widget: node.addWidget('combo', inputName, defaultValue, () => {}, {
+    const widget = node.addWidget(
+      'combo',
+      inputSpec.name,
+      defaultValue,
+      () => {},
+      {
         values: comboOptions
-      }) as IComboWidget
-    }
+      }
+    ) as IComboWidget
 
-    if (inputOptions.remote) {
+    if (inputSpec.remote) {
       const remoteWidget = useRemoteWidget({
-        inputData,
+        remoteConfig: inputSpec.remote,
         defaultValue,
         node,
-        widget: res.widget
+        widget
       })
-      if (inputOptions.remote.refresh_button) remoteWidget.addRefreshButton()
+      if (inputSpec.remote.refresh_button) remoteWidget.addRefreshButton()
 
-      const origOptions = res.widget.options
-      res.widget.options = new Proxy(
-        origOptions as Record<string | symbol, any>,
-        {
-          get(target, prop: string | symbol) {
-            if (prop !== 'values') return target[prop]
-            return remoteWidget.getValue()
-          }
+      const origOptions = widget.options
+      widget.options = new Proxy(origOptions as Record<string | symbol, any>, {
+        get(target, prop: string | symbol) {
+          if (prop !== 'values') return target[prop]
+          return remoteWidget.getValue()
         }
-      )
+      })
     }
 
-    if (inputOptions.control_after_generate) {
-      res.widget.linkedWidgets = addValueControlWidgets(
+    if (inputSpec.control_after_generate) {
+      widget.linkedWidgets = addValueControlWidgets(
         node,
-        res.widget,
+        widget,
         undefined,
         undefined,
-        inputData
+        transformInputSpecV2ToV1(inputSpec)
       )
     }
-    return res
+    return widget
   }
 
   return widgetConstructor

src/composables/widgets/useRemoteWidget.ts

@@ -2,7 +2,7 @@ import { LGraphNode } from '@comfyorg/litegraph'
 import { IWidget } from '@comfyorg/litegraph'
 import axios from 'axios'
 
-import type { InputSpec, RemoteWidgetConfig } from '@/schemas/nodeDefSchema'
+import type { RemoteWidgetConfig } from '@/schemas/nodeDefSchema'
 
 const MAX_RETRIES = 5
 const TIMEOUT = 4096
@@ -67,16 +67,15 @@ const fetchData = async (
 export function useRemoteWidget<
   T extends string | number | boolean | object
 >(options: {
-  inputData: InputSpec
+  remoteConfig: RemoteWidgetConfig
   defaultValue: T
   node: LGraphNode
   widget: IWidget
 }) {
-  const { inputData, defaultValue, node, widget } = options
-  const config = (inputData[1]?.remote ?? {}) as RemoteWidgetConfig
-  const { refresh = 0, max_retries = MAX_RETRIES } = config
+  const { remoteConfig, defaultValue, node, widget } = options
+  const { refresh = 0, max_retries = MAX_RETRIES } = remoteConfig
   const isPermanent = refresh <= 0
-  const cacheKey = createCacheKey(config)
+  const cacheKey = createCacheKey(remoteConfig)
   let isLoaded = false
   let refreshQueued = false
 
@@ -131,7 +130,10 @@ export function useRemoteWidget<
 
     try {
       currentEntry.controller = new AbortController()
-      currentEntry.fetchPromise = fetchData(config, currentEntry.controller)
+      currentEntry.fetchPromise = fetchData(
+        remoteConfig,
+        currentEntry.controller
+      )
       const data = await currentEntry.fetchPromise
 
       setSuccess(currentEntry, data)
@@ -146,11 +148,11 @@ export function useRemoteWidget<
   }
 
   const onRefresh = () => {
-    if (config.control_after_refresh) {
+    if (remoteConfig.control_after_refresh) {
       const data = getCachedValue()
       if (!Array.isArray(data)) return // control_after_refresh is only supported for array values
 
-      switch (config.control_after_refresh) {
+      switch (remoteConfig.control_after_refresh) {
         case 'first':
           widget.value = data[0] ?? defaultValue
           break

src/schemas/nodeDefSchema.ts

@@ -71,8 +71,8 @@ export const zComboInputOptions = zBaseInputOptions.extend({
   image_folder: z.enum(['input', 'output', 'temp']).optional(),
   allow_batch: z.boolean().optional(),
   video_upload: z.boolean().optional(),
-  remote: zRemoteWidgetConfig.optional(),
-  options: z.array(zComboOption).optional()
+  options: z.array(zComboOption).optional(),
+  remote: zRemoteWidgetConfig.optional()
 })
 
 const zIntInputSpec = z.tuple([z.literal('INT'), zIntInputOptions.optional()])

src/scripts/widgets.ts

@@ -293,6 +293,6 @@ export const ComfyWidgets: Record<string, ComfyWidgetConstructor> = {
   BOOLEAN: transformWidgetConstructorV2ToV1(useBooleanWidget()),
   STRING: transformWidgetConstructorV2ToV1(useStringWidget()),
   MARKDOWN: transformWidgetConstructorV2ToV1(useMarkdownWidget()),
-  COMBO: useComboWidget(),
+  COMBO: transformWidgetConstructorV2ToV1(useComboWidget()),
   IMAGEUPLOAD: useImageUploadWidget()
 }

src/stores/widgetStore.ts

@@ -2,11 +2,6 @@ import { defineStore } from 'pinia'
 import { computed, ref } from 'vue'
 
 import type { InputSpec as InputSpecV2 } from '@/schemas/nodeDef/nodeDefSchemaV2'
-import {
-  type ComboInputSpecV2,
-  type InputSpec,
-  isComboInputSpecV1
-} from '@/schemas/nodeDefSchema'
 import { ComfyWidgetConstructor, ComfyWidgets } from '@/scripts/widgets'
 
 export const useWidgetStore = defineStore('widget', () => {
@@ -45,42 +40,10 @@ export const useWidgetStore = defineStore('widget', () => {
     }
   }
 
-  function getDefaultValue(inputData: InputSpec) {
-    if (Array.isArray(inputData[0]))
-      return getDefaultValue(transformComboInput(inputData))
-
-    // @ts-expect-error InputSpec is not typed correctly
-    const widgetType = getWidgetType(inputData[0], inputData[1]?.name)
-
-    const [_, props] = inputData
-
-    if (!props) return undefined
-    if (props.default) return props.default
-
-    // @ts-expect-error InputSpec is not typed correctly
-    if (widgetType === 'COMBO' && props.options?.length) return props.options[0]
-    if (props.remote) return 'Loading...'
-    return undefined
-  }
-
-  const transformComboInput = (inputData: InputSpec): ComboInputSpecV2 => {
-    // @ts-expect-error InputSpec is not typed correctly
-    return isComboInputSpecV1(inputData)
-      ? [
-          'COMBO',
-          {
-            options: inputData[0],
-            ...Object(inputData[1] || {})
-          }
-        ]
-      : inputData
-  }
-
   return {
     widgets,
     getWidgetType,
     inputIsWidget,
-    registerCustomWidgets,
-    getDefaultValue
+    registerCustomWidgets
   }
 })

tests-ui/tests/composables/widgets/useComboWidget.test.ts

@@ -1,22 +1,14 @@
-import { createPinia, setActivePinia } from 'pinia'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 import { useComboWidget } from '@/composables/widgets/useComboWidget'
-import type { InputSpec } from '@/schemas/nodeDefSchema'
-
-vi.mock('@/stores/widgetStore', () => ({
-  useWidgetStore: () => ({
-    getDefaultValue: vi.fn()
-  })
-}))
+import type { InputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
 
 vi.mock('@/scripts/widgets', () => ({
   addValueControlWidgets: vi.fn()
 }))
 
 describe('useComboWidget', () => {
   beforeEach(() => {
-    setActivePinia(createPinia())
     vi.clearAllMocks()
   })
 
@@ -26,14 +18,12 @@ describe('useComboWidget', () => {
       addWidget: vi.fn().mockReturnValue({ options: {} } as any)
     }
 
-    const inputSpec: InputSpec = ['COMBO', undefined]
+    const inputSpec: InputSpec = {
+      type: 'COMBO',
+      name: 'inputName'
+    }
 
-    const widget = constructor(
-      mockNode as any,
-      'inputName',
-      inputSpec,
-      undefined as any
-    )
+    const widget = constructor(mockNode as any, inputSpec)
 
     expect(mockNode.addWidget).toHaveBeenCalledWith(
       'combo',
@@ -44,8 +34,6 @@ describe('useComboWidget', () => {
         values: []
       })
     )
-    expect(widget).toEqual({
-      widget: { options: {} }
-    })
+    expect(widget).toEqual({ options: {} })
   })
 })

tests-ui/tests/composables/widgets/useRemoteWidget.test.ts

@@ -2,7 +2,7 @@ import axios from 'axios'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
 import { useRemoteWidget } from '@/composables/widgets/useRemoteWidget'
-import type { ComboInputSpecV2 } from '@/schemas/nodeDefSchema'
+import { RemoteWidgetConfig } from '@/schemas/nodeDefSchema'
 
 vi.mock('axios', () => {
   return {
@@ -29,22 +29,16 @@ vi.mock('@/stores/settingStore', () => ({
 const FIRST_BACKOFF = 1000 // backoff is 1s on first retry
 const DEFAULT_VALUE = 'Loading...'
 
-function createMockInputData(overrides = {}): ComboInputSpecV2 {
-  return [
-    'COMBO',
-    {
-      name: 'test_widget',
-      remote: {
-        route: `/api/test/${Date.now()}${Math.random().toString(36).substring(2, 15)}`,
-        refresh: 0,
-        ...overrides
-      }
-    }
-  ]
+function createMockConfig(overrides = {}): RemoteWidgetConfig {
+  return {
+    route: `/api/test/${Date.now()}${Math.random().toString(36).substring(2, 15)}`,
+    refresh: 0,
+    ...overrides
+  }
 }
 
 const createMockOptions = (inputOverrides = {}) => ({
-  inputData: createMockInputData(inputOverrides),
+  remoteConfig: createMockConfig(inputOverrides),
   defaultValue: DEFAULT_VALUE,
   node: {} as any,
   widget: {} as any
@@ -81,7 +75,7 @@ async function getResolvedValue(hook: ReturnType<typeof useRemoteWidget>) {
 }
 
 describe('useRemoteWidget', () => {
-  let mockInputData: ComboInputSpecV2
+  let mockConfig: RemoteWidgetConfig
 
   beforeEach(() => {
     vi.clearAllMocks()
@@ -92,7 +86,7 @@ describe('useRemoteWidget', () => {
     vi.spyOn(Map.prototype, 'set').mockClear()
     vi.spyOn(Map.prototype, 'delete').mockClear()
 
-    mockInputData = createMockInputData()
+    mockConfig = createMockConfig()
   })
 
   afterEach(() => {