diff --git a/src/composables/widgets/useComboWidget.ts b/src/composables/widgets/useComboWidget.ts index 684cf1c8f..7607288a8 100644 --- a/src/composables/widgets/useComboWidget.ts +++ b/src/composables/widgets/useComboWidget.ts @@ -3,8 +3,8 @@ import type { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets' import { type InputSpec, - isComboInputSpec, - isComboInputSpecV2 + getComboSpecComboOptions, + isComboInputSpec } from '@/schemas/nodeDefSchema' import { addValueControlWidgets } from '@/scripts/widgets' import type { ComfyWidgetConstructor } from '@/scripts/widgets' @@ -24,9 +24,7 @@ export const useComboWidget = () => { const widgetStore = useWidgetStore() const inputOptions = inputData[1] ?? {} - const comboOptions = - (isComboInputSpecV2(inputData) ? inputOptions.options : inputData[0]) ?? - [] + const comboOptions = getComboSpecComboOptions(inputData) const defaultValue = widgetStore.getDefaultValue(inputData) diff --git a/src/extensions/core/widgetInputs.ts b/src/extensions/core/widgetInputs.ts index 1ab38203c..1d22bf0ac 100644 --- a/src/extensions/core/widgetInputs.ts +++ b/src/extensions/core/widgetInputs.ts @@ -14,10 +14,10 @@ import type { CanvasMouseEvent } from '@comfyorg/litegraph/dist/types/events' import type { InputSpec } from '@/schemas/nodeDefSchema' import { app } from '@/scripts/app' -import { clone } from '@/scripts/utils' import { ComfyWidgets, addValueControlWidgets } from '@/scripts/widgets' import { useNodeDefStore } from '@/stores/nodeDefStore' import { useSettingStore } from '@/stores/settingStore' +import { mergeInputSpec } from '@/utils/nodeDefUtil' import { applyTextReplacements } from '@/utils/searchAndReplace' import { isPrimitiveNode } from '@/utils/typeGuardUtil' @@ -603,113 +603,17 @@ export function mergeIfValid( config2: InputSpec, forceUpdate?: boolean, recreateWidget?: () => void, - config1?: unknown -): { customConfig: Record } { + config1?: InputSpec +): { customConfig: InputSpec[1] } { if (!config1) { config1 = getWidgetConfig(output) } - if (config1[0] instanceof Array) { - if (!isValidCombo(config1[0], config2[0])) return - } else if (config1[0] !== config2[0]) { - // Types dont match - console.log(`connection rejected: types dont match`, config1[0], config2[0]) - return - } + const customSpec = mergeInputSpec(config1, config2) - const keys = new Set([ - ...Object.keys(config1[1] ?? {}), - ...Object.keys(config2[1] ?? {}) - ]) - - let customConfig: Record | undefined - const getCustomConfig = () => { - if (!customConfig) { - customConfig = clone(config1[1] ?? {}) - } - return customConfig - } - - const isNumber = config1[0] === 'INT' || config1[0] === 'FLOAT' - for (const k of keys.values()) { - if ( - k !== 'default' && - k !== 'forceInput' && - k !== 'defaultInput' && - k !== 'control_after_generate' && - k !== 'multiline' && - k !== 'tooltip' && - k !== 'dynamicPrompts' - ) { - let v1 = config1[1][k] - let v2 = config2[1]?.[k] - - if (v1 === v2 || (!v1 && !v2)) continue - - if (isNumber) { - if (k === 'min') { - const theirMax = config2[1]?.['max'] - if (theirMax != null && v1 > theirMax) { - console.log('connection rejected: min > max', v1, theirMax) - return - } - getCustomConfig()[k] = - // @ts-expect-error InputSpec is not typed correctly - v1 == null ? v2 : v2 == null ? v1 : Math.max(v1, v2) - continue - } else if (k === 'max') { - const theirMin = config2[1]?.['min'] - if (theirMin != null && v1 < theirMin) { - console.log('connection rejected: max < min', v1, theirMin) - return - } - getCustomConfig()[k] = - // @ts-expect-error InputSpec is not typed correctly - v1 == null ? v2 : v2 == null ? v1 : Math.min(v1, v2) - continue - } else if (k === 'step') { - let step - if (v1 == null) { - // No current step - step = v2 - } else if (v2 == null) { - // No new step - step = v1 - } else { - if (v1 < v2) { - // Ensure v1 is larger for the mod - const a = v2 - v2 = v1 - v1 = a - } - // @ts-expect-error InputSpec is not typed correctly - if (v1 % v2) { - console.log( - 'connection rejected: steps not divisible', - 'current:', - v1, - 'new:', - v2 - ) - return - } - - step = v1 - } - - getCustomConfig()[k] = step - continue - } - } - - console.log(`connection rejected: config ${k} values dont match`, v1, v2) - return - } - } - - if (customConfig || forceUpdate) { - if (customConfig) { - output.widget[CONFIG] = [config1[0], customConfig] + if (customSpec || forceUpdate) { + if (customSpec) { + output.widget[CONFIG] = customSpec } const widget = recreateWidget?.call(this) @@ -723,7 +627,7 @@ export function mergeIfValid( } } - return { customConfig } + return { customConfig: customSpec[1] } } app.registerExtension({ diff --git a/src/schemas/nodeDefSchema.ts b/src/schemas/nodeDefSchema.ts index c00f983da..b88d4331d 100644 --- a/src/schemas/nodeDefSchema.ts +++ b/src/schemas/nodeDefSchema.ts @@ -15,6 +15,7 @@ const zRemoteWidgetConfig = z.object({ const zBaseInputOptions = z .object({ default: z.any().optional(), + /** @deprecated Group node uses this field. Remove when group node feature is removed. */ defaultInput: z.boolean().optional(), forceInput: z.boolean().optional(), tooltip: z.string().optional(), @@ -143,6 +144,30 @@ export function isComboInputSpec( return isComboInputSpecV1(inputSpec) || isComboInputSpecV2(inputSpec) } +/** + * Get the type of an input spec. + * + * @param inputSpec - The input spec to get the type of. + * @returns The type of the input spec. + */ +export function getInputSpecType(inputSpec: InputSpec): string { + return isComboInputSpec(inputSpec) ? 'COMBO' : inputSpec[0] +} + +/** + * Get the combo options from a combo input spec. + * + * @param inputSpec - The input spec to get the combo options from. + * @returns The combo options. + */ +export function getComboSpecComboOptions( + inputSpec: ComboInputSpec | ComboInputSpecV2 +): (number | string)[] { + return ( + (isComboInputSpecV2(inputSpec) ? inputSpec[1]?.options : inputSpec[0]) ?? [] + ) +} + const excludedLiterals = new Set(['INT', 'FLOAT', 'BOOLEAN', 'STRING', 'COMBO']) const zCustomInputSpec = z.tuple([ z.string().refine((value) => !excludedLiterals.has(value)), @@ -201,6 +226,8 @@ export type FloatInputOptions = z.infer export type BooleanInputOptions = z.infer export type StringInputOptions = z.infer export type ComboInputOptions = z.infer +export type BaseInputOptions = z.infer +export type NumericInputOptions = z.infer export type IntInputSpec = z.infer export type FloatInputSpec = z.infer diff --git a/src/utils/mathUtil.ts b/src/utils/mathUtil.ts new file mode 100644 index 000000000..c8645363f --- /dev/null +++ b/src/utils/mathUtil.ts @@ -0,0 +1,21 @@ +/** + * Finds the greatest common divisor (GCD) for two numbers. + * + * @param a - The first number. + * @param b - The second number. + * @returns The GCD of the two numbers. + */ +export const gcd = (a: number, b: number): number => { + return b === 0 ? a : gcd(b, a % b) +} + +/** + * Finds the least common multiple (LCM) for two numbers. + * + * @param a - The first number. + * @param b - The second number. + * @returns The LCM of the two numbers. + */ +export const lcm = (a: number, b: number): number => { + return Math.abs(a * b) / gcd(a, b) +} diff --git a/src/utils/nodeDefUtil.ts b/src/utils/nodeDefUtil.ts new file mode 100644 index 000000000..836cd5515 --- /dev/null +++ b/src/utils/nodeDefUtil.ts @@ -0,0 +1,142 @@ +import _ from 'lodash' + +import type { + ComboInputSpec, + ComboInputSpecV2, + FloatInputSpec, + InputSpec, + IntInputSpec, + NumericInputOptions +} from '@/schemas/nodeDefSchema' +import { + getComboSpecComboOptions, + getInputSpecType, + isComboInputSpec, + isFloatInputSpec, + isIntInputSpec +} from '@/schemas/nodeDefSchema' + +import { lcm } from './mathUtil' + +const IGNORE_KEYS = new Set([ + 'default', + 'forceInput', + 'defaultInput', + 'control_after_generate', + 'multiline', + 'tooltip', + 'dynamicPrompts' +]) + +const getRange = (options: NumericInputOptions) => { + const min = options.min ?? -Infinity + const max = options.max ?? Infinity + return { min, max } +} + +const mergeNumericInputSpec = ( + spec1: T, + spec2: T +): T | null => { + const type = spec1[0] + const options1 = spec1[1] ?? {} + const options2 = spec2[1] ?? {} + + const range1 = getRange(options1) + const range2 = getRange(options2) + + // If the ranges do not overlap, return null + if (range1.min > range2.max || range1.max < range2.min) { + return null + } + + const step1 = options1.step ?? 1 + const step2 = options2.step ?? 1 + + const mergedOptions = { + // Take intersection of ranges + min: Math.max(range1.min, range2.min), + max: Math.min(range1.max, range2.max), + step: lcm(step1, step2) + } + + return mergeCommonInputSpec( + [type, { ...options1, ...mergedOptions }] as unknown as T, + [type, { ...options2, ...mergedOptions }] as unknown as T + ) +} + +const mergeComboInputSpec = ( + spec1: T, + spec2: T +): T | null => { + const options1 = spec1[1] ?? {} + const options2 = spec2[1] ?? {} + + const comboOptions1 = getComboSpecComboOptions(spec1) + const comboOptions2 = getComboSpecComboOptions(spec2) + + const intersection = _.intersection(comboOptions1, comboOptions2) + + // If the intersection is empty, return null + if (intersection.length === 0) { + return null + } + + return mergeCommonInputSpec( + ['COMBO', { ...options1, options: intersection }] as unknown as T, + ['COMBO', { ...options2, options: intersection }] as unknown as T + ) +} + +const mergeCommonInputSpec = ( + spec1: T, + spec2: T +): T | null => { + const type = getInputSpecType(spec1) + const options1 = spec1[1] ?? {} + const options2 = spec2[1] ?? {} + + const compareKeys = _.union(_.keys(options1), _.keys(options2)).filter( + (key) => !IGNORE_KEYS.has(key) + ) + + const mergeIsValid = compareKeys.every((key) => { + const value1 = options1[key] + const value2 = options2[key] + return value1 === value2 || (_.isNil(value1) && _.isNil(value2)) + }) + + return mergeIsValid + ? ([type, { ...options1, ...options2 }] as unknown as T) + : null +} + +/** + * Merges two input specs. + * + * @param spec1 - The first input spec. + * @param spec2 - The second input spec. + * @returns The merged input spec, or null if the specs are not mergeable. + */ +export const mergeInputSpec = ( + spec1: InputSpec, + spec2: InputSpec +): InputSpec | null => { + const type1 = getInputSpecType(spec1) + const type2 = getInputSpecType(spec2) + + if (type1 !== type2) { + return null + } + + if (isIntInputSpec(spec1) || isFloatInputSpec(spec1)) { + return mergeNumericInputSpec(spec1, spec2 as typeof spec1) + } + + if (isComboInputSpec(spec1)) { + return mergeComboInputSpec(spec1, spec2 as typeof spec1) + } + + return mergeCommonInputSpec(spec1, spec2) +} diff --git a/tests-ui/tests/utils/nodeDefUtil.test.ts b/tests-ui/tests/utils/nodeDefUtil.test.ts new file mode 100644 index 000000000..1e8b39813 --- /dev/null +++ b/tests-ui/tests/utils/nodeDefUtil.test.ts @@ -0,0 +1,206 @@ +// @ts-strict-ignore +import { describe, expect, it } from 'vitest' + +import type { + ComboInputSpec, + ComboInputSpecV2, + FloatInputSpec, + InputSpec, + IntInputSpec +} from '@/schemas/nodeDefSchema' +import { mergeInputSpec } from '@/utils/nodeDefUtil' + +describe('nodeDefUtil', () => { + describe('mergeInputSpec', () => { + // Test numeric input specs (INT and FLOAT) + describe('numeric input specs', () => { + it('should merge INT specs with overlapping ranges', () => { + const spec1: IntInputSpec = ['INT', { min: 0, max: 10 }] + const spec2: IntInputSpec = ['INT', { min: 5, max: 15 }] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).not.toBeNull() + expect(result?.[0]).toBe('INT') + expect(result?.[1].min).toBe(5) + expect(result?.[1].max).toBe(10) + }) + + it('should return null for INT specs with non-overlapping ranges', () => { + const spec1: IntInputSpec = ['INT', { min: 0, max: 5 }] + const spec2: IntInputSpec = ['INT', { min: 10, max: 15 }] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).toBeNull() + }) + + it('should merge FLOAT specs with overlapping ranges', () => { + const spec1: FloatInputSpec = ['FLOAT', { min: 0.5, max: 10.5 }] + const spec2: FloatInputSpec = ['FLOAT', { min: 5.5, max: 15.5 }] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).not.toBeNull() + expect(result?.[0]).toBe('FLOAT') + expect(result?.[1].min).toBe(5.5) + expect(result?.[1].max).toBe(10.5) + }) + + it('should handle specs with undefined min/max values', () => { + const spec1: FloatInputSpec = ['FLOAT', { min: 0.5 }] + const spec2: FloatInputSpec = ['FLOAT', { max: 15.5 }] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).not.toBeNull() + expect(result?.[0]).toBe('FLOAT') + expect(result?.[1].min).toBe(0.5) + expect(result?.[1].max).toBe(15.5) + }) + + it('should merge step values using least common multiple', () => { + const spec1: IntInputSpec = ['INT', { min: 0, max: 10, step: 2 }] + const spec2: IntInputSpec = ['INT', { min: 0, max: 10, step: 3 }] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).not.toBeNull() + expect(result?.[0]).toBe('INT') + expect(result?.[1].step).toBe(6) // LCM of 2 and 3 is 6 + }) + + it('should use default step of 1 when step is not specified', () => { + const spec1: IntInputSpec = ['INT', { min: 0, max: 10 }] + const spec2: IntInputSpec = ['INT', { min: 0, max: 10, step: 4 }] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).not.toBeNull() + expect(result?.[0]).toBe('INT') + expect(result?.[1].step).toBe(4) // LCM of 1 and 4 is 4 + }) + + it('should handle step values for FLOAT specs', () => { + const spec1: FloatInputSpec = ['FLOAT', { min: 0, max: 10, step: 0.5 }] + const spec2: FloatInputSpec = ['FLOAT', { min: 0, max: 10, step: 0.25 }] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).not.toBeNull() + expect(result?.[0]).toBe('FLOAT') + expect(result?.[1].step).toBe(0.5) + }) + }) + + // Test combo input specs + describe('combo input specs', () => { + it('should merge COMBO specs with overlapping options', () => { + const spec1: ComboInputSpecV2 = ['COMBO', { options: ['A', 'B', 'C'] }] + const spec2: ComboInputSpecV2 = ['COMBO', { options: ['B', 'C', 'D'] }] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).not.toBeNull() + expect(result?.[0]).toBe('COMBO') + expect(result?.[1].options).toEqual(['B', 'C']) + }) + + it('should return null for COMBO specs with no overlapping options', () => { + const spec1: ComboInputSpecV2 = ['COMBO', { options: ['A', 'B'] }] + const spec2: ComboInputSpecV2 = ['COMBO', { options: ['C', 'D'] }] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).toBeNull() + }) + + it('should handle COMBO specs with additional properties', () => { + const spec1: ComboInputSpecV2 = [ + 'COMBO', + { + options: ['A', 'B', 'C'], + default: 'A', + tooltip: 'Select an option' + } + ] + const spec2: ComboInputSpecV2 = [ + 'COMBO', + { + options: ['B', 'C', 'D'], + default: 'B', + multiline: true + } + ] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).not.toBeNull() + expect(result?.[0]).toBe('COMBO') + expect(result?.[1].options).toEqual(['B', 'C']) + expect(result?.[1].default).toBe('B') + expect(result?.[1].tooltip).toBe('Select an option') + expect(result?.[1].multiline).toBe(true) + }) + + it('should handle v1 and v2 combo specs', () => { + const spec1: ComboInputSpec = [['A', 'B', 'C', 'D'], {}] + const spec2: ComboInputSpecV2 = ['COMBO', { options: ['C', 'D'] }] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).not.toBeNull() + expect(result?.[0]).toBe('COMBO') + expect(result?.[1].options).toEqual(['C', 'D']) + }) + }) + + // Test common input spec behavior + describe('common input spec behavior', () => { + it('should return null for specs with different types', () => { + const spec1: IntInputSpec = ['INT', { min: 0, max: 10 }] + const spec2: ComboInputSpecV2 = ['COMBO', { options: ['A', 'B'] }] + + const result = mergeInputSpec(spec1, spec2 as unknown as IntInputSpec) + + expect(result).toBeNull() + }) + + it('should ignore specified keys when comparing specs', () => { + const spec1: InputSpec = [ + 'STRING', + { + default: 'value1', + tooltip: 'Tooltip 1', + step: 1 + } + ] + const spec2: InputSpec = [ + 'STRING', + { + default: 'value2', + tooltip: 'Tooltip 2', + step: 1 + } + ] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).not.toBeNull() + expect(result?.[0]).toBe('STRING') + expect(result?.[1].default).toBe('value2') + expect(result?.[1].tooltip).toBe('Tooltip 2') + expect(result?.[1].step).toBe(1) + }) + + it('should return null if non-ignored properties differ', () => { + const spec1: InputSpec = ['STRING', { step: 1 }] + const spec2: InputSpec = ['STRING', { step: 2 }] + + const result = mergeInputSpec(spec1, spec2) + + expect(result).toBeNull() + }) + }) + }) +})