Use V2 schema in widget constructors (Part 1) (#2860)

src/composables/widgets/useBooleanWidget.ts

@@ -1,28 +1,33 @@
 import type { LGraphNode } from '@comfyorg/litegraph'
 
-import { type InputSpec, isBooleanInputSpec } from '@/schemas/nodeDefSchema'
-import type { ComfyWidgetConstructor } from '@/scripts/widgets'
+import {
+  type InputSpec,
+  isBooleanInputSpec
+} from '@/schemas/nodeDef/nodeDefSchemaV2'
+import { type ComfyWidgetConstructorV2 } from '@/scripts/widgets'
 
 export const useBooleanWidget = () => {
-  const widgetConstructor: ComfyWidgetConstructor = (
+  const widgetConstructor: ComfyWidgetConstructorV2 = (
     node: LGraphNode,
-    inputName: string,
-    inputData: InputSpec
+    inputSpec: InputSpec
   ) => {
-    if (!isBooleanInputSpec(inputData)) {
-      throw new Error(`Invalid input data: ${inputData}`)
+    if (!isBooleanInputSpec(inputSpec)) {
+      throw new Error(`Invalid input data: ${inputSpec}`)
     }
 
-    const inputOptions = inputData[1] ?? {}
-    const defaultVal = inputOptions?.default ?? false
+    const defaultVal = inputSpec.default ?? false
     const options = {
-      on: inputOptions?.label_on,
-      off: inputOptions?.label_off
+      on: inputSpec.label_on,
+      off: inputSpec.label_off
     }
 
-    return {
-      widget: node.addWidget('toggle', inputName, defaultVal, () => {}, options)
-    }
+    return node.addWidget(
+      'toggle',
+      inputSpec.name,
+      defaultVal,
+      () => {},
+      options
+    )
   }
 
   return widgetConstructor

src/composables/widgets/useFloatWidget.ts

@@ -2,8 +2,11 @@ import type { LGraphNode } from '@comfyorg/litegraph'
 import type { INumericWidget } from '@comfyorg/litegraph/dist/types/widgets'
 import _ from 'lodash'
 
-import { type InputSpec, isFloatInputSpec } from '@/schemas/nodeDefSchema'
-import type { ComfyWidgetConstructor } from '@/scripts/widgets'
+import {
+  type InputSpec,
+  isFloatInputSpec
+} from '@/schemas/nodeDef/nodeDefSchemaV2'
+import { type ComfyWidgetConstructorV2 } from '@/scripts/widgets'
 import { useSettingStore } from '@/stores/settingStore'
 
 function onFloatValueChange(this: INumericWidget, v: number) {
@@ -22,23 +25,18 @@ export const _for_testing = {
 }
 
 export const useFloatWidget = () => {
-  const widgetConstructor: ComfyWidgetConstructor = (
+  const widgetConstructor: ComfyWidgetConstructorV2 = (
     node: LGraphNode,
-    inputName: string,
-    inputData: InputSpec
+    inputSpec: InputSpec
   ) => {
-    if (!isFloatInputSpec(inputData)) {
-      throw new Error(`Invalid input data: ${inputData}`)
+    if (!isFloatInputSpec(inputSpec)) {
+      throw new Error(`Invalid input data: ${inputSpec}`)
     }
 
-    // TODO: Move to outer scope to avoid re-initializing on every call
-    // Blocked on ComfyWidgets lazy initialization.
-
     const settingStore = useSettingStore()
     const sliderEnabled = !settingStore.get('Comfy.DisableSliders')
-    const inputOptions = inputData[1] ?? {}
 
-    const display_type = inputOptions?.display
+    const display_type = inputSpec.display
     const widgetType =
       sliderEnabled && display_type == 'slider'
         ? 'slider'
@@ -46,33 +44,31 @@ export const useFloatWidget = () => {
           ? 'knob'
           : 'number'
 
-    const step = inputOptions.step ?? 0.5
+    const step = inputSpec.step ?? 0.5
     const precision =
       settingStore.get('Comfy.FloatRoundingPrecision') ||
       Math.max(0, -Math.floor(Math.log10(step)))
     const enableRounding = !settingStore.get('Comfy.DisableFloatRounding')
 
-    const defaultValue = inputOptions.default ?? 0
-    return {
-      widget: node.addWidget(
-        widgetType,
-        inputName,
-        defaultValue,
-        onFloatValueChange,
-        {
-          min: inputOptions.min ?? 0,
-          max: inputOptions.max ?? 2048,
-          round:
-            enableRounding && precision && !inputOptions.round
-              ? (1_000_000 * Math.pow(0.1, precision)) / 1_000_000
-              : (inputOptions.round as number),
-          /** @deprecated Use step2 instead. The 10x value is a legacy implementation. */
-          step: step * 10.0,
-          step2: step,
-          precision
-        }
-      )
-    }
+    const defaultValue = inputSpec.default ?? 0
+    return node.addWidget(
+      widgetType,
+      inputSpec.name,
+      defaultValue,
+      onFloatValueChange,
+      {
+        min: inputSpec.min ?? 0,
+        max: inputSpec.max ?? 2048,
+        round:
+          enableRounding && precision && !inputSpec.round
+            ? (1_000_000 * Math.pow(0.1, precision)) / 1_000_000
+            : (inputSpec.round as number),
+        /** @deprecated Use step2 instead. The 10x value is a legacy implementation. */
+        step: step * 10.0,
+        step2: step,
+        precision
+      }
+    )
   }
 
   return widgetConstructor

src/composables/widgets/useIntWidget.ts

@@ -1,10 +1,13 @@
 import type { LGraphNode } from '@comfyorg/litegraph'
 import type { INumericWidget } from '@comfyorg/litegraph/dist/types/widgets'
 
-import { type InputSpec, isIntInputSpec } from '@/schemas/nodeDefSchema'
-import type { ComfyApp } from '@/scripts/app'
+import { transformInputSpecV2ToV1 } from '@/schemas/nodeDef/migration'
 import {
-  type ComfyWidgetConstructor,
+  type InputSpec,
+  isIntInputSpec
+} from '@/schemas/nodeDef/nodeDefSchemaV2'
+import {
+  type ComfyWidgetConstructorV2,
   addValueControlWidget
 } from '@/scripts/widgets'
 import { useSettingStore } from '@/stores/settingStore'
@@ -33,21 +36,17 @@ export const _for_testing = {
 }
 
 export const useIntWidget = () => {
-  const widgetConstructor: ComfyWidgetConstructor = (
+  const widgetConstructor: ComfyWidgetConstructorV2 = (
     node: LGraphNode,
-    inputName: string,
-    inputData: InputSpec,
-    app: ComfyApp,
-    widgetName?: string
+    inputSpec: InputSpec
   ) => {
-    if (!isIntInputSpec(inputData)) {
-      throw new Error(`Invalid input data: ${inputData}`)
+    if (!isIntInputSpec(inputSpec)) {
+      throw new Error(`Invalid input data: ${inputSpec}`)
     }
 
     const settingStore = useSettingStore()
     const sliderEnabled = !settingStore.get('Comfy.DisableSliders')
-    const inputOptions = inputData[1] ?? {}
-    const display_type = inputOptions?.display
+    const display_type = inputSpec.display
     const widgetType =
       sliderEnabled && display_type == 'slider'
         ? 'slider'
@@ -55,38 +54,36 @@ export const useIntWidget = () => {
           ? 'knob'
           : 'number'
 
-    const step = inputOptions.step ?? 1
-    const defaultValue = inputOptions.default ?? 0
-    const result = {
-      widget: node.addWidget(
-        widgetType,
-        inputName,
-        defaultValue,
-        onValueChange,
-        {
-          min: inputOptions.min ?? 0,
-          max: inputOptions.max ?? 2048,
-          /** @deprecated Use step2 instead. The 10x value is a legacy implementation. */
-          step: step * 10,
-          step2: step,
-          precision: 0
-        }
-      )
-    }
+    const step = inputSpec.step ?? 1
+    const defaultValue = inputSpec.default ?? 0
+    const widget = node.addWidget(
+      widgetType,
+      inputSpec.name,
+      defaultValue,
+      onValueChange,
+      {
+        min: inputSpec.min ?? 0,
+        max: inputSpec.max ?? 2048,
+        /** @deprecated Use step2 instead. The 10x value is a legacy implementation. */
+        step: step * 10,
+        step2: step,
+        precision: 0
+      }
+    )
 
-    if (inputOptions.control_after_generate) {
+    if (inputSpec.control_after_generate) {
       const seedControl = addValueControlWidget(
         node,
-        result.widget,
+        widget,
         'randomize',
         undefined,
-        widgetName,
-        inputData
+        undefined,
+        transformInputSpecV2ToV1(inputSpec)
       )
-      result.widget.linkedWidgets = [seedControl]
+      widget.linkedWidgets = [seedControl]
     }
 
-    return result
+    return widget
   }
   return widgetConstructor
 }

src/composables/widgets/useSeedWidget.ts

@@ -1,39 +0,0 @@
-import type { LGraphNode } from '@comfyorg/litegraph'
-
-import { type InputSpec, isIntInputSpec } from '@/schemas/nodeDefSchema'
-import type { ComfyWidgetConstructor } from '@/scripts/widgets'
-import type { ComfyApp } from '@/types'
-
-import { useIntWidget } from './useIntWidget'
-
-export const useSeedWidget = () => {
-  const IntWidget = useIntWidget()
-
-  const widgetConstructor: ComfyWidgetConstructor = (
-    node: LGraphNode,
-    inputName: string,
-    inputData: InputSpec,
-    app: ComfyApp,
-    widgetName?: string
-  ) => {
-    if (!isIntInputSpec(inputData)) {
-      throw new Error(`Invalid input data: ${inputData}`)
-    }
-
-    return IntWidget(
-      node,
-      inputName,
-      [
-        'INT',
-        {
-          ...inputData[1],
-          control_after_generate: true
-        }
-      ],
-      app,
-      widgetName
-    )
-  }
-
-  return widgetConstructor
-}

src/schemas/nodeDef/migration.ts

@@ -25,14 +25,20 @@ export function transformNodeDefV1ToV2(
   // Process required inputs
   if (nodeDefV1.input?.required) {
     Object.entries(nodeDefV1.input.required).forEach(([name, inputSpecV1]) => {
-      inputs[name] = transformInputSpec(inputSpecV1, name, false)
+      inputs[name] = transformInputSpecV1ToV2(inputSpecV1, {
+        name,
+        isOptional: false
+      })
     })
   }
 
   // Process optional inputs
   if (nodeDefV1.input?.optional) {
     Object.entries(nodeDefV1.input.optional).forEach(([name, inputSpecV1]) => {
-      inputs[name] = transformInputSpec(inputSpecV1, name, true)
+      inputs[name] = transformInputSpecV1ToV2(inputSpecV1, {
+        name,
+        isOptional: true
+      })
     })
   }
 
@@ -81,11 +87,15 @@ export function transformNodeDefV1ToV2(
  * @param isOptional Whether the input is optional
  * @returns The transformed V2 input specification
  */
-function transformInputSpec(
+export function transformInputSpecV1ToV2(
   inputSpecV1: InputSpecV1,
-  name: string,
-  isOptional: boolean
+  kwargs: {
+    name: string
+    isOptional?: boolean
+  }
 ): InputSpecV2 {
+  const { name, isOptional = false } = kwargs
+
   // Extract options from the input spec
   const options = inputSpecV1[1] || {}
 
@@ -119,3 +129,9 @@ function transformInputSpec(
     ...baseProps
   }
 }
+
+export function transformInputSpecV2ToV1(
+  inputSpecV2: InputSpecV2
+): InputSpecV1 {
+  return [inputSpecV2.type, inputSpecV2]
+}

src/scripts/widgets.ts

@@ -11,15 +11,21 @@ import { useFloatWidget } from '@/composables/widgets/useFloatWidget'
 import { useImageUploadWidget } from '@/composables/widgets/useImageUploadWidget'
 import { useIntWidget } from '@/composables/widgets/useIntWidget'
 import { useMarkdownWidget } from '@/composables/widgets/useMarkdownWidget'
-import { useSeedWidget } from '@/composables/widgets/useSeedWidget'
 import { useStringWidget } from '@/composables/widgets/useStringWidget'
 import { t } from '@/i18n'
+import { transformInputSpecV1ToV2 } from '@/schemas/nodeDef/migration'
+import type { InputSpec as InputSpecV2 } from '@/schemas/nodeDef/nodeDefSchemaV2'
 import type { InputSpec } from '@/schemas/nodeDefSchema'
 import { useSettingStore } from '@/stores/settingStore'
 
 import type { ComfyApp } from './app'
 import './domWidget'
 
+export type ComfyWidgetConstructorV2 = (
+  node: LGraphNode,
+  inputSpec: InputSpecV2
+) => IWidget
+
 export type ComfyWidgetConstructor = (
   node: LGraphNode,
   inputName: string,
@@ -28,6 +34,24 @@ export type ComfyWidgetConstructor = (
   widgetName?: string
 ) => { widget: IWidget; minWidth?: number; minHeight?: number }
 
+/**
+ * Transforms a V2 widget constructor to a V1 widget constructor.
+ * @param widgetConstructorV2 The V2 widget constructor to transform.
+ * @returns The transformed V1 widget constructor.
+ */
+const transformWidgetConstructorV2ToV1 = (
+  widgetConstructorV2: ComfyWidgetConstructorV2
+): ComfyWidgetConstructor => {
+  return (node, inputName, inputData) => {
+    const inputSpec = transformInputSpecV1ToV2(inputData, {
+      name: inputName
+    })
+    return {
+      widget: widgetConstructorV2(node, inputSpec)
+    }
+  }
+}
+
 function controlValueRunBefore() {
   return useSettingStore().get('Comfy.WidgetControlMode') === 'before'
 }
@@ -251,14 +275,19 @@ export function addValueControlWidgets(
   return widgets
 }
 
-const SeedWidget = useSeedWidget()
+const seedWidget = transformWidgetConstructorV2ToV1((node, inputSpec) => {
+  return useIntWidget()(node, {
+    ...inputSpec,
+    control_after_generate: true
+  })
+})
 
 export const ComfyWidgets: Record<string, ComfyWidgetConstructor> = {
-  'INT:seed': SeedWidget,
-  'INT:noise_seed': SeedWidget,
-  INT: useIntWidget(),
-  FLOAT: useFloatWidget(),
-  BOOLEAN: useBooleanWidget(),
+  'INT:seed': seedWidget,
+  'INT:noise_seed': seedWidget,
+  INT: transformWidgetConstructorV2ToV1(useIntWidget()),
+  FLOAT: transformWidgetConstructorV2ToV1(useFloatWidget()),
+  BOOLEAN: transformWidgetConstructorV2ToV1(useBooleanWidget()),
   STRING: useStringWidget(),
   MARKDOWN: useMarkdownWidget(),
   COMBO: useComboWidget(),