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)
This commit is contained in:
Arjan Singh
2025-11-05 11:33:00 -08:00
committed by GitHub
parent 265f1257e7
commit a2ef569b9c
9 changed files with 1804 additions and 106 deletions

View File

@@ -4,13 +4,61 @@ import { LGraphNode } from '@/lib/litegraph/src/litegraph'
import type { IBaseWidget } from '@/lib/litegraph/src/types/widgets'
import { useAssetBrowserDialog } from '@/platform/assets/composables/useAssetBrowserDialog'
import { assetService } from '@/platform/assets/services/assetService'
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
import { useComboWidget } from '@/renderer/extensions/vueNodes/widgets/composables/useComboWidget'
import type { InputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
// Mock factory using actual type
function createMockAssetItem(overrides: Partial<AssetItem> = {}): AssetItem {
return {
id: 'test-asset-id',
name: 'test-image.png',
asset_hash: 'hash123',
size: 1024,
mime_type: 'image/png',
tags: ['input'],
created_at: new Date().toISOString(),
updated_at: new Date().toISOString(),
last_access_time: new Date().toISOString(),
...overrides
}
}
// Use vi.hoisted() to ensure mock state is initialized before mocks
const mockDistributionState = vi.hoisted(() => ({ isCloud: false }))
const mockUpdateInputs = vi.hoisted(() => vi.fn(() => Promise.resolve()))
const mockGetInputName = vi.hoisted(() => vi.fn((hash: string) => hash))
const mockAssetsStoreState = vi.hoisted(() => {
const inputAssets: AssetItem[] = []
return {
inputAssets,
inputLoading: false
}
})
vi.mock('@/scripts/widgets', () => ({
addValueControlWidgets: vi.fn()
}))
vi.mock('@/platform/distribution/types', () => ({
get isCloud() {
return mockDistributionState.isCloud
}
}))
vi.mock('@/stores/assetsStore', () => ({
useAssetsStore: vi.fn(() => ({
get inputAssets() {
return mockAssetsStoreState.inputAssets
},
get inputLoading() {
return mockAssetsStoreState.inputLoading
},
updateInputs: mockUpdateInputs,
getInputName: mockGetInputName
}))
}))
const mockSettingStoreGet = vi.fn(() => false)
vi.mock('@/platform/settings/settingStore', () => ({
useSettingStore: vi.fn(() => ({
@@ -42,7 +90,7 @@ vi.mock('@/platform/assets/composables/useAssetBrowserDialog', () => {
// Test factory functions
function createMockWidget(overrides: Partial<IBaseWidget> = {}): IBaseWidget {
const mockCallback = vi.fn()
return {
const widget: IBaseWidget = {
type: 'combo',
options: {},
name: 'testWidget',
@@ -50,7 +98,8 @@ function createMockWidget(overrides: Partial<IBaseWidget> = {}): IBaseWidget {
callback: mockCallback,
y: 0,
...overrides
} as IBaseWidget
}
return widget
}
function createMockNode(comfyClass = 'TestNode'): LGraphNode {
@@ -73,11 +122,12 @@ function createMockNode(comfyClass = 'TestNode'): LGraphNode {
}
function createMockInputSpec(overrides: Partial<InputSpec> = {}): InputSpec {
return {
const inputSpec: InputSpec = {
type: 'COMBO',
name: 'testInput',
...overrides
} as InputSpec
}
return inputSpec
}
describe('useComboWidget', () => {
@@ -86,6 +136,10 @@ describe('useComboWidget', () => {
mockSettingStoreGet.mockReturnValue(false)
vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(false)
vi.mocked(useAssetBrowserDialog).mockClear()
mockDistributionState.isCloud = false
mockAssetsStoreState.inputAssets = []
mockAssetsStoreState.inputLoading = false
mockUpdateInputs.mockClear()
})
it('should handle undefined spec', () => {
@@ -110,6 +164,7 @@ describe('useComboWidget', () => {
})
it('should create normal combo widget when asset API is disabled', () => {
mockDistributionState.isCloud = true
mockSettingStoreGet.mockReturnValue(false) // Asset API disabled
vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(true) // Widget is eligible
@@ -137,6 +192,7 @@ describe('useComboWidget', () => {
})
it('should create asset browser widget when API enabled', () => {
mockDistributionState.isCloud = true
mockSettingStoreGet.mockReturnValue(true)
vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(true)
@@ -169,36 +225,8 @@ describe('useComboWidget', () => {
expect(widget).toBe(mockWidget)
})
it('should create asset browser widget with options when API enabled', () => {
mockSettingStoreGet.mockReturnValue(true)
vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(true)
const constructor = useComboWidget()
const mockWidget = createMockWidget({
type: 'asset',
name: 'ckpt_name',
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(widget).toBe(mockWidget)
})
it('should use asset browser widget even when inputSpec has a default value but no options', () => {
it('should create asset browser widget when default value provided without options', () => {
mockDistributionState.isCloud = true
mockSettingStoreGet.mockReturnValue(true)
vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(true)
@@ -229,6 +257,7 @@ describe('useComboWidget', () => {
})
it('should show Select model when asset widget has undefined current value', () => {
mockDistributionState.isCloud = true
mockSettingStoreGet.mockReturnValue(true)
vi.mocked(assetService.isAssetBrowserEligible).mockReturnValue(true)
@@ -256,4 +285,203 @@ describe('useComboWidget', () => {
expect(mockSettingStoreGet).toHaveBeenCalledWith('Comfy.Assets.UseAssetAPI')
expect(widget).toBe(mockWidget)
})
describe('cloud input asset mapping', () => {
const HASH_FILENAME =
'72e786ff2a44d682c4294db0b7098e569832bc394efc6dad644e6ec85a78efb7.png'
const HASH_FILENAME_2 =
'a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456.jpg'
it.each([
{ nodeClass: 'LoadImage', inputName: 'image' },
{ nodeClass: 'LoadVideo', inputName: 'video' },
{ nodeClass: 'LoadAudio', inputName: 'audio' }
])(
'should create combo widget with getOptionLabel for $nodeClass in cloud',
({ nodeClass, inputName }) => {
mockDistributionState.isCloud = true
const constructor = useComboWidget()
const mockWidget = createMockWidget({
type: 'combo',
name: inputName,
value: HASH_FILENAME
})
const mockNode = createMockNode(nodeClass)
vi.mocked(mockNode.addWidget).mockReturnValue(mockWidget)
const inputSpec = createMockInputSpec({
name: inputName,
options: [HASH_FILENAME, HASH_FILENAME_2]
})
const widget = constructor(mockNode, inputSpec)
expect(mockNode.addWidget).toHaveBeenCalledWith(
'combo',
inputName,
HASH_FILENAME,
expect.any(Function),
expect.objectContaining({
values: [], // Empty initially, populated dynamically by Proxy
getOptionLabel: expect.any(Function)
})
)
expect(widget).toBe(mockWidget)
}
)
it("should format option labels using store's getInputName function", () => {
mockDistributionState.isCloud = true
mockGetInputName.mockReturnValue('Beautiful Sunset.png')
const constructor = useComboWidget()
const mockWidget = createMockWidget({
type: 'combo',
name: 'image',
value: HASH_FILENAME
})
const mockNode = createMockNode('LoadImage')
vi.mocked(mockNode.addWidget).mockReturnValue(mockWidget)
const inputSpec = createMockInputSpec({
name: 'image',
options: [HASH_FILENAME]
})
constructor(mockNode, inputSpec)
// Extract the injected getOptionLabel function with type narrowing
const addWidgetCall = vi.mocked(mockNode.addWidget).mock.calls[0]
const options = addWidgetCall[4]
if (typeof options !== 'object' || !options) {
throw new Error('Expected options to be an object')
}
if (!('getOptionLabel' in options)) {
throw new Error('Expected options to have getOptionLabel property')
}
if (typeof options.getOptionLabel !== 'function') {
throw new Error('Expected getOptionLabel to be a function')
}
// Test that the injected function calls getInputName
const result = options.getOptionLabel(HASH_FILENAME)
expect(mockGetInputName).toHaveBeenCalledWith(HASH_FILENAME)
expect(result).toBe('Beautiful Sunset.png')
})
it('should create normal combo widget for non-input nodes in cloud', () => {
mockDistributionState.isCloud = true
const constructor = useComboWidget()
const mockWidget = createMockWidget()
const mockNode = createMockNode('SomeOtherNode')
vi.mocked(mockNode.addWidget).mockReturnValue(mockWidget)
const inputSpec = createMockInputSpec({
name: 'option',
options: [HASH_FILENAME, HASH_FILENAME_2]
})
const widget = constructor(mockNode, inputSpec)
expect(mockNode.addWidget).toHaveBeenCalledWith(
'combo',
'option',
HASH_FILENAME,
expect.any(Function),
{ values: [HASH_FILENAME, HASH_FILENAME_2] }
)
expect(widget).toBe(mockWidget)
})
it('should create normal combo widget for LoadImage in OSS', () => {
mockDistributionState.isCloud = false
const constructor = useComboWidget()
const mockWidget = createMockWidget()
const mockNode = createMockNode('LoadImage')
vi.mocked(mockNode.addWidget).mockReturnValue(mockWidget)
const inputSpec = createMockInputSpec({
name: 'image',
options: [HASH_FILENAME, HASH_FILENAME_2]
})
const widget = constructor(mockNode, inputSpec)
expect(mockNode.addWidget).toHaveBeenCalledWith(
'combo',
'image',
HASH_FILENAME,
expect.any(Function),
{
values: [HASH_FILENAME, HASH_FILENAME_2]
}
)
expect(widget).toBe(mockWidget)
})
it('should trigger lazy load for cloud input nodes', () => {
mockDistributionState.isCloud = true
mockAssetsStoreState.inputAssets = []
mockAssetsStoreState.inputLoading = false
const constructor = useComboWidget()
const mockWidget = createMockWidget({ type: 'combo' })
const mockNode = createMockNode('LoadImage')
vi.mocked(mockNode.addWidget).mockReturnValue(mockWidget)
const inputSpec = createMockInputSpec({
name: 'image',
options: [HASH_FILENAME]
})
constructor(mockNode, inputSpec)
expect(mockUpdateInputs).toHaveBeenCalledTimes(1)
})
it('should not trigger lazy load if assets already loading', () => {
mockDistributionState.isCloud = true
mockAssetsStoreState.inputAssets = []
mockAssetsStoreState.inputLoading = true
const constructor = useComboWidget()
const mockWidget = createMockWidget({ type: 'combo' })
const mockNode = createMockNode('LoadImage')
vi.mocked(mockNode.addWidget).mockReturnValue(mockWidget)
const inputSpec = createMockInputSpec({
name: 'image',
options: [HASH_FILENAME]
})
constructor(mockNode, inputSpec)
expect(mockUpdateInputs).not.toHaveBeenCalled()
})
it('should not trigger lazy load if assets already loaded', () => {
mockDistributionState.isCloud = true
mockAssetsStoreState.inputAssets = [
createMockAssetItem({
id: 'asset-123',
name: 'image1.png',
asset_hash: HASH_FILENAME
})
]
mockAssetsStoreState.inputLoading = false
const constructor = useComboWidget()
const mockWidget = createMockWidget({ type: 'combo' })
const mockNode = createMockNode('LoadImage')
vi.mocked(mockNode.addWidget).mockReturnValue(mockWidget)
const inputSpec = createMockInputSpec({
name: 'image',
options: [HASH_FILENAME]
})
constructor(mockNode, inputSpec)
expect(mockUpdateInputs).not.toHaveBeenCalled()
})
})
})