From 4582c71583108ca82432080e6c55ad7445d06af8 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Mon, 28 Oct 2024 21:23:53 -0400
Subject: [PATCH] [Refactor] Rework modelStore and ModelLibrarySidebarTab
 (#1350)

* nit

* Rename

* nit

* Move load model folders to app level

* Various fixes

* nit

* nit

* wip

* nit

* nit

* nit

* Split

* nit

* Add back spinner

* nit

* nit

* Add refresh button

* nit

* nit

* Preserve model folder order

* Avoid order change on folder open
---
 .../sidebar/tabs/ModelLibrarySidebarTab.vue   | 134 +++++++++---------
 .../tabs/modelLibrary/ModelTreeLeaf.vue       |  11 +-
 src/i18n.ts                                   |   6 +
 src/scripts/api.ts                            |   9 +-
 src/scripts/app.ts                            |   5 +-
 src/stores/modelStore.ts                      | 107 +++++++-------
 src/views/GraphView.vue                       |   4 +
 tests-ui/tests/store/modelStore.test.ts       |  18 ++-
 8 files changed, 154 insertions(+), 140 deletions(-)

diff --git a/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
index 9d7231fd4..5dc561ad0 100644
--- a/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
@@ -3,6 +3,22 @@
     :title="$t('sideToolbar.modelLibrary')"
     class="bg-[var(--p-tree-background)]"
   >
+    <template #tool-buttons>
+      <Button
+        icon="pi pi-refresh"
+        @click="modelStore.loadModelFolders"
+        severity="secondary"
+        text
+        v-tooltip="$t('refresh')"
+      />
+      <Button
+        icon="pi pi-cloud-download"
+        @click="modelStore.loadModels"
+        severity="secondary"
+        text
+        v-tooltip="$t('loadAllFolders')"
+      />
+    </template>
     <template #header>
       <SearchBox
         class="model-lib-search-box p-4"
@@ -26,12 +42,17 @@
 </template>
 
 <script setup lang="ts">
+import Button from 'primevue/button'
 import SearchBox from '@/components/common/SearchBox.vue'
-import { useI18n } from 'vue-i18n'
 import TreeExplorer from '@/components/common/TreeExplorer.vue'
 import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue'
 import ModelTreeLeaf from '@/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue'
-import { ComfyModelDef, useModelStore } from '@/stores/modelStore'
+import {
+  ComfyModelDef,
+  ModelFolder,
+  ResourceState,
+  useModelStore
+} from '@/stores/modelStore'
 import { useModelToNodeStore } from '@/stores/modelToNodeStore'
 import { useSettingStore } from '@/stores/settingStore'
 import { useTreeExpansion } from '@/hooks/treeHooks'
@@ -39,11 +60,10 @@ import type {
   RenderedTreeExplorerNode,
   TreeExplorerNode
 } from '@/types/treeExplorerTypes'
-import { computed, ref, type ComputedRef, watch, toRef } from 'vue'
+import { computed, ref, watch, toRef, onMounted } from 'vue'
 import type { TreeNode } from 'primevue/treenode'
 import { app } from '@/scripts/app'
 import { buildTree } from '@/utils/treeUtil'
-const { t } = useI18n()
 const modelStore = useModelStore()
 const modelToNodeStore = useModelToNodeStore()
 const settingStore = useSettingStore()
@@ -51,65 +71,36 @@ const searchQuery = ref<string>('')
 const expandedKeys = ref<Record<string, boolean>>({})
 const { toggleNodeOnEvent } = useTreeExpansion(expandedKeys)
 
-const root: ComputedRef<TreeNode> = computed(() => {
-  let modelList: ComfyModelDef[] = []
-  if (!modelStore.modelFolders.length) {
-    modelStore.getModelFolders()
-  }
-  if (settingStore.get('Comfy.ModelLibrary.AutoLoadAll')) {
-    for (let folder of modelStore.modelFolders) {
-      modelStore.getModelsInFolderCached(folder)
-    }
-  }
-  for (let folder of modelStore.modelFolders) {
-    const models = modelStore.modelStoreMap[folder]
-    if (models) {
-      if (Object.values(models.models).length) {
-        modelList.push(...Object.values(models.models))
-      } else {
-        // ModelDef with key 'folder/a/b/c/' is treated as empty folder
-        const fakeModel = new ComfyModelDef('', folder)
-        fakeModel.is_fake_object = true
-        modelList.push(fakeModel)
-      }
-    } else {
-      const fakeModel = new ComfyModelDef('Loading', folder)
-      fakeModel.is_fake_object = true
-      modelList.push(fakeModel)
-    }
-  }
+const filteredModels = computed<ComfyModelDef[]>(() => {
   if (searchQuery.value) {
     const search = searchQuery.value.toLocaleLowerCase()
-    modelList = modelList.filter((model: ComfyModelDef) => {
+    return modelStore.models.filter((model: ComfyModelDef) => {
       return model.searchable.includes(search)
     })
   }
-  const tree: TreeNode = buildTree(modelList, (model: ComfyModelDef) =>
-    model.key.split('/')
+  return modelStore.models
+})
+
+type ModelOrFolder = ComfyModelDef | ModelFolder
+
+const root = computed<TreeNode>(() => {
+  const allNodes: ModelOrFolder[] = [
+    ...modelStore.modelFolders,
+    ...filteredModels.value
+  ]
+  return buildTree(allNodes, (modelOrFolder: ModelOrFolder) =>
+    modelOrFolder.key.split('/')
   )
-  return tree
 })
 
-const renderedRoot = computed<TreeExplorerNode<ComfyModelDef>>(() => {
+const renderedRoot = computed<TreeExplorerNode<ModelOrFolder>>(() => {
   const nameFormat = settingStore.get('Comfy.ModelLibrary.NameFormat')
-  const fillNodeInfo = (node: TreeNode): TreeExplorerNode<ComfyModelDef> => {
+  const fillNodeInfo = (node: TreeNode): TreeExplorerNode<ModelOrFolder> => {
     const children = node.children?.map(fillNodeInfo)
     const model: ComfyModelDef | null =
       node.leaf && node.data ? node.data : null
-    if (model?.is_fake_object) {
-      if (model.file_name === 'Loading') {
-        return {
-          key: node.key,
-          label: t('loading') + '...',
-          leaf: true,
-          data: node.data,
-          getIcon: (node: TreeExplorerNode<ComfyModelDef>) => {
-            return 'pi pi-spin pi-spinner'
-          },
-          children: []
-        }
-      }
-    }
+    const folder: ModelFolder | null =
+      !node.leaf && node.data ? node.data : null
 
     return {
       key: node.key,
@@ -120,32 +111,29 @@ const renderedRoot = computed<TreeExplorerNode<ComfyModelDef>>(() => {
         : node.label,
       leaf: node.leaf,
       data: node.data,
-      getIcon: (node: TreeExplorerNode<ComfyModelDef>) => {
-        if (node.leaf) {
-          if (node.data && node.data.image) {
-            return 'pi pi-fake-spacer'
-          }
-          return 'pi pi-file'
+      getIcon: () => {
+        if (model) {
+          return model.image ? 'pi pi-image' : 'pi pi-file'
         }
+        if (folder) {
+          return folder.state === ResourceState.Loading
+            ? 'pi pi-spin pi-spinner'
+            : 'pi pi-folder'
+        }
+        return 'pi pi-folder'
       },
-      getBadgeText: (node: TreeExplorerNode<ComfyModelDef>) => {
-        if (node.leaf) {
+      getBadgeText: () => {
+        // Return null to apply default badge text
+        // Return empty string to hide badge
+        if (!folder) {
           return null
         }
-        if (node.children?.length === 1) {
-          const onlyChild = node.children[0]
-          if (onlyChild.data?.is_fake_object) {
-            if (onlyChild.data.file_name === 'Loading') {
-              return ''
-            }
-          }
-        }
-        return null
+        return folder.state === ResourceState.Loaded ? null : ''
       },
       children,
       draggable: node.leaf,
       handleClick: (
-        node: RenderedTreeExplorerNode<ComfyModelDef>,
+        node: RenderedTreeExplorerNode<ModelOrFolder>,
         e: MouseEvent
       ) => {
         if (node.leaf) {
@@ -178,13 +166,19 @@ watch(
         const folderPath = key.split('/').slice(1).join('/')
         if (folderPath && !folderPath.includes('/')) {
           // Trigger (async) load of model data for this folder
-          modelStore.getModelsInFolderCached(folderPath)
+          modelStore.getLoadedModelFolder(folderPath)
         }
       }
     })
   },
   { deep: true }
 )
+
+onMounted(async () => {
+  if (settingStore.get('Comfy.ModelLibrary.AutoLoadAll')) {
+    await modelStore.loadModels()
+  }
+})
 </script>
 
 <style scoped>
diff --git a/src/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue b/src/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue
index 42465815a..d8b6988e1 100644
--- a/src/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue
+++ b/src/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue
@@ -57,9 +57,6 @@ const sidebarLocation = computed<'left' | 'right'>(() =>
 )
 
 const handleModelHover = async () => {
-  if (modelDef.value.is_fake_object) {
-    return
-  }
   const hoverTarget = modelContentElement.value
   const targetRect = hoverTarget.getBoundingClientRect()
 
@@ -87,7 +84,6 @@ const showPreview = computed(() => {
   return (
     isHovered.value &&
     modelDef.value &&
-    !modelDef.value.is_fake_object &&
     modelDef.value.has_loaded_metadata &&
     (modelDef.value.author ||
       modelDef.value.simplified_file_name != modelDef.value.title ||
@@ -99,9 +95,6 @@ const showPreview = computed(() => {
 })
 
 const handleMouseEnter = async () => {
-  if (modelDef.value.is_fake_object) {
-    return
-  }
   isHovered.value = true
   await nextTick()
   handleModelHover()
@@ -113,9 +106,7 @@ onMounted(() => {
   modelContentElement.value = container.value?.closest('.p-tree-node-content')
   modelContentElement.value?.addEventListener('mouseenter', handleMouseEnter)
   modelContentElement.value?.addEventListener('mouseleave', handleMouseLeave)
-  if (!modelDef.value.is_fake_object) {
-    modelDef.value.load()
-  }
+  modelDef.value.load()
 })
 
 onUnmounted(() => {
diff --git a/src/i18n.ts b/src/i18n.ts
index 989fe6c8e..c4cd0f019 100644
--- a/src/i18n.ts
+++ b/src/i18n.ts
@@ -2,6 +2,8 @@ import { createI18n } from 'vue-i18n'
 
 const messages = {
   en: {
+    loadAllFolders: 'Load All Folders',
+    refresh: 'Refresh',
     terminal: 'Terminal',
     videoFailedToLoad: 'Video failed to load',
     extensionName: 'Extension Name',
@@ -114,6 +116,8 @@ const messages = {
     }
   },
   zh: {
+    loadAllFolders: '加载所有文件夹',
+    refresh: '刷新',
     terminal: '终端',
     videoFailedToLoad: '视频加载失败',
     extensionName: '扩展名称',
@@ -225,6 +229,8 @@ const messages = {
     }
   },
   ru: {
+    refresh: 'Обновить',
+    loadAllFolders: 'Загрузить все папки',
     terminal: 'Терминал',
     videoFailedToLoad: 'Видео не удалось загрузить',
     extensionName: 'Название расширения',
diff --git a/src/scripts/api.ts b/src/scripts/api.ts
index 487ec70a0..18a54978a 100644
--- a/src/scripts/api.ts
+++ b/src/scripts/api.ts
@@ -366,7 +366,10 @@ class ComfyApi extends EventTarget {
     if (res.status === 404) {
       return []
     }
-    return await res.json()
+    const folderBlacklist = ['configs', 'custom_nodes']
+    return (await res.json()).filter(
+      (folder: string) => !folderBlacklist.includes(folder)
+    )
   }
 
   /**
@@ -374,10 +377,10 @@ class ComfyApi extends EventTarget {
    * @param {string} folder The folder to list models from, such as 'checkpoints'
    * @returns The list of model filenames within the specified folder
    */
-  async getModels(folder: string) {
+  async getModels(folder: string): Promise<string[]> {
     const res = await this.fetchApi(`/models/${folder}`)
     if (res.status === 404) {
-      return null
+      return []
     }
     return await res.json()
   }
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index c3b2cdad6..9a32bb9f2 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -2293,8 +2293,8 @@ export class ComfyApp {
       graphData.models &&
       useSettingStore().get('Comfy.Workflow.ShowMissingModelsWarning')
     ) {
-      for (let m of graphData.models) {
-        const models_available = await useModelStore().getModelsInFolderCached(
+      for (const m of graphData.models) {
+        const models_available = await useModelStore().getLoadedModelFolder(
           m.directory
         )
         if (models_available === null) {
@@ -2929,7 +2929,6 @@ export class ComfyApp {
     }
     if (this.vueAppReady) {
       useToastStore().add(requestToastMessage)
-      useModelStore().clearCache()
     }
 
     const defs = await api.getNodeDefs({
diff --git a/src/stores/modelStore.ts b/src/stores/modelStore.ts
index c8f29757c..f862ab454 100644
--- a/src/stores/modelStore.ts
+++ b/src/stores/modelStore.ts
@@ -1,4 +1,4 @@
-import { ref } from 'vue'
+import { computed, ref } from 'vue'
 import { defineStore } from 'pinia'
 import { api } from '@/scripts/api'
 
@@ -51,8 +51,6 @@ export class ComfyModelDef {
   has_loaded_metadata: boolean = false
   /** If true, a metadata load request has been triggered, but may or may not yet have finished loading */
   is_load_requested: boolean = false
-  /** If true, this is a fake model object used as a placeholder for something (eg a loading icon) */
-  is_fake_object: boolean = false
   /** A string full of auto-computed lowercase-only searchable text for this model */
   searchable: string = ''
 
@@ -140,70 +138,83 @@ export class ComfyModelDef {
   }
 }
 
+export enum ResourceState {
+  Uninitialized,
+  Loading,
+  Loaded
+}
+
 export class ModelFolder {
+  /** Models in this folder */
   models: Record<string, ComfyModelDef> = {}
+  state: ResourceState = ResourceState.Uninitialized
 
-  constructor(directory: string, models: string[]) {
-    for (const model of models) {
-      this.models[model] = new ComfyModelDef(model, directory)
-    }
+  constructor(public directory: string) {}
+
+  get key(): string {
+    return this.directory + '/'
   }
 
-  async loadModelMetadata(modelName: string) {
-    if (this.models[modelName]) {
-      await this.models[modelName].load()
+  /**
+   * Loads the models in this folder from the server
+   */
+  async load() {
+    if (this.state !== ResourceState.Uninitialized) {
+      return this
+    }
+    this.state = ResourceState.Loading
+    const models = await api.getModels(this.directory)
+    for (const model of models) {
+      this.models[model] = new ComfyModelDef(model, this.directory)
     }
+    this.state = ResourceState.Loaded
+    return this
   }
 }
 
-const folderBlacklist = ['configs', 'custom_nodes']
-
 /** Model store handler, wraps individual per-folder model stores */
-export const useModelStore = defineStore('modelStore', () => {
-  const modelStoreMap = ref<Record<string, ModelFolder | null>>({})
-  const isLoading = ref<Record<string, Promise<ModelFolder | null> | null>>({})
-  const modelFolders = ref<string[]>([])
+export const useModelStore = defineStore('models', () => {
+  const modelFolderNames = ref<string[]>([])
+  const modelFolderByName = ref<Record<string, ModelFolder>>({})
+  const modelFolders = computed<ModelFolder[]>(() =>
+    modelFolderNames.value.map(
+      (folderName) => modelFolderByName.value[folderName]
+    )
+  )
+  const models = computed<ComfyModelDef[]>(() =>
+    modelFolders.value.flatMap((folder) => Object.values(folder.models))
+  )
 
-  async function getModelsInFolderCached(
-    folder: string
-  ): Promise<ModelFolder | null> {
-    if (folder in modelStoreMap.value) {
-      return modelStoreMap.value[folder]
-    }
-    if (isLoading.value[folder]) {
-      return isLoading.value[folder]
+  /**
+   * Loads the model folders from the server
+   */
+  async function loadModelFolders() {
+    modelFolderNames.value = await api.getModelFolders()
+    modelFolderByName.value = {}
+    for (const folderName of modelFolderNames.value) {
+      modelFolderByName.value[folderName] = new ModelFolder(folderName)
     }
-    const promise = api.getModels(folder).then((models) => {
-      if (!models) {
-        return null
-      }
-      const store = new ModelFolder(folder, models)
-      modelStoreMap.value[folder] = store
-      isLoading.value[folder] = null
-      return store
-    })
-    isLoading.value[folder] = promise
-    return promise
   }
 
-  function clearCache() {
-    Object.keys(modelStoreMap.value).forEach((key) => {
-      delete modelStoreMap.value[key]
-    })
+  async function getLoadedModelFolder(
+    folderName: string
+  ): Promise<ModelFolder | null> {
+    const folder = modelFolderByName.value[folderName]
+    return folder ? await folder.load() : null
   }
 
-  async function getModelFolders() {
-    modelFolders.value = (await api.getModelFolders()).filter(
-      (folder) => !folderBlacklist.includes(folder)
-    )
+  /**
+   * Loads all model folders' contents from the server
+   */
+  async function loadModels() {
+    return Promise.all(modelFolders.value.map((folder) => folder.load()))
   }
 
   return {
-    modelStoreMap,
-    isLoading,
+    models,
     modelFolders,
-    getModelsInFolderCached,
-    clearCache,
-    getModelFolders
+    loadModelFolders,
+    loadModels,
+    getLoadedModelFolder
   }
 })
diff --git a/src/views/GraphView.vue b/src/views/GraphView.vue
index 2d592dd1a..054dd439d 100644
--- a/src/views/GraphView.vue
+++ b/src/views/GraphView.vue
@@ -37,6 +37,7 @@ import { useSidebarTabStore } from '@/stores/workspace/sidebarTabStore'
 import { useNodeBookmarkStore } from '@/stores/nodeBookmarkStore'
 import { useNodeDefStore, useNodeFrequencyStore } from '@/stores/nodeDefStore'
 import { useBottomPanelStore } from '@/stores/workspace/bottomPanelStore'
+import { useModelStore } from '@/stores/modelStore'
 
 setupAutoQueueHandler()
 
@@ -159,6 +160,9 @@ const onGraphReady = () => {
       // Load keybindings.
       useKeybindingStore().loadUserKeybindings()
 
+      // Load model folders
+      useModelStore().loadModelFolders()
+
       // Migrate legacy bookmarks
       useNodeBookmarkStore().migrateLegacyBookmarks()
 
diff --git a/tests-ui/tests/store/modelStore.test.ts b/tests-ui/tests/store/modelStore.test.ts
index 5af4d5884..db05ad31b 100644
--- a/tests-ui/tests/store/modelStore.test.ts
+++ b/tests-ui/tests/store/modelStore.test.ts
@@ -6,6 +6,7 @@ import { api } from '@/scripts/api'
 jest.mock('@/scripts/api', () => ({
   api: {
     getModels: jest.fn(),
+    getModelFolders: jest.fn(),
     viewMetadata: jest.fn()
   }
 }))
@@ -16,6 +17,7 @@ function enableMocks() {
     'sdv15.safetensors',
     'noinfo.safetensors'
   ])
+  ;(api.getModelFolders as jest.Mock).mockResolvedValue(['checkpoints', 'vae'])
   ;(api.viewMetadata as jest.Mock).mockImplementation((_, model) => {
     if (model === 'noinfo.safetensors') {
       return Promise.resolve({})
@@ -37,14 +39,15 @@ function enableMocks() {
 describe('useModelStore', () => {
   let store: ReturnType<typeof useModelStore>
 
-  beforeEach(() => {
+  beforeEach(async () => {
     setActivePinia(createPinia())
     store = useModelStore()
   })
 
   it('should load models', async () => {
     enableMocks()
-    const folderStore = await store.getModelsInFolderCached('checkpoints')
+    await store.loadModelFolders()
+    const folderStore = await store.getLoadedModelFolder('checkpoints')
     expect(folderStore).not.toBeNull()
     if (!folderStore) return
     expect(Object.keys(folderStore.models).length).toBe(3)
@@ -52,7 +55,8 @@ describe('useModelStore', () => {
 
   it('should load model metadata', async () => {
     enableMocks()
-    const folderStore = await store.getModelsInFolderCached('checkpoints')
+    await store.loadModelFolders()
+    const folderStore = await store.getLoadedModelFolder('checkpoints')
     expect(folderStore).not.toBeNull()
     if (!folderStore) return
     const model = folderStore.models['sdxl.safetensors']
@@ -69,7 +73,8 @@ describe('useModelStore', () => {
 
   it('should handle no metadata', async () => {
     enableMocks()
-    const folderStore = await store.getModelsInFolderCached('checkpoints')
+    await store.loadModelFolders()
+    const folderStore = await store.getLoadedModelFolder('checkpoints')
     expect(folderStore).not.toBeNull()
     if (!folderStore) return
     const model = folderStore.models['noinfo.safetensors']
@@ -84,8 +89,9 @@ describe('useModelStore', () => {
 
   it('should cache model information', async () => {
     enableMocks()
-    const folderStore1 = await store.getModelsInFolderCached('checkpoints')
-    const folderStore2 = await store.getModelsInFolderCached('checkpoints')
+    await store.loadModelFolders()
+    await store.getLoadedModelFolder('checkpoints')
+    await store.getLoadedModelFolder('checkpoints')
     expect(api.getModels).toHaveBeenCalledTimes(1)
   })
 })