[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
2026-04-27 01:39:47 +00:00 · 2024-10-28 21:23:53 -04:00
parent 757f0ced81
commit 4582c71583
8 changed files with 162 additions and 148 deletions
--- 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 tree
+  return modelStore.models
 })

-const renderedRoot = computed<TreeExplorerNode<ComfyModelDef>>(() => {
+type ModelOrFolder = ComfyModelDef | ModelFolder
+
+const root = computed<TreeNode>(() => {
+  const allNodes: ModelOrFolder[] = [
+    ...modelStore.modelFolders,
+    ...filteredModels.value
+  ]
+  return buildTree(allNodes, (modelOrFolder: ModelOrFolder) =>
+    modelOrFolder.key.split('/')
+  )
+})
+
+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>
--- 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(() => {
--- 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: 'Название расширения',
--- 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()
  }
--- 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({
--- 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 class ModelFolder {
-  models: Record<string, ComfyModelDef> = {}
+export enum ResourceState {
+  Uninitialized,
+  Loading,
+  Loaded
+}

-  constructor(directory: string, models: string[]) {
-    for (const model of models) {
-      this.models[model] = new ComfyModelDef(model, directory)
-    }
+export class ModelFolder {
+  /** Models in this folder */
+  models: Record<string, ComfyModelDef> = {}
+  state: ResourceState = ResourceState.Uninitialized
+
+  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[]>([])
-
-  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]
-    }
-    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 getModelFolders() {
-    modelFolders.value = (await api.getModelFolders()).filter(
-      (folder) => !folderBlacklist.includes(folder)
+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))
+  )
+
+  /**
+   * 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)
+    }
+  }
+
+  async function getLoadedModelFolder(
+    folderName: string
+  ): Promise<ModelFolder | null> {
+    const folder = modelFolderByName.value[folderName]
+    return folder ? await folder.load() : null
+  }
+
+  /**
+   * 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
  }
 })
--- 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()