[Refactor] Rework modelStore and ModelLibrarySidebarTab (Comfy-Org#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

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,90 +42,65 @@
 </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'
 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()
 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>

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(() => {

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: 'Название расширения',

src/scripts/api.ts

@@ -366,18 +366,21 @@ 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)
+    )
   }
 
   /**
    * Gets a list of models in the specified folder
    * @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()
   }

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({