Skip to content

Commit a2ef569

Browse files
authored
feat(ComboWidget): add ability to have mapped inputs (#6585)
## Summary 1. Add a `getOptionLabel` option to `ComboWidget` so users of it can map of custom labels to widget values. (e.g., `"My Photo" -> "my_photo_1235.png"`). 2. Utilize this ability in Cloud environment to map user uploaded filenames to their corresponding input asset. 3. Copious unit tests to make sure I didn't (AFAIK) break anything during the refactoring portion of development. 4. Bonus: Scope model browser to only show in cloud distributions until it's released elsewhere; should prevent some undesired UI behavior if a user accidentally enables the assetAPI. ## Review Focus Widget code: please double check the work there. ## Screenshots (if applicable) https://github.com/user-attachments/assets/a94b3203-c87f-4285-b692-479996859a5a ┆Issue is synchronized with this [Notion page](https://www.notion.so/PR-6585-Feat-input-mapping-2a26d73d365081faa667e49892c8d45a) by [Unito](https://www.unito.io)
1 parent 265f125 commit a2ef569

File tree

9 files changed

+1801
-103
lines changed

9 files changed

+1801
-103
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ export interface IWidgetOptions<TValues = unknown[]> {
2929
canvasOnly?: boolean
3030

3131
values?: TValues
32+
/** Optional function to format values for display (e.g., hash → human-readable name) */
33+
getOptionLabel?: (value?: string | null) => string
3234
callback?: IWidget['callback']
3335
}
3436

src/lib/litegraph/src/widgets/ComboWidget.ts

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,18 @@ export class ComboWidget
3434

3535
override get _displayValue() {
3636
if (this.computedDisabled) return ''
37+
38+
if (this.options.getOptionLabel) {
39+
try {
40+
return this.options.getOptionLabel(
41+
this.value ? String(this.value) : null
42+
)
43+
} catch (e) {
44+
console.error('Failed to map value:', e)
45+
return this.value ? String(this.value) : ''
46+
}
47+
}
48+
3749
const { values: rawValues } = this.options
3850
if (rawValues) {
3951
const values = typeof rawValues === 'function' ? rawValues() : rawValues
@@ -131,7 +143,31 @@ export class ComboWidget
131143
const values = this.getValues(node)
132144
const values_list = toArray(values)
133145

134-
// Handle center click - show dropdown menu
146+
// Use addItem to solve duplicate filename issues
147+
if (this.options.getOptionLabel) {
148+
const menuOptions = {
149+
scale: Math.max(1, canvas.ds.scale),
150+
event: e,
151+
className: 'dark',
152+
callback: (value: string) => {
153+
this.setValue(value, { e, node, canvas })
154+
}
155+
}
156+
const menu = new LiteGraph.ContextMenu([], menuOptions)
157+
158+
for (const value of values_list) {
159+
try {
160+
const label = this.options.getOptionLabel(String(value))
161+
menu.addItem(label, value, menuOptions)
162+
} catch (err) {
163+
console.error('Failed to map value:', err)
164+
menu.addItem(String(value), value, menuOptions)
165+
}
166+
}
167+
return
168+
}
169+
170+
// Show dropdown menu when user clicks on widget label
135171
const text_values = values != values_list ? Object.values(values) : values
136172
new LiteGraph.ContextMenu(text_values, {
137173
scale: Math.max(1, canvas.ds.scale),

src/platform/assets/services/assetService.ts

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -194,20 +194,31 @@ function createAssetService() {
194194
/**
195195
* Gets assets filtered by a specific tag
196196
*
197-
* @param tag - The tag to filter by (e.g., 'models')
197+
* @param tag - The tag to filter by (e.g., 'models', 'input')
198198
* @param includePublic - Whether to include public assets (default: true)
199+
* @param options - Pagination options
200+
* @param options.limit - Maximum number of assets to return (default: 500)
201+
* @param options.offset - Number of assets to skip (default: 0)
199202
* @returns Promise<AssetItem[]> - Full asset objects filtered by tag, excluding missing assets
200203
*/
201204
async function getAssetsByTag(
202205
tag: string,
203-
includePublic: boolean = true
206+
includePublic: boolean = true,
207+
{
208+
limit = DEFAULT_LIMIT,
209+
offset = 0
210+
}: { limit?: number; offset?: number } = {}
204211
): Promise<AssetItem[]> {
205212
const queryParams = new URLSearchParams({
206213
include_tags: tag,
207-
limit: DEFAULT_LIMIT.toString(),
214+
limit: limit.toString(),
208215
include_public: includePublic ? 'true' : 'false'
209216
})
210217

218+
if (offset > 0) {
219+
queryParams.set('offset', offset.toString())
220+
}
221+
211222
const data = await handleAssetRequest(
212223
`${ASSETS_ENDPOINT}?${queryParams.toString()}`,
213224
`assets for tag ${tag}`

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

Lines changed: 163 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
assetItemSchema
1212
} from '@/platform/assets/schemas/assetSchema'
1313
import { assetService } from '@/platform/assets/services/assetService'
14+
import { isCloud } from '@/platform/distribution/types'
1415
import { useSettingStore } from '@/platform/settings/settingStore'
1516
import { transformInputSpecV2ToV1 } from '@/schemas/nodeDef/migration'
1617
import { isComboInputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
@@ -22,6 +23,8 @@ import { ComponentWidgetImpl, addWidget } from '@/scripts/domWidget'
2223
import type { BaseDOMWidget } from '@/scripts/domWidget'
2324
import { addValueControlWidgets } from '@/scripts/widgets'
2425
import type { ComfyWidgetConstructorV2 } from '@/scripts/widgets'
26+
import { useAssetsStore } from '@/stores/assetsStore'
27+
import { getMediaTypeFromFilename } from '@/utils/formatUtil'
2528

2629
import { useRemoteWidget } from './useRemoteWidget'
2730

@@ -32,6 +35,20 @@ const getDefaultValue = (inputSpec: ComboInputSpec) => {
3235
return undefined
3336
}
3437

38+
// Map node types to expected media types
39+
const NODE_MEDIA_TYPE_MAP: Record<string, 'image' | 'video' | 'audio'> = {
40+
LoadImage: 'image',
41+
LoadVideo: 'video',
42+
LoadAudio: 'audio'
43+
}
44+
45+
// Map node types to placeholder i18n keys
46+
const NODE_PLACEHOLDER_MAP: Record<string, string> = {
47+
LoadImage: 'widgets.uploadSelect.placeholderImage',
48+
LoadVideo: 'widgets.uploadSelect.placeholderVideo',
49+
LoadAudio: 'widgets.uploadSelect.placeholderAudio'
50+
}
51+
3552
const addMultiSelectWidget = (
3653
node: LGraphNode,
3754
inputSpec: ComboInputSpec
@@ -55,94 +72,176 @@ const addMultiSelectWidget = (
5572
return widget
5673
}
5774

58-
const addComboWidget = (
75+
const createAssetBrowserWidget = (
5976
node: LGraphNode,
60-
inputSpec: ComboInputSpec
77+
inputSpec: ComboInputSpec,
78+
defaultValue: string | undefined
6179
): IBaseWidget => {
62-
const settingStore = useSettingStore()
63-
const isUsingAssetAPI = settingStore.get('Comfy.Assets.UseAssetAPI')
64-
const isEligible = assetService.isAssetBrowserEligible(
65-
node.comfyClass,
66-
inputSpec.name
67-
)
80+
const currentValue = defaultValue
81+
const displayLabel = currentValue ?? t('widgets.selectModel')
82+
const assetBrowserDialog = useAssetBrowserDialog()
83+
84+
const widget = node.addWidget(
85+
'asset',
86+
inputSpec.name,
87+
displayLabel,
88+
async function (this: IBaseWidget) {
89+
if (!isAssetWidget(widget)) {
90+
throw new Error(`Expected asset widget but received ${widget.type}`)
91+
}
92+
await assetBrowserDialog.show({
93+
nodeType: node.comfyClass || '',
94+
inputName: inputSpec.name,
95+
currentValue: widget.value,
96+
onAssetSelected: (asset) => {
97+
const validatedAsset = assetItemSchema.safeParse(asset)
6898

69-
if (isUsingAssetAPI && isEligible) {
70-
const currentValue = getDefaultValue(inputSpec)
71-
const displayLabel = currentValue ?? t('widgets.selectModel')
99+
if (!validatedAsset.success) {
100+
console.error(
101+
'Invalid asset item:',
102+
validatedAsset.error.errors,
103+
'Received:',
104+
asset
105+
)
106+
return
107+
}
72108

73-
const assetBrowserDialog = useAssetBrowserDialog()
109+
const filename = validatedAsset.data.user_metadata?.filename
110+
const validatedFilename = assetFilenameSchema.safeParse(filename)
74111

75-
const widget = node.addWidget(
76-
'asset',
77-
inputSpec.name,
78-
displayLabel,
79-
async function (this: IBaseWidget) {
80-
if (!isAssetWidget(widget)) {
81-
throw new Error(`Expected asset widget but received ${widget.type}`)
82-
}
83-
await assetBrowserDialog.show({
84-
nodeType: node.comfyClass || '',
85-
inputName: inputSpec.name,
86-
currentValue: widget.value,
87-
onAssetSelected: (asset) => {
88-
const validatedAsset = assetItemSchema.safeParse(asset)
89-
90-
if (!validatedAsset.success) {
91-
console.error(
92-
'Invalid asset item:',
93-
validatedAsset.error.errors,
94-
'Received:',
95-
asset
96-
)
97-
return
98-
}
99-
100-
const filename = validatedAsset.data.user_metadata?.filename
101-
const validatedFilename = assetFilenameSchema.safeParse(filename)
102-
103-
if (!validatedFilename.success) {
104-
console.error(
105-
'Invalid asset filename:',
106-
validatedFilename.error.errors,
107-
'for asset:',
108-
validatedAsset.data.id
109-
)
110-
return
111-
}
112-
113-
const oldValue = widget.value
114-
this.value = validatedFilename.data
115-
node.onWidgetChanged?.(
116-
widget.name,
117-
validatedFilename.data,
118-
oldValue,
119-
widget
112+
if (!validatedFilename.success) {
113+
console.error(
114+
'Invalid asset filename:',
115+
validatedFilename.error.errors,
116+
'for asset:',
117+
validatedAsset.data.id
120118
)
119+
return
121120
}
122-
})
121+
122+
const oldValue = widget.value
123+
this.value = validatedFilename.data
124+
node.onWidgetChanged?.(
125+
widget.name,
126+
validatedFilename.data,
127+
oldValue,
128+
widget
129+
)
130+
}
131+
})
132+
}
133+
)
134+
135+
return widget
136+
}
137+
138+
const createInputMappingWidget = (
139+
node: LGraphNode,
140+
inputSpec: ComboInputSpec,
141+
defaultValue: string | undefined
142+
): IBaseWidget => {
143+
const assetsStore = useAssetsStore()
144+
145+
const widget = node.addWidget(
146+
'combo',
147+
inputSpec.name,
148+
defaultValue ?? '',
149+
() => {},
150+
{
151+
values: [],
152+
getOptionLabel: (value?: string | null) => {
153+
if (!value) {
154+
const placeholderKey =
155+
NODE_PLACEHOLDER_MAP[node.comfyClass ?? ''] ??
156+
'widgets.uploadSelect.placeholder'
157+
return t(placeholderKey)
158+
}
159+
return assetsStore.getInputName(value)
123160
}
124-
)
161+
}
162+
)
125163

126-
return widget
164+
if (assetsStore.inputAssets.length === 0 && !assetsStore.inputLoading) {
165+
void assetsStore.updateInputs().then(() => {
166+
// edge for users using nodes with 0 prior inputs
167+
// force canvas refresh the first time they add an asset
168+
// so they see filenames instead of hashes.
169+
node.setDirtyCanvas(true, false)
170+
})
127171
}
128172

129-
// Create normal combo widget
173+
const origOptions = widget.options
174+
widget.options = new Proxy(origOptions, {
175+
get(target, prop) {
176+
if (prop !== 'values') {
177+
return target[prop as keyof typeof target]
178+
}
179+
return assetsStore.inputAssets
180+
.filter(
181+
(asset) =>
182+
getMediaTypeFromFilename(asset.name) ===
183+
NODE_MEDIA_TYPE_MAP[node.comfyClass ?? '']
184+
)
185+
.map((asset) => asset.asset_hash)
186+
.filter((hash): hash is string => !!hash)
187+
}
188+
})
189+
190+
if (inputSpec.control_after_generate) {
191+
if (!isComboWidget(widget)) {
192+
throw new Error(`Expected combo widget but received ${widget.type}`)
193+
}
194+
widget.linkedWidgets = addValueControlWidgets(
195+
node,
196+
widget,
197+
undefined,
198+
undefined,
199+
transformInputSpecV2ToV1(inputSpec)
200+
)
201+
}
202+
203+
return widget
204+
}
205+
206+
const addComboWidget = (
207+
node: LGraphNode,
208+
inputSpec: ComboInputSpec
209+
): IBaseWidget => {
130210
const defaultValue = getDefaultValue(inputSpec)
131-
const comboOptions = inputSpec.options ?? []
211+
212+
if (isCloud) {
213+
const settingStore = useSettingStore()
214+
const isUsingAssetAPI = settingStore.get('Comfy.Assets.UseAssetAPI')
215+
const isEligible = assetService.isAssetBrowserEligible(
216+
node.comfyClass,
217+
inputSpec.name
218+
)
219+
220+
if (isUsingAssetAPI && isEligible) {
221+
return createAssetBrowserWidget(node, inputSpec, defaultValue)
222+
}
223+
224+
if (NODE_MEDIA_TYPE_MAP[node.comfyClass ?? '']) {
225+
return createInputMappingWidget(node, inputSpec, defaultValue)
226+
}
227+
}
228+
229+
// Standard combo widget
132230
const widget = node.addWidget(
133231
'combo',
134232
inputSpec.name,
135233
defaultValue,
136234
() => {},
137235
{
138-
values: comboOptions
236+
values: inputSpec.options ?? []
139237
}
140238
)
141239

142240
if (inputSpec.remote) {
143241
if (!isComboWidget(widget)) {
144242
throw new Error(`Expected combo widget but received ${widget.type}`)
145243
}
244+
146245
const remoteWidget = useRemoteWidget({
147246
remoteConfig: inputSpec.remote,
148247
defaultValue,
@@ -166,6 +265,7 @@ const addComboWidget = (
166265
if (!isComboWidget(widget)) {
167266
throw new Error(`Expected combo widget but received ${widget.type}`)
168267
}
268+
169269
widget.linkedWidgets = addValueControlWidgets(
170270
node,
171271
widget,

0 commit comments

Comments
 (0)