From 739ebd3d04722c95ea5e892e531260c9df40b678 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Tue, 29 Oct 2024 09:24:56 -0400
Subject: [PATCH] Auto-expand model library tree on search (#1357)

---
 .../sidebar/tabs/ModelLibrarySidebarTab.vue   | 30 ++++++++++++-------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
index 5dc561ad0..a325a9921 100644
--- a/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
@@ -24,6 +24,7 @@
         class="model-lib-search-box p-4"
         v-model:modelValue="searchQuery"
         :placeholder="$t('searchModels') + '...'"
+        @search="handleSearch"
       />
     </template>
     <template #body>
@@ -60,7 +61,7 @@ import type {
   RenderedTreeExplorerNode,
   TreeExplorerNode
 } from '@/types/treeExplorerTypes'
-import { computed, ref, watch, toRef, onMounted } from 'vue'
+import { computed, ref, watch, toRef, onMounted, nextTick } from 'vue'
 import type { TreeNode } from 'primevue/treenode'
 import { app } from '@/scripts/app'
 import { buildTree } from '@/utils/treeUtil'
@@ -69,17 +70,26 @@ const modelToNodeStore = useModelToNodeStore()
 const settingStore = useSettingStore()
 const searchQuery = ref<string>('')
 const expandedKeys = ref<Record<string, boolean>>({})
-const { toggleNodeOnEvent } = useTreeExpansion(expandedKeys)
+const { expandNode, toggleNodeOnEvent } = useTreeExpansion(expandedKeys)
 
-const filteredModels = computed<ComfyModelDef[]>(() => {
-  if (searchQuery.value) {
-    const search = searchQuery.value.toLocaleLowerCase()
-    return modelStore.models.filter((model: ComfyModelDef) => {
-      return model.searchable.includes(search)
-    })
+const filteredModels = ref<ComfyModelDef[]>([])
+const handleSearch = async (query: string) => {
+  if (!query) {
+    filteredModels.value = modelStore.models
+    expandedKeys.value = {}
+    return
   }
-  return modelStore.models
-})
+  // Load all models to ensure we have the latest data
+  await modelStore.loadModels()
+  const search = query.toLocaleLowerCase()
+  filteredModels.value = modelStore.models.filter((model: ComfyModelDef) => {
+    return model.searchable.includes(search)
+  })
+
+  nextTick(() => {
+    expandNode(root.value)
+  })
+}
 
 type ModelOrFolder = ComfyModelDef | ModelFolder