From cb1d66f9deb10474ffb51dde7189c885c1a36b49 Mon Sep 17 00:00:00 2001 From: AustinMroz Date: Thu, 11 Sep 2025 12:00:34 -0700 Subject: [PATCH] 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 --- src/lib/litegraph/src/types/widgets.ts | 7 ++ src/lib/litegraph/src/widgets/AssetWidget.ts | 41 +++++++++ src/lib/litegraph/src/widgets/widgetMap.ts | 4 + .../widgets/composables/useComboWidget.ts | 20 ++-- .../composables/useComboWidget.test.ts | 91 ++++++++++++++++--- 5 files changed, 139 insertions(+), 24 deletions(-) create mode 100644 src/lib/litegraph/src/widgets/AssetWidget.ts diff --git a/src/lib/litegraph/src/types/widgets.ts b/src/lib/litegraph/src/types/widgets.ts index fa47c6a93..850fe9bcb 100644 --- a/src/lib/litegraph/src/types/widgets.ts +++ b/src/lib/litegraph/src/types/widgets.ts @@ -76,6 +76,7 @@ export type IWidget = | IImageCompareWidget | ISelectButtonWidget | ITextareaWidget + | IAssetWidget export interface IBooleanWidget extends IBaseWidget { type: 'toggle' @@ -224,6 +225,12 @@ export interface ITextareaWidget extends IBaseWidget { value: string } +export interface IAssetWidget + extends IBaseWidget> { + type: 'asset' + value: string +} + /** * Valid widget types. TS cannot provide easily extensible type safety for this at present. * Override linkedWidgets[] diff --git a/src/lib/litegraph/src/widgets/AssetWidget.ts b/src/lib/litegraph/src/widgets/AssetWidget.ts new file mode 100644 index 000000000..f8a8e1209 --- /dev/null +++ b/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 + 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) + } +} diff --git a/src/lib/litegraph/src/widgets/widgetMap.ts b/src/lib/litegraph/src/widgets/widgetMap.ts index 42a8f5663..02cdb5597 100644 --- a/src/lib/litegraph/src/widgets/widgetMap.ts +++ b/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( 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) } diff --git a/src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts b/src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts index 586cfb27e..102642816 100644 --- a/src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts +++ b/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 = () => { diff --git a/tests-ui/tests/renderer/extensions/vueNodes/widgets/composables/useComboWidget.test.ts b/tests-ui/tests/renderer/extensions/vueNodes/widgets/composables/useComboWidget.test.ts index ce3acffc5..1ec9bcf1a 100644 --- a/tests-ui/tests/renderer/extensions/vueNodes/widgets/composables/useComboWidget.test.ts +++ b/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', - 'Select model', + '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', // Should fallback to this instead of undefined expect.any(Function) ) expect(mockSettingStoreGet).toHaveBeenCalledWith('Comfy.Assets.UseAssetAPI')