[feat] Implement AssetService behind settings flag (#5404)

* [feat] add Comfy.Assets.UseAssetAPI to CORE_SETTINGS

* [feat] create AssetService

1. Add service for accessing new Asset API
2. Add fallback model paths logic so empty model directories appear for
   the user.
3. Copious tests for them all.

Co-Authored-By: Claude <noreply@anthropic.com>

* [feat] switch between assets and file paths for model data

* [feat] ignore assets with "missing" tag

* [fix] formatting and style

* [fix] call assets API with the correct filters

* [feat] elminate unused modelPath code

* [fix] remove stray comment

* [fix] model manager api was not parsed correctly

---------

Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Arjan Singh
2025-09-08 11:49:55 -07:00
committed by snomiao
parent 94ad3a1df1
commit 3c5ec3e63f
8 changed files with 455 additions and 17 deletions

View File

@@ -2,7 +2,9 @@ import { createPinia, setActivePinia } from 'pinia'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { api } from '@/scripts/api'
import { assetService } from '@/services/assetService'
import { useModelStore } from '@/stores/modelStore'
import { useSettingStore } from '@/stores/settingStore'
// Mock the api
vi.mock('@/scripts/api', () => ({
@@ -13,7 +15,32 @@ vi.mock('@/scripts/api', () => ({
}
}))
function enableMocks() {
// Mock the assetService
vi.mock('@/services/assetService', () => ({
assetService: {
getAssetModelFolders: vi.fn(),
getAssetModels: vi.fn()
}
}))
// Mock the settingStore
vi.mock('@/stores/settingStore', () => ({
useSettingStore: vi.fn()
}))
function enableMocks(useAssetAPI = false) {
// Mock settingStore to return the useAssetAPI setting
const mockSettingStore = {
get: vi.fn().mockImplementation((key: string) => {
if (key === 'Comfy.Assets.UseAssetAPI') {
return useAssetAPI
}
return false
})
}
vi.mocked(useSettingStore).mockReturnValue(mockSettingStore as any)
// Mock experimental API - returns objects with name and folders properties
vi.mocked(api.getModels).mockResolvedValue([
{ name: 'sdxl.safetensors', pathIndex: 0 },
{ name: 'sdv15.safetensors', pathIndex: 0 },
@@ -23,6 +50,18 @@ function enableMocks() {
{ name: 'checkpoints', folders: ['/path/to/checkpoints'] },
{ name: 'vae', folders: ['/path/to/vae'] }
])
// Mock asset API - also returns objects with name and folders properties
vi.mocked(assetService.getAssetModelFolders).mockResolvedValue([
{ name: 'checkpoints', folders: ['/path/to/checkpoints'] },
{ name: 'vae', folders: ['/path/to/vae'] }
])
vi.mocked(assetService.getAssetModels).mockResolvedValue([
{ name: 'sdxl.safetensors', pathIndex: 0 },
{ name: 'sdv15.safetensors', pathIndex: 0 },
{ name: 'noinfo.safetensors', pathIndex: 0 }
])
vi.mocked(api.viewMetadata).mockImplementation((_, model) => {
if (model === 'noinfo.safetensors') {
return Promise.resolve({})
@@ -46,26 +85,25 @@ describe('useModelStore', () => {
beforeEach(async () => {
setActivePinia(createPinia())
store = useModelStore()
vi.resetAllMocks()
})
it('should load models', async () => {
enableMocks()
store = useModelStore()
await store.loadModelFolders()
const folderStore = await store.getLoadedModelFolder('checkpoints')
expect(folderStore).not.toBeNull()
if (!folderStore) return
expect(Object.keys(folderStore.models).length).toBe(3)
expect(folderStore).toBeDefined()
expect(Object.keys(folderStore!.models)).toHaveLength(3)
})
it('should load model metadata', async () => {
enableMocks()
store = useModelStore()
await store.loadModelFolders()
const folderStore = await store.getLoadedModelFolder('checkpoints')
expect(folderStore).not.toBeNull()
if (!folderStore) return
const model = folderStore.models['0/sdxl.safetensors']
expect(folderStore).toBeDefined()
const model = folderStore!.models['0/sdxl.safetensors']
await model.load()
expect(model.title).toBe('Title of sdxl.safetensors')
expect(model.architecture_id).toBe('stable-diffusion-xl-base-v1')
@@ -79,11 +117,11 @@ describe('useModelStore', () => {
it('should handle no metadata', async () => {
enableMocks()
store = useModelStore()
await store.loadModelFolders()
const folderStore = await store.getLoadedModelFolder('checkpoints')
expect(folderStore).not.toBeNull()
if (!folderStore) return
const model = folderStore.models['0/noinfo.safetensors']
expect(folderStore).toBeDefined()
const model = folderStore!.models['0/noinfo.safetensors']
await model.load()
expect(model.file_name).toBe('noinfo.safetensors')
expect(model.title).toBe('noinfo')
@@ -95,6 +133,7 @@ describe('useModelStore', () => {
it('should cache model information', async () => {
enableMocks()
store = useModelStore()
await store.loadModelFolders()
expect(api.getModels).toHaveBeenCalledTimes(0)
await store.getLoadedModelFolder('checkpoints')
@@ -102,4 +141,36 @@ describe('useModelStore', () => {
await store.getLoadedModelFolder('checkpoints')
expect(api.getModels).toHaveBeenCalledTimes(1)
})
describe('API switching functionality', () => {
it('should use experimental API for complete workflow when UseAssetAPI setting is false', async () => {
enableMocks(false) // useAssetAPI = false
store = useModelStore()
await store.loadModelFolders()
const folderStore = await store.getLoadedModelFolder('checkpoints')
// Both APIs return objects with .name property, modelStore extracts folder.name in both cases
expect(api.getModelFolders).toHaveBeenCalledTimes(1)
expect(api.getModels).toHaveBeenCalledWith('checkpoints')
expect(assetService.getAssetModelFolders).toHaveBeenCalledTimes(0)
expect(assetService.getAssetModels).toHaveBeenCalledTimes(0)
expect(folderStore).toBeDefined()
expect(Object.keys(folderStore!.models)).toHaveLength(3)
})
it('should use asset API for complete workflow when UseAssetAPI setting is true', async () => {
enableMocks(true) // useAssetAPI = true
store = useModelStore()
await store.loadModelFolders()
const folderStore = await store.getLoadedModelFolder('checkpoints')
// Both APIs return objects with .name property, modelStore extracts folder.name in both cases
expect(assetService.getAssetModelFolders).toHaveBeenCalledTimes(1)
expect(assetService.getAssetModels).toHaveBeenCalledWith('checkpoints')
expect(api.getModelFolders).toHaveBeenCalledTimes(0)
expect(api.getModels).toHaveBeenCalledTimes(0)
expect(folderStore).toBeDefined()
expect(Object.keys(folderStore!.models)).toHaveLength(3)
})
})
})