[Refactor] Add util to merge input spec (#2834)

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)
 

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<string, unknown> } {
+  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<string, unknown> | 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({

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<typeof zFloatInputOptions>
 export type BooleanInputOptions = z.infer<typeof zBooleanInputOptions>
 export type StringInputOptions = z.infer<typeof zStringInputOptions>
 export type ComboInputOptions = z.infer<typeof zComboInputOptions>
+export type BaseInputOptions = z.infer<typeof zBaseInputOptions>
+export type NumericInputOptions = z.infer<typeof zNumericInputOptions>
 
 export type IntInputSpec = z.infer<typeof zIntInputSpec>
 export type FloatInputSpec = z.infer<typeof zFloatInputSpec>

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)
+}

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<string>([
+  '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 = <T extends IntInputSpec | FloatInputSpec>(
+  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 = <T extends ComboInputSpec | ComboInputSpecV2>(
+  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 = <T extends InputSpec>(
+  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)
+}