Add Asset Widget (#5475)

* [feat] carve out path to call asset browser in combo widget

* Add Asset Widget

* [feat] add fallback "Select model" label

---------

Co-authored-by: Arjan Singh <[email protected]>

Author: AustinMroz

src/lib/litegraph/src/types/widgets.ts

@@ -76,6 +76,7 @@ export type IWidget =
   | IImageCompareWidget
   | ISelectButtonWidget
   | ITextareaWidget
+  | IAssetWidget
 
 export interface IBooleanWidget extends IBaseWidget<boolean, 'toggle'> {
   type: 'toggle'
@@ -224,6 +225,12 @@ export interface ITextareaWidget extends IBaseWidget<string, 'textarea'> {
   value: string
 }
 
+export interface IAssetWidget
+  extends IBaseWidget<string, 'asset', IWidgetOptions<string[]>> {
+  type: 'asset'
+  value: string
+}
+
 /**
  * Valid widget types.  TS cannot provide easily extensible type safety for this at present.
  * Override linkedWidgets[]

src/lib/litegraph/src/widgets/AssetWidget.ts

@@ -0,0 +1,41 @@
+import type { LGraphNode } from '@/lib/litegraph/src/LGraphNode'
+import type { IAssetWidget } from '@/lib/litegraph/src/types/widgets'
+
+import { BaseWidget, type DrawWidgetOptions } from './BaseWidget'
+
+export class AssetWidget
+  extends BaseWidget<IAssetWidget>
+  implements IAssetWidget
+{
+  constructor(widget: IAssetWidget, node: LGraphNode) {
+    super(widget, node)
+    this.type ??= 'asset'
+    this.value = widget.value?.toString() ?? ''
+  }
+
+  override get _displayValue(): string {
+    return String(this.value) //FIXME: Resolve asset name
+  }
+
+  override drawWidget(
+    ctx: CanvasRenderingContext2D,
+    { width, showText = true }: DrawWidgetOptions
+  ) {
+    // Store original context attributes
+    const { fillStyle, strokeStyle, textAlign } = ctx
+
+    this.drawWidgetShape(ctx, { width, showText })
+
+    if (showText) {
+      this.drawTruncatingText({ ctx, width, leftPadding: 0, rightPadding: 0 })
+    }
+
+    // Restore original context attributes
+    Object.assign(ctx, { textAlign, strokeStyle, fillStyle })
+  }
+
+  override onClick() {
+    //Open Modal
+    this.callback?.(this.value)
+  }
+}

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

@@ -7,6 +7,7 @@ import type {
 } from '@/lib/litegraph/src/types/widgets'
 import { toClass } from '@/lib/litegraph/src/utils/type'
 
+import { AssetWidget } from './AssetWidget'
 import { BaseWidget } from './BaseWidget'
 import { BooleanWidget } from './BooleanWidget'
 import { ButtonWidget } from './ButtonWidget'
@@ -47,6 +48,7 @@ export type WidgetTypeMap = {
   imagecompare: ImageCompareWidget
   selectbutton: SelectButtonWidget
   textarea: TextareaWidget
+  asset: AssetWidget
   [key: string]: BaseWidget
 }
 
@@ -115,6 +117,8 @@ export function toConcreteWidget<TWidget extends IWidget | IBaseWidget>(
       return toClass(SelectButtonWidget, narrowedWidget, node)
     case 'textarea':
       return toClass(TextareaWidget, narrowedWidget, node)
+    case 'asset':
+      return toClass(AssetWidget, narrowedWidget, node)
     default: {
       if (wrapLegacyWidgets) return toClass(LegacyWidget, widget, node)
     }

src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts

@@ -69,19 +69,15 @@ const addComboWidget = (
   )
 
   if (isUsingAssetAPI && isEligible) {
-    // Create button widget for Asset Browser
+    // Get the default value for the button text (currently selected model)
     const currentValue = getDefaultValue(inputSpec)
+    const displayLabel = currentValue ?? t('widgets.selectModel')
 
-    const widget = node.addWidget(
-      'button',
-      inputSpec.name,
-      t('widgets.selectModel'),
-      () => {
-        console.log(
-          `Asset Browser would open here for:\nNode: ${node.type}\nWidget: ${inputSpec.name}\nCurrent Value:${currentValue}`
-        )
-      }
-    )
+    const widget = node.addWidget('asset', inputSpec.name, displayLabel, () => {
+      console.log(
+        `Asset Browser would open here for:\nNode: ${node.type}\nWidget: ${inputSpec.name}\nCurrent Value:${currentValue}`
+      )
+    })
 
     return widget
   }
@@ -129,7 +125,7 @@ const addComboWidget = (
     )
   }
 
-  return widget
+  return widget as IBaseWidget
 }
 
 export const useComboWidget = () => {

tests-ui/tests/renderer/extensions/vueNodes/widgets/composables/useComboWidget.test.ts

@@ -88,8 +88,8 @@ describe('useComboWidget', () => {
   })
 
   it('should create normal combo widget when asset API is disabled', () => {
-    mockSettingStoreGet.mockReturnValue(false)
-    vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(true)
+    mockSettingStoreGet.mockReturnValue(false) // Asset API disabled
+    vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(true) // Widget is eligible
 
     const constructor = useComboWidget()
     const mockWidget = createMockWidget()
@@ -101,6 +101,7 @@ describe('useComboWidget', () => {
     })
 
     const widget = constructor(mockNode, inputSpec)
+    expect(widget).toBe(mockWidget)
 
     expect(mockNode.addWidget).toHaveBeenCalledWith(
       'combo',
@@ -142,15 +143,15 @@ describe('useComboWidget', () => {
     expect(widget).toBe(mockWidget)
   })
 
-  it('should create asset browser button widget when API enabled and widget eligible', () => {
+  it('should create asset browser widget when API enabled and widget eligible', () => {
     mockSettingStoreGet.mockReturnValue(true)
     vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(true)
 
     const constructor = useComboWidget()
     const mockWidget = createMockWidget({
-      type: 'button',
+      type: 'asset',
       name: 'ckpt_name',
-      value: 'Select model'
+      value: 'model1.safetensors'
     })
     const mockNode = createMockNode('CheckpointLoaderSimple')
     vi.mocked(mockNode.addWidget).mockReturnValue(mockWidget)
@@ -162,9 +163,9 @@ describe('useComboWidget', () => {
     const widget = constructor(mockNode, inputSpec)
 
     expect(mockNode.addWidget).toHaveBeenCalledWith(
-      'button',
+      'asset',
       'ckpt_name',
-      'Select model',
+      'model1.safetensors',
       expect.any(Function)
     )
     expect(mockSettingStoreGet).toHaveBeenCalledWith('Comfy.Assets.UseAssetAPI')
@@ -175,15 +176,48 @@ describe('useComboWidget', () => {
     expect(widget).toBe(mockWidget)
   })
 
-  it('should use asset browser button even when inputSpec has a default value but no options', () => {
+  it('should create asset browser widget with options when API enabled and widget eligible', () => {
     mockSettingStoreGet.mockReturnValue(true)
     vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(true)
 
     const constructor = useComboWidget()
     const mockWidget = createMockWidget({
-      type: 'button',
+      type: 'asset',
       name: 'ckpt_name',
-      value: 'Select model'
+      value: 'model1.safetensors'
+    })
+    const mockNode = createMockNode('CheckpointLoaderSimple')
+    vi.mocked(mockNode.addWidget).mockReturnValue(mockWidget)
+    const inputSpec = createMockInputSpec({
+      name: 'ckpt_name',
+      options: ['model1.safetensors', 'model2.safetensors']
+    })
+
+    const widget = constructor(mockNode, inputSpec)
+
+    expect(mockNode.addWidget).toHaveBeenCalledWith(
+      'asset',
+      'ckpt_name',
+      'model1.safetensors',
+      expect.any(Function)
+    )
+    expect(mockSettingStoreGet).toHaveBeenCalledWith('Comfy.Assets.UseAssetAPI')
+    expect(vi.mocked(assetService.isAssetBrowserEligible)).toHaveBeenCalledWith(
+      'ckpt_name',
+      'CheckpointLoaderSimple'
+    )
+    expect(widget).toBe(mockWidget)
+  })
+
+  it('should use asset browser widget even when inputSpec has a default value but no options', () => {
+    mockSettingStoreGet.mockReturnValue(true)
+    vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(true)
+
+    const constructor = useComboWidget()
+    const mockWidget = createMockWidget({
+      type: 'asset',
+      name: 'ckpt_name',
+      value: 'fallback.safetensors'
     })
     const mockNode = createMockNode('CheckpointLoaderSimple')
     vi.mocked(mockNode.addWidget).mockReturnValue(mockWidget)
@@ -196,9 +230,42 @@ describe('useComboWidget', () => {
     const widget = constructor(mockNode, inputSpec)
 
     expect(mockNode.addWidget).toHaveBeenCalledWith(
-      'button',
+      'asset',
+      'ckpt_name',
+      'fallback.safetensors',
+      expect.any(Function)
+    )
+    expect(mockSettingStoreGet).toHaveBeenCalledWith('Comfy.Assets.UseAssetAPI')
+    expect(vi.mocked(assetService.isAssetBrowserEligible)).toHaveBeenCalledWith(
+      'ckpt_name',
+      'CheckpointLoaderSimple'
+    )
+    expect(widget).toBe(mockWidget)
+  })
+
+  it('should show Select model when asset widget has undefined current value', () => {
+    mockSettingStoreGet.mockReturnValue(true)
+    vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(true)
+
+    const constructor = useComboWidget()
+    const mockWidget = createMockWidget({
+      type: 'asset',
+      name: 'ckpt_name',
+      value: 'Select model'
+    })
+    const mockNode = createMockNode('CheckpointLoaderSimple')
+    vi.mocked(mockNode.addWidget).mockReturnValue(mockWidget)
+    const inputSpec = createMockInputSpec({
+      name: 'ckpt_name'
+      // Note: no default, no options, not remote - getDefaultValue returns undefined
+    })
+
+    const widget = constructor(mockNode, inputSpec)
+
+    expect(mockNode.addWidget).toHaveBeenCalledWith(
+      'asset',
       'ckpt_name',
-      'Select model',
+      'Select model', // Should fallback to this instead of undefined
       expect.any(Function)
     )
     expect(mockSettingStoreGet).toHaveBeenCalledWith('Comfy.Assets.UseAssetAPI')