[Refactor] useSeedWidget composable (#2508)

src/composables/widgets/useIntWidget.ts

@@ -1,7 +1,11 @@
 import type { LGraphNode } from '@comfyorg/litegraph'
 import type { INumericWidget } from '@comfyorg/litegraph/dist/types/widgets'
 
-import type { ComfyWidgetConstructor } from '@/scripts/widgets'
+import type { ComfyApp } from '@/scripts/app'
+import {
+  type ComfyWidgetConstructor,
+  addValueControlWidget
+} from '@/scripts/widgets'
 import { useSettingStore } from '@/stores/settingStore'
 import { InputSpec } from '@/types/apiTypes'
 import { getNumberDefaults } from '@/utils/mathUtil'
@@ -10,7 +14,9 @@ export const useIntWidget = () => {
   const widgetConstructor: ComfyWidgetConstructor = (
     node: LGraphNode,
     inputName: string,
-    inputData: InputSpec
+    inputData: InputSpec,
+    app?: ComfyApp,
+    widgetName?: string
   ) => {
     const settingStore = useSettingStore()
     const sliderEnabled = !settingStore.get('Comfy.DisableSliders')
@@ -28,7 +34,7 @@ export const useIntWidget = () => {
     })
     config.precision = 0
 
-    return {
+    const result = {
       widget: node.addWidget(
         widgetType,
         inputName,
@@ -44,6 +50,20 @@ export const useIntWidget = () => {
         config
       )
     }
+
+    if (inputData[1]?.control_after_generate) {
+      const seedControl = addValueControlWidget(
+        node,
+        result.widget,
+        'randomize',
+        undefined,
+        widgetName,
+        inputData
+      )
+      result.widget.linkedWidgets = [seedControl]
+    }
+
+    return result
   }
   return widgetConstructor
 }

src/composables/widgets/useSeedWidget.ts

@@ -0,0 +1,27 @@
+import type { LGraphNode } from '@comfyorg/litegraph'
+
+import type { ComfyWidgetConstructor } from '@/scripts/widgets'
+import type { ComfyApp } from '@/types'
+import type { InputSpec } from '@/types/apiTypes'
+
+import { useIntWidget } from './useIntWidget'
+
+export const useSeedWidget = () => {
+  const IntWidget = useIntWidget()
+
+  const widgetConstructor: ComfyWidgetConstructor = (
+    node: LGraphNode,
+    inputName: string,
+    inputData: InputSpec,
+    app?: ComfyApp,
+    widgetName?: string
+  ) => {
+    inputData[1] = {
+      ...inputData[1],
+      control_after_generate: true
+    }
+    return IntWidget(node, inputName, inputData, app, widgetName)
+  }
+
+  return widgetConstructor
+}

src/scripts/widgets.ts

@@ -17,6 +17,7 @@ import { Markdown as TiptapMarkdown } from 'tiptap-markdown'
 import { useFloatWidget } from '@/composables/widgets/useFloatWidget'
 import { useIntWidget } from '@/composables/widgets/useIntWidget'
 import { useRemoteWidget } from '@/composables/widgets/useRemoteWidget'
+import { useSeedWidget } from '@/composables/widgets/useSeedWidget'
 import { useStringWidget } from '@/composables/widgets/useStringWidget'
 import { useSettingStore } from '@/stores/settingStore'
 import { useToastStore } from '@/stores/toastStore'
@@ -51,36 +52,6 @@ export function updateControlWidgetLabel(widget) {
 export const IS_CONTROL_WIDGET = Symbol()
 const HAS_EXECUTED = Symbol()
 
-function getNumberDefaults(
-  inputData: InputSpec,
-  defaultStep,
-  precision,
-  enable_rounding
-) {
-  let defaultVal = inputData[1]['default']
-  let { min, max, step, round } = inputData[1]
-
-  if (defaultVal == undefined) defaultVal = 0
-  if (min == undefined) min = 0
-  if (max == undefined) max = 2048
-  if (step == undefined) step = defaultStep
-  // precision is the number of decimal places to show.
-  // by default, display the the smallest number of decimal places such that changes of size step are visible.
-  if (precision == undefined) {
-    precision = Math.max(-Math.floor(Math.log10(step)), 0)
-  }
-
-  if (enable_rounding && (round == undefined || round === true)) {
-    // by default, round the value to those decimal places shown.
-    round = Math.round(1000000 * Math.pow(0.1, precision)) / 1000000
-  }
-
-  return {
-    val: defaultVal,
-    config: { min, max, step: 10.0 * step, round, precision }
-  }
-}
-
 export function addValueControlWidget(
   node,
   targetWidget,
@@ -284,60 +255,6 @@ export function addValueControlWidgets(
   return widgets
 }
 
-function seedWidget(node, inputName, inputData: InputSpec, app, widgetName) {
-  const seed = createIntWidget(node, inputName, inputData, app, true)
-  const seedControl = addValueControlWidget(
-    node,
-    seed.widget,
-    'randomize',
-    undefined,
-    widgetName,
-    inputData
-  )
-
-  seed.widget.linkedWidgets = [seedControl]
-  return seed
-}
-
-function createIntWidget(
-  node,
-  inputName,
-  inputData: InputSpec,
-  app,
-  isSeedInput: boolean = false
-) {
-  const control = inputData[1]?.control_after_generate
-  if (!isSeedInput && control) {
-    return seedWidget(
-      node,
-      inputName,
-      inputData,
-      app,
-      typeof control === 'string' ? control : undefined
-    )
-  }
-
-  let widgetType = isSlider(inputData[1]['display'], app)
-  const { val, config } = getNumberDefaults(inputData, 1, 0, true)
-  Object.assign(config, { precision: 0 })
-  return {
-    widget: node.addWidget(
-      widgetType,
-      inputName,
-      val,
-      function (v) {
-        const s = this.options.step / 10
-        let sh = this.options.min % s
-        if (isNaN(sh)) {
-          sh = 0
-        }
-        this.value = Math.round((v - sh) / s) * s + sh
-      },
-      config
-    )
-  }
-}
-
 function addMarkdownWidget(node, name: string, opts, app: ComfyApp) {
   TiptapMarkdown.configure({
     html: false,
@@ -425,19 +342,13 @@ function addMarkdownWidget(node, name: string, opts, app: ComfyApp) {
   return { minWidth: 400, minHeight: 200, widget }
 }
 
-function isSlider(display, app) {
-  if (app.ui.settings.getSettingValue('Comfy.DisableSliders')) {
-    return 'number'
-  }
-
-  return display === 'slider' ? 'slider' : 'number'
-}
+const SeedWidget = useSeedWidget()
 
 export const ComfyWidgets: Record<string, ComfyWidgetConstructor> = {
-  'INT:seed': seedWidget,
-  'INT:noise_seed': seedWidget,
-  FLOAT: useFloatWidget(),
+  'INT:seed': SeedWidget,
+  'INT:noise_seed': SeedWidget,
   INT: useIntWidget(),
+  FLOAT: useFloatWidget(),
   BOOLEAN(node, inputName, inputData) {
     let defaultVal = false
     let options = {}

src/types/apiTypes.ts

@@ -297,7 +297,12 @@ const zIntInputSpec = inputSpec([
     step: z.number().optional(),
     // Note: Many node authors are using INT to pass list of INT.
     // TODO: Add list of ints type.
-    default: z.union([z.number(), z.array(z.number())]).optional()
+    default: z.union([z.number(), z.array(z.number())]).optional(),
+    /**
+     * If true, a linked widget will be added to the node to select the mode
+     * of `control_after_generate`.
+     */
+    control_after_generate: z.boolean().optional()
   })
 ])