Add multi-select widget

Author: Chenlei Hu

src/components/graph/widgets/MultiSelectWidget.vue

@@ -0,0 +1,31 @@
+<template>
+  <div>
+    <MultiSelect
+      v-model="selectedItems"
+      :options="options"
+      filter
+      :placeholder="placeholder"
+      :maxSelectedLabels="3"
+      class="w-full"
+    />
+  </div>
+</template>
+
+<script setup lang="ts">
+import MultiSelect from 'primevue/multiselect'
+import { computed, defineModel } from 'vue'
+
+import type { ComboInputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
+import type { ComponentWidget } from '@/scripts/domWidget'
+
+const selectedItems = defineModel<string[]>({ required: true })
+const { widget } = defineProps<{
+  widget: ComponentWidget<string[]>
+}>()
+const options = computed(
+  () => (widget.inputSpec as ComboInputSpec).options ?? []
+)
+const placeholder = computed(
+  () => (widget.inputSpec as ComboInputSpec).placeholder ?? 'Select items'
+)
+</script>

src/composables/widgets/useComboWidget.ts

@@ -1,16 +1,24 @@
 import type { LGraphNode } from '@comfyorg/litegraph'
 import type { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets'
+import { ref } from 'vue'
 
+import MultiSelectWidget from '@/components/graph/widgets/MultiSelectWidget.vue'
 import { transformInputSpecV2ToV1 } from '@/schemas/nodeDef/migration'
 import {
   ComboInputSpec,
   type InputSpec,
   isComboInputSpec
 } from '@/schemas/nodeDef/nodeDefSchemaV2'
+import {
+  type BaseDOMWidget,
+  ComponentWidgetImpl,
+  addWidget
+} from '@/scripts/domWidget'
 import {
   type ComfyWidgetConstructorV2,
   addValueControlWidgets
 } from '@/scripts/widgets'
+import { generateUUID } from '@/utils/formatUtil'
 
 import { useRemoteWidget } from './useRemoteWidget'
 
@@ -21,56 +29,81 @@ const getDefaultValue = (inputSpec: ComboInputSpec) => {
   return undefined
 }
 
-export const useComboWidget = () => {
-  const widgetConstructor: ComfyWidgetConstructorV2 = (
-    node: LGraphNode,
-    inputSpec: InputSpec
-  ) => {
-    if (!isComboInputSpec(inputSpec)) {
-      throw new Error(`Invalid input data: ${inputSpec}`)
+const addMultiSelectWidget = (node: LGraphNode, inputSpec: ComboInputSpec) => {
+  const widgetValue = ref<string[]>([])
+  const widget = new ComponentWidgetImpl({
+    id: generateUUID(),
+    node,
+    name: inputSpec.name,
+    component: MultiSelectWidget,
+    inputSpec,
+    options: {
+      getValue: () => widgetValue.value,
+      setValue: (value: string[]) => {
+        widgetValue.value = value
+      }
     }
+  })
+  addWidget(node, widget as BaseDOMWidget<object | string>)
+  // TODO: Add remote support to multi-select widget
+  return widget
+}
 
-    const comboOptions = inputSpec.options ?? []
-    const defaultValue = getDefaultValue(inputSpec)
+const addComboWidget = (node: LGraphNode, inputSpec: ComboInputSpec) => {
+  const defaultValue = getDefaultValue(inputSpec)
+  const comboOptions = inputSpec.options ?? []
+  const widget = node.addWidget(
+    'combo',
+    inputSpec.name,
+    defaultValue,
+    () => {},
+    {
+      values: comboOptions
+    }
+  ) as IComboWidget
 
-    const widget = node.addWidget(
-      'combo',
-      inputSpec.name,
+  if (inputSpec.remote) {
+    const remoteWidget = useRemoteWidget({
+      remoteConfig: inputSpec.remote,
       defaultValue,
-      () => {},
-      {
-        values: comboOptions
+      node,
+      widget
+    })
+    if (inputSpec.remote.refresh_button) remoteWidget.addRefreshButton()
+
+    const origOptions = widget.options
+    widget.options = new Proxy(origOptions as Record<string | symbol, any>, {
+      get(target, prop: string | symbol) {
+        if (prop !== 'values') return target[prop]
+        return remoteWidget.getValue()
       }
-    ) as IComboWidget
+    })
+  }
 
-    if (inputSpec.remote) {
-      const remoteWidget = useRemoteWidget({
-        remoteConfig: inputSpec.remote,
-        defaultValue,
-        node,
-        widget
-      })
-      if (inputSpec.remote.refresh_button) remoteWidget.addRefreshButton()
+  if (inputSpec.control_after_generate) {
+    widget.linkedWidgets = addValueControlWidgets(
+      node,
+      widget,
+      undefined,
+      undefined,
+      transformInputSpecV2ToV1(inputSpec)
+    )
+  }
 
-      const origOptions = widget.options
-      widget.options = new Proxy(origOptions as Record<string | symbol, any>, {
-        get(target, prop: string | symbol) {
-          if (prop !== 'values') return target[prop]
-          return remoteWidget.getValue()
-        }
-      })
-    }
+  return widget
+}
 
-    if (inputSpec.control_after_generate) {
-      widget.linkedWidgets = addValueControlWidgets(
-        node,
-        widget,
-        undefined,
-        undefined,
-        transformInputSpecV2ToV1(inputSpec)
-      )
+export const useComboWidget = () => {
+  const widgetConstructor: ComfyWidgetConstructorV2 = (
+    node: LGraphNode,
+    inputSpec: InputSpec
+  ) => {
+    if (!isComboInputSpec(inputSpec)) {
+      throw new Error(`Invalid input data: ${inputSpec}`)
     }
-    return widget
+    return inputSpec.multi_select
+      ? addMultiSelectWidget(node, inputSpec)
+      : addComboWidget(node, inputSpec)
   }
 
   return widgetConstructor

src/schemas/nodeDefSchema.ts

@@ -72,7 +72,11 @@ export const zComboInputOptions = zBaseInputOptions.extend({
   allow_batch: z.boolean().optional(),
   video_upload: z.boolean().optional(),
   options: z.array(zComboOption).optional(),
-  remote: zRemoteWidgetConfig.optional()
+  remote: zRemoteWidgetConfig.optional(),
+  /** Whether the widget is a multi-select widget. */
+  multi_select: z.boolean().optional(),
+  /** Placeholder when no item is selected in multi-select widget. */
+  placeholder: z.string().optional()
 })
 
 const zIntInputSpec = z.tuple([z.literal('INT'), zIntInputOptions.optional()])

src/scripts/domWidget.ts

@@ -8,6 +8,7 @@ import _ from 'lodash'
 import type { Component } from 'vue'
 
 import { useChainCallback } from '@/composables/functional/useChainCallback'
+import type { InputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
 import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { generateUUID } from '@/utils/formatUtil'
 
@@ -21,9 +22,9 @@ export interface BaseDOMWidget<V extends object | string>
 
   // BaseDOMWidget properties
   /** The unique ID of the widget. */
-  id: string
+  readonly id: string
   /** The node that the widget belongs to. */
-  node: LGraphNode
+  readonly node: LGraphNode
   /** Whether the widget is visible. */
   isVisible(): boolean
 }
@@ -47,7 +48,8 @@ export interface DOMWidget<T extends HTMLElement, V extends object | string>
  */
 export interface ComponentWidget<V extends object | string>
   extends BaseDOMWidget<V> {
-  component: Component
+  readonly component: Component
+  readonly inputSpec: InputSpec
 }
 
 export interface DOMWidgetOptions<V extends object | string>
@@ -203,19 +205,22 @@ export class ComponentWidgetImpl<V extends object | string>
   implements ComponentWidget<V>
 {
   readonly component: Component
+  readonly inputSpec: InputSpec
 
   constructor(obj: {
     id: string
     node: LGraphNode
     name: string
     component: Component
+    inputSpec: InputSpec
     options: DOMWidgetOptions<V>
   }) {
     super({
       ...obj,
       type: 'custom'
     })
     this.component = obj.component
+    this.inputSpec = obj.inputSpec
   }
 
   computeLayoutSize() {
@@ -229,6 +234,23 @@ export class ComponentWidgetImpl<V extends object | string>
   }
 }
 
+export const addWidget = <W extends BaseDOMWidget<object | string>>(
+  node: LGraphNode,
+  widget: W
+) => {
+  node.addCustomWidget(widget)
+  node.onRemoved = useChainCallback(node.onRemoved, () => {
+    widget.onRemove?.()
+  })
+
+  node.onResize = useChainCallback(node.onResize, () => {
+    widget.options.beforeResize?.call(widget, node)
+    widget.options.afterResize?.call(widget, node)
+  })
+
+  useDomWidgetStore().registerWidget(widget)
+}
+
 LGraphNode.prototype.addDOMWidget = function <
   T extends HTMLElement,
   V extends object | string
@@ -239,20 +261,16 @@ LGraphNode.prototype.addDOMWidget = function <
   element: T,
   options: DOMWidgetOptions<V> = {}
 ): DOMWidget<T, V> {
+  const widget = new DOMWidgetImpl({
+    id: generateUUID(),
+    node: this,
+    name,
+    type,
+    element,
+    options: { hideOnZoom: true, ...options }
+  })
   // Note: Before `LGraphNode.configure` is called, `this.id` is always `-1`.
-  const widget = this.addCustomWidget(
-    new DOMWidgetImpl({
-      id: generateUUID(),
-      node: this,
-      name,
-      type,
-      element,
-      options: {
-        hideOnZoom: true,
-        ...options
-      }
-    })
-  )
+  addWidget(this, widget as unknown as BaseDOMWidget<object | string>)
 
   // Workaround for https://github.com/Comfy-Org/ComfyUI_frontend/issues/2493
   // Some custom nodes are explicitly expecting getter and setter of `value`
@@ -267,16 +285,5 @@ LGraphNode.prototype.addDOMWidget = function <
     }
   })
 
-  this.onRemoved = useChainCallback(this.onRemoved, () => {
-    widget.onRemove()
-  })
-
-  this.onResize = useChainCallback(this.onResize, () => {
-    options.beforeResize?.call(widget, this)
-    options.afterResize?.call(widget, this)
-  })
-
-  useDomWidgetStore().registerWidget(widget)
-
   return widget
 }