feat(cloud): add asset widget support for PrimitiveNode model selection

src/extensions/core/widgetInputs.ts

@@ -12,17 +12,20 @@ import type {
   IBaseWidget,
   TWidgetValue
 } from '@/lib/litegraph/src/types/widgets'
+import { assetService } from '@/platform/assets/services/assetService'
+import { createAssetWidget } from '@/platform/assets/utils/createAssetWidget'
+import { isCloud } from '@/platform/distribution/types'
 import type { InputSpec } from '@/schemas/nodeDefSchema'
 import { app } from '@/scripts/app'
 import {
   ComfyWidgets,
   addValueControlWidgets,
   isValidWidgetType
 } from '@/scripts/widgets'
+import { isPrimitiveNode } from '@/renderer/utils/nodeTypeGuards'
 import { CONFIG, GET_CONFIG } from '@/services/litegraphService'
 import { mergeInputSpec } from '@/utils/nodeDefUtil'
 import { applyTextReplacements } from '@/utils/searchAndReplace'
-import { isPrimitiveNode } from '@/renderer/utils/nodeTypeGuards'
 
 const replacePropertyName = 'Run widget replace on values'
 export class PrimitiveNode extends LGraphNode {
@@ -103,7 +106,7 @@ export class PrimitiveNode extends LGraphNode {
 
   override onAfterGraphConfigured() {
     if (this.outputs[0].links?.length && !this.widgets?.length) {
-      this.#onFirstConnection()
+      this._onFirstConnection()
 
       // Populate widget values from config data
       if (this.widgets && this.widgets_values) {
@@ -116,7 +119,7 @@ export class PrimitiveNode extends LGraphNode {
       }
 
       // Merge values if required
-      this.#mergeWidgetConfig()
+      this._mergeWidgetConfig()
     }
   }
 
@@ -133,11 +136,11 @@ export class PrimitiveNode extends LGraphNode {
     const links = this.outputs[0].links
     if (connected) {
       if (links?.length && !this.widgets?.length) {
-        this.#onFirstConnection()
+        this._onFirstConnection()
       }
     } else {
       // We may have removed a link that caused the constraints to change
-      this.#mergeWidgetConfig()
+      this._mergeWidgetConfig()
 
       if (!links?.length) {
         this.onLastDisconnect()
@@ -159,7 +162,7 @@ export class PrimitiveNode extends LGraphNode {
     }
 
     if (this.outputs[slot].links?.length) {
-      const valid = this.#isValidConnection(input)
+      const valid = this._isValidConnection(input)
       if (valid) {
         // On connect of additional outputs, copy our value to their widget
         this.applyToGraph([{ target_id: target_node.id, target_slot } as LLink])
@@ -170,7 +173,7 @@ export class PrimitiveNode extends LGraphNode {
     return true
   }
 
-  #onFirstConnection(recreating?: boolean) {
+  private _onFirstConnection(recreating?: boolean) {
     // First connection can fire before the graph is ready on initial load so random things can be missing
     if (!this.outputs[0].links || !this.graph) {
       this.onLastDisconnect()
@@ -204,7 +207,7 @@ export class PrimitiveNode extends LGraphNode {
     this.outputs[0].name = type
     this.outputs[0].widget = widget
 
-    this.#createWidget(
+    this._createWidget(
       widget[CONFIG] ?? config,
       theirNode,
       widget.name,
@@ -213,7 +216,7 @@ export class PrimitiveNode extends LGraphNode {
     )
   }
 
-  #createWidget(
+  private _createWidget(
     inputData: InputSpec,
     node: LGraphNode,
     widgetName: string,
@@ -228,6 +231,20 @@ export class PrimitiveNode extends LGraphNode {
     // Store current size as addWidget resizes the node
     const [oldWidth, oldHeight] = this.size
     let widget: IBaseWidget
+
+    // Cloud: Use asset widget for model-eligible inputs
+    if (isCloud && type === 'COMBO') {
+      const isEligible = assetService.isAssetBrowserEligible(
+        node.comfyClass,
+        widgetName
+      )
+      if (isEligible) {
+        widget = this._createAssetWidget(node, widgetName, inputData)
+        this._finalizeWidget(widget, oldWidth, oldHeight, recreating)
+        return
+      }
+    }
+
     if (isValidWidgetType(type)) {
       widget = (ComfyWidgets[type](this, 'value', inputData, app) || {}).widget
     } else {
@@ -277,20 +294,49 @@ export class PrimitiveNode extends LGraphNode {
       }
     }
 
-    // When our value changes, update other widgets to reflect our changes
-    // e.g. so LoadImage shows correct image
+    this._finalizeWidget(widget, oldWidth, oldHeight, recreating)
+  }
+
+  private _createAssetWidget(
+    targetNode: LGraphNode,
+    _widgetName: string,
+    inputData: InputSpec
+  ): IBaseWidget {
+    const defaultValue = inputData[1]?.default as string | undefined
+    return createAssetWidget({
+      node: this,
+      widgetName: 'value',
+      nodeTypeForBrowser: targetNode.comfyClass ?? '',
+      defaultValue,
+      onValueChange: (widget, newValue, oldValue) => {
+        widget.callback?.(
+          widget.value,
+          app.canvas,
+          this,
+          app.canvas.graph_mouse,
+          {} as CanvasPointerEvent
+        )
+        this.onWidgetChanged?.(widget.name, newValue, oldValue, widget)
+      }
+    })
+  }
+
+  private _finalizeWidget(
+    widget: IBaseWidget,
+    oldWidth: number,
+    oldHeight: number,
+    recreating: boolean
+  ) {
     widget.callback = useChainCallback(widget.callback, () => {
       this.applyToGraph()
     })
 
-    // Use the biggest dimensions in case the widgets caused the node to grow
     this.setSize([
       Math.max(this.size[0], oldWidth),
       Math.max(this.size[1], oldHeight)
     ])
 
     if (!recreating) {
-      // Grow our node more if required
       const sz = this.computeSize()
       if (this.size[0] < sz[0]) {
         this.size[0] = sz[0]
@@ -307,16 +353,16 @@ export class PrimitiveNode extends LGraphNode {
 
   recreateWidget() {
     const values = this.widgets?.map((w) => w.value)
-    this.#removeWidgets()
-    this.#onFirstConnection(true)
+    this._removeWidgets()
+    this._onFirstConnection(true)
     if (values?.length && this.widgets) {
       for (let i = 0; i < this.widgets.length; i++)
         this.widgets[i].value = values[i]
     }
     return this.widgets?.[0]
   }
 
-  #mergeWidgetConfig() {
+  private _mergeWidgetConfig() {
     // Merge widget configs if the node has multiple outputs
     const output = this.outputs[0]
     const links = output.links ?? []
@@ -348,11 +394,11 @@ export class PrimitiveNode extends LGraphNode {
       const theirInput = theirNode.inputs[link.target_slot]
 
       // Call is valid connection so it can merge the configs when validating
-      this.#isValidConnection(theirInput, hasConfig)
+      this._isValidConnection(theirInput, hasConfig)
     }
   }
 
-  #isValidConnection(input: INodeInputSlot, forceUpdate?: boolean) {
+  private _isValidConnection(input: INodeInputSlot, forceUpdate?: boolean) {
     // Only allow connections where the configs match
     const output = this.outputs?.[0]
     const config2 = (input.widget?.[GET_CONFIG] as () => InputSpec)?.()
@@ -367,7 +413,7 @@ export class PrimitiveNode extends LGraphNode {
     )
   }
 
-  #removeWidgets() {
+  private _removeWidgets() {
     if (this.widgets) {
       // Allow widgets to cleanup
       for (const w of this.widgets) {
@@ -398,7 +444,7 @@ export class PrimitiveNode extends LGraphNode {
     this.outputs[0].name = 'connect to widget input'
     delete this.outputs[0].widget
 
-    this.#removeWidgets()
+    this._removeWidgets()
   }
 }
 

src/lib/litegraph/src/litegraph.ts

@@ -150,7 +150,9 @@ export { BaseWidget } from './widgets/BaseWidget'
 
 export { LegacyWidget } from './widgets/LegacyWidget'
 
-export { isComboWidget, isAssetWidget } from './widgets/widgetMap'
+export { isComboWidget } from './widgets/widgetMap'
+/** @knipIgnoreUnusedButUsedByCustomNodes */
+export { isAssetWidget } from './widgets/widgetMap'
 // Additional test-specific exports
 export { LGraphButton } from './LGraphButton'
 export { MovingOutputLink } from './canvas/MovingOutputLink'

src/lib/litegraph/src/widgets/widgetMap.ts

@@ -141,7 +141,10 @@ export function isComboWidget(widget: IBaseWidget): widget is IComboWidget {
   return widget.type === 'combo'
 }
 
-/** Type guard: Narrow **from {@link IBaseWidget}** to {@link IAssetWidget}. */
+/**
+ * Type guard: Narrow **from {@link IBaseWidget}** to {@link IAssetWidget}.
+ * @knipIgnoreUnusedButUsedByCustomNodes
+ */
 export function isAssetWidget(widget: IBaseWidget): widget is IAssetWidget {
   return widget.type === 'asset'
 }

src/platform/assets/utils/createAssetWidget.ts

@@ -0,0 +1,88 @@
+import { t } from '@/i18n'
+import type { LGraphNode } from '@/lib/litegraph/src/litegraph'
+import type {
+  IBaseWidget,
+  IWidgetAssetOptions
+} from '@/lib/litegraph/src/types/widgets'
+import { useAssetBrowserDialog } from '@/platform/assets/composables/useAssetBrowserDialog'
+import {
+  assetFilenameSchema,
+  assetItemSchema
+} from '@/platform/assets/schemas/assetSchema'
+import { getAssetFilename } from '@/platform/assets/utils/assetMetadataUtils'
+
+interface CreateAssetWidgetParams {
+  /** The node to add the widget to */
+  node: LGraphNode
+  /** The widget name */
+  widgetName: string
+  /** The node type to show in asset browser (may differ from node.comfyClass for PrimitiveNode) */
+  nodeTypeForBrowser: string
+  /** Default value for the widget */
+  defaultValue?: string
+  /** Callback when widget value changes */
+  onValueChange?: (
+    widget: IBaseWidget,
+    newValue: string,
+    oldValue: unknown
+  ) => void
+}
+
+/**
+ * Creates an asset widget that opens the Asset Browser dialog for model selection.
+ * Used by both regular nodes (via useComboWidget) and PrimitiveNode.
+ *
+ * @param params - Configuration for the asset widget
+ * @returns The created asset widget
+ */
+export function createAssetWidget(
+  params: CreateAssetWidgetParams
+): IBaseWidget {
+  const { node, widgetName, nodeTypeForBrowser, defaultValue, onValueChange } =
+    params
+
+  const displayLabel = defaultValue ?? t('widgets.selectModel')
+  const assetBrowserDialog = useAssetBrowserDialog()
+
+  async function openModal(widget: IBaseWidget) {
+    await assetBrowserDialog.show({
+      nodeType: nodeTypeForBrowser,
+      inputName: widgetName,
+      currentValue: widget.value as string,
+      onAssetSelected: (asset) => {
+        const validatedAsset = assetItemSchema.safeParse(asset)
+
+        if (!validatedAsset.success) {
+          console.error(
+            'Invalid asset item:',
+            validatedAsset.error.errors,
+            'Received:',
+            asset
+          )
+          return
+        }
+
+        const filename = getAssetFilename(validatedAsset.data)
+        const validatedFilename = assetFilenameSchema.safeParse(filename)
+
+        if (!validatedFilename.success) {
+          console.error(
+            'Invalid asset filename:',
+            validatedFilename.error.errors,
+            'for asset:',
+            validatedAsset.data.id
+          )
+          return
+        }
+
+        const oldValue = widget.value
+        widget.value = validatedFilename.data
+        onValueChange?.(widget, validatedFilename.data, oldValue)
+      }
+    })
+  }
+
+  const options: IWidgetAssetOptions = { openModal }
+
+  return node.addWidget('asset', widgetName, displayLabel, () => {}, options)
+}