diff --git a/src/extensions/core/widgetInputs.ts b/src/extensions/core/widgetInputs.ts index 889cb4cbe..cab43004e 100644 --- a/src/extensions/core/widgetInputs.ts +++ b/src/extensions/core/widgetInputs.ts @@ -12,6 +12,10 @@ 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 { useSettingStore } from '@/platform/settings/settingStore' import type { InputSpec } from '@/schemas/nodeDefSchema' import { app } from '@/scripts/app' import { @@ -19,10 +23,10 @@ import { 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 { @@ -228,6 +232,24 @@ 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 when asset API is enabled + if (isCloud && type === 'COMBO') { + const settingStore = useSettingStore() + const isUsingAssetAPI = settingStore.get('Comfy.Assets.UseAssetAPI') + const isEligible = assetService.isAssetBrowserEligible( + node.comfyClass, + widgetName + ) + if (isUsingAssetAPI && isEligible) { + widget = this._createAssetWidget(node, widgetName, inputData) + const theirWidget = node.widgets?.find((w) => w.name === widgetName) + if (theirWidget) widget.value = theirWidget.value + this._finalizeWidget(widget, oldWidth, oldHeight, recreating) + return + } + } + if (isValidWidgetType(type)) { widget = (ComfyWidgets[type](this, 'value', inputData, app) || {}).widget } else { @@ -277,20 +299,50 @@ 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, + targetInputName: string, + inputData: InputSpec + ): IBaseWidget { + const defaultValue = inputData[1]?.default as string | undefined + return createAssetWidget({ + node: this, + widgetName: 'value', + nodeTypeForBrowser: targetNode.comfyClass ?? '', + inputNameForBrowser: targetInputName, + 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] diff --git a/src/lib/litegraph/src/litegraph.ts b/src/lib/litegraph/src/litegraph.ts index 59d62a84f..960dc190b 100644 --- a/src/lib/litegraph/src/litegraph.ts +++ b/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' diff --git a/src/lib/litegraph/src/widgets/widgetMap.ts b/src/lib/litegraph/src/widgets/widgetMap.ts index 37b906efb..7b2eaea6a 100644 --- a/src/lib/litegraph/src/widgets/widgetMap.ts +++ b/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' } diff --git a/src/locales/en/main.json b/src/locales/en/main.json index e4d93d2ba..de6628d77 100644 --- a/src/locales/en/main.json +++ b/src/locales/en/main.json @@ -2633,6 +2633,10 @@ "errorUploadFailed": "Failed to import asset. Please try again.", "errorUserTokenAccessDenied": "Your API token does not have access to this resource. Please check your token permissions.", "errorUserTokenInvalid": "Your stored API token is invalid or expired. Please update your token in settings.", + "invalidAsset": "Invalid Asset", + "invalidAssetDetail": "The selected asset could not be validated. Please try again.", + "invalidFilename": "Invalid Filename", + "invalidFilenameDetail": "The asset filename could not be determined. Please try again.", "failedToCreateNode": "Failed to create node. Please try again or check console for details.", "fileFormats": "File formats", "fileName": "File Name", diff --git a/src/platform/assets/utils/createAssetWidget.ts b/src/platform/assets/utils/createAssetWidget.ts new file mode 100644 index 000000000..ed48ab2b7 --- /dev/null +++ b/src/platform/assets/utils/createAssetWidget.ts @@ -0,0 +1,111 @@ +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 { useToastStore } from '@/platform/updates/common/toastStore' +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 + /** Input name for asset browser filtering (defaults to widgetName if not provided) */ + inputNameForBrowser?: 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, + inputNameForBrowser, + defaultValue, + onValueChange + } = params + + const displayLabel = defaultValue ?? t('widgets.selectModel') + const assetBrowserDialog = useAssetBrowserDialog() + + async function openModal(widget: IBaseWidget) { + const toastStore = useToastStore() + + await assetBrowserDialog.show({ + nodeType: nodeTypeForBrowser, + inputName: inputNameForBrowser ?? 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 + ) + toastStore.add({ + severity: 'error', + summary: t('assetBrowser.invalidAsset'), + detail: t('assetBrowser.invalidAssetDetail'), + life: 5000 + }) + 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 + ) + toastStore.add({ + severity: 'error', + summary: t('assetBrowser.invalidFilename'), + detail: t('assetBrowser.invalidFilenameDetail'), + life: 5000 + }) + 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) +} diff --git a/src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts b/src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts index f398f9b0b..d57f10427 100644 --- a/src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts +++ b/src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts @@ -3,18 +3,10 @@ import { ref } from 'vue' import MultiSelectWidget from '@/components/graph/widgets/MultiSelectWidget.vue' import { t } from '@/i18n' import type { LGraphNode } from '@/lib/litegraph/src/litegraph' -import { isAssetWidget, isComboWidget } 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' +import { isComboWidget } from '@/lib/litegraph/src/litegraph' +import type { IBaseWidget } 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 { useSettingStore } from '@/platform/settings/settingStore' import type { @@ -90,71 +82,20 @@ const addMultiSelectWidget = ( return widget } -const createAssetBrowserWidget = ( +function createAssetBrowserWidget( node: LGraphNode, inputSpec: ComboInputSpec, defaultValue: string | undefined -): IBaseWidget => { - const currentValue = defaultValue - const displayLabel = currentValue ?? t('widgets.selectModel') - const assetBrowserDialog = useAssetBrowserDialog() - - async function openModal(widget: IBaseWidget) { - if (!isAssetWidget(widget)) { - throw new Error(`Expected asset widget but received ${widget.type}`) +): IBaseWidget { + return createAssetWidget({ + node, + widgetName: inputSpec.name, + nodeTypeForBrowser: node.comfyClass ?? '', + defaultValue, + onValueChange: (widget, newValue, oldValue) => { + node.onWidgetChanged?.(widget.name, newValue, oldValue, widget) } - await assetBrowserDialog.show({ - nodeType: node.comfyClass || '', - inputName: inputSpec.name, - currentValue: widget.value, - 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 - node.onWidgetChanged?.( - widget.name, - validatedFilename.data, - oldValue, - widget - ) - } - }) - } - const options: IWidgetAssetOptions = { openModal } - - const widget = node.addWidget( - 'asset', - inputSpec.name, - displayLabel, - () => undefined, - options - ) - - return widget + }) } const createInputMappingWidget = (