From 35e6cabfe7575d846504c27a8f3189087ae500e3 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Wed, 5 Mar 2025 13:12:51 -0500 Subject: [PATCH] Use v2 input spec for combo widget (#2878) --- src/composables/widgets/useComboWidget.ts | 80 ++++++++++--------- src/composables/widgets/useRemoteWidget.ts | 20 ++--- src/schemas/nodeDefSchema.ts | 4 +- src/scripts/widgets.ts | 2 +- src/stores/widgetStore.ts | 39 +-------- .../widgets/useComboWidget.test.ts | 26 ++---- .../widgets/useRemoteWidget.test.ts | 26 +++--- 7 files changed, 75 insertions(+), 122 deletions(-) diff --git a/src/composables/widgets/useComboWidget.ts b/src/composables/widgets/useComboWidget.ts index 7607288a87..9fa194f0d5 100644 --- a/src/composables/widgets/useComboWidget.ts +++ b/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 comboOptions = inputSpec.options ?? [] + const defaultValue = getDefaultValue(inputSpec) - const defaultValue = widgetStore.getDefaultValue(inputData) - - 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, - { - 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, { + 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 diff --git a/src/composables/widgets/useRemoteWidget.ts b/src/composables/widgets/useRemoteWidget.ts index 19c5cee3a6..0edfbde29e 100644 --- a/src/composables/widgets/useRemoteWidget.ts +++ b/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 diff --git a/src/schemas/nodeDefSchema.ts b/src/schemas/nodeDefSchema.ts index 8d1a4467d4..7c9dd39b62 100644 --- a/src/schemas/nodeDefSchema.ts +++ b/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()]) diff --git a/src/scripts/widgets.ts b/src/scripts/widgets.ts index 86eea8f427..052991d452 100644 --- a/src/scripts/widgets.ts +++ b/src/scripts/widgets.ts @@ -293,6 +293,6 @@ export const ComfyWidgets: Record = { BOOLEAN: transformWidgetConstructorV2ToV1(useBooleanWidget()), STRING: transformWidgetConstructorV2ToV1(useStringWidget()), MARKDOWN: transformWidgetConstructorV2ToV1(useMarkdownWidget()), - COMBO: useComboWidget(), + COMBO: transformWidgetConstructorV2ToV1(useComboWidget()), IMAGEUPLOAD: useImageUploadWidget() } diff --git a/src/stores/widgetStore.ts b/src/stores/widgetStore.ts index beda023e60..51ff4373b6 100644 --- a/src/stores/widgetStore.ts +++ b/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 } }) diff --git a/tests-ui/tests/composables/widgets/useComboWidget.test.ts b/tests-ui/tests/composables/widgets/useComboWidget.test.ts index 7cf6e60993..cbdd74bab8 100644 --- a/tests-ui/tests/composables/widgets/useComboWidget.test.ts +++ b/tests-ui/tests/composables/widgets/useComboWidget.test.ts @@ -1,14 +1,7 @@ -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() @@ -16,7 +9,6 @@ vi.mock('@/scripts/widgets', () => ({ 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: {} }) }) }) diff --git a/tests-ui/tests/composables/widgets/useRemoteWidget.test.ts b/tests-ui/tests/composables/widgets/useRemoteWidget.test.ts index cc17ad902d..f50decfefa 100644 --- a/tests-ui/tests/composables/widgets/useRemoteWidget.test.ts +++ b/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) { } 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(() => {