mirror of
https://github.com/Comfy-Org/ComfyUI_frontend.git
synced 2026-07-16 08:49:09 +00:00
Compare commits
49 Commits
fix/deflak
...
synap5e/as
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
84491c3c60 | ||
|
|
e8f5617da7 | ||
|
|
65f80ecee7 | ||
|
|
f6bbaf4b9e | ||
|
|
1e7d48623a | ||
|
|
c8afd4cf8a | ||
|
|
abfb89b990 | ||
|
|
259fd9a62f | ||
|
|
724e31d235 | ||
|
|
9a7793ccc8 | ||
|
|
317cc81196 | ||
|
|
64740d9d4a | ||
|
|
9788ed2439 | ||
|
|
7b841548ce | ||
|
|
fcb7d838ef | ||
|
|
873a85e59e | ||
|
|
ca33752569 | ||
|
|
e6916bb665 | ||
|
|
fdc4651934 | ||
|
|
62123c4c0d | ||
|
|
1cb04bef92 | ||
|
|
f5e221b955 | ||
|
|
fd1f2726a5 | ||
|
|
8d51d933fc | ||
|
|
be564b232b | ||
|
|
b62405a6ca | ||
|
|
33c1806673 | ||
|
|
9342caad0a | ||
|
|
48aed9e0d9 | ||
|
|
dd3dceeaca | ||
|
|
3605fc1d75 | ||
|
|
f97195e392 | ||
|
|
0488d0d8a4 | ||
|
|
aea28a06de | ||
|
|
5847b3d148 | ||
|
|
187239712c | ||
|
|
e53b7b4ba8 | ||
|
|
ab56cf9e82 | ||
|
|
f772fc0123 | ||
|
|
a5f4559df8 | ||
|
|
6103dd164e | ||
|
|
f7d672e3a5 | ||
|
|
6f0eaefe1b | ||
|
|
55ceec0a16 | ||
|
|
5ac1b15266 | ||
|
|
df826415ca | ||
|
|
ec25e874a2 | ||
|
|
c8e4029f96 | ||
|
|
5ebeb580ca |
@@ -418,28 +418,26 @@ export class AssetsSidebarTab extends SidebarTab {
|
||||
async openSettingsMenu() {
|
||||
await this.dismissToasts()
|
||||
await this.settingsButton.click()
|
||||
// Wait for the popover content to render. Use the default timeout so slower
|
||||
// (e.g. cloud) app inits don't burst-fail a tight explicit timeout.
|
||||
await expect(
|
||||
this.listViewOption.or(this.gridViewOption).first()
|
||||
).toBeVisible()
|
||||
// Wait for popover content to render
|
||||
await this.listViewOption
|
||||
.or(this.gridViewOption)
|
||||
.first()
|
||||
.waitFor({ state: 'visible', timeout: 3000 })
|
||||
}
|
||||
|
||||
async openFilterMenu() {
|
||||
await this.dismissToasts()
|
||||
await this.filterButton.click()
|
||||
// Wait for the filter popover to open. Use the default timeout so slower
|
||||
// (e.g. cloud) app inits don't burst-fail a tight explicit timeout.
|
||||
await expect(this.filterCheckbox('Image')).toBeVisible()
|
||||
await this.filterCheckbox('Image').waitFor({
|
||||
state: 'visible',
|
||||
timeout: 3000
|
||||
})
|
||||
}
|
||||
|
||||
async toggleMediaTypeFilter(
|
||||
filter: MediaFilterKind | MediaFilterLabel
|
||||
): Promise<void> {
|
||||
const checkbox = this.filterCheckbox(filter)
|
||||
// Ensure the popover has finished opening before reading its state; a stale
|
||||
// read here races the reka-ui slide/fade animation.
|
||||
await expect(checkbox).toBeVisible()
|
||||
const before = await checkbox.getAttribute('aria-checked')
|
||||
await checkbox.click()
|
||||
const expected = before === 'true' ? 'false' : 'true'
|
||||
|
||||
@@ -422,17 +422,9 @@ test.describe('Vue Node Moving', { tag: '@vue-nodes' }, () => {
|
||||
loadCheckpointHeaderPos
|
||||
)
|
||||
|
||||
// Poll the header position so the assertion retries until the touch pan
|
||||
// has settled, instead of reading a single mid-animation bounding box.
|
||||
// A screen bounding box read is pixel-quantized, so assert to the nearest
|
||||
// pixel (precision 0 == within 0.5px) rather than the default precision 2
|
||||
// (within 0.005px), which sub-pixel canvas rounding cannot satisfy.
|
||||
await expect
|
||||
.poll(() => getHeaderPos(comfyPage, 'Load Checkpoint').then((p) => p.x))
|
||||
.toBeCloseTo(loadCheckpointHeaderPos.x + 64, 0)
|
||||
await expect
|
||||
.poll(() => getHeaderPos(comfyPage, 'Load Checkpoint').then((p) => p.y))
|
||||
.toBeCloseTo(loadCheckpointHeaderPos.y + 64, 0)
|
||||
const newHeaderPos = await getLoadCheckpointHeaderPos(comfyPage)
|
||||
expect(newHeaderPos.x).toBeCloseTo(loadCheckpointHeaderPos.x + 64)
|
||||
expect(newHeaderPos.y).toBeCloseTo(loadCheckpointHeaderPos.y + 64)
|
||||
}
|
||||
)
|
||||
})
|
||||
|
||||
@@ -426,7 +426,7 @@ describe('shouldPreventRekaDismiss', () => {
|
||||
expect(event.defaultPrevented).toBe(true)
|
||||
})
|
||||
|
||||
it.for(['p-dialog', 'p-select-overlay'])(
|
||||
it.for(['p-dialog', 'p-select-overlay', 'p-toast'])(
|
||||
'focus-outside on a sibling %s portal does not dismiss the parent',
|
||||
(className) => {
|
||||
const overlay = document.createElement('div')
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
// PrimeVue overlays (Select, ColorPicker, Popover, Autocomplete, stacked
|
||||
// PrimeVue Dialogs) teleport to body. Reka treats clicks on body-portaled
|
||||
// elements as outside its dialog and would auto-dismiss on the first
|
||||
// interaction, tearing the overlay down mid-interaction. Treat any
|
||||
// PrimeVue overlay click as inside.
|
||||
// PrimeVue Dialogs, Toasts) teleport to body. Reka treats clicks on
|
||||
// body-portaled elements as outside its dialog and would auto-dismiss on the
|
||||
// first interaction, tearing the overlay down mid-interaction. Treat any
|
||||
// PrimeVue overlay click as inside. Toasts matter for focus-outside: when a
|
||||
// button disables itself mid-action (e.g. a confirm entering its loading
|
||||
// state), the browser drops focus and recovery can land on the toast's close
|
||||
// button, which must not dismiss the dialog underneath.
|
||||
const PRIMEVUE_OVERLAY_SELECTORS =
|
||||
'.p-select-overlay, .p-colorpicker-panel, .p-popover, .p-autocomplete-overlay, .p-overlay, .p-overlay-mask, .p-dialog'
|
||||
'.p-select-overlay, .p-colorpicker-panel, .p-popover, .p-autocomplete-overlay, .p-overlay, .p-overlay-mask, .p-dialog, .p-toast'
|
||||
|
||||
// Reka portals its own dialogs / popovers / menus into the body too. When a
|
||||
// nested Reka layer opens on top of a non-modal parent, the parent's
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import { createTestingPinia } from '@pinia/testing'
|
||||
import { fromPartial } from '@total-typescript/shoehorn'
|
||||
import userEvent from '@testing-library/user-event'
|
||||
import { render, screen } from '@testing-library/vue'
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
import { nextTick } from 'vue'
|
||||
@@ -18,7 +19,10 @@ const {
|
||||
mockGetNodeProvider,
|
||||
mockToggleNodeOnEvent,
|
||||
mockRefreshModelFolder,
|
||||
downloadStoreState
|
||||
mockLoadModels,
|
||||
downloadStoreState,
|
||||
settingState,
|
||||
modelsState
|
||||
} = vi.hoisted(() => {
|
||||
let capturedRoot: TreeExplorerNode<unknown> | null = null
|
||||
return {
|
||||
@@ -33,7 +37,13 @@ const {
|
||||
mockGetNodeProvider: vi.fn(),
|
||||
mockToggleNodeOnEvent: vi.fn(),
|
||||
mockRefreshModelFolder: vi.fn().mockResolvedValue(undefined),
|
||||
downloadStoreState: { setLastCompleted: (_: unknown) => {} }
|
||||
mockLoadModels: vi.fn().mockResolvedValue([]),
|
||||
downloadStoreState: { setLastCompleted: (_: unknown) => {} },
|
||||
settingState: { useAssetAPI: false, autoLoadAll: false },
|
||||
modelsState: {
|
||||
push: (_: unknown) => {},
|
||||
reset: () => {}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
@@ -54,20 +64,30 @@ const mockModel = fromPartial<ComfyModelDef>({
|
||||
searchable: 'checkpoints/model.safetensors'
|
||||
})
|
||||
|
||||
vi.mock('@/stores/modelStore', () => ({
|
||||
ResourceState: {
|
||||
Loading: 'loading',
|
||||
Loaded: 'loaded'
|
||||
},
|
||||
useModelStore: () => ({
|
||||
modelFolders: [],
|
||||
models: [mockModel],
|
||||
loadModels: vi.fn().mockResolvedValue([]),
|
||||
loadModelFolders: vi.fn().mockResolvedValue([]),
|
||||
refresh: vi.fn().mockResolvedValue(undefined),
|
||||
refreshModelFolder: mockRefreshModelFolder
|
||||
})
|
||||
}))
|
||||
vi.mock('@/stores/modelStore', async () => {
|
||||
const { reactive } = await import('vue')
|
||||
const models = reactive<ComfyModelDef[]>([])
|
||||
modelsState.push = (model: unknown) => {
|
||||
models.push(model as ComfyModelDef)
|
||||
}
|
||||
modelsState.reset = () => {
|
||||
models.splice(0, models.length, mockModel)
|
||||
}
|
||||
return {
|
||||
ResourceState: {
|
||||
Loading: 'loading',
|
||||
Loaded: 'loaded'
|
||||
},
|
||||
useModelStore: () => ({
|
||||
modelFolders: [],
|
||||
models,
|
||||
loadModels: mockLoadModels,
|
||||
loadModelFolders: vi.fn().mockResolvedValue([]),
|
||||
refresh: vi.fn().mockResolvedValue(undefined),
|
||||
refreshModelFolder: mockRefreshModelFolder
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
vi.mock('@/stores/assetDownloadStore', async () => {
|
||||
const { ref } = await import('vue')
|
||||
@@ -92,6 +112,10 @@ vi.mock('@/platform/settings/settingStore', () => ({
|
||||
useSettingStore: () => ({
|
||||
get: vi.fn((key: string) => {
|
||||
if (key === 'Comfy.ModelLibrary.NameFormat') return 'filename'
|
||||
if (key === 'Comfy.Assets.UseAssetAPI') return settingState.useAssetAPI
|
||||
if (key === 'Comfy.ModelLibrary.AutoLoadAll') {
|
||||
return settingState.autoLoadAll
|
||||
}
|
||||
return false
|
||||
})
|
||||
})
|
||||
@@ -104,26 +128,45 @@ vi.mock('@/composables/useTreeExpansion', () => ({
|
||||
})
|
||||
}))
|
||||
|
||||
vi.mock('@/components/common/TreeExplorer.vue', () => ({
|
||||
default: {
|
||||
name: 'TreeExplorer',
|
||||
template: '<div data-testid="tree-explorer" />',
|
||||
props: ['root', 'expandedKeys'],
|
||||
setup(props: { root: TreeExplorerNode<unknown> }) {
|
||||
captureRoot(props.root)
|
||||
vi.mock('@/components/common/TreeExplorer.vue', async () => {
|
||||
const { watchEffect } = await import('vue')
|
||||
return {
|
||||
default: {
|
||||
name: 'TreeExplorer',
|
||||
template: '<div data-testid="tree-explorer" />',
|
||||
props: ['root', 'expandedKeys'],
|
||||
setup(props: { root: TreeExplorerNode<unknown> }) {
|
||||
watchEffect(() => captureRoot(props.root))
|
||||
}
|
||||
}
|
||||
}
|
||||
}))
|
||||
})
|
||||
|
||||
vi.mock('@/components/ui/search-input/SearchInput.vue', () => ({
|
||||
default: {
|
||||
name: 'SearchInput',
|
||||
template: '<input data-testid="search-input" />',
|
||||
template: '<input data-testid="search-input" @input="onInput" />',
|
||||
props: ['modelValue', 'placeholder'],
|
||||
setup() {
|
||||
return { focus: vi.fn() }
|
||||
},
|
||||
expose: ['focus']
|
||||
emits: ['update:modelValue', 'search'],
|
||||
setup(
|
||||
_props: unknown,
|
||||
{
|
||||
emit,
|
||||
expose
|
||||
}: {
|
||||
emit: (event: 'update:modelValue' | 'search', value: string) => void
|
||||
expose: (exposed: Record<string, unknown>) => void
|
||||
}
|
||||
) {
|
||||
expose({ focus: vi.fn() })
|
||||
return {
|
||||
onInput: (event: Event) => {
|
||||
const value = (event.target as HTMLInputElement).value
|
||||
emit('update:modelValue', value)
|
||||
emit('search', value)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}))
|
||||
|
||||
@@ -134,7 +177,8 @@ vi.mock('./SidebarTopArea.vue', () => ({
|
||||
vi.mock('./SidebarTabTemplate.vue', () => ({
|
||||
default: {
|
||||
name: 'SidebarTabTemplate',
|
||||
template: '<div><slot name="header" /><slot name="body" /></div>'
|
||||
template:
|
||||
'<div><slot name="tool-buttons" /><slot name="header" /><slot name="body" /></div>'
|
||||
}
|
||||
}))
|
||||
|
||||
@@ -157,13 +201,17 @@ describe('ModelLibrarySidebarTab', () => {
|
||||
vi.clearAllMocks()
|
||||
resetRoot()
|
||||
downloadStoreState.setLastCompleted(null)
|
||||
settingState.useAssetAPI = false
|
||||
settingState.autoLoadAll = false
|
||||
modelsState.reset()
|
||||
})
|
||||
|
||||
function renderComponent() {
|
||||
return render(ModelLibrarySidebarTab, {
|
||||
global: {
|
||||
plugins: [createTestingPinia({ stubActions: false }), i18n],
|
||||
stubs: { teleport: true }
|
||||
stubs: { teleport: true },
|
||||
directives: { tooltip: {} }
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -236,4 +284,67 @@ describe('ModelLibrarySidebarTab', () => {
|
||||
|
||||
expect(mockRefreshModelFolder).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
describe('search', () => {
|
||||
it('updates active search results when a reload adds a matching model', async () => {
|
||||
const user = userEvent.setup()
|
||||
renderComponent()
|
||||
await nextTick()
|
||||
|
||||
await user.type(screen.getByTestId('search-input'), 'model')
|
||||
await nextTick()
|
||||
|
||||
expect(mockLoadModels).toHaveBeenCalled()
|
||||
const leafLabels = () => {
|
||||
const { children: folders = [] } = getRoot()
|
||||
return folders.flatMap(({ children: leaves = [] }) =>
|
||||
leaves.map((leaf) => leaf.label)
|
||||
)
|
||||
}
|
||||
expect(leafLabels()).toEqual(['model'])
|
||||
|
||||
// A completed scan reloads the store while the search is still active.
|
||||
modelsState.push(
|
||||
fromPartial<ComfyModelDef>({
|
||||
key: 'checkpoints/model-new.safetensors',
|
||||
file_name: 'model-new.safetensors',
|
||||
simplified_file_name: 'model-new',
|
||||
title: 'Model New',
|
||||
directory: 'checkpoints',
|
||||
searchable: 'checkpoints/model-new.safetensors'
|
||||
})
|
||||
)
|
||||
await nextTick()
|
||||
|
||||
expect(leafLabels()).toEqual(['model', 'model-new'])
|
||||
})
|
||||
})
|
||||
|
||||
describe('asset mode', () => {
|
||||
it('hides the load-all button and eager-loads models on mount', async () => {
|
||||
settingState.useAssetAPI = true
|
||||
renderComponent()
|
||||
await nextTick()
|
||||
|
||||
expect(screen.queryByLabelText('g.loadAllFolders')).toBeNull()
|
||||
expect(screen.getByLabelText('g.refresh')).toBeInTheDocument()
|
||||
expect(mockLoadModels).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
it('legacy mode keeps the load-all button and stays lazy by default', async () => {
|
||||
renderComponent()
|
||||
await nextTick()
|
||||
|
||||
expect(screen.getByLabelText('g.loadAllFolders')).toBeInTheDocument()
|
||||
expect(mockLoadModels).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('legacy mode still honors AutoLoadAll', async () => {
|
||||
settingState.autoLoadAll = true
|
||||
renderComponent()
|
||||
await nextTick()
|
||||
|
||||
expect(mockLoadModels).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
@@ -11,6 +11,7 @@
|
||||
<i class="icon-[lucide--refresh-cw] size-4" />
|
||||
</Button>
|
||||
<Button
|
||||
v-if="!usesAssetAPI"
|
||||
v-tooltip.bottom="$t('g.loadAllFolders')"
|
||||
variant="muted-textonly"
|
||||
size="icon"
|
||||
@@ -77,28 +78,28 @@ import { buildTree } from '@/utils/treeUtil'
|
||||
const modelStore = useModelStore()
|
||||
const modelToNodeStore = useModelToNodeStore()
|
||||
const settingStore = useSettingStore()
|
||||
const usesAssetAPI = computed(() =>
|
||||
settingStore.get('Comfy.Assets.UseAssetAPI')
|
||||
)
|
||||
const assetDownloadStore = useAssetDownloadStore()
|
||||
const searchBoxRef = ref()
|
||||
const searchQuery = ref<string>('')
|
||||
const expandedKeys = ref<Record<string, boolean>>({})
|
||||
const { expandNode, toggleNodeOnEvent } = useTreeExpansion(expandedKeys)
|
||||
|
||||
const filteredModels = ref<ComfyModelDef[]>([])
|
||||
const filteredModels = computed<ComfyModelDef[]>(() => {
|
||||
const search = searchQuery.value.toLocaleLowerCase()
|
||||
if (!search) return []
|
||||
return modelStore.models.filter((model) => model.searchable.includes(search))
|
||||
})
|
||||
|
||||
const handleSearch = async (query: string) => {
|
||||
if (!query) {
|
||||
filteredModels.value = []
|
||||
expandedKeys.value = {}
|
||||
return
|
||||
}
|
||||
// Load all models to ensure we have the latest data
|
||||
// Load all models to ensure results cover folders not yet opened
|
||||
await modelStore.loadModels()
|
||||
const search = query.toLocaleLowerCase()
|
||||
filteredModels.value = modelStore.models.filter((model: ComfyModelDef) => {
|
||||
return model.searchable.includes(search)
|
||||
})
|
||||
|
||||
await nextTick()
|
||||
expandNode(root.value)
|
||||
}
|
||||
|
||||
type ModelOrFolder = ComfyModelDef | ModelFolder
|
||||
@@ -112,6 +113,12 @@ const root = computed<TreeNode>(() => {
|
||||
)
|
||||
})
|
||||
|
||||
watch(root, async (newRoot) => {
|
||||
if (!searchQuery.value) return
|
||||
await nextTick()
|
||||
expandNode(newRoot)
|
||||
})
|
||||
|
||||
const renderedRoot = computed<TreeExplorerNode<ModelOrFolder>>(() => {
|
||||
const nameFormat = settingStore.get('Comfy.ModelLibrary.NameFormat')
|
||||
const fillNodeInfo = (node: TreeNode): TreeExplorerNode<ModelOrFolder> => {
|
||||
@@ -193,7 +200,13 @@ watch(
|
||||
|
||||
onMounted(async () => {
|
||||
searchBoxRef.value?.focus()
|
||||
if (settingStore.get('Comfy.ModelLibrary.AutoLoadAll')) {
|
||||
// In asset mode the whole library resolves from one cached walk, so eager
|
||||
// loading is cheap and keeps search and folder badges complete from the
|
||||
// start; AutoLoadAll remains the opt-in for the request-per-folder legacy path.
|
||||
if (
|
||||
usesAssetAPI.value ||
|
||||
settingStore.get('Comfy.ModelLibrary.AutoLoadAll')
|
||||
) {
|
||||
await modelStore.loadModels()
|
||||
}
|
||||
})
|
||||
|
||||
@@ -26,6 +26,7 @@ import { computed, nextTick, onMounted, onUnmounted, ref } from 'vue'
|
||||
import TreeExplorerTreeNode from '@/components/common/TreeExplorerTreeNode.vue'
|
||||
import { useSettingStore } from '@/platform/settings/settingStore'
|
||||
import type { ComfyModelDef } from '@/stores/modelStore'
|
||||
import { getModelPreviewUrl } from '@/stores/modelStore'
|
||||
import type { RenderedTreeExplorerNode } from '@/types/treeExplorerTypes'
|
||||
|
||||
import ModelPreview from './ModelPreview.vue'
|
||||
@@ -37,17 +38,7 @@ const props = defineProps<{
|
||||
// Note: The leaf node should always have a model definition on node.data.
|
||||
const modelDef = computed<ComfyModelDef>(() => props.node.data!)
|
||||
|
||||
const modelPreviewUrl = computed(() => {
|
||||
if (modelDef.value.image) {
|
||||
return modelDef.value.image
|
||||
}
|
||||
const folder = modelDef.value.directory
|
||||
const path_index = modelDef.value.path_index
|
||||
const extension = modelDef.value.file_name.split('.').pop()
|
||||
const filename = modelDef.value.file_name.replace(`.${extension}`, '.webp')
|
||||
const encodedFilename = encodeURIComponent(filename).replace(/%2F/g, '/')
|
||||
return `/api/experiment/models/preview/${folder}/${path_index}/${encodedFilename}`
|
||||
})
|
||||
const modelPreviewUrl = computed(() => getModelPreviewUrl(modelDef.value))
|
||||
|
||||
const previewRef = ref<InstanceType<typeof ModelPreview> | null>(null)
|
||||
const modelPreviewStyle = ref<CSSProperties>({
|
||||
|
||||
@@ -33,7 +33,8 @@ export enum ServerFeatureFlag {
|
||||
SHOW_SIGNIN_BUTTON = 'show_signin_button',
|
||||
UNIFIED_CLOUD_AUTH = 'unified_cloud_auth',
|
||||
CONSOLIDATED_BILLING_ENABLED = 'consolidated_billing_enabled',
|
||||
SIGNUP_TURNSTILE = 'signup_turnstile'
|
||||
SIGNUP_TURNSTILE = 'signup_turnstile',
|
||||
SUPPORTS_MODEL_TYPE_TAGS = 'supports_model_type_tags'
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -208,6 +209,12 @@ export function useFeatureFlags() {
|
||||
remoteConfig.value.signup_turnstile,
|
||||
'off'
|
||||
)
|
||||
},
|
||||
get supportsModelTypeTags() {
|
||||
return api.getServerFeature(
|
||||
ServerFeatureFlag.SUPPORTS_MODEL_TYPE_TAGS,
|
||||
false
|
||||
)
|
||||
}
|
||||
})
|
||||
|
||||
|
||||
@@ -9,6 +9,20 @@ import { useAssetsStore } from '@/stores/assetsStore'
|
||||
|
||||
const mockAssetsByKey = vi.hoisted(() => new Map<string, AssetItem[]>())
|
||||
const mockLoadingByKey = vi.hoisted(() => new Map<string, boolean>())
|
||||
const mockSupportsModelTypeTags = vi.hoisted(() => ({ value: false }))
|
||||
|
||||
vi.mock('@/composables/useFeatureFlags', () => ({
|
||||
useFeatureFlags: () => ({
|
||||
flags: {
|
||||
get supportsModelTypeTags() {
|
||||
return mockSupportsModelTypeTags.value
|
||||
},
|
||||
get modelUploadButtonEnabled() {
|
||||
return false
|
||||
}
|
||||
}
|
||||
})
|
||||
}))
|
||||
|
||||
vi.mock('@/i18n', () => ({
|
||||
t: (key: string, params?: Record<string, string>) =>
|
||||
@@ -214,6 +228,7 @@ describe('AssetBrowserModal', () => {
|
||||
vi.resetAllMocks()
|
||||
mockAssetsByKey.clear()
|
||||
mockLoadingByKey.clear()
|
||||
mockSupportsModelTypeTags.value = false
|
||||
})
|
||||
|
||||
describe('Integration with useAssetBrowser', () => {
|
||||
@@ -420,5 +435,20 @@ describe('AssetBrowserModal', () => {
|
||||
'assetBrowser.allCategory:{"category":"Checkpoints"}'
|
||||
)
|
||||
})
|
||||
|
||||
it('strips the model_type: prefix from the title when the flag is on', async () => {
|
||||
mockSupportsModelTypeTags.value = true
|
||||
const assets = [
|
||||
createTestAsset('asset1', 'Model A', 'model_type:checkpoints')
|
||||
]
|
||||
mockAssetsByKey.set('CheckpointLoaderSimple', assets)
|
||||
|
||||
renderModal({ nodeType: 'CheckpointLoaderSimple' })
|
||||
await flushPromises()
|
||||
|
||||
expect(screen.getByTestId('modal-title').textContent).toBe(
|
||||
'assetBrowser.allCategory:{"category":"Checkpoints"}'
|
||||
)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
@@ -100,6 +100,7 @@ import SearchInput from '@/components/ui/search-input/SearchInput.vue'
|
||||
import Button from '@/components/ui/button/Button.vue'
|
||||
import BaseModalLayout from '@/components/widget/layout/BaseModalLayout.vue'
|
||||
import LeftSidePanel from '@/components/widget/panel/LeftSidePanel.vue'
|
||||
import { useFeatureFlags } from '@/composables/useFeatureFlags'
|
||||
import { usePrimeVueOverlayChildStyle } from '@/composables/usePopoverSizing'
|
||||
import AssetFilterBar from '@/platform/assets/components/AssetFilterBar.vue'
|
||||
import AssetGrid from '@/platform/assets/components/AssetGrid.vue'
|
||||
@@ -109,12 +110,14 @@ import { useAssetBrowser } from '@/platform/assets/composables/useAssetBrowser'
|
||||
import { useModelTypes } from '@/platform/assets/composables/useModelTypes'
|
||||
import { useModelUpload } from '@/platform/assets/composables/useModelUpload'
|
||||
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
|
||||
import { getPrimaryCategoryTag } from '@/platform/assets/utils/assetMetadataUtils'
|
||||
import { formatCategoryLabel } from '@/platform/assets/utils/categoryLabel'
|
||||
import { useAssetsStore } from '@/stores/assetsStore'
|
||||
import { useModelToNodeStore } from '@/stores/modelToNodeStore'
|
||||
import { OnCloseKey } from '@/types/widgetTypes'
|
||||
|
||||
const { t } = useI18n()
|
||||
const { flags } = useFeatureFlags()
|
||||
const assetStore = useAssetsStore()
|
||||
const modelToNodeStore = useModelToNodeStore()
|
||||
const breakpoints = useBreakpoints(breakpointsTailwind)
|
||||
@@ -191,9 +194,21 @@ const focusedAsset = ref<AssetDisplayItem | null>(null)
|
||||
const isRightPanelOpen = ref(false)
|
||||
|
||||
const primaryCategoryTag = computed(() => {
|
||||
const modelTypeMode = flags.supportsModelTypeTags
|
||||
// A node-typed picker is FOR a category; title off that category rather
|
||||
// than guessing from the first asset, whose first model_type value may be
|
||||
// a different category it shares a root with.
|
||||
if (modelTypeMode && props.nodeType) {
|
||||
const mapped = modelToNodeStore.getCategoryForNodeType(props.nodeType)
|
||||
if (mapped) return mapped
|
||||
}
|
||||
|
||||
const assets = fetchedAssets.value ?? []
|
||||
// Covered assets title off the model_type value they group under (so title
|
||||
// and grouping cannot diverge); uncovered assets keep the legacy verbatim
|
||||
// first tag.
|
||||
const tagFromAssets = assets
|
||||
.map((asset) => asset.tags?.find((tag) => tag !== 'models'))
|
||||
.map((asset) => getPrimaryCategoryTag(asset, modelTypeMode))
|
||||
.find((tag): tag is string => typeof tag === 'string' && tag.length > 0)
|
||||
|
||||
if (tagFromAssets) return tagFromAssets
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import { createTestingPinia } from '@pinia/testing'
|
||||
import { describe, expect, it, vi } from 'vitest'
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest'
|
||||
import { createI18n } from 'vue-i18n'
|
||||
|
||||
import { render, screen } from '@testing-library/vue'
|
||||
@@ -14,6 +14,14 @@ vi.mock('@/composables/useCopyToClipboard', () => ({
|
||||
})
|
||||
}))
|
||||
|
||||
const mockDistribution = vi.hoisted(() => ({ isCloud: false }))
|
||||
vi.mock('@/platform/distribution/types', async (importOriginal) => ({
|
||||
...(await importOriginal<object>()),
|
||||
get isCloud() {
|
||||
return mockDistribution.isCloud
|
||||
}
|
||||
}))
|
||||
|
||||
const i18n = createI18n({
|
||||
legacy: false,
|
||||
locale: 'en',
|
||||
@@ -41,6 +49,10 @@ describe('ModelInfoPanel', () => {
|
||||
...overrides
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
mockDistribution.isCloud = false
|
||||
})
|
||||
|
||||
function renderPanel(asset: AssetDisplayItem) {
|
||||
return render(ModelInfoPanel, {
|
||||
props: { asset },
|
||||
@@ -138,6 +150,18 @@ describe('ModelInfoPanel', () => {
|
||||
).toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('shows an editable model type dropdown for a mutable asset on cloud', () => {
|
||||
mockDistribution.isCloud = true
|
||||
renderPanel(createMockAsset({ is_immutable: false }))
|
||||
expect(screen.getByRole('combobox')).toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('keeps the model type read-only on core even for a mutable asset', () => {
|
||||
mockDistribution.isCloud = false
|
||||
renderPanel(createMockAsset({ is_immutable: false }))
|
||||
expect(screen.queryByRole('combobox')).not.toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('renders base models field', () => {
|
||||
const asset = createMockAsset({
|
||||
user_metadata: { base_model: ['SDXL'] }
|
||||
|
||||
@@ -71,7 +71,7 @@
|
||||
</span>
|
||||
</template>
|
||||
<ModelInfoField :label="t('assetBrowser.modelInfo.modelType')">
|
||||
<Select v-if="!isImmutable" v-model="selectedModelType">
|
||||
<Select v-if="isModelTypeEditable" v-model="selectedModelType">
|
||||
<SelectTrigger class="w-full">
|
||||
<SelectValue
|
||||
:placeholder="t('assetBrowser.modelInfo.selectModelType')"
|
||||
@@ -215,6 +215,7 @@ import { useI18n } from 'vue-i18n'
|
||||
|
||||
import EditableText from '@/components/common/EditableText.vue'
|
||||
import { useCopyToClipboard } from '@/composables/useCopyToClipboard'
|
||||
import { useFeatureFlags } from '@/composables/useFeatureFlags'
|
||||
import PropertiesAccordionItem from '@/components/rightSidePanel/layout/PropertiesAccordionItem.vue'
|
||||
import Button from '@/components/ui/button/Button.vue'
|
||||
import Select from '@/components/ui/select/Select.vue'
|
||||
@@ -229,17 +230,19 @@ import TagsInputItemDelete from '@/components/ui/tags-input/TagsInputItemDelete.
|
||||
import TagsInputItemText from '@/components/ui/tags-input/TagsInputItemText.vue'
|
||||
import type { AssetDisplayItem } from '@/platform/assets/composables/useAssetBrowser'
|
||||
import { useModelTypes } from '@/platform/assets/composables/useModelTypes'
|
||||
import { isCloud } from '@/platform/distribution/types'
|
||||
import type { AssetUserMetadata } from '@/platform/assets/schemas/assetSchema'
|
||||
import {
|
||||
buildModelTypeTagUpdate,
|
||||
getAssetAdditionalTags,
|
||||
getAssetBaseModels,
|
||||
getAssetDescription,
|
||||
getAssetDisplayName,
|
||||
getAssetFilename,
|
||||
getAssetModelType,
|
||||
getAssetSourceUrl,
|
||||
getAssetTriggerPhrases,
|
||||
getAssetUserDescription,
|
||||
getEditableModelType,
|
||||
getSourceName
|
||||
} from '@/platform/assets/utils/assetMetadataUtils'
|
||||
import { useAssetsStore } from '@/stores/assetsStore'
|
||||
@@ -265,6 +268,7 @@ const { asset, cacheKey, selectContentStyle } = defineProps<{
|
||||
}>()
|
||||
|
||||
const assetsStore = useAssetsStore()
|
||||
const { flags } = useFeatureFlags()
|
||||
const { modelTypes } = useModelTypes()
|
||||
|
||||
const pendingUpdates = ref<AssetUserMetadata>({})
|
||||
@@ -272,6 +276,9 @@ const pendingModelType = ref<string | undefined>(undefined)
|
||||
const isEditingDisplayName = ref(false)
|
||||
|
||||
const isImmutable = computed(() => asset.is_immutable ?? true)
|
||||
// Retagging a model rewrites its asset tags; core is filesystem-backed and does
|
||||
// not yet move the file to match, so the model type is read-only off-cloud.
|
||||
const isModelTypeEditable = computed(() => !isImmutable.value && isCloud)
|
||||
const displayName = computed(
|
||||
() => pendingUpdates.value.name ?? getAssetDisplayName(asset)
|
||||
)
|
||||
@@ -318,12 +325,17 @@ function handleDisplayNameEdit(newName: string) {
|
||||
}
|
||||
|
||||
const debouncedSaveModelType = useDebounceFn((newModelType: string) => {
|
||||
if (isImmutable.value) return
|
||||
const currentModelType = getAssetModelType(asset)
|
||||
if (!isModelTypeEditable.value) return
|
||||
const currentModelType = getEditableModelType(
|
||||
asset,
|
||||
flags.supportsModelTypeTags
|
||||
)
|
||||
if (currentModelType === newModelType) return
|
||||
const newTags = asset.tags
|
||||
.filter((tag) => tag !== currentModelType)
|
||||
.concat(newModelType)
|
||||
const newTags = buildModelTypeTagUpdate(
|
||||
asset,
|
||||
newModelType,
|
||||
flags.supportsModelTypeTags
|
||||
)
|
||||
assetsStore.updateAssetTags(asset, newTags, cacheKey)
|
||||
}, 500)
|
||||
|
||||
@@ -345,7 +357,10 @@ const userDescription = computed({
|
||||
})
|
||||
|
||||
const selectedModelType = computed({
|
||||
get: () => pendingModelType.value ?? getAssetModelType(asset) ?? undefined,
|
||||
get: () =>
|
||||
pendingModelType.value ??
|
||||
getEditableModelType(asset, flags.supportsModelTypeTags) ??
|
||||
undefined,
|
||||
set: (value: string | undefined) => {
|
||||
if (!value) return
|
||||
pendingModelType.value = value
|
||||
|
||||
@@ -25,10 +25,22 @@ vi.mock('@/i18n', () => ({
|
||||
d: (date: Date) => date.toLocaleDateString()
|
||||
}))
|
||||
|
||||
const mockSupportsModelTypeTags = vi.hoisted(() => ({ value: false }))
|
||||
vi.mock('@/composables/useFeatureFlags', () => ({
|
||||
useFeatureFlags: () => ({
|
||||
flags: {
|
||||
get supportsModelTypeTags() {
|
||||
return mockSupportsModelTypeTags.value
|
||||
}
|
||||
}
|
||||
})
|
||||
}))
|
||||
|
||||
describe('useAssetBrowser', () => {
|
||||
beforeEach(() => {
|
||||
setActivePinia(createPinia())
|
||||
vi.restoreAllMocks()
|
||||
mockSupportsModelTypeTags.value = false
|
||||
})
|
||||
|
||||
// Test fixtures - minimal data focused on functionality being tested
|
||||
@@ -138,6 +150,25 @@ describe('useAssetBrowser', () => {
|
||||
})
|
||||
})
|
||||
|
||||
it('strips the model_type: prefix from the badge when the flag is on', () => {
|
||||
mockSupportsModelTypeTags.value = true
|
||||
const apiAsset = createApiAsset({
|
||||
tags: ['models', 'model_type:checkpoints', 'sdxl']
|
||||
})
|
||||
|
||||
const { filteredAssets } = useAssetBrowser(ref([apiAsset]))
|
||||
const result = filteredAssets.value[0]
|
||||
|
||||
expect(result.badges).toContainEqual({
|
||||
label: 'checkpoints',
|
||||
type: 'type'
|
||||
})
|
||||
expect(result.badges).not.toContainEqual({
|
||||
label: 'model_type:checkpoints',
|
||||
type: 'type'
|
||||
})
|
||||
})
|
||||
|
||||
it('handles tags with multiple slashes in badges', () => {
|
||||
const apiAsset = createApiAsset({
|
||||
tags: ['models', 'checkpoint/subfolder/model-name']
|
||||
@@ -668,6 +699,34 @@ describe('useAssetBrowser', () => {
|
||||
])
|
||||
})
|
||||
|
||||
it('groups by model_type:* value and ignores other tags when the flag is on', () => {
|
||||
mockSupportsModelTypeTags.value = true
|
||||
const assets = [
|
||||
createApiAsset({ tags: ['models', 'model_type:checkpoints', 'sdxl'] }),
|
||||
createApiAsset({ tags: ['models', 'model_type:LLM'] })
|
||||
]
|
||||
|
||||
const { navItems } = useAssetBrowser(ref(assets))
|
||||
|
||||
const typeGroup = navItems.value[2] as { items: { id: string }[] }
|
||||
expect(typeGroup.items.map((i) => i.id)).toEqual(['LLM', 'checkpoints'])
|
||||
})
|
||||
|
||||
it('ignores model_type: and groups by bare tags when the flag is off', () => {
|
||||
const assets = [
|
||||
createApiAsset({ tags: ['models', 'model_type:checkpoints'] }),
|
||||
createApiAsset({ tags: ['models', 'model_type:LLM'] })
|
||||
]
|
||||
|
||||
const { navItems } = useAssetBrowser(ref(assets))
|
||||
|
||||
const typeGroup = navItems.value[2] as { items: { id: string }[] }
|
||||
expect(typeGroup.items.map((i) => i.id)).toEqual([
|
||||
'model_type:LLM',
|
||||
'model_type:checkpoints'
|
||||
])
|
||||
})
|
||||
|
||||
it('handles assets with no category tag', () => {
|
||||
const assets = [
|
||||
createApiAsset({ tags: ['models'] }), // No second tag
|
||||
|
||||
@@ -19,10 +19,13 @@ import {
|
||||
} from '@/platform/assets/utils/assetFilterUtils'
|
||||
import {
|
||||
getAssetBaseModels,
|
||||
getAssetFilename
|
||||
getAssetCategories,
|
||||
getAssetFilename,
|
||||
getAssetTypeBadges
|
||||
} from '@/platform/assets/utils/assetMetadataUtils'
|
||||
import { MODELS_TAG } from '@/platform/assets/services/assetService'
|
||||
import { sortAssets } from '@/platform/assets/utils/assetSortUtils'
|
||||
import { useFeatureFlags } from '@/composables/useFeatureFlags'
|
||||
import { useAssetDownloadStore } from '@/stores/assetDownloadStore'
|
||||
import type { NavGroupData, NavItemData } from '@/types/navTypes'
|
||||
|
||||
@@ -43,18 +46,19 @@ export interface AssetDisplayItem extends AssetItem {
|
||||
}
|
||||
}
|
||||
|
||||
const displayItemCache = new WeakMap<AssetItem, AssetDisplayItem>()
|
||||
const displayItemCache = new WeakMap<
|
||||
AssetItem,
|
||||
{ modelTypeMode: boolean; item: AssetDisplayItem }
|
||||
>()
|
||||
|
||||
function buildDisplayItem(asset: AssetItem): AssetDisplayItem {
|
||||
function buildDisplayItem(
|
||||
asset: AssetItem,
|
||||
modelTypeMode: boolean
|
||||
): AssetDisplayItem {
|
||||
const badges: AssetBadge[] = []
|
||||
|
||||
const typeTag = asset.tags.find((tag) => tag !== 'models')
|
||||
if (typeTag) {
|
||||
const badgeLabel = typeTag.includes('/')
|
||||
? typeTag.substring(typeTag.indexOf('/') + 1)
|
||||
: typeTag
|
||||
|
||||
badges.push({ label: badgeLabel, type: 'type' })
|
||||
for (const typeBadge of getAssetTypeBadges(asset, modelTypeMode)) {
|
||||
badges.push({ label: typeBadge, type: 'type' })
|
||||
}
|
||||
|
||||
for (const model of getAssetBaseModels(asset)) {
|
||||
@@ -75,12 +79,15 @@ function buildDisplayItem(asset: AssetItem): AssetDisplayItem {
|
||||
}
|
||||
}
|
||||
|
||||
function transformAssetForDisplay(asset: AssetItem): AssetDisplayItem {
|
||||
function transformAssetForDisplay(
|
||||
asset: AssetItem,
|
||||
modelTypeMode: boolean
|
||||
): AssetDisplayItem {
|
||||
const cached = displayItemCache.get(asset)
|
||||
if (cached) return cached
|
||||
const built = buildDisplayItem(asset)
|
||||
displayItemCache.set(asset, built)
|
||||
return built
|
||||
if (cached && cached.modelTypeMode === modelTypeMode) return cached.item
|
||||
const item = buildDisplayItem(asset, modelTypeMode)
|
||||
displayItemCache.set(asset, { modelTypeMode, item })
|
||||
return item
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -93,6 +100,7 @@ export function useAssetBrowser(
|
||||
const assets = computed<AssetItem[]>(() => assetsSource.value ?? [])
|
||||
const assetDownloadStore = useAssetDownloadStore()
|
||||
const { sessionDownloadCount } = storeToRefs(assetDownloadStore)
|
||||
const { flags } = useFeatureFlags()
|
||||
|
||||
// State
|
||||
const searchQuery = ref('')
|
||||
@@ -122,12 +130,10 @@ export function useAssetBrowser(
|
||||
})
|
||||
|
||||
const typeCategories = computed<NavItemData[]>(() => {
|
||||
const modelTypeMode = flags.supportsModelTypeTags
|
||||
const categories = assets.value
|
||||
.filter((asset) => asset.tags.includes(MODELS_TAG))
|
||||
.flatMap((asset) =>
|
||||
asset.tags.filter((tag) => tag !== MODELS_TAG && tag.length > 0)
|
||||
)
|
||||
.map((tag) => tag.split('/')[0])
|
||||
.flatMap((asset) => getAssetCategories(asset, modelTypeMode))
|
||||
|
||||
return Array.from(new Set(categories))
|
||||
.sort()
|
||||
@@ -191,7 +197,9 @@ export function useAssetBrowser(
|
||||
|
||||
// Category-filtered assets for filter options (before search/format/base model filters)
|
||||
const categoryFilteredAssets = computed(() => {
|
||||
return assets.value.filter(filterByCategory(selectedCategory.value))
|
||||
return assets.value.filter(
|
||||
filterByCategory(selectedCategory.value, flags.supportsModelTypeTags)
|
||||
)
|
||||
})
|
||||
|
||||
const { availableFileFormats, availableBaseModels } = useAssetFilterOptions(
|
||||
@@ -248,7 +256,10 @@ export function useAssetBrowser(
|
||||
const sortedAssets = sortAssets(filtered, filters.value.sortBy)
|
||||
|
||||
// Transform to display format
|
||||
return sortedAssets.map(transformAssetForDisplay)
|
||||
const modelTypeMode = flags.supportsModelTypeTags
|
||||
return sortedAssets.map((asset) =>
|
||||
transformAssetForDisplay(asset, modelTypeMode)
|
||||
)
|
||||
})
|
||||
|
||||
function updateFilters(newFilters: AssetFilterState) {
|
||||
|
||||
@@ -38,7 +38,10 @@ vi.mock('@/scripts/api', () => ({
|
||||
api: {
|
||||
fetchApi: vi.fn(),
|
||||
addEventListener: vi.fn(),
|
||||
apiURL: vi.fn((path: string) => path)
|
||||
apiURL: vi.fn((path: string) => path),
|
||||
getServerFeature: vi.fn(
|
||||
(_name: string, defaultValue?: unknown) => defaultValue
|
||||
)
|
||||
}
|
||||
}))
|
||||
|
||||
@@ -279,6 +282,43 @@ describe('useUploadModelWizard', () => {
|
||||
expect(result?.modelType).toBe('checkpoints')
|
||||
})
|
||||
|
||||
it('namespaces the tag but keeps user_metadata.model_type bare when the backend supports it', async () => {
|
||||
const { assetService } =
|
||||
await import('@/platform/assets/services/assetService')
|
||||
const { api } = await import('@/scripts/api')
|
||||
vi.mocked(assetService.uploadAssetAsync).mockResolvedValue({
|
||||
type: 'sync',
|
||||
asset: {
|
||||
id: 'asset-1',
|
||||
name: 'model.safetensors',
|
||||
tags: ['models', 'model_type:checkpoints']
|
||||
}
|
||||
})
|
||||
vi.mocked(api.getServerFeature).mockImplementation((name, defaultValue) =>
|
||||
name === 'supports_model_type_tags' ? true : defaultValue
|
||||
)
|
||||
|
||||
try {
|
||||
const wizard = setupUploadModelWizard(modelTypes, {
|
||||
requiredModelType: 'checkpoints'
|
||||
})
|
||||
wizard.wizardData.value.url = 'https://civitai.com/models/12345'
|
||||
|
||||
await wizard.uploadModel()
|
||||
|
||||
const uploadArg = vi.mocked(assetService.uploadAssetAsync).mock
|
||||
.calls[0][0]
|
||||
expect(uploadArg.tags).toEqual(['models', 'model_type:checkpoints'])
|
||||
expect(uploadArg.user_metadata?.model_type).toBe('checkpoints')
|
||||
// The namespaced returned tag must not trip the required-type guard.
|
||||
expect(wizard.uploadTypeMismatch.value).toBeNull()
|
||||
} finally {
|
||||
vi.mocked(api.getServerFeature).mockImplementation(
|
||||
(_name, defaultValue) => defaultValue
|
||||
)
|
||||
}
|
||||
})
|
||||
|
||||
it('returns the synced asset filename for sync imports', async () => {
|
||||
const { assetService } =
|
||||
await import('@/platform/assets/services/assetService')
|
||||
@@ -347,6 +387,65 @@ describe('useUploadModelWizard', () => {
|
||||
})
|
||||
})
|
||||
|
||||
it('treats a namespaced model_type: tag as satisfying the required type', async () => {
|
||||
const { assetService } =
|
||||
await import('@/platform/assets/services/assetService')
|
||||
vi.mocked(assetService.uploadAssetAsync).mockResolvedValue({
|
||||
type: 'sync',
|
||||
asset: {
|
||||
id: 'asset-1',
|
||||
name: 'model.safetensors',
|
||||
tags: ['models', 'model_type:checkpoints']
|
||||
}
|
||||
})
|
||||
|
||||
const wizard = setupUploadModelWizard(
|
||||
ref([
|
||||
{ name: 'Checkpoint', value: 'checkpoints' },
|
||||
{ name: 'LoRA', value: 'loras' }
|
||||
]),
|
||||
{ requiredModelType: 'checkpoints' }
|
||||
)
|
||||
wizard.wizardData.value.url = 'https://civitai.com/models/12345'
|
||||
|
||||
const result = await wizard.uploadModel()
|
||||
|
||||
expect(result).not.toBeNull()
|
||||
expect(wizard.uploadTypeMismatch.value).toBeNull()
|
||||
})
|
||||
|
||||
it('strips the model_type: prefix from the imported-type label on a real mismatch', async () => {
|
||||
const { assetService } =
|
||||
await import('@/platform/assets/services/assetService')
|
||||
vi.mocked(assetService.uploadAssetAsync).mockResolvedValue({
|
||||
type: 'sync',
|
||||
asset: {
|
||||
id: 'asset-lora',
|
||||
name: 'model.safetensors',
|
||||
tags: ['models', 'model_type:loras']
|
||||
}
|
||||
})
|
||||
|
||||
const wizard = setupUploadModelWizard(
|
||||
ref([
|
||||
{ name: 'Checkpoint', value: 'checkpoints' },
|
||||
{ name: 'LoRA', value: 'loras' }
|
||||
]),
|
||||
{ requiredModelType: 'checkpoints' }
|
||||
)
|
||||
wizard.wizardData.value.url = 'https://civitai.com/models/12345'
|
||||
|
||||
const result = await wizard.uploadModel()
|
||||
|
||||
expect(result).toBeNull()
|
||||
expect(wizard.uploadTypeMismatch.value).toEqual({
|
||||
importedModelType: 'loras',
|
||||
importedModelTypeLabel: 'LoRA',
|
||||
requiredModelType: 'checkpoints',
|
||||
requiredModelTypeLabel: 'Checkpoint'
|
||||
})
|
||||
})
|
||||
|
||||
it('does not block sync imports as mismatches without a required model type', async () => {
|
||||
const { assetService } =
|
||||
await import('@/platform/assets/services/assetService')
|
||||
|
||||
@@ -2,6 +2,7 @@ import type { Ref } from 'vue'
|
||||
import { computed, ref, watch } from 'vue'
|
||||
import { useI18n } from 'vue-i18n'
|
||||
|
||||
import { useFeatureFlags } from '@/composables/useFeatureFlags'
|
||||
import { st } from '@/i18n'
|
||||
import { civitaiImportSource } from '@/platform/assets/importSources/civitaiImportSource'
|
||||
import { huggingfaceImportSource } from '@/platform/assets/importSources/huggingfaceImportSource'
|
||||
@@ -11,7 +12,11 @@ import type {
|
||||
} from '@/platform/assets/schemas/assetSchema'
|
||||
import { assetService } from '@/platform/assets/services/assetService'
|
||||
import type { ImportSource } from '@/platform/assets/types/importSource'
|
||||
import { getAssetFilename } from '@/platform/assets/utils/assetMetadataUtils'
|
||||
import {
|
||||
getAssetFilename,
|
||||
stripModelTypePrefix,
|
||||
toModelTypeTag
|
||||
} from '@/platform/assets/utils/assetMetadataUtils'
|
||||
import { validateSourceUrl } from '@/platform/assets/utils/importSourceUtil'
|
||||
import { useAssetDownloadStore } from '@/stores/assetDownloadStore'
|
||||
import { useAssetsStore } from '@/stores/assetsStore'
|
||||
@@ -68,6 +73,7 @@ export function useUploadModelWizard(
|
||||
options: UploadModelWizardOptions = {}
|
||||
) {
|
||||
const { t } = useI18n()
|
||||
const { flags } = useFeatureFlags()
|
||||
const assetsStore = useAssetsStore()
|
||||
const assetDownloadStore = useAssetDownloadStore()
|
||||
const modelToNodeStore = useModelToNodeStore()
|
||||
@@ -271,19 +277,22 @@ export function useUploadModelWizard(
|
||||
}
|
||||
|
||||
function getImportedModelType(asset: AssetItem): string | undefined {
|
||||
const knownType = asset.tags.find(
|
||||
(tag) =>
|
||||
tag !== MODEL_ROOT_TAG &&
|
||||
const subtypeTags = asset.tags
|
||||
.filter((tag) => tag !== MODEL_ROOT_TAG)
|
||||
.map(stripModelTypePrefix)
|
||||
return (
|
||||
subtypeTags.find((tag) =>
|
||||
modelTypes.value.some((type) => type.value === tag)
|
||||
) ?? subtypeTags[0]
|
||||
)
|
||||
return knownType ?? asset.tags.find((tag) => tag !== MODEL_ROOT_TAG)
|
||||
}
|
||||
|
||||
function blockMismatchedImportedModel(
|
||||
asset: AssetItem,
|
||||
requiredType: string
|
||||
): boolean {
|
||||
if (asset.tags.includes(requiredType)) return false
|
||||
if (asset.tags.map(stripModelTypePrefix).includes(requiredType))
|
||||
return false
|
||||
|
||||
const importedType = getImportedModelType(asset)
|
||||
uploadStatus.value = 'error'
|
||||
@@ -317,7 +326,11 @@ export function useUploadModelWizard(
|
||||
|
||||
try {
|
||||
const modelType = resolvedModelType.value
|
||||
const tags = modelType ? ['models', modelType] : ['models']
|
||||
const subtypeTag =
|
||||
modelType && flags.supportsModelTypeTags
|
||||
? toModelTypeTag(modelType)
|
||||
: modelType
|
||||
const tags = subtypeTag ? [MODEL_ROOT_TAG, subtypeTag] : [MODEL_ROOT_TAG]
|
||||
const filename =
|
||||
wizardData.value.metadata?.filename ||
|
||||
wizardData.value.metadata?.name ||
|
||||
|
||||
@@ -11,6 +11,8 @@ const zAsset = z.object({
|
||||
tags: z.array(z.string()).optional().default([]),
|
||||
preview_id: z.string().nullable().optional(),
|
||||
display_name: z.string().optional(),
|
||||
/** Path within the model's category folder, i.e. the value a loader widget expects. */
|
||||
loader_path: z.string().nullish(),
|
||||
preview_url: z.string().optional(),
|
||||
thumbnail_url: z.string().optional(),
|
||||
created_at: z.string().optional(),
|
||||
@@ -27,11 +29,6 @@ const zAssetResponse = zListAssetsResponse
|
||||
assets: z.array(zAsset)
|
||||
})
|
||||
|
||||
const zModelFolder = z.object({
|
||||
name: z.string(),
|
||||
folders: z.array(z.string())
|
||||
})
|
||||
|
||||
// Zod schema for ModelFile to align with interface
|
||||
const zModelFile = z.object({
|
||||
name: z.string(),
|
||||
@@ -100,7 +97,6 @@ export type AssetItem = z.infer<typeof zAsset>
|
||||
export type AssetResponse = z.infer<typeof zAssetResponse>
|
||||
export type AssetMetadata = z.infer<typeof zAssetMetadata>
|
||||
export type AsyncUploadResponse = z.infer<typeof zAsyncUploadResponse>
|
||||
export type ModelFolder = z.infer<typeof zModelFolder>
|
||||
export type ModelFile = z.infer<typeof zModelFile>
|
||||
|
||||
/** Payload for updating an asset via PUT /assets/:id */
|
||||
@@ -132,4 +128,10 @@ export type TagsOperationResult = z.infer<typeof tagsOperationResultSchema>
|
||||
export interface ModelFolderInfo {
|
||||
name: string
|
||||
folders: string[]
|
||||
/**
|
||||
* The folder's raw registered extension allowlist from
|
||||
* `/experiment/models`. An empty array means match-all; absent on older
|
||||
* backends.
|
||||
*/
|
||||
extensions?: string[]
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@ import { api } from '@/scripts/api'
|
||||
|
||||
const mockDistributionState = vi.hoisted(() => ({ isCloud: false }))
|
||||
const mockSettingStoreGet = vi.hoisted(() => vi.fn(() => false))
|
||||
const mockSupportsModelTypeTags = vi.hoisted(() => ({ value: true }))
|
||||
|
||||
vi.mock('@/platform/distribution/types', () => ({
|
||||
get isCloud() {
|
||||
@@ -19,6 +20,16 @@ vi.mock('@/platform/distribution/types', () => ({
|
||||
}
|
||||
}))
|
||||
|
||||
vi.mock('@/composables/useFeatureFlags', () => ({
|
||||
useFeatureFlags: () => ({
|
||||
flags: {
|
||||
get supportsModelTypeTags() {
|
||||
return mockSupportsModelTypeTags.value
|
||||
}
|
||||
}
|
||||
})
|
||||
}))
|
||||
|
||||
vi.mock('@/platform/settings/settingStore', () => ({
|
||||
useSettingStore: vi.fn(() => ({
|
||||
get: mockSettingStoreGet
|
||||
@@ -40,7 +51,9 @@ vi.mock('@/stores/modelToNodeStore', () => {
|
||||
|
||||
vi.mock('@/scripts/api', () => ({
|
||||
api: {
|
||||
fetchApi: vi.fn()
|
||||
fetchApi: vi.fn(),
|
||||
addCustomEventListener: vi.fn(),
|
||||
removeCustomEventListener: vi.fn()
|
||||
}
|
||||
}))
|
||||
|
||||
@@ -87,6 +100,7 @@ function validAsset(overrides: Partial<AssetItem> = {}): AssetItem {
|
||||
return {
|
||||
id: 'asset-1',
|
||||
name: 'model.safetensors',
|
||||
loader_path: overrides.name ?? 'model.safetensors',
|
||||
tags: ['models'],
|
||||
...overrides
|
||||
}
|
||||
@@ -416,32 +430,329 @@ describe(assetService.deleteAsset, () => {
|
||||
})
|
||||
})
|
||||
|
||||
describe(assetService.getAssetModelFolders, () => {
|
||||
describe(assetService.getAssetModels, () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
assetService.invalidateModelBuckets()
|
||||
mockSupportsModelTypeTags.value = true
|
||||
})
|
||||
|
||||
it('walks the models tag once, excluding missing assets', async () => {
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildAssetListResponse([
|
||||
validAsset({ id: 'a', tags: ['models', 'model_type:checkpoints'] })
|
||||
])
|
||||
)
|
||||
|
||||
await assetService.getAssetModels('checkpoints')
|
||||
|
||||
expect(fetchApiMock).toHaveBeenCalledTimes(1)
|
||||
const requestedUrl = fetchApiMock.mock.calls[0]?.[0] as string
|
||||
const params = new URL(requestedUrl, 'http://localhost').searchParams
|
||||
expect(params.get('include_tags')).toBe('models')
|
||||
expect(params.get('exclude_tags')).toBe(MISSING_TAG)
|
||||
})
|
||||
|
||||
it('buckets by bare tags when model_type tags are unsupported', async () => {
|
||||
mockSupportsModelTypeTags.value = false
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildAssetListResponse([
|
||||
validAsset({
|
||||
id: 'a',
|
||||
name: 'a.safetensors',
|
||||
tags: ['models', 'checkpoints']
|
||||
})
|
||||
])
|
||||
)
|
||||
|
||||
const models = await assetService.getAssetModels('checkpoints')
|
||||
|
||||
expect(models).toEqual([{ name: 'a.safetensors', pathIndex: 0 }])
|
||||
})
|
||||
|
||||
it('drops uncategorized model assets with a warning', async () => {
|
||||
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildAssetListResponse([
|
||||
validAsset({
|
||||
id: 'ok',
|
||||
name: 'ok.safetensors',
|
||||
tags: ['models', 'model_type:loras']
|
||||
}),
|
||||
validAsset({
|
||||
id: 'uncat',
|
||||
name: 'orphan.safetensors',
|
||||
tags: ['models']
|
||||
})
|
||||
])
|
||||
)
|
||||
|
||||
const loras = await assetService.getAssetModels('loras')
|
||||
|
||||
expect(loras).toEqual([{ name: 'ok.safetensors', pathIndex: 0 }])
|
||||
expect(warn).toHaveBeenCalledWith(
|
||||
expect.stringContaining('orphan.safetensors')
|
||||
)
|
||||
warn.mockRestore()
|
||||
})
|
||||
|
||||
it('maps loader_path and drops unloadable assets without one', async () => {
|
||||
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildAssetListResponse([
|
||||
validAsset({
|
||||
id: 'nested',
|
||||
name: 'model.safetensors',
|
||||
loader_path: 'sdxl/model.safetensors',
|
||||
tags: ['models', 'model_type:checkpoints']
|
||||
}),
|
||||
validAsset({
|
||||
id: 'orphan',
|
||||
name: 'orphan.safetensors',
|
||||
loader_path: null,
|
||||
tags: ['models', 'model_type:checkpoints']
|
||||
}),
|
||||
validAsset({
|
||||
id: 'other-folder',
|
||||
name: 'lora.safetensors',
|
||||
tags: ['models', 'model_type:loras']
|
||||
})
|
||||
])
|
||||
)
|
||||
|
||||
const models = await assetService.getAssetModels('checkpoints')
|
||||
|
||||
expect(models).toEqual([{ name: 'sdxl/model.safetensors', pathIndex: 0 }])
|
||||
expect(warn).toHaveBeenCalledWith(
|
||||
expect.stringContaining('orphan.safetensors')
|
||||
)
|
||||
warn.mockRestore()
|
||||
})
|
||||
|
||||
it('drops assets whose loader path is traversal-shaped', async () => {
|
||||
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildAssetListResponse([
|
||||
validAsset({
|
||||
id: 'evil',
|
||||
name: 'evil.safetensors',
|
||||
loader_path: '../../secrets/evil.safetensors',
|
||||
tags: ['models', 'model_type:checkpoints']
|
||||
}),
|
||||
validAsset({
|
||||
id: 'ok',
|
||||
name: 'fine.safetensors',
|
||||
tags: ['models', 'model_type:checkpoints']
|
||||
})
|
||||
])
|
||||
)
|
||||
|
||||
const models = await assetService.getAssetModels('checkpoints')
|
||||
|
||||
expect(models).toEqual([{ name: 'fine.safetensors', pathIndex: 0 }])
|
||||
expect(warn).toHaveBeenCalledWith(expect.stringContaining('unsafe'))
|
||||
warn.mockRestore()
|
||||
})
|
||||
|
||||
it('groups slashed bare tags by their top-level segment', async () => {
|
||||
mockSupportsModelTypeTags.value = false
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildAssetListResponse([
|
||||
validAsset({
|
||||
id: 'slashed',
|
||||
name: 'model1.safetensors',
|
||||
tags: ['models', 'Chatterbox/subfolder1/model1']
|
||||
})
|
||||
])
|
||||
)
|
||||
|
||||
const models = await assetService.getAssetModels('Chatterbox')
|
||||
|
||||
expect(models).toEqual([{ name: 'model1.safetensors', pathIndex: 0 }])
|
||||
})
|
||||
|
||||
it('falls back to filename metadata then name on bare-tag backends', async () => {
|
||||
mockSupportsModelTypeTags.value = false
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildAssetListResponse([
|
||||
validAsset({
|
||||
id: 'cloud-hash',
|
||||
name: 'blake3-content-hash',
|
||||
loader_path: null,
|
||||
user_metadata: { filename: 'sdxl/cloud-model.safetensors' },
|
||||
tags: ['models', 'checkpoints']
|
||||
}),
|
||||
validAsset({
|
||||
id: 'bare',
|
||||
name: 'plain.safetensors',
|
||||
loader_path: null,
|
||||
tags: ['models', 'checkpoints']
|
||||
})
|
||||
])
|
||||
)
|
||||
|
||||
const models = await assetService.getAssetModels('checkpoints')
|
||||
|
||||
expect(models).toEqual([
|
||||
{ name: 'sdxl/cloud-model.safetensors', pathIndex: 0 },
|
||||
{ name: 'plain.safetensors', pathIndex: 0 }
|
||||
])
|
||||
})
|
||||
|
||||
it('orders each folder subdirectories-first then files, alphabetically', async () => {
|
||||
const checkpointAsset = (id: string, loaderPath: string) =>
|
||||
validAsset({
|
||||
id,
|
||||
name: loaderPath.split('/').pop()!,
|
||||
loader_path: loaderPath,
|
||||
tags: ['models', 'model_type:checkpoints']
|
||||
})
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildAssetListResponse([
|
||||
checkpointAsset('1', 'sdxl/base.safetensors'),
|
||||
checkpointAsset('2', 'v1-5.safetensors'),
|
||||
checkpointAsset('3', 'sdxl/refiner.safetensors'),
|
||||
checkpointAsset('4', 'anything.safetensors'),
|
||||
checkpointAsset('5', 'dynamicrafter/model.safetensors')
|
||||
])
|
||||
)
|
||||
|
||||
const models = await assetService.getAssetModels('checkpoints')
|
||||
|
||||
expect(models.map((m) => m.name)).toEqual([
|
||||
'dynamicrafter/model.safetensors',
|
||||
'sdxl/base.safetensors',
|
||||
'sdxl/refiner.safetensors',
|
||||
'anything.safetensors',
|
||||
'v1-5.safetensors'
|
||||
])
|
||||
})
|
||||
|
||||
it('does not let a stale in-flight walk overwrite an invalidated cache', async () => {
|
||||
let resolveStaleWalk!: (response: Response) => void
|
||||
fetchApiMock.mockReturnValueOnce(
|
||||
new Promise<Response>((resolve) => {
|
||||
resolveStaleWalk = resolve
|
||||
})
|
||||
)
|
||||
const staleRead = assetService.getAssetModels('checkpoints')
|
||||
|
||||
assetService.invalidateModelBuckets()
|
||||
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildAssetListResponse([
|
||||
validAsset({
|
||||
id: 'fresh',
|
||||
name: 'fresh.safetensors',
|
||||
tags: ['models', 'model_type:checkpoints']
|
||||
})
|
||||
])
|
||||
)
|
||||
const freshModels = await assetService.getAssetModels('checkpoints')
|
||||
expect(freshModels.map((m) => m.name)).toEqual(['fresh.safetensors'])
|
||||
|
||||
resolveStaleWalk(
|
||||
buildAssetListResponse([
|
||||
validAsset({
|
||||
id: 'stale',
|
||||
name: 'stale.safetensors',
|
||||
tags: ['models', 'model_type:checkpoints']
|
||||
})
|
||||
])
|
||||
)
|
||||
await staleRead
|
||||
|
||||
const cachedModels = await assetService.getAssetModels('checkpoints')
|
||||
expect(cachedModels.map((m) => m.name)).toEqual(['fresh.safetensors'])
|
||||
expect(fetchApiMock).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
|
||||
it('places multi-category assets in every folder from a single walk', async () => {
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildAssetListResponse([
|
||||
validAsset({
|
||||
id: 'shared',
|
||||
name: 'dual_use.safetensors',
|
||||
loader_path: 'dual_use.safetensors',
|
||||
tags: [
|
||||
'models',
|
||||
'model_type:checkpoints',
|
||||
'model_type:diffusion_models'
|
||||
]
|
||||
})
|
||||
])
|
||||
)
|
||||
|
||||
const checkpoints = await assetService.getAssetModels('checkpoints')
|
||||
const diffusion = await assetService.getAssetModels('diffusion_models')
|
||||
|
||||
expect(checkpoints).toEqual([
|
||||
{ name: 'dual_use.safetensors', pathIndex: 0 }
|
||||
])
|
||||
expect(diffusion).toEqual([{ name: 'dual_use.safetensors', pathIndex: 0 }])
|
||||
// Both folder reads resolve from a single memoized models walk.
|
||||
expect(fetchApiMock).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
})
|
||||
|
||||
describe(assetService.onModelsScanned, () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
it('requests missing-tag exclusion and returns alphabetical unique folders without include_public', async () => {
|
||||
it('invokes the callback when the scan event fires and unsubscribes cleanly', () => {
|
||||
const callback = vi.fn()
|
||||
|
||||
const unsubscribe = assetService.onModelsScanned(callback)
|
||||
|
||||
const [eventType, handler] = vi.mocked(api.addCustomEventListener).mock
|
||||
.calls[0]!
|
||||
expect(eventType).toBe('assets.seed.fast_complete')
|
||||
|
||||
handler!(new CustomEvent(eventType))
|
||||
expect(callback).toHaveBeenCalledOnce()
|
||||
|
||||
unsubscribe()
|
||||
expect(api.removeCustomEventListener).toHaveBeenCalledWith(
|
||||
eventType,
|
||||
handler
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe(assetService.seedModelAssets, () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
it('POSTs the models root to the seed endpoint', async () => {
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildAssetListResponse([
|
||||
validAsset({ id: 'a', tags: ['models', 'loras'] }),
|
||||
validAsset({ id: 'b', tags: ['models', 'checkpoints'] }),
|
||||
validAsset({ id: 'c', tags: ['models', 'configs'] }),
|
||||
validAsset({ id: 'e', tags: ['models', 'loras'] })
|
||||
])
|
||||
buildResponse({ status: 'started' }, { status: 202 })
|
||||
)
|
||||
|
||||
const folders = await assetService.getAssetModelFolders()
|
||||
await assetService.seedModelAssets()
|
||||
|
||||
expect(folders).toEqual([
|
||||
{ name: 'checkpoints', folders: [] },
|
||||
{ name: 'loras', folders: [] }
|
||||
])
|
||||
expect(fetchApiMock).toHaveBeenCalledWith('/assets/seed', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ roots: ['models'] })
|
||||
})
|
||||
})
|
||||
|
||||
const requestedUrl = fetchApiMock.mock.calls[0]?.[0] as string
|
||||
const params = new URL(requestedUrl, 'http://localhost').searchParams
|
||||
expect(params.has('include_public')).toBe(false)
|
||||
expect(params.get('exclude_tags')).toBe(MISSING_TAG)
|
||||
it('treats an already-running scan (409) as success', async () => {
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildResponse({ status: 'already_running' }, { ok: false, status: 409 })
|
||||
)
|
||||
|
||||
await expect(assetService.seedModelAssets()).resolves.toBeUndefined()
|
||||
})
|
||||
|
||||
it('throws on other error statuses', async () => {
|
||||
fetchApiMock.mockResolvedValueOnce(
|
||||
buildResponse({}, { ok: false, status: 500 })
|
||||
)
|
||||
|
||||
await expect(assetService.seedModelAssets()).rejects.toThrow('500')
|
||||
})
|
||||
})
|
||||
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
import { fromZodError } from 'zod-validation-error'
|
||||
import { z } from 'zod'
|
||||
|
||||
import { useFeatureFlags } from '@/composables/useFeatureFlags'
|
||||
import { st } from '@/i18n'
|
||||
|
||||
import {
|
||||
assetFilenameSchema,
|
||||
assetItemSchema,
|
||||
assetResponseSchema,
|
||||
asyncUploadResponseSchema,
|
||||
@@ -17,9 +19,9 @@ import type {
|
||||
AssetUpdatePayload,
|
||||
AsyncUploadResponse,
|
||||
ModelFile,
|
||||
ModelFolder,
|
||||
TagsOperationResult
|
||||
} from '@/platform/assets/schemas/assetSchema'
|
||||
import { getAssetFilename } from '@/platform/assets/utils/assetMetadataUtils'
|
||||
import { isCloud } from '@/platform/distribution/types'
|
||||
import { useSettingStore } from '@/platform/settings/settingStore'
|
||||
import { api } from '@/scripts/api'
|
||||
@@ -180,6 +182,7 @@ function getLocalizedErrorMessage(errorCode: string): string {
|
||||
}
|
||||
|
||||
const ASSETS_ENDPOINT = '/assets'
|
||||
const ASSETS_SEED_ENDPOINT = '/assets/seed'
|
||||
const ASSETS_DOWNLOAD_ENDPOINT = '/assets/download'
|
||||
const ASSETS_EXPORT_ENDPOINT = '/assets/export'
|
||||
const EXPERIMENTAL_WARNING = `EXPERIMENTAL: If you are seeing this please make sure "Comfy.Assets.UseAssetAPI" is set to "false" in your ComfyUI Settings.\n`
|
||||
@@ -187,6 +190,8 @@ const DEFAULT_LIMIT = 500
|
||||
const INPUT_ASSETS_WITH_PUBLIC_LIMIT = 500
|
||||
|
||||
export const MODELS_TAG = 'models'
|
||||
/** Prefix for the namespaced tag that carries a model's folder category, e.g. `model_type:checkpoints`. */
|
||||
const MODEL_TYPE_TAG_PREFIX = 'model_type:'
|
||||
export const INPUT_TAG = 'input'
|
||||
export const OUTPUT_TAG = 'output'
|
||||
/** Asset tag used by the backend for placeholder records that are not installed. */
|
||||
@@ -209,6 +214,48 @@ function normalizeAssetTags(tags: string[]): string[] {
|
||||
return tags.map((tag) => tag.trim()).filter(Boolean)
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves the model folder a tag represents, or undefined when the tag is not
|
||||
* a folder category. `supports_model_type_tags` backends carry the category as
|
||||
* a namespaced `model_type:<folder>` tag; older backends mint bare tags, which
|
||||
* may carry subfolder paths (e.g. `Chatterbox/sub/model`) and group by their
|
||||
* top-level segment, matching the asset browser's legacy grouping.
|
||||
*/
|
||||
function modelFolderFromTag(
|
||||
tag: string,
|
||||
modelTypeMode: boolean
|
||||
): string | undefined {
|
||||
if (modelTypeMode) {
|
||||
return tag.startsWith(MODEL_TYPE_TAG_PREFIX)
|
||||
? tag.slice(MODEL_TYPE_TAG_PREFIX.length)
|
||||
: undefined
|
||||
}
|
||||
if (tag === MODELS_TAG || tag.length === 0) return undefined
|
||||
return tag.split('/')[0]
|
||||
}
|
||||
|
||||
/**
|
||||
* Orders loader paths as subdirectories before files at every level,
|
||||
* alphabetical within each group. The asset API returns models in storage
|
||||
* order, which would otherwise interleave root-level files with folder
|
||||
* contents in the sidebar tree.
|
||||
*/
|
||||
function compareLoaderPaths(a: string, b: string): number {
|
||||
const aSegments = a.split('/')
|
||||
const bSegments = b.split('/')
|
||||
const sharedDepth = Math.min(aSegments.length, bSegments.length)
|
||||
for (let i = 0; i < sharedDepth; i++) {
|
||||
const aIsFile = i === aSegments.length - 1
|
||||
const bIsFile = i === bSegments.length - 1
|
||||
if (aIsFile !== bIsFile) return aIsFile ? 1 : -1
|
||||
const order = aSegments[i].localeCompare(bSegments[i], undefined, {
|
||||
numeric: true
|
||||
})
|
||||
if (order !== 0) return order
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
async function withCallerAbort<T>(
|
||||
promise: Promise<T>,
|
||||
signal?: AbortSignal
|
||||
@@ -269,6 +316,26 @@ function createAssetService() {
|
||||
let inputAssetsIncludingPublicRequestId = 0
|
||||
let pendingInputAssetsIncludingPublic: Promise<AssetItem[]> | null = null
|
||||
|
||||
/**
|
||||
* Model assets bucketed by folder category, built from a single walk of the
|
||||
* `models` tag rather than a fetch per category. Shared by the folder list
|
||||
* and per-folder listings so the sidebar loads every model in one pass.
|
||||
*/
|
||||
let modelBuckets: Map<string, AssetItem[]> | null = null
|
||||
let modelBucketsRequestId = 0
|
||||
let pendingModelBuckets: Promise<Map<string, AssetItem[]>> | null = null
|
||||
|
||||
/**
|
||||
* Discards the cached model buckets so the next read re-walks the models
|
||||
* tag. Bumping the request id keeps a walk that was already in flight from
|
||||
* repopulating the cache with pre-invalidation data.
|
||||
*/
|
||||
function invalidateModelBuckets(): void {
|
||||
modelBucketsRequestId++
|
||||
modelBuckets = null
|
||||
pendingModelBuckets = null
|
||||
}
|
||||
|
||||
/** Invalidates the cached public-inclusive input assets without aborting in-flight readers. */
|
||||
function invalidateInputAssetsIncludingPublic(): void {
|
||||
inputAssetsIncludingPublicRequestId++
|
||||
@@ -330,51 +397,156 @@ function createAssetService() {
|
||||
return validateAssetResponse(data)
|
||||
}
|
||||
/**
|
||||
* Gets a list of model folder keys from the asset API
|
||||
*
|
||||
* Logic:
|
||||
* 1. Extract directory names directly from asset tags
|
||||
* 2. Filter out blacklisted directories
|
||||
* 3. Return alphabetically sorted directories with assets
|
||||
*
|
||||
* @returns The list of model folder keys
|
||||
* Walks every `models`-tagged asset once and buckets each into the folder
|
||||
* categories carried by its `model_type:` tags. A single asset lands in every
|
||||
* category it is tagged with (e.g. a shared-root model in both `checkpoints`
|
||||
* and `diffusion_models`). Which folders are actually shown is decided by
|
||||
* `/experiment/models`; models with no category tag are dropped with a warning
|
||||
* rather than hidden silently.
|
||||
*/
|
||||
async function getAssetModelFolders(): Promise<ModelFolder[]> {
|
||||
const data = await handleAssetRequest(
|
||||
{ includeTags: [MODELS_TAG] },
|
||||
'model folders'
|
||||
)
|
||||
async function buildModelBuckets(): Promise<Map<string, AssetItem[]>> {
|
||||
const assets = await getAllAssetsByTag(MODELS_TAG, true)
|
||||
const modelTypeMode = useFeatureFlags().flags.supportsModelTypeTags
|
||||
const buckets = new Map<string, AssetItem[]>()
|
||||
|
||||
// Blacklist directories we don't want to show
|
||||
const blacklistedDirectories = new Set(['configs'])
|
||||
for (const asset of assets) {
|
||||
const folders = asset.tags
|
||||
.map((tag) => modelFolderFromTag(tag, modelTypeMode))
|
||||
.filter((folder): folder is string => folder !== undefined)
|
||||
|
||||
const folderTags = data.assets
|
||||
.flatMap((asset) => asset.tags)
|
||||
.filter((tag) => tag !== MODELS_TAG && !blacklistedDirectories.has(tag))
|
||||
const discoveredFolders = new Set<string>(folderTags)
|
||||
if (folders.length === 0) {
|
||||
console.warn(
|
||||
`Asset ${asset.id} (${asset.name}) is tagged '${MODELS_TAG}' but has no model category; skipping.`
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
// Return only discovered folders in alphabetical order
|
||||
const sortedFolders = Array.from(discoveredFolders).toSorted()
|
||||
return sortedFolders.map((name) => ({ name, folders: [] }))
|
||||
// On loader_path-contract backends a null loader_path marks an
|
||||
// unloadable asset (e.g. an orphan): it must not mint a widget value,
|
||||
// and `name` is deprecated for path semantics.
|
||||
if (modelTypeMode && !asset.loader_path) {
|
||||
console.warn(
|
||||
`Asset ${asset.id} (${asset.name}) has no loader_path; skipping.`
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
// The loader value flows into viewMetadata URLs and widget values, so a
|
||||
// traversal-shaped path must not pass through even if the backend's own
|
||||
// validation ever regresses.
|
||||
const loaderValue = asset.loader_path ?? getAssetFilename(asset)
|
||||
if (!assetFilenameSchema.safeParse(loaderValue).success) {
|
||||
console.warn(
|
||||
`Asset ${asset.id} (${asset.name}) has an unsafe loader path ('${loaderValue}'); skipping.`
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
for (const folder of folders) {
|
||||
const bucket = buckets.get(folder)
|
||||
if (bucket) bucket.push(asset)
|
||||
else buckets.set(folder, [asset])
|
||||
}
|
||||
}
|
||||
|
||||
for (const bucket of buckets.values()) {
|
||||
bucket.sort((a, b) =>
|
||||
compareLoaderPaths(
|
||||
a.loader_path ?? getAssetFilename(a),
|
||||
b.loader_path ?? getAssetFilename(b)
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
return buckets
|
||||
}
|
||||
|
||||
/** Returns the memoized model buckets, walking the models tag on first read. */
|
||||
async function loadModelBuckets(): Promise<Map<string, AssetItem[]>> {
|
||||
if (modelBuckets) return modelBuckets
|
||||
if (pendingModelBuckets) return pendingModelBuckets
|
||||
|
||||
const requestId = ++modelBucketsRequestId
|
||||
const walk = async () => {
|
||||
try {
|
||||
const buckets = await buildModelBuckets()
|
||||
if (requestId === modelBucketsRequestId) {
|
||||
modelBuckets = buckets
|
||||
}
|
||||
return buckets
|
||||
} finally {
|
||||
if (requestId === modelBucketsRequestId) {
|
||||
pendingModelBuckets = null
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pendingModelBuckets = walk()
|
||||
return pendingModelBuckets
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets a list of models in the specified folder from the asset API
|
||||
* Gets the models in the specified folder from the single models walk.
|
||||
* @param folder The folder to list models from, such as 'checkpoints'
|
||||
* @returns The list of model filenames within the specified folder
|
||||
*/
|
||||
async function getAssetModels(folder: string): Promise<ModelFile[]> {
|
||||
const data = await handleAssetRequest(
|
||||
{ includeTags: [MODELS_TAG, folder] },
|
||||
`models for ${folder}`
|
||||
)
|
||||
|
||||
return data.assets.map((asset) => ({
|
||||
name: asset.name,
|
||||
const buckets = await loadModelBuckets()
|
||||
return (buckets.get(folder) ?? []).map((asset) => ({
|
||||
// `loader_path` is the category-relative path the loader widget expects
|
||||
// and the source for the sidebar tree. Backends that predate it (bare-tag
|
||||
// mode; today's cloud) fall back to the filename metadata — the same
|
||||
// value the asset browser serializes — rather than `name`, which is a
|
||||
// content hash on cloud.
|
||||
name: asset.loader_path ?? getAssetFilename(asset),
|
||||
// Asset records carry no root identity, so every model reports root 0.
|
||||
// Known limitation on multi-root categories (extra_model_paths.yaml):
|
||||
// preview reads target root 0 (wrong file or 404 for secondary-root
|
||||
// files), and same-relative-path files in different roots collapse
|
||||
// onto one sidebar row. Metadata is unaffected unless relative paths
|
||||
// collide (/view_metadata searches roots in order without an index),
|
||||
// as are loader widget values; lifting this needs the backend to carry
|
||||
// root identity on assets.
|
||||
pathIndex: 0
|
||||
}))
|
||||
}
|
||||
|
||||
/**
|
||||
* Asks the backend to rescan the model roots on disk so newly added files
|
||||
* become assets. Fire-and-forget: the scan's fast (insert) phase already
|
||||
* writes the category tags and filenames the sidebar needs and is announced
|
||||
* by an `assets.seed.fast_complete` websocket event. A 409 means a scan is
|
||||
* already running, which will emit the same event, so it is not an error.
|
||||
*/
|
||||
async function seedModelAssets(): Promise<void> {
|
||||
const res = await api.fetchApi(ASSETS_SEED_ENDPOINT, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ roots: ['models'] })
|
||||
})
|
||||
if (!res.ok && res.status !== 409) {
|
||||
throw new Error(
|
||||
`Unable to start asset scan: Server returned ${res.status}`
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Subscribes to the backend's scan fast-phase completion broadcast — the
|
||||
* moment newly scanned files' tags and loader paths become queryable. The
|
||||
* wire-level event (`assets.seed.fast_complete`) is owned here; consumers
|
||||
* receive a callback and an unsubscribe function.
|
||||
*/
|
||||
function onModelsScanned(callback: () => void | Promise<void>): () => void {
|
||||
const handler = () => {
|
||||
void callback()
|
||||
}
|
||||
api.addCustomEventListener('assets.seed.fast_complete', handler)
|
||||
return () => {
|
||||
api.removeCustomEventListener('assets.seed.fast_complete', handler)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Checks if a widget input should use the asset browser based on both input name and node comfyClass
|
||||
*
|
||||
@@ -971,8 +1143,10 @@ function createAssetService() {
|
||||
}
|
||||
|
||||
return {
|
||||
getAssetModelFolders,
|
||||
getAssetModels,
|
||||
invalidateModelBuckets,
|
||||
onModelsScanned,
|
||||
seedModelAssets,
|
||||
isAssetAPIEnabled,
|
||||
isAssetBrowserEligible,
|
||||
shouldUseAssetBrowser,
|
||||
|
||||
@@ -45,7 +45,7 @@ describe('assetFilterUtils properties', () => {
|
||||
it('filterByCategory("all") accepts every asset', () => {
|
||||
fc.assert(
|
||||
fc.property(arbAssetItem, (asset) => {
|
||||
expect(filterByCategory('all')(asset)).toBe(true)
|
||||
expect(filterByCategory('all', false)(asset)).toBe(true)
|
||||
})
|
||||
)
|
||||
})
|
||||
@@ -56,7 +56,7 @@ describe('assetFilterUtils properties', () => {
|
||||
fc.array(arbAssetItem, { maxLength: 30 }),
|
||||
fc.stringMatching(/^[a-z]{1,8}$/),
|
||||
(assets, category) => {
|
||||
const filter = filterByCategory(category)
|
||||
const filter = filterByCategory(category, false)
|
||||
const result = assets.filter(filter)
|
||||
expect(result.length).toBeLessThanOrEqual(assets.length)
|
||||
for (const item of result) {
|
||||
|
||||
@@ -26,19 +26,54 @@ function createAsset(
|
||||
|
||||
describe('filterByCategory', () => {
|
||||
it.for([
|
||||
{ category: 'all', tags: ['checkpoint'], expected: true },
|
||||
{ category: 'checkpoint', tags: ['checkpoint'], expected: true },
|
||||
{ category: 'lora', tags: ['checkpoint'], expected: false },
|
||||
{ category: 'all', tags: ['checkpoint'], mode: false, expected: true },
|
||||
{
|
||||
category: 'checkpoint',
|
||||
tags: ['checkpoint'],
|
||||
mode: false,
|
||||
expected: true
|
||||
},
|
||||
{ category: 'lora', tags: ['checkpoint'], mode: false, expected: false },
|
||||
{
|
||||
category: 'checkpoint',
|
||||
tags: ['models', 'checkpoint/xl'],
|
||||
mode: false,
|
||||
expected: true
|
||||
},
|
||||
{ category: 'xl', tags: ['models', 'checkpoint/xl'], expected: false }
|
||||
{
|
||||
category: 'xl',
|
||||
tags: ['models', 'checkpoint/xl'],
|
||||
mode: false,
|
||||
expected: false
|
||||
},
|
||||
{
|
||||
category: 'checkpoints',
|
||||
tags: ['models', 'model_type:checkpoints'],
|
||||
mode: true,
|
||||
expected: true
|
||||
},
|
||||
{
|
||||
category: 'LLM',
|
||||
tags: ['models', 'model_type:LLM'],
|
||||
mode: true,
|
||||
expected: true
|
||||
},
|
||||
{
|
||||
category: 'sdxl',
|
||||
tags: ['models', 'model_type:checkpoints', 'sdxl'],
|
||||
mode: true,
|
||||
expected: false
|
||||
},
|
||||
{
|
||||
category: 'checkpoints',
|
||||
tags: ['models', 'model_type:checkpoints'],
|
||||
mode: false,
|
||||
expected: false
|
||||
}
|
||||
])(
|
||||
'category=$category with tags=$tags returns $expected',
|
||||
({ category, tags, expected }) => {
|
||||
const filter = filterByCategory(category)
|
||||
'category=$category tags=$tags mode=$mode returns $expected',
|
||||
({ category, tags, mode, expected }) => {
|
||||
const filter = filterByCategory(category, mode)
|
||||
const asset = createAsset('model.safetensors', { tags })
|
||||
expect(filter(asset)).toBe(expected)
|
||||
}
|
||||
|
||||
@@ -1,21 +1,14 @@
|
||||
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
|
||||
import type { OwnershipOption } from '@/platform/assets/types/filterTypes'
|
||||
import { getAssetBaseModels } from '@/platform/assets/utils/assetMetadataUtils'
|
||||
import {
|
||||
getAssetBaseModels,
|
||||
getAssetCategories
|
||||
} from '@/platform/assets/utils/assetMetadataUtils'
|
||||
|
||||
export function filterByCategory(category: string) {
|
||||
export function filterByCategory(category: string, modelTypeMode: boolean) {
|
||||
return (asset: AssetItem) => {
|
||||
if (category === 'all') return true
|
||||
|
||||
// Check if any tag matches the category (for exact matches)
|
||||
if (asset.tags.includes(category)) return true
|
||||
|
||||
// Check if any tag's top-level folder matches the category
|
||||
return asset.tags.some((tag) => {
|
||||
if (typeof tag === 'string' && tag.includes('/')) {
|
||||
return tag.split('/')[0] === category
|
||||
}
|
||||
return false
|
||||
})
|
||||
return getAssetCategories(asset, modelTypeMode).includes(category)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2,22 +2,34 @@ import { afterEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
|
||||
import {
|
||||
MISSING_TAG,
|
||||
MODELS_TAG
|
||||
} from '@/platform/assets/services/assetService'
|
||||
import {
|
||||
buildModelTypeTagUpdate,
|
||||
getAssetAdditionalTags,
|
||||
getAssetBaseModel,
|
||||
getAssetBaseModels,
|
||||
getAssetCardTitle,
|
||||
getAssetCategories,
|
||||
getAssetDescription,
|
||||
getAssetDisplayFilename,
|
||||
getAssetDisplayName,
|
||||
getAssetFilename,
|
||||
getAssetMetadataDimensions,
|
||||
getAssetModelType,
|
||||
getAssetNodeCategoryCandidates,
|
||||
getAssetSourceUrl,
|
||||
getPrimaryCategoryTag,
|
||||
getAssetStoredFilename,
|
||||
getAssetTriggerPhrases,
|
||||
getAssetTypeBadges,
|
||||
getAssetUserDescription,
|
||||
getEditableModelType,
|
||||
getSourceName,
|
||||
resolveDisplayImageDimensions
|
||||
resolveDisplayImageDimensions,
|
||||
stripModelTypePrefix,
|
||||
toModelTypeTag
|
||||
} from '@/platform/assets/utils/assetMetadataUtils'
|
||||
|
||||
const { isCloudRef } = vi.hoisted(() => ({
|
||||
@@ -274,7 +286,17 @@ describe('assetMetadataUtils', () => {
|
||||
tags: ['models'],
|
||||
expected: null
|
||||
},
|
||||
{ name: 'returns null when tags empty', tags: [], expected: null }
|
||||
{ name: 'returns null when tags empty', tags: [], expected: null },
|
||||
{
|
||||
name: 'never returns a raw model_type: literal (no round-trip into edit widgets)',
|
||||
tags: ['models', 'model_type:checkpoints'],
|
||||
expected: null
|
||||
},
|
||||
{
|
||||
name: 'skips model_type: tags in favour of the bare twin',
|
||||
tags: ['models', 'model_type:checkpoints', 'checkpoints'],
|
||||
expected: 'checkpoints'
|
||||
}
|
||||
])('$name', ({ tags, expected }) => {
|
||||
const asset = { ...mockAsset, tags }
|
||||
expect(getAssetModelType(asset)).toBe(expected)
|
||||
@@ -540,3 +562,353 @@ describe('assetMetadataUtils', () => {
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe('getAssetCategories', () => {
|
||||
const asset = (tags: string[]): AssetItem => ({
|
||||
id: 'a',
|
||||
name: 'model.safetensors',
|
||||
tags
|
||||
})
|
||||
|
||||
it('uses model_type:* values as the group and disregards other tags in model_type mode', () => {
|
||||
expect(
|
||||
getAssetCategories(
|
||||
asset(['models', 'model_type:checkpoints', 'sdxl']),
|
||||
true
|
||||
)
|
||||
).toEqual(['checkpoints'])
|
||||
})
|
||||
|
||||
it('preserves the model_type value casing', () => {
|
||||
expect(
|
||||
getAssetCategories(asset(['models', 'model_type:LLM']), true)
|
||||
).toEqual(['LLM'])
|
||||
})
|
||||
|
||||
it('routes an uncovered asset by its bare tags in model_type mode', () => {
|
||||
expect(getAssetCategories(asset(['models', 'checkpoints']), true)).toEqual([
|
||||
'checkpoints'
|
||||
])
|
||||
})
|
||||
|
||||
it('ignores model_type: and uses bare-tag grouping when mode is off', () => {
|
||||
expect(
|
||||
getAssetCategories(
|
||||
asset(['models', 'model_type:checkpoints', 'sdxl']),
|
||||
false
|
||||
)
|
||||
).toEqual(['model_type:checkpoints', 'sdxl'])
|
||||
})
|
||||
|
||||
it('never surfaces namespace residue as a category for uncovered assets in mode', () => {
|
||||
expect(
|
||||
getAssetCategories(asset(['models', 'model_type:', 'sdxl']), true)
|
||||
).toEqual(['sdxl'])
|
||||
})
|
||||
})
|
||||
|
||||
describe('getPrimaryCategoryTag', () => {
|
||||
const asset = (tags: string[]): AssetItem => ({
|
||||
id: 'a',
|
||||
name: 'model.safetensors',
|
||||
tags
|
||||
})
|
||||
|
||||
it('uses the model_type value a covered asset groups under', () => {
|
||||
expect(
|
||||
getPrimaryCategoryTag(asset(['models', 'sdxl', 'model_type:vae']), true)
|
||||
).toBe('vae')
|
||||
})
|
||||
|
||||
it('keeps the legacy verbatim tag for an uncovered hierarchical asset', () => {
|
||||
expect(
|
||||
getPrimaryCategoryTag(asset(['models', 'Chatterbox/sub/model']), true)
|
||||
).toBe('Chatterbox/sub/model')
|
||||
})
|
||||
|
||||
it('skips namespace residue instead of titling off a raw model_type: tag', () => {
|
||||
expect(
|
||||
getPrimaryCategoryTag(asset(['models', 'model_type:']), true)
|
||||
).toBeUndefined()
|
||||
})
|
||||
|
||||
it('returns the legacy first non-models tag when mode is off', () => {
|
||||
expect(
|
||||
getPrimaryCategoryTag(asset(['models', 'model_type:vae']), false)
|
||||
).toBe('model_type:vae')
|
||||
})
|
||||
})
|
||||
|
||||
describe('getAssetNodeCategoryCandidates', () => {
|
||||
const asset = (tags: string[]): AssetItem => ({
|
||||
id: 'a',
|
||||
name: 'model.safetensors',
|
||||
tags
|
||||
})
|
||||
|
||||
it('orders the most specific (deepest) tag ahead of a flat model_type value', () => {
|
||||
expect(
|
||||
getAssetNodeCategoryCandidates(
|
||||
asset(['models', 'model_type:LLM', 'LLM/Qwen-VL/Qwen3-0.6B']),
|
||||
true
|
||||
)
|
||||
).toEqual(['LLM/Qwen-VL/Qwen3-0.6B', 'LLM'])
|
||||
})
|
||||
|
||||
it('strips the model_type: prefix when it is the only candidate', () => {
|
||||
expect(
|
||||
getAssetNodeCategoryCandidates(asset(['models', 'model_type:vae']), true)
|
||||
).toEqual(['vae'])
|
||||
})
|
||||
|
||||
it('keeps a model_type value ahead of an equally-deep bare tag', () => {
|
||||
expect(
|
||||
getAssetNodeCategoryCandidates(
|
||||
asset(['models', 'model_type:checkpoints', 'sdxl']),
|
||||
true
|
||||
)
|
||||
).toEqual(['checkpoints', 'sdxl'])
|
||||
})
|
||||
|
||||
it('demotes an unrelated deeper bare tag below the model_type value', () => {
|
||||
expect(
|
||||
getAssetNodeCategoryCandidates(
|
||||
asset(['models', 'model_type:vae', 'foo/bar']),
|
||||
true
|
||||
)
|
||||
).toEqual(['vae', 'foo/bar'])
|
||||
})
|
||||
|
||||
it('keeps unrelated bare tags as trailing fallbacks rather than dropping them', () => {
|
||||
expect(
|
||||
getAssetNodeCategoryCandidates(
|
||||
asset(['models', 'model_type:LLM', 'LLM/Qwen-VL', 'foo/bar/baz']),
|
||||
true
|
||||
)
|
||||
).toEqual(['LLM/Qwen-VL', 'LLM', 'foo/bar/baz'])
|
||||
})
|
||||
|
||||
it('keeps a hierarchical tag intact', () => {
|
||||
expect(
|
||||
getAssetNodeCategoryCandidates(
|
||||
asset(['models', 'chatterbox/chatterbox_vc']),
|
||||
true
|
||||
)
|
||||
).toEqual(['chatterbox/chatterbox_vc'])
|
||||
})
|
||||
|
||||
it('returns no candidates when only reserved tags are present', () => {
|
||||
expect(
|
||||
getAssetNodeCategoryCandidates(asset(['models', 'missing']), true)
|
||||
).toEqual([])
|
||||
})
|
||||
|
||||
it('uses the first non-reserved tag verbatim when mode is off', () => {
|
||||
expect(
|
||||
getAssetNodeCategoryCandidates(asset(['models', 'model_type:vae']), false)
|
||||
).toEqual(['model_type:vae'])
|
||||
expect(
|
||||
getAssetNodeCategoryCandidates(asset(['models', 'checkpoints']), false)
|
||||
).toEqual(['checkpoints'])
|
||||
})
|
||||
})
|
||||
|
||||
describe('getAssetTypeBadges', () => {
|
||||
const asset = (tags: string[]): AssetItem => ({
|
||||
id: 'a',
|
||||
name: 'model.safetensors',
|
||||
tags
|
||||
})
|
||||
|
||||
it('strips the model_type: prefix in model_type mode (no raw leak)', () => {
|
||||
expect(
|
||||
getAssetTypeBadges(
|
||||
asset(['models', 'model_type:checkpoints', 'sdxl']),
|
||||
true
|
||||
)
|
||||
).toEqual(['checkpoints'])
|
||||
})
|
||||
|
||||
it('badges the model_type value even when a bare tag comes first, matching the grouping', () => {
|
||||
expect(
|
||||
getAssetTypeBadges(asset(['models', 'foo', 'model_type:bar']), true)
|
||||
).toEqual(['bar'])
|
||||
})
|
||||
|
||||
it('badges every category a shared multi-type asset groups under', () => {
|
||||
expect(
|
||||
getAssetTypeBadges(
|
||||
asset([
|
||||
'models',
|
||||
'model_type:checkpoints',
|
||||
'model_type:diffusion_models'
|
||||
]),
|
||||
true
|
||||
)
|
||||
).toEqual(['checkpoints', 'diffusion_models'])
|
||||
})
|
||||
|
||||
it('falls back to the bare tag for an uncovered asset in model_type mode', () => {
|
||||
expect(getAssetTypeBadges(asset(['models', 'sdxl']), true)).toEqual([
|
||||
'sdxl'
|
||||
])
|
||||
})
|
||||
|
||||
it('returns no badge rather than a blank one for a malformed empty model_type: tag', () => {
|
||||
expect(getAssetTypeBadges(asset(['models', 'model_type:']), true)).toEqual(
|
||||
[]
|
||||
)
|
||||
})
|
||||
|
||||
it('leaks the literal model_type: tag when mode is off', () => {
|
||||
expect(
|
||||
getAssetTypeBadges(asset(['models', 'model_type:checkpoints']), false)
|
||||
).toEqual(['model_type:checkpoints'])
|
||||
})
|
||||
|
||||
it('shows the segment after the first slash for a bare hierarchical tag', () => {
|
||||
expect(
|
||||
getAssetTypeBadges(asset(['models', 'checkpoint/xl']), false)
|
||||
).toEqual(['xl'])
|
||||
})
|
||||
|
||||
it('returns no badges when only the models tag is present', () => {
|
||||
expect(getAssetTypeBadges(asset(['models']), true)).toEqual([])
|
||||
})
|
||||
})
|
||||
|
||||
describe('stripModelTypePrefix', () => {
|
||||
it('removes the model_type: prefix when present', () => {
|
||||
expect(stripModelTypePrefix('model_type:checkpoints')).toBe('checkpoints')
|
||||
})
|
||||
|
||||
it('leaves a tag without the prefix unchanged', () => {
|
||||
expect(stripModelTypePrefix('checkpoints')).toBe('checkpoints')
|
||||
expect(stripModelTypePrefix('checkpoint/xl')).toBe('checkpoint/xl')
|
||||
})
|
||||
})
|
||||
|
||||
describe('toModelTypeTag', () => {
|
||||
it('prefixes a folder_name with the model_type namespace', () => {
|
||||
expect(toModelTypeTag('checkpoints')).toBe('model_type:checkpoints')
|
||||
expect(toModelTypeTag('ultralytics_bbox')).toBe(
|
||||
'model_type:ultralytics_bbox'
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe('getEditableModelType', () => {
|
||||
const asset = (tags: string[]): AssetItem => ({
|
||||
id: 'a',
|
||||
name: 'model.safetensors',
|
||||
tags
|
||||
})
|
||||
|
||||
it('returns the stripped model_type value in model_type mode', () => {
|
||||
expect(
|
||||
getEditableModelType(
|
||||
asset(['models', 'checkpoints', 'model_type:checkpoints']),
|
||||
true
|
||||
)
|
||||
).toBe('checkpoints')
|
||||
})
|
||||
|
||||
it('falls back to the bare tag for an uncovered asset in model_type mode', () => {
|
||||
expect(getEditableModelType(asset(['models', 'sam2']), true)).toBe('sam2')
|
||||
})
|
||||
|
||||
it('uses the legacy first-non-models tag when mode is off', () => {
|
||||
expect(
|
||||
getEditableModelType(asset(['models', 'checkpoints', 'sdxl']), false)
|
||||
).toBe('checkpoints')
|
||||
})
|
||||
|
||||
it('returns null when only the models tag is present', () => {
|
||||
expect(getEditableModelType(asset(['models']), true)).toBeNull()
|
||||
})
|
||||
})
|
||||
|
||||
describe('buildModelTypeTagUpdate', () => {
|
||||
const asset = (tags: string[]): AssetItem => ({
|
||||
id: 'a',
|
||||
name: 'model.safetensors',
|
||||
tags
|
||||
})
|
||||
|
||||
it('swaps the bare subtype tag when mode is off', () => {
|
||||
expect(
|
||||
buildModelTypeTagUpdate(asset(['models', 'checkpoints']), 'loras', false)
|
||||
).toEqual(['models', 'loras'])
|
||||
})
|
||||
|
||||
it('preserves user labels and swaps only the subtype tag when mode is off', () => {
|
||||
expect(
|
||||
buildModelTypeTagUpdate(
|
||||
asset(['models', 'checkpoints', 'sdxl']),
|
||||
'loras',
|
||||
false
|
||||
)
|
||||
).toEqual(['models', 'sdxl', 'loras'])
|
||||
})
|
||||
|
||||
it('writes only the model_type form for a covered asset, leaving the bare twin for the backend', () => {
|
||||
expect(
|
||||
buildModelTypeTagUpdate(
|
||||
asset(['models', 'checkpoints', 'model_type:checkpoints']),
|
||||
'loras',
|
||||
true
|
||||
)
|
||||
).toEqual(['models', 'checkpoints', 'model_type:loras'])
|
||||
})
|
||||
|
||||
it('replaces every existing model_type form for a shared-path dual-tagged asset', () => {
|
||||
expect(
|
||||
buildModelTypeTagUpdate(
|
||||
asset([
|
||||
'models',
|
||||
'diffusion_models',
|
||||
'model_type:diffusion_models',
|
||||
'model_type:unet_gguf'
|
||||
]),
|
||||
'loras',
|
||||
true
|
||||
)
|
||||
).toEqual(['models', 'diffusion_models', 'model_type:loras'])
|
||||
})
|
||||
|
||||
it('drops the bare current type for an uncovered asset in model_type mode', () => {
|
||||
expect(
|
||||
buildModelTypeTagUpdate(asset(['models', 'sam2']), 'loras', true)
|
||||
).toEqual(['models', 'model_type:loras'])
|
||||
})
|
||||
|
||||
it('keeps user labels untouched in model_type mode', () => {
|
||||
expect(
|
||||
buildModelTypeTagUpdate(
|
||||
asset(['models', 'checkpoints', 'model_type:checkpoints', 'sdxl']),
|
||||
'loras',
|
||||
true
|
||||
)
|
||||
).toEqual(['models', 'checkpoints', 'sdxl', 'model_type:loras'])
|
||||
})
|
||||
})
|
||||
|
||||
describe('reserved tag mirrors', () => {
|
||||
const asset = (tags: string[]): AssetItem => ({
|
||||
id: 'a',
|
||||
name: 'model.safetensors',
|
||||
tags
|
||||
})
|
||||
|
||||
it("treats assetService's canonical reserved tags as reserved (locals must not drift)", () => {
|
||||
expect(getAssetCategories(asset([MODELS_TAG, 'x']), false)).toEqual(['x'])
|
||||
expect(
|
||||
getAssetNodeCategoryCandidates(
|
||||
asset([MODELS_TAG, MISSING_TAG, 'x']),
|
||||
true
|
||||
)
|
||||
).toEqual(['x'])
|
||||
expect(getAssetTypeBadges(asset([MODELS_TAG, 'x']), false)).toEqual(['x'])
|
||||
expect(getAssetModelType(asset([MODELS_TAG]))).toBeNull()
|
||||
})
|
||||
})
|
||||
|
||||
@@ -2,6 +2,11 @@ import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
|
||||
import { isCloud } from '@/platform/distribution/types'
|
||||
import { isCivitaiUrl } from '@/utils/formatUtil'
|
||||
|
||||
// Reserved tag literals (mirror assetService's MODELS_TAG/MISSING_TAG). Kept
|
||||
// local so this leaf util doesn't pull the heavier assetService -> i18n chain.
|
||||
const MODELS_TAG = 'models'
|
||||
const MISSING_TAG = 'missing'
|
||||
|
||||
/**
|
||||
* Type-safe utilities for extracting metadata from assets.
|
||||
* These utilities check user_metadata first, then metadata, then fallback.
|
||||
@@ -140,16 +145,236 @@ export function getSourceName(url: string): string {
|
||||
return 'Source'
|
||||
}
|
||||
|
||||
export const MODEL_TYPE_TAG_PREFIX = 'model_type:'
|
||||
|
||||
/**
|
||||
* Extracts the model type from asset tags
|
||||
* Extracts the model type from asset tags as a bare (non-namespaced) value.
|
||||
* Never returns a raw `model_type:*` literal: this value feeds edit widgets
|
||||
* whose save path writes tags back verbatim, so a namespaced tag leaking
|
||||
* through here would round-trip the prefixed literal into the tag set.
|
||||
* @param asset - The asset to extract model type from
|
||||
* @returns The model type string or null if not present
|
||||
*/
|
||||
export function getAssetModelType(asset: AssetItem): string | null {
|
||||
const typeTag = asset.tags?.find((tag) => tag && tag !== 'models')
|
||||
const typeTag = asset.tags?.find(
|
||||
(tag) => tag && tag !== MODELS_TAG && !tag.startsWith(MODEL_TYPE_TAG_PREFIX)
|
||||
)
|
||||
return typeTag ?? null
|
||||
}
|
||||
|
||||
/**
|
||||
* Builds the namespaced subtype tag the backend stores in `model_type:` mode.
|
||||
* The argument is a discovery folder_name (e.g. `checkpoints`,
|
||||
* `ultralytics_bbox`); the backend keeps the bare directory-path twin in sync.
|
||||
*/
|
||||
export function toModelTypeTag(folderName: string): string {
|
||||
return `${MODEL_TYPE_TAG_PREFIX}${folderName}`
|
||||
}
|
||||
|
||||
/** Strips the `model_type:` prefix off each namespaced tag, dropping non-`model_type:` tags. */
|
||||
function getModelTypeTagValues(asset: AssetItem): string[] {
|
||||
return asset.tags
|
||||
.filter((tag) => tag.startsWith(MODEL_TYPE_TAG_PREFIX))
|
||||
.map((tag) => tag.slice(MODEL_TYPE_TAG_PREFIX.length))
|
||||
.filter((tag) => tag.length > 0)
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves the folder_name shown as the asset's current model type in the edit
|
||||
* dropdown. In `modelTypeMode` the stripped `model_type:` value is authoritative
|
||||
* (covered assets); an uncovered asset with no `model_type:` tag falls back to
|
||||
* its bare subtype tag, mirroring the read-side grouping. Outside the mode this
|
||||
* is the legacy first-non-`models` tag.
|
||||
*/
|
||||
export function getEditableModelType(
|
||||
asset: AssetItem,
|
||||
modelTypeMode: boolean
|
||||
): string | null {
|
||||
if (modelTypeMode) {
|
||||
const [modelType] = getModelTypeTagValues(asset)
|
||||
if (modelType) return modelType
|
||||
}
|
||||
return getAssetModelType(asset)
|
||||
}
|
||||
|
||||
/**
|
||||
* Computes the tag set for re-typing a model asset to `newFolderName`. In
|
||||
* `modelTypeMode` only the `model_type:` form is written — the backend keeps the
|
||||
* bare directory-path twin in sync, so existing `model_type:` tags are dropped
|
||||
* (covered assets) or the bare current type is dropped (uncovered assets) and
|
||||
* the new `model_type:<folder_name>` is added. Outside the mode it swaps the
|
||||
* legacy bare subtype tag, preserving the pre-namespace behavior.
|
||||
*/
|
||||
export function buildModelTypeTagUpdate(
|
||||
asset: AssetItem,
|
||||
newFolderName: string,
|
||||
modelTypeMode: boolean
|
||||
): string[] {
|
||||
if (!modelTypeMode) {
|
||||
const currentType = getAssetModelType(asset)
|
||||
return asset.tags.filter((tag) => tag !== currentType).concat(newFolderName)
|
||||
}
|
||||
|
||||
const modelTypeTags = asset.tags.filter((tag) =>
|
||||
tag.startsWith(MODEL_TYPE_TAG_PREFIX)
|
||||
)
|
||||
const currentBareType = getAssetModelType(asset)
|
||||
const tagsToRemove =
|
||||
modelTypeTags.length > 0
|
||||
? new Set(modelTypeTags)
|
||||
: new Set(currentBareType ? [currentBareType] : [])
|
||||
|
||||
return asset.tags
|
||||
.filter((tag) => !tagsToRemove.has(tag))
|
||||
.concat(toModelTypeTag(newFolderName))
|
||||
}
|
||||
|
||||
/** Legacy grouping: each non-`models` tag's top-level path segment. */
|
||||
function getBareTagCategories(asset: AssetItem): string[] {
|
||||
return asset.tags
|
||||
.filter((tag) => tag !== MODELS_TAG && tag.length > 0)
|
||||
.map((tag) => tag.split('/')[0])
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves the category keys a model asset is grouped under.
|
||||
*
|
||||
* `modelTypeMode` reflects whether the backend declares the `model_type:` tag
|
||||
* scheme (the `supports_model_type_tags` capability). When true, an asset's
|
||||
* `model_type:*` values are authoritative; an asset with no `model_type:` tag
|
||||
* still routes by its bare tags. When false (the default) categories come from
|
||||
* the legacy bare-tag top-level grouping and `model_type:` is ignored.
|
||||
*/
|
||||
export function getAssetCategories(
|
||||
asset: AssetItem,
|
||||
modelTypeMode: boolean
|
||||
): string[] {
|
||||
if (modelTypeMode) {
|
||||
const modelTypes = getModelTypeTagValues(asset)
|
||||
if (modelTypes.length > 0) return modelTypes
|
||||
// Uncovered assets route by bare tags, but namespace residue (e.g. a
|
||||
// malformed empty `model_type:`) must not surface as a raw category.
|
||||
return getBareTagCategories(asset).filter(
|
||||
(category) => !category.startsWith(MODEL_TYPE_TAG_PREFIX)
|
||||
)
|
||||
}
|
||||
|
||||
return getBareTagCategories(asset)
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves the primary tag a browser surface titles itself after. In
|
||||
* `modelTypeMode` a covered asset uses its first `model_type:*` value — the
|
||||
* key it groups under — while an uncovered asset keeps the legacy selection
|
||||
* (first non-`models` tag, verbatim, hierarchical paths intact). Outside the
|
||||
* mode this is exactly the legacy selection.
|
||||
*/
|
||||
export function getPrimaryCategoryTag(
|
||||
asset: AssetItem,
|
||||
modelTypeMode: boolean
|
||||
): string | undefined {
|
||||
if (modelTypeMode) {
|
||||
const [modelType] = getModelTypeTagValues(asset)
|
||||
if (modelType) return modelType
|
||||
return asset.tags.find(
|
||||
(tag) => tag !== MODELS_TAG && !tag.startsWith(MODEL_TYPE_TAG_PREFIX)
|
||||
)
|
||||
}
|
||||
return asset.tags.find((tag) => tag !== MODELS_TAG)
|
||||
}
|
||||
|
||||
/** Number of `parent/child` segments in a tag, used to pick the most specific. */
|
||||
function pathDepth(tag: string): number {
|
||||
return tag.split('/').length
|
||||
}
|
||||
|
||||
/** Removes the `model_type:` namespace prefix from a tag when present. */
|
||||
export function stripModelTypePrefix(tag: string): string {
|
||||
return tag.startsWith(MODEL_TYPE_TAG_PREFIX)
|
||||
? tag.slice(MODEL_TYPE_TAG_PREFIX.length)
|
||||
: tag
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves the labels shown as an asset card's type badges.
|
||||
*
|
||||
* In `modelTypeMode` a covered asset badges every `model_type:*` value — the
|
||||
* same keys it groups under (`getAssetCategories`) — so a shared-root asset
|
||||
* tagged with several categories carries each of them; whichever category
|
||||
* view the card appears in is represented on the card. Uncovered assets (and
|
||||
* legacy mode) keep the original single selection: first non-`models` tag,
|
||||
* with bare hierarchical tags showing the segment after the first `/`.
|
||||
*/
|
||||
export function getAssetTypeBadges(
|
||||
asset: AssetItem,
|
||||
modelTypeMode: boolean
|
||||
): string[] {
|
||||
if (modelTypeMode) {
|
||||
const modelTypes = getModelTypeTagValues(asset)
|
||||
if (modelTypes.length > 0) return modelTypes
|
||||
}
|
||||
const typeTag = asset.tags.find(
|
||||
(tag) =>
|
||||
tag !== MODELS_TAG &&
|
||||
!(modelTypeMode && tag.startsWith(MODEL_TYPE_TAG_PREFIX))
|
||||
)
|
||||
if (!typeTag) return []
|
||||
return [
|
||||
typeTag.includes('/') ? typeTag.slice(typeTag.indexOf('/') + 1) : typeTag
|
||||
]
|
||||
}
|
||||
|
||||
/**
|
||||
* Ordered node-category candidates for an asset, most specific first.
|
||||
*
|
||||
* Callers resolve a node provider by trying each candidate in order and taking
|
||||
* the first that maps to a provider. The full (possibly hierarchical) value is
|
||||
* kept so `modelToNodeStore`'s `parent/child` fallback still works.
|
||||
*
|
||||
* In `modelTypeMode` (backend declares `supports_model_type_tags`) candidates
|
||||
* come in two tiers. Tier 1: the stripped `model_type:*` values plus bare tags
|
||||
* *related* to one of them (equal to it, or extending it as a `parent/child`
|
||||
* path), ordered by descending depth — so a resolvable `LLM/Qwen-VL/...` twin
|
||||
* wins over a flat `model_type:LLM`, while ties keep `model_type:*` values
|
||||
* ahead of bare tags. Tier 2: unrelated bare tags (e.g. a user-added
|
||||
* `foo/bar`), tried only when nothing authoritative resolves — they can no
|
||||
* longer pre-empt a resolvable `model_type:*` value however deep they are.
|
||||
* An uncovered asset (no `model_type:` tag) routes by all its bare tags,
|
||||
* deepest first. Outside `modelTypeMode` the legacy first-non-reserved tag is
|
||||
* used verbatim.
|
||||
*/
|
||||
export function getAssetNodeCategoryCandidates(
|
||||
asset: AssetItem,
|
||||
modelTypeMode: boolean
|
||||
): string[] {
|
||||
if (!modelTypeMode) {
|
||||
const legacy = asset.tags.find(
|
||||
(tag) => tag !== MODELS_TAG && tag !== MISSING_TAG
|
||||
)
|
||||
return legacy ? [legacy] : []
|
||||
}
|
||||
|
||||
const bareTags = asset.tags.filter(
|
||||
(tag) =>
|
||||
tag !== MODELS_TAG &&
|
||||
tag !== MISSING_TAG &&
|
||||
!tag.startsWith(MODEL_TYPE_TAG_PREFIX)
|
||||
)
|
||||
|
||||
const byDepthDesc = (a: string, b: string) => pathDepth(b) - pathDepth(a)
|
||||
|
||||
const modelTypes = getModelTypeTagValues(asset)
|
||||
if (modelTypes.length === 0) return bareTags.toSorted(byDepthDesc)
|
||||
|
||||
const isRelated = (tag: string) =>
|
||||
modelTypes.some((type) => tag === type || tag.startsWith(`${type}/`))
|
||||
|
||||
return [
|
||||
...[...modelTypes, ...bareTags.filter(isRelated)].sort(byDepthDesc),
|
||||
...bareTags.filter((tag) => !isRelated(tag)).sort(byDepthDesc)
|
||||
]
|
||||
}
|
||||
|
||||
/**
|
||||
* Extracts user description from asset user_metadata
|
||||
* @param asset - The asset to extract user description from
|
||||
|
||||
@@ -1,14 +1,25 @@
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
|
||||
import { resolveModelNodeFromAsset } from '@/platform/assets/utils/resolveModelNodeFromAsset'
|
||||
|
||||
const mockGetNodeProvider = vi.hoisted(() => vi.fn())
|
||||
const mockSupportsModelTypeTags = vi.hoisted(() => ({ value: false }))
|
||||
|
||||
vi.mock('@/stores/modelToNodeStore', () => ({
|
||||
useModelToNodeStore: () => ({ getNodeProvider: mockGetNodeProvider })
|
||||
}))
|
||||
|
||||
vi.mock('@/composables/useFeatureFlags', () => ({
|
||||
useFeatureFlags: () => ({
|
||||
flags: {
|
||||
get supportsModelTypeTags() {
|
||||
return mockSupportsModelTypeTags.value
|
||||
}
|
||||
}
|
||||
})
|
||||
}))
|
||||
|
||||
function createMockAsset(overrides: Partial<AssetItem> = {}): AssetItem {
|
||||
return {
|
||||
id: 'asset-123',
|
||||
@@ -49,6 +60,11 @@ describe('resolveModelNodeFromAsset', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
vi.spyOn(console, 'error').mockImplementation(() => {})
|
||||
mockSupportsModelTypeTags.value = false
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks()
|
||||
})
|
||||
|
||||
describe('valid assets', () => {
|
||||
@@ -68,6 +84,48 @@ describe('resolveModelNodeFromAsset', () => {
|
||||
}
|
||||
})
|
||||
|
||||
it('strips the model_type: prefix when resolving the provider in model_type mode', () => {
|
||||
mockSupportsModelTypeTags.value = true
|
||||
mockProvider(createMockNodeProvider())
|
||||
const result = resolveModelNodeFromAsset(
|
||||
createMockAsset({ tags: ['models', 'model_type:vae'] })
|
||||
)
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
expect(mockGetNodeProvider).toHaveBeenCalledWith('vae')
|
||||
})
|
||||
|
||||
it('skips an unresolvable incidental tag and resolves via the model_type value', () => {
|
||||
mockSupportsModelTypeTags.value = true
|
||||
mockGetNodeProvider.mockImplementation((category: string) =>
|
||||
category === 'vae' ? createMockNodeProvider() : undefined
|
||||
)
|
||||
const result = resolveModelNodeFromAsset(
|
||||
createMockAsset({ tags: ['models', 'model_type:vae', 'foo/bar'] })
|
||||
)
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
expect(mockGetNodeProvider).toHaveBeenCalledWith('foo/bar')
|
||||
expect(mockGetNodeProvider).toHaveBeenCalledWith('vae')
|
||||
})
|
||||
|
||||
it('prefers the deepest resolvable path over a flat model_type value', () => {
|
||||
mockSupportsModelTypeTags.value = true
|
||||
mockGetNodeProvider.mockImplementation((category: string) =>
|
||||
category === 'LLM/Qwen-VL/Qwen3-0.6B'
|
||||
? createMockNodeProvider()
|
||||
: undefined
|
||||
)
|
||||
const result = resolveModelNodeFromAsset(
|
||||
createMockAsset({
|
||||
tags: ['models', 'model_type:LLM', 'LLM/Qwen-VL/Qwen3-0.6B']
|
||||
})
|
||||
)
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
expect(mockGetNodeProvider).toHaveBeenCalledWith('LLM/Qwen-VL/Qwen3-0.6B')
|
||||
})
|
||||
|
||||
it('falls back to metadata.filename when user_metadata.filename missing', () => {
|
||||
mockProvider(createMockNodeProvider())
|
||||
const result = resolveModelNodeFromAsset(
|
||||
@@ -201,7 +259,7 @@ describe('resolveModelNodeFromAsset', () => {
|
||||
if (!result.success) {
|
||||
expect(result.error.code).toBe('NO_PROVIDER')
|
||||
expect(result.error.message).toContain('checkpoints')
|
||||
expect(result.error.details?.category).toBe('checkpoints')
|
||||
expect(result.error.details?.candidates).toEqual(['checkpoints'])
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
@@ -4,7 +4,11 @@ import {
|
||||
MISSING_TAG,
|
||||
MODELS_TAG
|
||||
} from '@/platform/assets/services/assetService'
|
||||
import { getAssetFilename } from '@/platform/assets/utils/assetMetadataUtils'
|
||||
import {
|
||||
getAssetFilename,
|
||||
getAssetNodeCategoryCandidates
|
||||
} from '@/platform/assets/utils/assetMetadataUtils'
|
||||
import { useFeatureFlags } from '@/composables/useFeatureFlags'
|
||||
import { useModelToNodeStore } from '@/stores/modelToNodeStore'
|
||||
import type { ModelNodeProvider } from '@/stores/modelToNodeStore'
|
||||
|
||||
@@ -81,10 +85,12 @@ export function resolveModelNodeFromAsset(
|
||||
}
|
||||
}
|
||||
|
||||
const category = validAsset.tags.find(
|
||||
(tag) => tag !== MODELS_TAG && tag !== MISSING_TAG
|
||||
const { flags } = useFeatureFlags()
|
||||
const candidates = getAssetNodeCategoryCandidates(
|
||||
validAsset,
|
||||
flags.supportsModelTypeTags
|
||||
)
|
||||
if (!category) {
|
||||
if (candidates.length === 0) {
|
||||
console.error(
|
||||
`Asset ${validAsset.id} has no valid category tag. Available tags: ${validAsset.tags.join(', ')} (expected tag other than '${MODELS_TAG}' or '${MISSING_TAG}')`
|
||||
)
|
||||
@@ -99,19 +105,31 @@ export function resolveModelNodeFromAsset(
|
||||
}
|
||||
}
|
||||
|
||||
const provider = useModelToNodeStore().getNodeProvider(category)
|
||||
if (!provider) {
|
||||
console.error(`No node provider registered for category: ${category}`)
|
||||
const modelToNodeStore = useModelToNodeStore()
|
||||
const resolved = candidates
|
||||
.map((category) => ({
|
||||
category,
|
||||
provider: modelToNodeStore.getNodeProvider(category)
|
||||
}))
|
||||
.find((candidate) => candidate.provider !== undefined)
|
||||
|
||||
if (!resolved?.provider) {
|
||||
// Known gap (out of scope for FE-1076): flat `model_type:LLM`-style tags
|
||||
// whose loaders are only registered hierarchically land here until the
|
||||
// backend emits a subtype-carrying tag.
|
||||
console.error(
|
||||
`No node provider registered for category: ${candidates.join(', ')}`
|
||||
)
|
||||
return {
|
||||
success: false,
|
||||
error: {
|
||||
code: 'NO_PROVIDER',
|
||||
message: `No node provider registered for category: ${category}`,
|
||||
message: `No node provider registered for category: ${candidates.join(', ')}`,
|
||||
assetId: validAsset.id,
|
||||
details: { category }
|
||||
details: { candidates }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return { success: true, value: { provider, filename } }
|
||||
return { success: true, value: { provider: resolved.provider, filename } }
|
||||
}
|
||||
|
||||
@@ -1215,6 +1215,15 @@ export const CORE_SETTINGS: SettingParams[] = [
|
||||
defaultValue: isCloud ? true : false,
|
||||
experimental: true
|
||||
},
|
||||
{
|
||||
id: 'Comfy.ModelLibrary.UseAssetBrowser',
|
||||
name: 'Use the asset browser for the model library',
|
||||
type: 'hidden',
|
||||
tooltip:
|
||||
'When enabled alongside the asset API, the model library opens the asset browser. Otherwise it opens the sidebar tree.',
|
||||
defaultValue: isCloud ? true : false,
|
||||
experimental: true
|
||||
},
|
||||
{
|
||||
id: 'Comfy.VersionCompatibility.DisableWarnings',
|
||||
name: 'Disable version compatibility warnings',
|
||||
|
||||
@@ -428,6 +428,7 @@ const zSettings = z.object({
|
||||
'Comfy.VueNodes.Enabled': z.boolean(),
|
||||
'Comfy.AppBuilder.VueNodeSwitchDismissed': z.boolean(),
|
||||
'Comfy.Assets.UseAssetAPI': z.boolean(),
|
||||
'Comfy.ModelLibrary.UseAssetBrowser': z.boolean(),
|
||||
'Comfy.Queue.QPOV2': z.boolean(),
|
||||
'Comfy.Queue.ShowRunProgressBar': z.boolean(),
|
||||
'Comfy-Desktop.AutoUpdate': z.boolean(),
|
||||
|
||||
@@ -875,6 +875,113 @@ describe('assetsStore - Model Assets Cache (Cloud)', () => {
|
||||
})
|
||||
})
|
||||
|
||||
describe('pagination safety', () => {
|
||||
it('stops instead of looping when the backend ignores offset', async () => {
|
||||
const store = useAssetsStore()
|
||||
const nodeType = 'CheckpointLoaderSimple'
|
||||
|
||||
// A backend that ignores offset returns the same full page every time.
|
||||
const fullPage = Array.from({ length: 500 }, (_, i) =>
|
||||
createMockAsset(`asset-${i}`)
|
||||
)
|
||||
vi.mocked(assetService.getAssetsForNodeType).mockResolvedValue(fullPage)
|
||||
|
||||
await store.updateModelsForNodeType(nodeType)
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(
|
||||
vi.mocked(assetService.getAssetsForNodeType)
|
||||
).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
|
||||
expect(
|
||||
vi.mocked(assetService.getAssetsForNodeType)
|
||||
).toHaveBeenCalledTimes(2)
|
||||
expect(store.getAssets(nodeType)).toHaveLength(500)
|
||||
})
|
||||
|
||||
it('continues past an all-duplicate page whose content differs from the previous page', async () => {
|
||||
const store = useAssetsStore()
|
||||
const nodeType = 'CheckpointLoaderSimple'
|
||||
|
||||
// Concurrent writes can shift pagination windows so a page is all
|
||||
// already-seen assets without the backend ignoring offset; later pages
|
||||
// can still hold unseen assets.
|
||||
const fullPage = Array.from({ length: 500 }, (_, i) =>
|
||||
createMockAsset(`asset-${i}`)
|
||||
)
|
||||
const samePageReordered = [...fullPage].reverse()
|
||||
const finalPage = [createMockAsset('late-arrival')]
|
||||
|
||||
let callCount = 0
|
||||
vi.mocked(assetService.getAssetsForNodeType).mockImplementation(
|
||||
async () => {
|
||||
callCount++
|
||||
if (callCount === 1) return fullPage
|
||||
if (callCount === 2) return samePageReordered
|
||||
return finalPage
|
||||
}
|
||||
)
|
||||
|
||||
await store.updateModelsForNodeType(nodeType)
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(
|
||||
vi.mocked(assetService.getAssetsForNodeType)
|
||||
).toHaveBeenCalledTimes(3)
|
||||
})
|
||||
expect(store.getAssets(nodeType).map((a) => a.id)).toContain(
|
||||
'late-arrival'
|
||||
)
|
||||
})
|
||||
|
||||
it('terminates when an offset-ignoring backend alternates page orderings', async () => {
|
||||
const store = useAssetsStore()
|
||||
const nodeType = 'CheckpointLoaderSimple'
|
||||
|
||||
// Same full page served forever with a nondeterministic ordering: no
|
||||
// page ever contributes a new ID, and no two consecutive pages are
|
||||
// identical. The walk must still stop.
|
||||
const fullPage = Array.from({ length: 500 }, (_, i) =>
|
||||
createMockAsset(`asset-${i}`)
|
||||
)
|
||||
const reversed = [...fullPage].reverse()
|
||||
let callCount = 0
|
||||
vi.mocked(assetService.getAssetsForNodeType).mockImplementation(
|
||||
async () => {
|
||||
callCount++
|
||||
return callCount % 2 === 1 ? fullPage : reversed
|
||||
}
|
||||
)
|
||||
|
||||
await store.updateModelsForNodeType(nodeType)
|
||||
|
||||
expect(callCount).toBeLessThanOrEqual(5)
|
||||
expect(store.getAssets(nodeType)).toHaveLength(500)
|
||||
})
|
||||
})
|
||||
|
||||
describe('refresh error surfacing', () => {
|
||||
it('surfaces a failed refresh on the committed state consumers read', async () => {
|
||||
const store = useAssetsStore()
|
||||
const nodeType = 'CheckpointLoaderSimple'
|
||||
|
||||
vi.mocked(assetService.getAssetsForNodeType).mockResolvedValueOnce([
|
||||
createMockAsset('existing')
|
||||
])
|
||||
await store.updateModelsForNodeType(nodeType)
|
||||
expect(store.getError(nodeType)).toBeUndefined()
|
||||
|
||||
vi.mocked(assetService.getAssetsForNodeType).mockRejectedValueOnce(
|
||||
new Error('backend down')
|
||||
)
|
||||
await store.updateModelsForNodeType(nodeType)
|
||||
|
||||
expect(store.getAssets(nodeType).map((a) => a.id)).toEqual(['existing'])
|
||||
expect(store.getError(nodeType)?.message).toBe('backend down')
|
||||
})
|
||||
})
|
||||
|
||||
describe('concurrent request handling', () => {
|
||||
it('should short-circuit concurrent calls to prevent duplicate work', async () => {
|
||||
const store = useAssetsStore()
|
||||
@@ -924,6 +1031,34 @@ describe('assetsStore - Model Assets Cache (Cloud)', () => {
|
||||
vi.mocked(assetService.getAssetsForNodeType)
|
||||
).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
|
||||
it('keeps a newer request single-flighted when a stale request finishes after invalidation', async () => {
|
||||
const store = useAssetsStore()
|
||||
const nodeType = 'CheckpointLoaderSimple'
|
||||
|
||||
let resolveFirst!: (assets: AssetItem[]) => void
|
||||
const firstFetch = new Promise<AssetItem[]>((resolve) => {
|
||||
resolveFirst = resolve
|
||||
})
|
||||
vi.mocked(assetService.getAssetsForNodeType)
|
||||
.mockReturnValueOnce(firstFetch)
|
||||
.mockReturnValue(new Promise<AssetItem[]>(() => {}))
|
||||
|
||||
const staleRequest = store.updateModelsForNodeType(nodeType)
|
||||
store.invalidateCategory('checkpoints')
|
||||
void store.updateModelsForNodeType(nodeType)
|
||||
|
||||
resolveFirst([createMockAsset('stale')])
|
||||
await staleRequest
|
||||
|
||||
// The stale request's teardown must not evict the newer request's
|
||||
// single-flight entry: a third call short-circuits instead of starting
|
||||
// a duplicate walk.
|
||||
void store.updateModelsForNodeType(nodeType)
|
||||
expect(
|
||||
vi.mocked(assetService.getAssetsForNodeType)
|
||||
).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
})
|
||||
|
||||
describe('shallowReactive state reactivity', () => {
|
||||
@@ -1435,6 +1570,35 @@ describe('assetsStore - Model Assets Cache (Cloud)', () => {
|
||||
})
|
||||
})
|
||||
|
||||
describe('assetsStore - Model Assets Cache (non-cloud)', () => {
|
||||
beforeEach(() => {
|
||||
setActivePinia(createTestingPinia({ stubActions: false }))
|
||||
mockIsCloud.value = false
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
it('caches model assets fetched by tag on non-cloud builds', async () => {
|
||||
const store = useAssetsStore()
|
||||
vi.mocked(assetService.getAssetsByTag).mockResolvedValue([
|
||||
{
|
||||
id: 'm1',
|
||||
name: 'sd_xl_base_1.0.safetensors',
|
||||
tags: ['checkpoints', 'models']
|
||||
},
|
||||
{ id: 'm2', name: 'lora.safetensors', tags: ['loras', 'models'] }
|
||||
])
|
||||
|
||||
await store.updateModelsForTag('models')
|
||||
|
||||
expect(assetService.getAssetsByTag).toHaveBeenCalledWith(
|
||||
'models',
|
||||
true,
|
||||
expect.anything()
|
||||
)
|
||||
expect(store.getAssets('tag:models')).toHaveLength(2)
|
||||
})
|
||||
})
|
||||
|
||||
describe('assetsStore - Deletion State and Input Mapping', () => {
|
||||
beforeEach(() => {
|
||||
setActivePinia(createTestingPinia({ stubActions: false }))
|
||||
|
||||
@@ -394,421 +394,444 @@ export const useAssetsStore = defineStore('assets', () => {
|
||||
* Multiple node types sharing the same category share the same cache entry.
|
||||
* Public API accepts nodeType for backwards compatibility but translates
|
||||
* to category internally using modelToNodeStore.getCategoryForNodeType().
|
||||
* Cloud-only feature - empty Maps in desktop builds
|
||||
*
|
||||
* Runs on every distribution; whether anything fetches through it is
|
||||
* decided by consumers via `assetService.isAssetAPIEnabled()`, which stays
|
||||
* the authoritative off-cloud gate.
|
||||
*/
|
||||
const getModelState = () => {
|
||||
if (isCloud) {
|
||||
const modelStateByCategory = ref(new Map<string, ModelPaginationState>())
|
||||
const modelStateByCategory = ref(new Map<string, ModelPaginationState>())
|
||||
|
||||
const assetsArrayCache = new Map<
|
||||
string,
|
||||
{ source: Map<string, AssetItem>; array: AssetItem[] }
|
||||
>()
|
||||
const assetsArrayCache = new Map<
|
||||
string,
|
||||
{ source: Map<string, AssetItem>; array: AssetItem[] }
|
||||
>()
|
||||
|
||||
const pendingRequestByCategory = new Map<string, ModelPaginationState>()
|
||||
const pendingPromiseByCategory = new Map<string, Promise<void>>()
|
||||
const pendingRequestByCategory = new Map<string, ModelPaginationState>()
|
||||
const pendingPromiseByCategory = new Map<string, Promise<void>>()
|
||||
|
||||
function createState(
|
||||
existingAssets?: Map<string, AssetItem>
|
||||
): ModelPaginationState {
|
||||
const assets = new Map(existingAssets)
|
||||
return reactive({
|
||||
assets,
|
||||
offset: 0,
|
||||
hasMore: true,
|
||||
isLoading: true
|
||||
})
|
||||
function createState(
|
||||
existingAssets?: Map<string, AssetItem>
|
||||
): ModelPaginationState {
|
||||
const assets = new Map(existingAssets)
|
||||
return reactive({
|
||||
assets,
|
||||
offset: 0,
|
||||
hasMore: true,
|
||||
isLoading: true
|
||||
})
|
||||
}
|
||||
|
||||
function isStale(category: string, state: ModelPaginationState): boolean {
|
||||
const committed = modelStateByCategory.value.get(category)
|
||||
const pending = pendingRequestByCategory.get(category)
|
||||
return committed !== state && pending !== state
|
||||
}
|
||||
|
||||
const EMPTY_ASSETS: AssetItem[] = []
|
||||
|
||||
/**
|
||||
* Resolve a key to a category. Handles both nodeType and tag:xxx formats.
|
||||
* @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models')
|
||||
* @returns The category or undefined if not resolvable
|
||||
*/
|
||||
function resolveCategory(key: string): string | undefined {
|
||||
if (key.startsWith('tag:')) {
|
||||
return key
|
||||
}
|
||||
return modelToNodeStore.getCategoryForNodeType(key)
|
||||
}
|
||||
|
||||
/**
|
||||
* Get assets by nodeType or tag key.
|
||||
* Translates nodeType to category internally for cache lookup.
|
||||
* @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models')
|
||||
*/
|
||||
function getAssets(key: string): AssetItem[] {
|
||||
const category = resolveCategory(key)
|
||||
if (!category) return EMPTY_ASSETS
|
||||
|
||||
const state = modelStateByCategory.value.get(category)
|
||||
const assetsMap = state?.assets
|
||||
if (!assetsMap) return EMPTY_ASSETS
|
||||
|
||||
const cached = assetsArrayCache.get(category)
|
||||
if (cached && cached.source === assetsMap) {
|
||||
return cached.array
|
||||
}
|
||||
|
||||
function isStale(category: string, state: ModelPaginationState): boolean {
|
||||
const committed = modelStateByCategory.value.get(category)
|
||||
const pending = pendingRequestByCategory.get(category)
|
||||
return committed !== state && pending !== state
|
||||
const array = Array.from(assetsMap.values())
|
||||
assetsArrayCache.set(category, { source: assetsMap, array })
|
||||
return array
|
||||
}
|
||||
|
||||
function isLoading(key: string): boolean {
|
||||
const category = resolveCategory(key)
|
||||
if (!category) return false
|
||||
return modelStateByCategory.value.get(category)?.isLoading ?? false
|
||||
}
|
||||
|
||||
function getError(key: string): Error | undefined {
|
||||
const category = resolveCategory(key)
|
||||
if (!category) return undefined
|
||||
return modelStateByCategory.value.get(category)?.error
|
||||
}
|
||||
|
||||
function hasMore(key: string): boolean {
|
||||
const category = resolveCategory(key)
|
||||
if (!category) return false
|
||||
return modelStateByCategory.value.get(category)?.hasMore ?? false
|
||||
}
|
||||
|
||||
function hasAssetKey(key: string): boolean {
|
||||
const category = resolveCategory(key)
|
||||
if (!category) return false
|
||||
return modelStateByCategory.value.has(category)
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a category exists in the cache.
|
||||
* Checks both direct category keys and tag-prefixed keys.
|
||||
* @param category The category to check (e.g., 'checkpoints', 'loras')
|
||||
*/
|
||||
function hasCategory(category: string): boolean {
|
||||
return (
|
||||
modelStateByCategory.value.has(category) ||
|
||||
modelStateByCategory.value.has(`tag:${category}`)
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Internal helper to fetch and cache assets for a category.
|
||||
* Loads first batch immediately, then progressively loads remaining batches.
|
||||
* Keeps existing data visible until new data is successfully fetched.
|
||||
*
|
||||
* Concurrent calls for the same category are short-circuited: if a request
|
||||
* is already in progress (tracked via pendingRequestByCategory), subsequent
|
||||
* calls return immediately to avoid redundant work.
|
||||
*/
|
||||
async function updateModelsForCategory(
|
||||
category: string,
|
||||
fetcher: (options: PaginationOptions) => Promise<AssetItem[]>
|
||||
): Promise<void> {
|
||||
if (pendingPromiseByCategory.has(category)) {
|
||||
return pendingPromiseByCategory.get(category)!
|
||||
}
|
||||
|
||||
const EMPTY_ASSETS: AssetItem[] = []
|
||||
const existingState = modelStateByCategory.value.get(category)
|
||||
const state = createState(existingState?.assets)
|
||||
|
||||
/**
|
||||
* Resolve a key to a category. Handles both nodeType and tag:xxx formats.
|
||||
* @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models')
|
||||
* @returns The category or undefined if not resolvable
|
||||
*/
|
||||
function resolveCategory(key: string): string | undefined {
|
||||
if (key.startsWith('tag:')) {
|
||||
return key
|
||||
}
|
||||
return modelToNodeStore.getCategoryForNodeType(key)
|
||||
const seenIds = new Set<string>()
|
||||
const seenPageSignatures = new Set<string>()
|
||||
let consecutiveNoProgressPages = 0
|
||||
|
||||
const hasExistingData = modelStateByCategory.value.has(category)
|
||||
if (hasExistingData) {
|
||||
pendingRequestByCategory.set(category, state)
|
||||
} else {
|
||||
// Also track in pending map for initial loads to prevent concurrent calls
|
||||
pendingRequestByCategory.set(category, state)
|
||||
modelStateByCategory.value.set(category, state)
|
||||
}
|
||||
|
||||
/**
|
||||
* Get assets by nodeType or tag key.
|
||||
* Translates nodeType to category internally for cache lookup.
|
||||
* @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models')
|
||||
*/
|
||||
function getAssets(key: string): AssetItem[] {
|
||||
const category = resolveCategory(key)
|
||||
if (!category) return EMPTY_ASSETS
|
||||
async function loadBatches(): Promise<void> {
|
||||
while (state.hasMore) {
|
||||
try {
|
||||
const newAssets = await fetcher({
|
||||
limit: MODEL_BATCH_SIZE,
|
||||
offset: state.offset
|
||||
})
|
||||
|
||||
const state = modelStateByCategory.value.get(category)
|
||||
const assetsMap = state?.assets
|
||||
if (!assetsMap) return EMPTY_ASSETS
|
||||
if (isStale(category, state)) return
|
||||
|
||||
const cached = assetsArrayCache.get(category)
|
||||
if (cached && cached.source === assetsMap) {
|
||||
return cached.array
|
||||
}
|
||||
|
||||
const array = Array.from(assetsMap.values())
|
||||
assetsArrayCache.set(category, { source: assetsMap, array })
|
||||
return array
|
||||
}
|
||||
|
||||
function isLoading(key: string): boolean {
|
||||
const category = resolveCategory(key)
|
||||
if (!category) return false
|
||||
return modelStateByCategory.value.get(category)?.isLoading ?? false
|
||||
}
|
||||
|
||||
function getError(key: string): Error | undefined {
|
||||
const category = resolveCategory(key)
|
||||
if (!category) return undefined
|
||||
return modelStateByCategory.value.get(category)?.error
|
||||
}
|
||||
|
||||
function hasMore(key: string): boolean {
|
||||
const category = resolveCategory(key)
|
||||
if (!category) return false
|
||||
return modelStateByCategory.value.get(category)?.hasMore ?? false
|
||||
}
|
||||
|
||||
function hasAssetKey(key: string): boolean {
|
||||
const category = resolveCategory(key)
|
||||
if (!category) return false
|
||||
return modelStateByCategory.value.has(category)
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a category exists in the cache.
|
||||
* Checks both direct category keys and tag-prefixed keys.
|
||||
* @param category The category to check (e.g., 'checkpoints', 'loras')
|
||||
*/
|
||||
function hasCategory(category: string): boolean {
|
||||
return (
|
||||
modelStateByCategory.value.has(category) ||
|
||||
modelStateByCategory.value.has(`tag:${category}`)
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Internal helper to fetch and cache assets for a category.
|
||||
* Loads first batch immediately, then progressively loads remaining batches.
|
||||
* Keeps existing data visible until new data is successfully fetched.
|
||||
*
|
||||
* Concurrent calls for the same category are short-circuited: if a request
|
||||
* is already in progress (tracked via pendingRequestByCategory), subsequent
|
||||
* calls return immediately to avoid redundant work.
|
||||
*/
|
||||
async function updateModelsForCategory(
|
||||
category: string,
|
||||
fetcher: (options: PaginationOptions) => Promise<AssetItem[]>
|
||||
): Promise<void> {
|
||||
if (pendingPromiseByCategory.has(category)) {
|
||||
return pendingPromiseByCategory.get(category)!
|
||||
}
|
||||
|
||||
const existingState = modelStateByCategory.value.get(category)
|
||||
const state = createState(existingState?.assets)
|
||||
|
||||
const seenIds = new Set<string>()
|
||||
|
||||
const hasExistingData = modelStateByCategory.value.has(category)
|
||||
if (hasExistingData) {
|
||||
pendingRequestByCategory.set(category, state)
|
||||
} else {
|
||||
// Also track in pending map for initial loads to prevent concurrent calls
|
||||
pendingRequestByCategory.set(category, state)
|
||||
modelStateByCategory.value.set(category, state)
|
||||
}
|
||||
|
||||
async function loadBatches(): Promise<void> {
|
||||
while (state.hasMore) {
|
||||
try {
|
||||
const newAssets = await fetcher({
|
||||
limit: MODEL_BATCH_SIZE,
|
||||
offset: state.offset
|
||||
})
|
||||
|
||||
if (isStale(category, state)) return
|
||||
|
||||
const isFirstBatch = state.offset === 0
|
||||
if (isFirstBatch) {
|
||||
assetsArrayCache.delete(category)
|
||||
if (hasExistingData) {
|
||||
pendingRequestByCategory.delete(category)
|
||||
modelStateByCategory.value.set(category, state)
|
||||
}
|
||||
const isFirstBatch = state.offset === 0
|
||||
if (isFirstBatch) {
|
||||
assetsArrayCache.delete(category)
|
||||
if (hasExistingData) {
|
||||
pendingRequestByCategory.delete(category)
|
||||
modelStateByCategory.value.set(category, state)
|
||||
}
|
||||
|
||||
// Merge new assets into existing map and track seen IDs
|
||||
for (const asset of newAssets) {
|
||||
seenIds.add(asset.id)
|
||||
state.assets.set(asset.id, asset)
|
||||
}
|
||||
state.assets = new Map(state.assets)
|
||||
|
||||
state.offset += newAssets.length
|
||||
state.hasMore = newAssets.length === MODEL_BATCH_SIZE
|
||||
|
||||
if (isFirstBatch) {
|
||||
state.isLoading = false
|
||||
}
|
||||
|
||||
if (state.hasMore) {
|
||||
await new Promise((resolve) => setTimeout(resolve, 50))
|
||||
}
|
||||
} catch (err) {
|
||||
if (isStale(category, state)) return
|
||||
console.error(`Error loading batch for ${category}:`, err)
|
||||
|
||||
state.error = err instanceof Error ? err : new Error(String(err))
|
||||
state.hasMore = false
|
||||
state.isLoading = false
|
||||
pendingRequestByCategory.delete(category)
|
||||
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
const staleIds = [...state.assets.keys()].filter(
|
||||
(id) => !seenIds.has(id)
|
||||
)
|
||||
for (const id of staleIds) {
|
||||
state.assets.delete(id)
|
||||
// Merge new assets into existing map and track seen IDs
|
||||
const uniqueIdsBefore = seenIds.size
|
||||
for (const asset of newAssets) {
|
||||
seenIds.add(asset.id)
|
||||
state.assets.set(asset.id, asset)
|
||||
}
|
||||
state.assets = new Map(state.assets)
|
||||
|
||||
// Termination guards for backends that do not honour `offset`.
|
||||
// A page whose exact ID sequence was already served means the
|
||||
// walk is cycling, however the pages are ordered — stop. A single
|
||||
// all-duplicate page with fresh content (concurrent writes
|
||||
// shifting pagination windows) keeps going, but a run of them
|
||||
// with no new IDs is treated as exhausted so reordered responses
|
||||
// can never loop forever.
|
||||
const batchSignature = newAssets.map((asset) => asset.id).join(',')
|
||||
const isRepeatedPage =
|
||||
newAssets.length > 0 && seenPageSignatures.has(batchSignature)
|
||||
seenPageSignatures.add(batchSignature)
|
||||
const madeProgress = seenIds.size > uniqueIdsBefore
|
||||
consecutiveNoProgressPages = madeProgress
|
||||
? 0
|
||||
: consecutiveNoProgressPages + 1
|
||||
state.offset += newAssets.length
|
||||
state.hasMore =
|
||||
newAssets.length === MODEL_BATCH_SIZE &&
|
||||
!isRepeatedPage &&
|
||||
consecutiveNoProgressPages < 3
|
||||
|
||||
if (isFirstBatch) {
|
||||
state.isLoading = false
|
||||
}
|
||||
|
||||
if (state.hasMore) {
|
||||
await new Promise((resolve) => setTimeout(resolve, 50))
|
||||
}
|
||||
} catch (err) {
|
||||
if (isStale(category, state)) return
|
||||
console.error(`Error loading batch for ${category}:`, err)
|
||||
|
||||
state.error = err instanceof Error ? err : new Error(String(err))
|
||||
state.hasMore = false
|
||||
state.isLoading = false
|
||||
// A refresh that fails before its first batch never replaces the
|
||||
// committed state, so mirror the error onto the state consumers
|
||||
// actually read (getError) instead of only the discarded one.
|
||||
const committed = modelStateByCategory.value.get(category)
|
||||
if (committed && committed !== state) {
|
||||
committed.error = state.error
|
||||
}
|
||||
if (pendingRequestByCategory.get(category) === state) {
|
||||
pendingRequestByCategory.delete(category)
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
assetsArrayCache.delete(category)
|
||||
}
|
||||
|
||||
const staleIds = [...state.assets.keys()].filter(
|
||||
(id) => !seenIds.has(id)
|
||||
)
|
||||
for (const id of staleIds) {
|
||||
state.assets.delete(id)
|
||||
}
|
||||
assetsArrayCache.delete(category)
|
||||
if (pendingRequestByCategory.get(category) === state) {
|
||||
pendingRequestByCategory.delete(category)
|
||||
}
|
||||
}
|
||||
|
||||
const promise = loadBatches().finally(() => {
|
||||
// Guard both cleanups: an invalidateCategory during an awaited fetch
|
||||
// lets a newer request register its own entries before this one's
|
||||
// teardown runs, and an unconditional delete would evict the newer
|
||||
// request's entry and break single-flighting.
|
||||
const promise = loadBatches().finally(() => {
|
||||
if (pendingPromiseByCategory.get(category) === promise) {
|
||||
pendingPromiseByCategory.delete(category)
|
||||
})
|
||||
pendingPromiseByCategory.set(category, promise)
|
||||
await promise
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetch and cache model assets for a specific node type.
|
||||
* Translates nodeType to category internally - multiple node types
|
||||
* sharing the same category will share the same cache entry.
|
||||
* @param nodeType The node type to fetch assets for (e.g., 'CheckpointLoaderSimple')
|
||||
*/
|
||||
async function updateModelsForNodeType(nodeType: string): Promise<void> {
|
||||
const category = modelToNodeStore.getCategoryForNodeType(nodeType)
|
||||
if (!category) return
|
||||
|
||||
// Use category as cache key but fetch using nodeType for API compatibility
|
||||
await updateModelsForCategory(category, (opts) =>
|
||||
assetService.getAssetsForNodeType(nodeType, opts)
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetch and cache model assets for a specific tag
|
||||
* @param tag The tag to fetch assets for (e.g., 'models')
|
||||
*/
|
||||
async function updateModelsForTag(tag: string): Promise<void> {
|
||||
const category = `tag:${tag}`
|
||||
await updateModelsForCategory(category, (opts) =>
|
||||
assetService.getAssetsByTag(tag, true, opts)
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Invalidate the cache for a specific category.
|
||||
* Forces a refetch on next access.
|
||||
* @param category The category to invalidate (e.g., 'checkpoints', 'loras')
|
||||
*/
|
||||
function invalidateCategory(category: string): void {
|
||||
modelStateByCategory.value.delete(category)
|
||||
assetsArrayCache.delete(category)
|
||||
pendingRequestByCategory.delete(category)
|
||||
pendingPromiseByCategory.delete(category)
|
||||
}
|
||||
|
||||
/**
|
||||
* Optimistically update an asset in the cache
|
||||
* @param assetId The asset ID to update
|
||||
* @param updates Partial asset data to merge
|
||||
* @param cacheKey Optional cache key to target (nodeType or 'tag:xxx')
|
||||
*/
|
||||
function updateAssetInCache(
|
||||
assetId: string,
|
||||
updates: Partial<AssetItem>,
|
||||
cacheKey?: string
|
||||
) {
|
||||
const category = cacheKey ? resolveCategory(cacheKey) : undefined
|
||||
if (cacheKey && !category) return
|
||||
|
||||
const categoriesToCheck = category
|
||||
? [category]
|
||||
: Array.from(modelStateByCategory.value.keys())
|
||||
|
||||
for (const cat of categoriesToCheck) {
|
||||
const state = modelStateByCategory.value.get(cat)
|
||||
if (!state?.assets) continue
|
||||
|
||||
const existingAsset = state.assets.get(assetId)
|
||||
if (existingAsset) {
|
||||
const updatedAsset = { ...existingAsset, ...updates }
|
||||
state.assets.set(assetId, updatedAsset)
|
||||
assetsArrayCache.delete(cat)
|
||||
if (cacheKey) return
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
pendingPromiseByCategory.set(category, promise)
|
||||
await promise
|
||||
}
|
||||
|
||||
/**
|
||||
* Update asset metadata with optimistic cache update
|
||||
* @param asset The asset to update
|
||||
* @param userMetadata The user_metadata to save
|
||||
* @param cacheKey Optional cache key to target for optimistic update
|
||||
*/
|
||||
async function updateAssetMetadata(
|
||||
asset: AssetItem,
|
||||
userMetadata: Record<string, unknown>,
|
||||
cacheKey?: string
|
||||
) {
|
||||
const originalMetadata = asset.user_metadata
|
||||
updateAssetInCache(asset.id, { user_metadata: userMetadata }, cacheKey)
|
||||
/**
|
||||
* Fetch and cache model assets for a specific node type.
|
||||
* Translates nodeType to category internally - multiple node types
|
||||
* sharing the same category will share the same cache entry.
|
||||
* @param nodeType The node type to fetch assets for (e.g., 'CheckpointLoaderSimple')
|
||||
*/
|
||||
async function updateModelsForNodeType(nodeType: string): Promise<void> {
|
||||
const category = modelToNodeStore.getCategoryForNodeType(nodeType)
|
||||
if (!category) return
|
||||
|
||||
try {
|
||||
const updatedAsset = await assetService.updateAsset(asset.id, {
|
||||
user_metadata: userMetadata
|
||||
})
|
||||
updateAssetInCache(asset.id, updatedAsset, cacheKey)
|
||||
} catch (error) {
|
||||
console.error('Failed to update asset metadata:', error)
|
||||
updateAssetInCache(
|
||||
asset.id,
|
||||
{ user_metadata: originalMetadata },
|
||||
cacheKey
|
||||
)
|
||||
// Use category as cache key but fetch using nodeType for API compatibility
|
||||
await updateModelsForCategory(category, (opts) =>
|
||||
assetService.getAssetsForNodeType(nodeType, opts)
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetch and cache model assets for a specific tag
|
||||
* @param tag The tag to fetch assets for (e.g., 'models')
|
||||
*/
|
||||
async function updateModelsForTag(tag: string): Promise<void> {
|
||||
const category = `tag:${tag}`
|
||||
await updateModelsForCategory(category, (opts) =>
|
||||
assetService.getAssetsByTag(tag, true, opts)
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Invalidate the cache for a specific category.
|
||||
* Forces a refetch on next access.
|
||||
* @param category The category to invalidate (e.g., 'checkpoints', 'loras')
|
||||
*/
|
||||
function invalidateCategory(category: string): void {
|
||||
modelStateByCategory.value.delete(category)
|
||||
assetsArrayCache.delete(category)
|
||||
pendingRequestByCategory.delete(category)
|
||||
pendingPromiseByCategory.delete(category)
|
||||
}
|
||||
|
||||
/**
|
||||
* Optimistically update an asset in the cache
|
||||
* @param assetId The asset ID to update
|
||||
* @param updates Partial asset data to merge
|
||||
* @param cacheKey Optional cache key to target (nodeType or 'tag:xxx')
|
||||
*/
|
||||
function updateAssetInCache(
|
||||
assetId: string,
|
||||
updates: Partial<AssetItem>,
|
||||
cacheKey?: string
|
||||
) {
|
||||
const category = cacheKey ? resolveCategory(cacheKey) : undefined
|
||||
if (cacheKey && !category) return
|
||||
|
||||
const categoriesToCheck = category
|
||||
? [category]
|
||||
: Array.from(modelStateByCategory.value.keys())
|
||||
|
||||
for (const cat of categoriesToCheck) {
|
||||
const state = modelStateByCategory.value.get(cat)
|
||||
if (!state?.assets) continue
|
||||
|
||||
const existingAsset = state.assets.get(assetId)
|
||||
if (existingAsset) {
|
||||
const updatedAsset = { ...existingAsset, ...updates }
|
||||
state.assets.set(assetId, updatedAsset)
|
||||
assetsArrayCache.delete(cat)
|
||||
if (cacheKey) return
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Update asset tags using add/remove endpoints
|
||||
* @param asset The asset to update (used to read current tags)
|
||||
* @param newTags The desired tags array
|
||||
* @param cacheKey Optional cache key to target for optimistic update
|
||||
*/
|
||||
async function updateAssetTags(
|
||||
asset: AssetItem,
|
||||
newTags: string[],
|
||||
cacheKey?: string
|
||||
) {
|
||||
const originalTags = asset.tags
|
||||
const tagsToAdd = difference(newTags, originalTags)
|
||||
const tagsToRemove = difference(originalTags, newTags)
|
||||
|
||||
if (tagsToAdd.length === 0 && tagsToRemove.length === 0) return
|
||||
|
||||
updateAssetInCache(asset.id, { tags: newTags }, cacheKey)
|
||||
|
||||
let removedTagsOnServer: string[] = []
|
||||
try {
|
||||
let removeResult: TagsOperationResult | undefined
|
||||
if (tagsToRemove.length > 0) {
|
||||
removeResult = await assetService.removeAssetTags(
|
||||
asset.id,
|
||||
tagsToRemove
|
||||
)
|
||||
removedTagsOnServer = removeResult.removed ?? tagsToRemove
|
||||
}
|
||||
|
||||
const addResult =
|
||||
tagsToAdd.length > 0
|
||||
? await assetService.addAssetTags(asset.id, tagsToAdd)
|
||||
: undefined
|
||||
|
||||
const finalTags = (addResult ?? removeResult)?.total_tags
|
||||
if (finalTags) {
|
||||
updateAssetInCache(asset.id, { tags: finalTags }, cacheKey)
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Failed to update asset tags:', error)
|
||||
updateAssetInCache(asset.id, { tags: originalTags }, cacheKey)
|
||||
|
||||
if (removedTagsOnServer.length > 0) {
|
||||
try {
|
||||
await assetService.addAssetTags(asset.id, removedTagsOnServer)
|
||||
} catch (compensationError) {
|
||||
console.error(
|
||||
'Failed to restore tags after partial failure; invalidating cache to force refetch:',
|
||||
compensationError
|
||||
)
|
||||
const categoriesToInvalidate = new Set<string>()
|
||||
const resolved = cacheKey ? resolveCategory(cacheKey) : undefined
|
||||
if (resolved) {
|
||||
categoriesToInvalidate.add(resolved)
|
||||
}
|
||||
for (const [
|
||||
category,
|
||||
state
|
||||
] of modelStateByCategory.value.entries()) {
|
||||
if (state.assets?.has(asset.id)) {
|
||||
categoriesToInvalidate.add(category)
|
||||
}
|
||||
}
|
||||
for (const category of categoriesToInvalidate) {
|
||||
invalidateCategory(category)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Invalidate model caches for a given category (e.g., 'checkpoints', 'loras')
|
||||
* Clears the category cache and tag-based caches so next access triggers refetch
|
||||
* @param category The model category to invalidate (e.g., 'checkpoints')
|
||||
*/
|
||||
function invalidateModelsForCategory(category: string): void {
|
||||
invalidateCategory(category)
|
||||
invalidateCategory(`tag:${category}`)
|
||||
invalidateCategory('tag:models')
|
||||
}
|
||||
|
||||
return {
|
||||
getAssets,
|
||||
isLoading,
|
||||
getError,
|
||||
hasMore,
|
||||
hasAssetKey,
|
||||
hasCategory,
|
||||
updateModelsForNodeType,
|
||||
updateModelsForTag,
|
||||
invalidateCategory,
|
||||
updateAssetMetadata,
|
||||
updateAssetTags,
|
||||
invalidateModelsForCategory
|
||||
}
|
||||
}
|
||||
|
||||
const emptyAssets: AssetItem[] = []
|
||||
/**
|
||||
* Update asset metadata with optimistic cache update
|
||||
* @param asset The asset to update
|
||||
* @param userMetadata The user_metadata to save
|
||||
* @param cacheKey Optional cache key to target for optimistic update
|
||||
*/
|
||||
async function updateAssetMetadata(
|
||||
asset: AssetItem,
|
||||
userMetadata: Record<string, unknown>,
|
||||
cacheKey?: string
|
||||
) {
|
||||
const originalMetadata = asset.user_metadata
|
||||
updateAssetInCache(asset.id, { user_metadata: userMetadata }, cacheKey)
|
||||
|
||||
try {
|
||||
const updatedAsset = await assetService.updateAsset(asset.id, {
|
||||
user_metadata: userMetadata
|
||||
})
|
||||
updateAssetInCache(asset.id, updatedAsset, cacheKey)
|
||||
} catch (error) {
|
||||
console.error('Failed to update asset metadata:', error)
|
||||
updateAssetInCache(
|
||||
asset.id,
|
||||
{ user_metadata: originalMetadata },
|
||||
cacheKey
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Update asset tags using add/remove endpoints
|
||||
* @param asset The asset to update (used to read current tags)
|
||||
* @param newTags The desired tags array
|
||||
* @param cacheKey Optional cache key to target for optimistic update
|
||||
*/
|
||||
async function updateAssetTags(
|
||||
asset: AssetItem,
|
||||
newTags: string[],
|
||||
cacheKey?: string
|
||||
) {
|
||||
const originalTags = asset.tags
|
||||
const tagsToAdd = difference(newTags, originalTags)
|
||||
const tagsToRemove = difference(originalTags, newTags)
|
||||
|
||||
if (tagsToAdd.length === 0 && tagsToRemove.length === 0) return
|
||||
|
||||
updateAssetInCache(asset.id, { tags: newTags }, cacheKey)
|
||||
|
||||
let removedTagsOnServer: string[] = []
|
||||
try {
|
||||
let removeResult: TagsOperationResult | undefined
|
||||
if (tagsToRemove.length > 0) {
|
||||
removeResult = await assetService.removeAssetTags(
|
||||
asset.id,
|
||||
tagsToRemove
|
||||
)
|
||||
removedTagsOnServer = removeResult.removed ?? tagsToRemove
|
||||
}
|
||||
|
||||
const addResult =
|
||||
tagsToAdd.length > 0
|
||||
? await assetService.addAssetTags(asset.id, tagsToAdd)
|
||||
: undefined
|
||||
|
||||
const finalTags = (addResult ?? removeResult)?.total_tags
|
||||
if (finalTags) {
|
||||
updateAssetInCache(asset.id, { tags: finalTags }, cacheKey)
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Failed to update asset tags:', error)
|
||||
updateAssetInCache(asset.id, { tags: originalTags }, cacheKey)
|
||||
|
||||
if (removedTagsOnServer.length > 0) {
|
||||
try {
|
||||
await assetService.addAssetTags(asset.id, removedTagsOnServer)
|
||||
} catch (compensationError) {
|
||||
console.error(
|
||||
'Failed to restore tags after partial failure; invalidating cache to force refetch:',
|
||||
compensationError
|
||||
)
|
||||
const categoriesToInvalidate = new Set<string>()
|
||||
const resolved = cacheKey ? resolveCategory(cacheKey) : undefined
|
||||
if (resolved) {
|
||||
categoriesToInvalidate.add(resolved)
|
||||
}
|
||||
for (const [
|
||||
category,
|
||||
state
|
||||
] of modelStateByCategory.value.entries()) {
|
||||
if (state.assets?.has(asset.id)) {
|
||||
categoriesToInvalidate.add(category)
|
||||
}
|
||||
}
|
||||
for (const category of categoriesToInvalidate) {
|
||||
invalidateCategory(category)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Invalidate model caches for a given category (e.g., 'checkpoints', 'loras')
|
||||
* Clears the category cache and tag-based caches so next access triggers refetch
|
||||
* @param category The model category to invalidate (e.g., 'checkpoints')
|
||||
*/
|
||||
function invalidateModelsForCategory(category: string): void {
|
||||
invalidateCategory(category)
|
||||
invalidateCategory(`tag:${category}`)
|
||||
invalidateCategory('tag:models')
|
||||
}
|
||||
|
||||
return {
|
||||
getAssets: () => emptyAssets,
|
||||
isLoading: () => false,
|
||||
getError: () => undefined,
|
||||
hasMore: () => false,
|
||||
hasAssetKey: () => false,
|
||||
hasCategory: () => false,
|
||||
updateModelsForNodeType: async () => {},
|
||||
invalidateCategory: () => {},
|
||||
updateModelsForTag: async () => {},
|
||||
updateAssetMetadata: async () => {},
|
||||
updateAssetTags: async () => {},
|
||||
invalidateModelsForCategory: () => {}
|
||||
getAssets,
|
||||
isLoading,
|
||||
getError,
|
||||
hasMore,
|
||||
hasAssetKey,
|
||||
hasCategory,
|
||||
updateModelsForNodeType,
|
||||
updateModelsForTag,
|
||||
invalidateCategory,
|
||||
updateAssetMetadata,
|
||||
updateAssetTags,
|
||||
invalidateModelsForCategory
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -5,7 +5,12 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
import { assetService } from '@/platform/assets/services/assetService'
|
||||
import { useSettingStore } from '@/platform/settings/settingStore'
|
||||
import { api } from '@/scripts/api'
|
||||
import { useModelStore } from '@/stores/modelStore'
|
||||
import {
|
||||
ResourceState,
|
||||
effectiveModelExtensions,
|
||||
matchesModelExtension,
|
||||
useModelStore
|
||||
} from '@/stores/modelStore'
|
||||
|
||||
// Mock the api
|
||||
vi.mock('@/scripts/api', () => ({
|
||||
@@ -15,6 +20,7 @@ vi.mock('@/scripts/api', () => ({
|
||||
viewMetadata: vi.fn(),
|
||||
apiURL: vi.fn((path: string) => `http://localhost:8188${path}`),
|
||||
addEventListener: vi.fn(),
|
||||
addCustomEventListener: vi.fn(),
|
||||
removeEventListener: vi.fn()
|
||||
}
|
||||
}))
|
||||
@@ -22,8 +28,10 @@ vi.mock('@/scripts/api', () => ({
|
||||
// Mock the assetService
|
||||
vi.mock('@/platform/assets/services/assetService', () => ({
|
||||
assetService: {
|
||||
getAssetModelFolders: vi.fn(),
|
||||
getAssetModels: vi.fn()
|
||||
getAssetModels: vi.fn(),
|
||||
invalidateModelBuckets: vi.fn(),
|
||||
onModelsScanned: vi.fn(),
|
||||
seedModelAssets: vi.fn()
|
||||
}
|
||||
}))
|
||||
|
||||
@@ -57,16 +65,15 @@ function enableMocks(useAssetAPI = false) {
|
||||
{ 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'] }
|
||||
])
|
||||
// Asset API supplies only the per-folder model contents; folders come from
|
||||
// api.getModelFolders in both paths.
|
||||
vi.mocked(assetService.getAssetModels).mockResolvedValue([
|
||||
{ name: 'sdxl.safetensors', pathIndex: 0 },
|
||||
{ name: 'sdv15.safetensors', pathIndex: 0 },
|
||||
{ name: 'noinfo.safetensors', pathIndex: 0 }
|
||||
])
|
||||
vi.mocked(assetService.seedModelAssets).mockResolvedValue(undefined)
|
||||
vi.mocked(assetService.onModelsScanned).mockReturnValue(() => {})
|
||||
|
||||
vi.mocked(api.viewMetadata).mockImplementation((_, model) => {
|
||||
if (model === 'noinfo.safetensors') {
|
||||
@@ -209,6 +216,193 @@ describe('useModelStore', () => {
|
||||
expect(api.getModelFolders).toHaveBeenCalledTimes(2)
|
||||
expect(api.getModels).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('kicks off a backend scan when models come from the asset API', async () => {
|
||||
enableMocks(true)
|
||||
store = useModelStore()
|
||||
|
||||
await store.refresh()
|
||||
|
||||
expect(assetService.seedModelAssets).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
it('does not scan on the legacy listing path', async () => {
|
||||
enableMocks(false)
|
||||
store = useModelStore()
|
||||
|
||||
await store.refresh()
|
||||
|
||||
expect(assetService.seedModelAssets).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe('concurrent folder loads', () => {
|
||||
it('does not let a stale folder response overwrite a fresher one', async () => {
|
||||
enableMocks()
|
||||
let resolveStale!: (value: { name: string; folders: string[] }[]) => void
|
||||
vi.mocked(api.getModelFolders).mockReturnValueOnce(
|
||||
new Promise((resolve) => {
|
||||
resolveStale = resolve
|
||||
})
|
||||
)
|
||||
store = useModelStore()
|
||||
const staleLoad = store.loadModelFolders()
|
||||
|
||||
vi.mocked(api.getModelFolders).mockResolvedValueOnce([
|
||||
{ name: 'fresh-folder', folders: ['/fresh'] }
|
||||
])
|
||||
await store.loadModelFolders()
|
||||
expect(store.modelFolders.map((f) => f.directory)).toEqual([
|
||||
'fresh-folder'
|
||||
])
|
||||
|
||||
resolveStale([{ name: 'stale-folder', folders: ['/stale'] }])
|
||||
await staleLoad
|
||||
|
||||
expect(store.modelFolders.map((f) => f.directory)).toEqual([
|
||||
'fresh-folder'
|
||||
])
|
||||
})
|
||||
})
|
||||
|
||||
it('eager-loading before boot loads the folder structure first', async () => {
|
||||
enableMocks()
|
||||
store = useModelStore()
|
||||
|
||||
await store.loadModels()
|
||||
|
||||
expect(api.getModelFolders).toHaveBeenCalledTimes(1)
|
||||
expect(api.getModels).toHaveBeenCalledWith('checkpoints')
|
||||
expect(api.getModels).toHaveBeenCalledWith('vae')
|
||||
})
|
||||
|
||||
describe('refreshModelFolder races', () => {
|
||||
it('keeps the newer refresh when an older one for the same folder finishes last', async () => {
|
||||
enableMocks()
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
await store.getLoadedModelFolder('checkpoints')
|
||||
|
||||
let resolveOld!: (value: { name: string; pathIndex: number }[]) => void
|
||||
vi.mocked(api.getModels).mockReturnValueOnce(
|
||||
new Promise((resolve) => {
|
||||
resolveOld = resolve
|
||||
})
|
||||
)
|
||||
const oldRefresh = store.refreshModelFolder('checkpoints')
|
||||
|
||||
vi.mocked(api.getModels).mockResolvedValueOnce([
|
||||
{ name: 'newer.safetensors', pathIndex: 0 }
|
||||
])
|
||||
await store.refreshModelFolder('checkpoints')
|
||||
|
||||
resolveOld([{ name: 'older.safetensors', pathIndex: 0 }])
|
||||
await oldRefresh
|
||||
|
||||
const folder = await store.getLoadedModelFolder('checkpoints')
|
||||
expect(folder!.models['0/newer.safetensors']).toBeDefined()
|
||||
expect(folder!.models['0/older.safetensors']).toBeUndefined()
|
||||
})
|
||||
|
||||
it('does not resurrect a stale folder over a fresher structure', async () => {
|
||||
enableMocks()
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
await store.getLoadedModelFolder('checkpoints')
|
||||
|
||||
let resolveStaleContents!: (
|
||||
value: { name: string; pathIndex: number }[]
|
||||
) => void
|
||||
vi.mocked(api.getModels).mockReturnValueOnce(
|
||||
new Promise((resolve) => {
|
||||
resolveStaleContents = resolve
|
||||
})
|
||||
)
|
||||
const staleRefresh = store.refreshModelFolder('checkpoints')
|
||||
|
||||
// A full reload rebuilds the folder structure mid-refresh.
|
||||
await store.loadModelFolders()
|
||||
const freshFolder = await store.getLoadedModelFolder('checkpoints')
|
||||
|
||||
resolveStaleContents([{ name: 'stale.safetensors', pathIndex: 0 }])
|
||||
await staleRefresh
|
||||
|
||||
const current = await store.getLoadedModelFolder('checkpoints')
|
||||
expect(current).toBe(freshFolder)
|
||||
expect(current!.models['0/stale.safetensors']).toBeUndefined()
|
||||
})
|
||||
})
|
||||
|
||||
describe('scan fast-phase completion', () => {
|
||||
it('re-loads folders whose eager load was still in flight when the reload fired', async () => {
|
||||
enableMocks(true)
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
|
||||
// Eager load starts but its fetch never lands before the scan event.
|
||||
let resolveEager!: (value: { name: string; pathIndex: number }[]) => void
|
||||
vi.mocked(assetService.getAssetModels).mockReturnValueOnce(
|
||||
new Promise((resolve) => {
|
||||
resolveEager = resolve
|
||||
})
|
||||
)
|
||||
const eagerLoad = store.getLoadedModelFolder('checkpoints')
|
||||
|
||||
const scanCallback = vi.mocked(assetService.onModelsScanned).mock
|
||||
.calls[0]?.[0]
|
||||
await scanCallback!()
|
||||
|
||||
// The rebuilt folder must have been re-loaded, not left uninitialized
|
||||
// while the original request finishes into a detached folder object.
|
||||
const folder = store.modelFolders.find(
|
||||
(f) => f.directory === 'checkpoints'
|
||||
)
|
||||
expect(folder!.state).toBe(ResourceState.Loaded)
|
||||
|
||||
resolveEager([{ name: 'detached.safetensors', pathIndex: 0 }])
|
||||
await eagerLoad
|
||||
const current = await store.getLoadedModelFolder('checkpoints')
|
||||
expect(current!.models['0/detached.safetensors']).toBeUndefined()
|
||||
})
|
||||
|
||||
it('re-loads previously loaded folders when the event fires', async () => {
|
||||
enableMocks(true)
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
await store.getLoadedModelFolder('checkpoints')
|
||||
expect(assetService.getAssetModels).toHaveBeenCalledTimes(1)
|
||||
|
||||
const scanCallback = vi.mocked(assetService.onModelsScanned).mock
|
||||
.calls[0]?.[0]
|
||||
expect(scanCallback).toBeDefined()
|
||||
await scanCallback!()
|
||||
await vi.waitFor(() => {
|
||||
expect(assetService.getAssetModels).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
|
||||
expect(assetService.invalidateModelBuckets).toHaveBeenCalled()
|
||||
expect(assetService.seedModelAssets).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('logs instead of rejecting when the post-scan reload fails', async () => {
|
||||
const error = vi.spyOn(console, 'error').mockImplementation(() => {})
|
||||
enableMocks(true)
|
||||
vi.mocked(api.getModelFolders).mockRejectedValue(
|
||||
new Error('transient network failure')
|
||||
)
|
||||
store = useModelStore()
|
||||
const scanCallback = vi.mocked(assetService.onModelsScanned).mock
|
||||
.calls[0]?.[0]
|
||||
|
||||
await scanCallback!()
|
||||
await vi.waitFor(() => {
|
||||
expect(error).toHaveBeenCalledWith(
|
||||
expect.stringContaining('reload'),
|
||||
expect.any(Error)
|
||||
)
|
||||
})
|
||||
error.mockRestore()
|
||||
})
|
||||
})
|
||||
|
||||
describe('API switching functionality', () => {
|
||||
@@ -218,28 +412,117 @@ describe('useModelStore', () => {
|
||||
await store.loadModelFolders()
|
||||
const folderStore = await store.getLoadedModelFolder('checkpoints')
|
||||
|
||||
// Both APIs return objects with .name property, modelStore extracts folder.name in both cases
|
||||
// Folders come from /experiment/models; legacy path also serves models.
|
||||
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 () => {
|
||||
it('should use asset API for model contents but /experiment/models for folders when UseAssetAPI 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)
|
||||
// Folders always come from /experiment/models; only contents use the asset API.
|
||||
expect(api.getModelFolders).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)
|
||||
})
|
||||
|
||||
it('filters asset-path folder contents by the folder extensions', async () => {
|
||||
enableMocks(true)
|
||||
vi.mocked(api.getModelFolders).mockResolvedValue([
|
||||
{ name: 'checkpoints', folders: ['/p'], extensions: ['.safetensors'] }
|
||||
])
|
||||
vi.mocked(assetService.getAssetModels).mockResolvedValue([
|
||||
{ name: 'keep.safetensors', pathIndex: 0 },
|
||||
{ name: 'notes.txt', pathIndex: 0 }
|
||||
])
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
const folder = await store.getLoadedModelFolder('checkpoints')
|
||||
|
||||
const names = Object.values(folder!.models).map((m) => m.file_name)
|
||||
expect(names).toEqual(['keep.safetensors'])
|
||||
})
|
||||
|
||||
it('hides non-model noise in match-all folders on the asset path', async () => {
|
||||
enableMocks(true)
|
||||
vi.mocked(api.getModelFolders).mockResolvedValue([
|
||||
{ name: 'LLM', folders: ['/p'], extensions: [] }
|
||||
])
|
||||
vi.mocked(assetService.getAssetModels).mockResolvedValue([
|
||||
{ name: 'model.safetensors', pathIndex: 0 },
|
||||
{ name: 'README.md', pathIndex: 0 }
|
||||
])
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
const folder = await store.getLoadedModelFolder('LLM')
|
||||
|
||||
const names = Object.values(folder!.models).map((m) => m.file_name)
|
||||
expect(names).toEqual(['model.safetensors'])
|
||||
})
|
||||
|
||||
it('leaves the legacy listing unfiltered', async () => {
|
||||
enableMocks(false)
|
||||
vi.mocked(api.getModelFolders).mockResolvedValue([
|
||||
{ name: 'checkpoints', folders: ['/p'], extensions: ['.safetensors'] }
|
||||
])
|
||||
vi.mocked(api.getModels).mockResolvedValue([
|
||||
{ name: 'keep.safetensors', pathIndex: 0 },
|
||||
{ name: 'legacy-visible.gguf', pathIndex: 0 }
|
||||
])
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
const folder = await store.getLoadedModelFolder('checkpoints')
|
||||
|
||||
const names = Object.values(folder!.models).map((m) => m.file_name)
|
||||
expect(names).toEqual(['keep.safetensors', 'legacy-visible.gguf'])
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe(matchesModelExtension, () => {
|
||||
it('keeps files whose extension is in the folder list', () => {
|
||||
expect(
|
||||
matchesModelExtension('a.safetensors', ['.safetensors', '.ckpt'])
|
||||
).toBe(true)
|
||||
expect(matchesModelExtension('a.txt', ['.safetensors'])).toBe(false)
|
||||
})
|
||||
|
||||
it('matches case-insensitively and on subpaths', () => {
|
||||
expect(
|
||||
matchesModelExtension('sub/dir/A.SAFETENSORS', ['.safetensors'])
|
||||
).toBe(true)
|
||||
})
|
||||
|
||||
it('is permissive when there are no real extensions', () => {
|
||||
// Unfiltered folders (empty) and the `folder`/`''` sentinels show everything.
|
||||
expect(matchesModelExtension('readme.md', [])).toBe(true)
|
||||
expect(matchesModelExtension('anything', ['folder'])).toBe(true)
|
||||
expect(matchesModelExtension('anything', [''])).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe(effectiveModelExtensions, () => {
|
||||
it('uses a registered allowlist verbatim', () => {
|
||||
expect(effectiveModelExtensions(['.gguf'])).toEqual(['.gguf'])
|
||||
})
|
||||
|
||||
it('substitutes the default list for match-all folders', () => {
|
||||
const effective = effectiveModelExtensions([])
|
||||
expect(effective).toContain('.safetensors')
|
||||
expect(matchesModelExtension('readme.md', effective)).toBe(false)
|
||||
})
|
||||
|
||||
it('treats an absent field (older backends) like match-all', () => {
|
||||
expect(effectiveModelExtensions(undefined)).toEqual(
|
||||
effectiveModelExtensions([])
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
import { defineStore } from 'pinia'
|
||||
import { computed, ref } from 'vue'
|
||||
import { computed, onScopeDispose, ref } from 'vue'
|
||||
|
||||
import type { ModelFile } from '@/platform/assets/schemas/assetSchema'
|
||||
import { assetService } from '@/platform/assets/services/assetService'
|
||||
import { isCloud } from '@/platform/distribution/types'
|
||||
import { useSettingStore } from '@/platform/settings/settingStore'
|
||||
import { api } from '@/scripts/api'
|
||||
|
||||
@@ -100,6 +101,11 @@ export class ComfyModelDef {
|
||||
if (this.has_loaded_metadata || this.is_load_requested) {
|
||||
return
|
||||
}
|
||||
// viewMetadata reads the safetensors header off local disk; on Cloud the
|
||||
// model bytes live in object storage so there is nothing to read.
|
||||
if (isCloud) {
|
||||
return
|
||||
}
|
||||
this.is_load_requested = true
|
||||
try {
|
||||
const metadata = await api.viewMetadata(this.directory, this.file_name)
|
||||
@@ -156,6 +162,66 @@ export enum ResourceState {
|
||||
Loaded
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves the preview image for a model: embedded metadata thumbnail when
|
||||
* loaded, otherwise the server-rendered `.webp` preview. The preview endpoint
|
||||
* reads a rendered thumbnail off local disk, which is unavailable on Cloud
|
||||
* (model bytes live in object storage), so Cloud resolves to no preview.
|
||||
*/
|
||||
export function getModelPreviewUrl(model: ComfyModelDef): string {
|
||||
if (model.image) return model.image
|
||||
if (isCloud) return ''
|
||||
const extension = model.file_name.split('.').pop()
|
||||
const filename = model.file_name.replace(`.${extension}`, '.webp')
|
||||
const encodedFilename = encodeURIComponent(filename).replace(/%2F/g, '/')
|
||||
return `/api/experiment/models/preview/${model.directory}/${model.path_index}/${encodedFilename}`
|
||||
}
|
||||
|
||||
/**
|
||||
* FE-owned copy of core's default `supported_pt_extensions`, applied to
|
||||
* match-all folders (empty registered allowlist) so they don't surface
|
||||
* README/config noise. Accepted to go stale across core version bumps; the
|
||||
* whole surface is expected to be short-lived.
|
||||
*/
|
||||
const DEFAULT_MODEL_EXTENSIONS = [
|
||||
'.ckpt',
|
||||
'.pt',
|
||||
'.pt2',
|
||||
'.bin',
|
||||
'.pth',
|
||||
'.safetensors',
|
||||
'.pkl',
|
||||
'.sft'
|
||||
]
|
||||
|
||||
/**
|
||||
* Resolves a folder's display allowlist from its raw registered `extensions`
|
||||
* (`/experiment/models`): non-empty is used verbatim; an empty array
|
||||
* (match-all) or an absent field (older backends) takes the FE default list,
|
||||
* reproducing the legacy sidebar's global-set behavior so nothing that used
|
||||
* to be hidden starts showing.
|
||||
*/
|
||||
export function effectiveModelExtensions(
|
||||
extensions: string[] | undefined
|
||||
): string[] {
|
||||
return extensions?.length ? extensions : DEFAULT_MODEL_EXTENSIONS
|
||||
}
|
||||
|
||||
/**
|
||||
* Whether a model file belongs in a folder given its display allowlist. An
|
||||
* empty list, or a list with no real (`.`-prefixed) extensions (the
|
||||
* `'folder'`/`''` sentinels), leaves the folder unfiltered.
|
||||
*/
|
||||
export function matchesModelExtension(
|
||||
fileName: string,
|
||||
extensions: string[]
|
||||
): boolean {
|
||||
const realExtensions = extensions.filter((ext) => ext.startsWith('.'))
|
||||
if (realExtensions.length === 0) return true
|
||||
const lower = fileName.toLowerCase()
|
||||
return realExtensions.some((ext) => lower.endsWith(ext.toLowerCase()))
|
||||
}
|
||||
|
||||
export class ModelFolder {
|
||||
/** Models in this folder */
|
||||
models: Record<string, ComfyModelDef> = {}
|
||||
@@ -163,7 +229,8 @@ export class ModelFolder {
|
||||
|
||||
constructor(
|
||||
public directory: string,
|
||||
private getModelsFunc: (folder: string) => Promise<ModelFile[]>
|
||||
private getModelsFunc: (folder: string) => Promise<ModelFile[]>,
|
||||
public readonly extensions: string[] = []
|
||||
) {}
|
||||
|
||||
get key(): string {
|
||||
@@ -180,6 +247,7 @@ export class ModelFolder {
|
||||
this.state = ResourceState.Loading
|
||||
const models = await this.getModelsFunc(this.directory)
|
||||
for (const model of models) {
|
||||
if (!matchesModelExtension(model.name, this.extensions)) continue
|
||||
this.models[`${model.pathIndex}/${model.name}`] = new ComfyModelDef(
|
||||
model.name,
|
||||
this.directory,
|
||||
@@ -212,22 +280,33 @@ export const useModelStore = defineStore('models', () => {
|
||||
: (folder) => api.getModels(folder)
|
||||
}
|
||||
|
||||
let modelFoldersRequestId = 0
|
||||
|
||||
/**
|
||||
* Loads the model folders from the server
|
||||
* Loads the model folders from the server.
|
||||
*
|
||||
* The folder list (and its registration order) always comes from
|
||||
* `/experiment/models`, the source of truth for which model folders exist;
|
||||
* only the per-folder contents differ between the asset API and legacy paths.
|
||||
* Concurrent loads (manual refresh racing the scan-complete reload) commit
|
||||
* only the newest request so a slow stale response cannot overwrite a
|
||||
* fresher folder structure.
|
||||
*/
|
||||
async function loadModelFolders() {
|
||||
const useAssetAPI: boolean = settingStore.get('Comfy.Assets.UseAssetAPI')
|
||||
|
||||
const resData = useAssetAPI
|
||||
? await assetService.getAssetModelFolders()
|
||||
: await api.getModelFolders()
|
||||
const requestId = ++modelFoldersRequestId
|
||||
const resData = await api.getModelFolders()
|
||||
if (requestId !== modelFoldersRequestId) return
|
||||
modelFolderNames.value = resData.map((folder) => folder.name)
|
||||
modelFolderByName.value = {}
|
||||
const useAssetAPI: boolean = settingStore.get('Comfy.Assets.UseAssetAPI')
|
||||
const getModelsFunc = createGetModelsFunc()
|
||||
for (const folderName of modelFolderNames.value) {
|
||||
modelFolderByName.value[folderName] = new ModelFolder(
|
||||
folderName,
|
||||
getModelsFunc
|
||||
for (const folder of resData) {
|
||||
modelFolderByName.value[folder.name] = new ModelFolder(
|
||||
folder.name,
|
||||
getModelsFunc,
|
||||
// Display filtering applies to the asset walk only; the legacy
|
||||
// listing keeps its historical server-side (global-set) filtering.
|
||||
useAssetAPI ? effectiveModelExtensions(folder.extensions) : []
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -240,9 +319,15 @@ export const useModelStore = defineStore('models', () => {
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads all model folders' contents from the server
|
||||
* Loads all model folders' contents from the server. Loads the folder
|
||||
* structure first when it has not arrived yet — eager loading can run
|
||||
* before app boot's own loadModelFolders call resolves, and iterating an
|
||||
* empty folder list would silently load nothing.
|
||||
*/
|
||||
async function loadModels() {
|
||||
if (modelFolderNames.value.length === 0) {
|
||||
await loadModelFolders()
|
||||
}
|
||||
return Promise.all(modelFolders.value.map((folder) => folder.load()))
|
||||
}
|
||||
|
||||
@@ -253,25 +338,45 @@ export const useModelStore = defineStore('models', () => {
|
||||
* a newly-introduced folder type is picked up without dropping other
|
||||
* folders' loaded contents.
|
||||
*/
|
||||
const folderRefreshIds = new Map<string, number>()
|
||||
|
||||
async function refreshModelFolder(folderName: string) {
|
||||
assetService.invalidateModelBuckets()
|
||||
if (!(folderName in modelFolderByName.value)) {
|
||||
await refresh()
|
||||
return
|
||||
}
|
||||
const folder = new ModelFolder(folderName, createGetModelsFunc())
|
||||
const requestId = modelFoldersRequestId
|
||||
const refreshId = (folderRefreshIds.get(folderName) ?? 0) + 1
|
||||
folderRefreshIds.set(folderName, refreshId)
|
||||
const folder = new ModelFolder(
|
||||
folderName,
|
||||
createGetModelsFunc(),
|
||||
modelFolderByName.value[folderName].extensions
|
||||
)
|
||||
await folder.load()
|
||||
// A full reload may have rebuilt the folder structure while this folder
|
||||
// refreshed, and a newer refresh of the same folder may have already
|
||||
// committed; committing then would resurrect a stale folder object.
|
||||
if (requestId !== modelFoldersRequestId) return
|
||||
if (folderRefreshIds.get(folderName) !== refreshId) return
|
||||
modelFolderByName.value[folderName] = folder
|
||||
}
|
||||
|
||||
/**
|
||||
* Refreshes the folder structure and re-loads any folder whose contents
|
||||
* had previously been loaded. Used by manual refresh actions ("r" key,
|
||||
* sidebar refresh button) to pick up on-disk changes without losing the
|
||||
* currently-visible contents.
|
||||
* Re-fetches the folder structure and re-loads any folder whose contents
|
||||
* had previously been loaded, picking up server-side changes without
|
||||
* losing the currently-visible contents.
|
||||
*/
|
||||
async function refresh() {
|
||||
async function reloadModels() {
|
||||
assetService.invalidateModelBuckets()
|
||||
// Loading counts as previously loaded: a scan-complete reload can land
|
||||
// while the eager load is still in flight, and replacing those folder
|
||||
// objects without re-loading them would strand the sidebar on
|
||||
// uninitialized folders whose original loads finish into detached
|
||||
// objects.
|
||||
const previouslyLoaded = modelFolders.value
|
||||
.filter((folder) => folder.state === ResourceState.Loaded)
|
||||
.filter((folder) => folder.state !== ResourceState.Uninitialized)
|
||||
.map((folder) => folder.directory)
|
||||
await loadModelFolders()
|
||||
await Promise.all(
|
||||
@@ -281,6 +386,43 @@ export const useModelStore = defineStore('models', () => {
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Asks the backend to rescan the model roots so files added on disk since
|
||||
* startup become assets. Skipped on Cloud (models are ingested via uploads,
|
||||
* not scanned from disk) and on the legacy listing path (which reads the
|
||||
* filesystem live on every request).
|
||||
*/
|
||||
async function requestModelScan() {
|
||||
if (isCloud) return
|
||||
if (!settingStore.get('Comfy.Assets.UseAssetAPI')) return
|
||||
try {
|
||||
await assetService.seedModelAssets()
|
||||
} catch (error) {
|
||||
console.warn('Unable to start model asset scan', error)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Manual refresh ("r" key, sidebar refresh button): kicks off a backend
|
||||
* rescan and immediately re-loads the currently known server state; the
|
||||
* scan completion subscription below re-loads again with whatever the
|
||||
* scan discovered. The scan is deliberately not awaited so it runs
|
||||
* concurrently with the reload.
|
||||
*/
|
||||
async function refresh() {
|
||||
void requestModelScan()
|
||||
await reloadModels()
|
||||
}
|
||||
|
||||
const unsubscribeModelsScanned = assetService.onModelsScanned(async () => {
|
||||
try {
|
||||
await reloadModels()
|
||||
} catch (error) {
|
||||
console.error('Failed to reload the model library after a scan', error)
|
||||
}
|
||||
})
|
||||
onScopeDispose(unsubscribeModelsScanned)
|
||||
|
||||
return {
|
||||
models,
|
||||
modelFolders,
|
||||
|
||||
@@ -5,12 +5,24 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
import { useSidebarTabStore } from '@/stores/workspace/sidebarTabStore'
|
||||
|
||||
const { mockGetSetting, mockRegisterCommand, mockRegisterCommands } =
|
||||
vi.hoisted(() => ({
|
||||
const {
|
||||
mockGetSetting,
|
||||
mockRegisterCommand,
|
||||
mockRegisterCommands,
|
||||
mockBrowseModelAssets,
|
||||
registeredCommands,
|
||||
commandStoreCommands
|
||||
} = vi.hoisted(() => {
|
||||
const registeredCommands: { id: string; function: () => unknown }[] = []
|
||||
return {
|
||||
mockGetSetting: vi.fn(),
|
||||
mockRegisterCommand: vi.fn(),
|
||||
mockRegisterCommands: vi.fn()
|
||||
}))
|
||||
mockRegisterCommand: vi.fn((command) => registeredCommands.push(command)),
|
||||
mockRegisterCommands: vi.fn(),
|
||||
mockBrowseModelAssets: vi.fn(),
|
||||
registeredCommands,
|
||||
commandStoreCommands: [] as { id: string; function: () => unknown }[]
|
||||
}
|
||||
})
|
||||
|
||||
vi.mock('@/platform/settings/settingStore', () => ({
|
||||
useSettingStore: () => ({
|
||||
@@ -21,7 +33,7 @@ vi.mock('@/platform/settings/settingStore', () => ({
|
||||
vi.mock('@/stores/commandStore', () => ({
|
||||
useCommandStore: () => ({
|
||||
registerCommand: mockRegisterCommand,
|
||||
commands: []
|
||||
commands: commandStoreCommands
|
||||
})
|
||||
}))
|
||||
|
||||
@@ -99,8 +111,18 @@ describe('useSidebarTabStore', () => {
|
||||
mockGetSetting.mockReset()
|
||||
mockRegisterCommand.mockClear()
|
||||
mockRegisterCommands.mockClear()
|
||||
mockBrowseModelAssets.mockClear()
|
||||
registeredCommands.length = 0
|
||||
commandStoreCommands.length = 0
|
||||
})
|
||||
|
||||
const toggleModelLibrary = async () => {
|
||||
const toggleCommand = registeredCommands.find(
|
||||
(command) => command.id === 'Workspace.ToggleSidebarTab.model-library'
|
||||
)
|
||||
await toggleCommand?.function()
|
||||
}
|
||||
|
||||
it('registers the job history tab when QPO V2 is enabled', () => {
|
||||
mockGetSetting.mockImplementation((key: string) =>
|
||||
key === 'Comfy.Queue.QPOV2' ? true : undefined
|
||||
@@ -160,4 +182,63 @@ describe('useSidebarTabStore', () => {
|
||||
])
|
||||
expect(mockRegisterCommand).toHaveBeenCalledTimes(6)
|
||||
})
|
||||
|
||||
describe('model library view selection', () => {
|
||||
it('toggles the sidebar tab when the asset view is disabled', async () => {
|
||||
mockGetSetting.mockImplementation((key: string) =>
|
||||
key === 'Comfy.ModelLibrary.UseAssetBrowser' ? false : undefined
|
||||
)
|
||||
commandStoreCommands.push({
|
||||
id: 'Comfy.BrowseModelAssets',
|
||||
function: mockBrowseModelAssets
|
||||
})
|
||||
|
||||
const store = useSidebarTabStore()
|
||||
store.registerCoreSidebarTabs()
|
||||
|
||||
await toggleModelLibrary()
|
||||
|
||||
expect(store.activeSidebarTabId).toBe('model-library')
|
||||
expect(mockBrowseModelAssets).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('opens the asset browser when the browser and asset API are enabled', async () => {
|
||||
mockGetSetting.mockImplementation((key: string) =>
|
||||
key === 'Comfy.ModelLibrary.UseAssetBrowser' ||
|
||||
key === 'Comfy.Assets.UseAssetAPI'
|
||||
? true
|
||||
: undefined
|
||||
)
|
||||
commandStoreCommands.push({
|
||||
id: 'Comfy.BrowseModelAssets',
|
||||
function: mockBrowseModelAssets
|
||||
})
|
||||
|
||||
const store = useSidebarTabStore()
|
||||
store.registerCoreSidebarTabs()
|
||||
|
||||
await toggleModelLibrary()
|
||||
|
||||
expect(mockBrowseModelAssets).toHaveBeenCalledOnce()
|
||||
expect(store.activeSidebarTabId).toBeNull()
|
||||
})
|
||||
|
||||
it('falls back to the sidebar tree when the asset API is disabled', async () => {
|
||||
mockGetSetting.mockImplementation((key: string) =>
|
||||
key === 'Comfy.ModelLibrary.UseAssetBrowser' ? true : false
|
||||
)
|
||||
commandStoreCommands.push({
|
||||
id: 'Comfy.BrowseModelAssets',
|
||||
function: mockBrowseModelAssets
|
||||
})
|
||||
|
||||
const store = useSidebarTabStore()
|
||||
store.registerCoreSidebarTabs()
|
||||
|
||||
await toggleModelLibrary()
|
||||
|
||||
expect(store.activeSidebarTabId).toBe('model-library')
|
||||
expect(mockBrowseModelAssets).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
@@ -76,8 +76,13 @@ export const useSidebarTabStore = defineStore('sidebarTab', () => {
|
||||
const settingStore = useSettingStore()
|
||||
const commandStore = useCommandStore()
|
||||
|
||||
// The asset browser cannot function without the asset API, so the
|
||||
// browser routing derives from both settings: with the API disabled
|
||||
// the browser setting is inert and the tab always opens the sidebar
|
||||
// tree, rather than prompt-correcting the combination.
|
||||
if (
|
||||
tab.id === 'model-library' &&
|
||||
settingStore.get('Comfy.ModelLibrary.UseAssetBrowser') &&
|
||||
settingStore.get('Comfy.Assets.UseAssetAPI')
|
||||
) {
|
||||
await commandStore.commands
|
||||
|
||||
Reference in New Issue
Block a user