From a32940bcf63056479aad6c199724d63adf781ed5 Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Fri, 14 Feb 2025 13:58:12 +0800
Subject: [PATCH] Enhance folder management in Lora API and UI; update folder
 tags dynamically and improve cache handling

---
 routes/api_routes.py                  |  6 +++-
 services/file_monitor.py              |  3 --
 static/js/api/loraApi.js              | 51 +++++++++++++++++++++++----
 static/js/managers/DownloadManager.js |  2 +-
 utils/model_utils.py                  |  3 +-
 5 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/routes/api_routes.py b/routes/api_routes.py
index 8f6f455c..e369644a 100644
--- a/routes/api_routes.py
+++ b/routes/api_routes.py
@@ -147,13 +147,17 @@ class ApiRoutes:
                 self._format_lora_response(item) 
                 for item in result['items']
             ]
+
+            # Get all available folders from cache
+            cache = await self.scanner.get_cached_data()
             
             return web.json_response({
                 'items': formatted_items,
                 'total': result['total'],
                 'page': result['page'],
                 'page_size': result['page_size'],
-                'total_pages': result['total_pages']
+                'total_pages': result['total_pages'],
+                'folders': cache.folders
             })
             
         except Exception as e:
diff --git a/services/file_monitor.py b/services/file_monitor.py
index f5045d4e..1e809935 100644
--- a/services/file_monitor.py
+++ b/services/file_monitor.py
@@ -25,14 +25,11 @@ class LoraFileHandler(FileSystemEventHandler):
 
     def _should_ignore(self, path: str) -> bool:
         """Check if path should be ignored"""
-        logger.info(f"Checking ignore for {path}")
-        logger.info(f"Ignore paths: {self._ignore_paths}")
         return path.replace(os.sep, '/') in self._ignore_paths
 
     def add_ignore_path(self, path: str, file_size: int = 0):
         """Add path to ignore list with dynamic timeout based on file size"""
         self._ignore_paths.add(path.replace(os.sep, '/'))
-        logger.info(f"Update ignore paths: {self._ignore_paths}")
         
         # Calculate timeout based on file size, with a minimum value
         # Assuming average download speed of 1MB/s
diff --git a/static/js/api/loraApi.js b/static/js/api/loraApi.js
index 457315df..4abeaac5 100644
--- a/static/js/api/loraApi.js
+++ b/static/js/api/loraApi.js
@@ -4,7 +4,7 @@ import { createLoraCard } from '../components/LoraCard.js';
 import { initializeInfiniteScroll } from '../utils/infiniteScroll.js';
 import { showDeleteModal } from '../utils/modalUtils.js';
 
-export async function loadMoreLoras() {
+export async function loadMoreLoras(boolUpdateFolders = false) {
     if (state.isLoading || !state.hasMore) return;
     
     state.isLoading = true;
@@ -56,6 +56,10 @@ export async function loadMoreLoras() {
         } else {
             state.hasMore = false;
         }
+
+        if (boolUpdateFolders && data.folders) {
+            updateFolderTags(data.folders);
+        }
         
     } catch (error) {
         console.error('Error loading loras:', error);
@@ -65,6 +69,39 @@ export async function loadMoreLoras() {
     }
 }
 
+function updateFolderTags(folders) {
+    const folderTagsContainer = document.querySelector('.folder-tags');
+    if (!folderTagsContainer) return;
+
+    // Keep track of currently selected folder
+    const currentFolder = state.activeFolder;
+
+    // Create HTML for folder tags
+    const tagsHTML = folders.map(folder => {
+        const isActive = folder === currentFolder;
+        return `<div class="tag ${isActive ? 'active' : ''}" data-folder="${folder}">${folder}</div>`;
+    }).join('');
+
+    // Update the container
+    folderTagsContainer.innerHTML = tagsHTML;
+
+    // Reattach click handlers
+    const tags = folderTagsContainer.querySelectorAll('.tag');
+    tags.forEach(tag => {
+        tag.addEventListener('click', function() {
+            const folder = this.dataset.folder;
+            // Remove active class from all tags
+            tags.forEach(t => t.classList.remove('active'));
+            // Add active class to clicked tag
+            this.classList.add('active');
+            // Update state and reload
+            state.activeFolder = folder;
+            state.currentPage = 1;
+            resetAndReload();
+        });
+    });
+}
+
 export async function fetchCivitai() {
     let ws = null;
     
@@ -224,7 +261,7 @@ export function appendLoraCards(loras) {
     });
 }
 
-export async function resetAndReload() {
+export async function resetAndReload(boolUpdateFolders = false) {
     console.log('Resetting with state:', { ...state });
     
     state.currentPage = 1;
@@ -240,14 +277,16 @@ export async function resetAndReload() {
     
     initializeInfiniteScroll();
     
-    await loadMoreLoras();
+    await loadMoreLoras(boolUpdateFolders);
 }
 
-export async function refreshLoras() {
+export async function refreshLoras(boolShowToast = true) {
     try {
         state.loadingManager.showSimpleLoading('Refreshing loras...');
-        await resetAndReload();
-        showToast('Refresh complete', 'success');
+        await resetAndReload(true);
+        if (boolShowToast){
+            showToast('Refresh complete', 'success');
+        } 
     } catch (error) {
         console.error('Refresh failed:', error);
         showToast('Failed to refresh loras', 'error');
diff --git a/static/js/managers/DownloadManager.js b/static/js/managers/DownloadManager.js
index bdf6d395..96726f3d 100644
--- a/static/js/managers/DownloadManager.js
+++ b/static/js/managers/DownloadManager.js
@@ -194,7 +194,7 @@ export class DownloadManager {
             modalManager.closeModal('downloadModal');
             
             // Refresh the grid to show new model
-            window.refreshLoras();
+            window.refreshLoras(false);
         } catch (error) {
             showToast(error.message, 'error');
         }
diff --git a/utils/model_utils.py b/utils/model_utils.py
index 893a0f5d..5a731dcf 100644
--- a/utils/model_utils.py
+++ b/utils/model_utils.py
@@ -7,7 +7,8 @@ BASE_MODEL_MAPPING = {
     "sdxl": "SDXL",
     "sd-v2": "SD2.0",
     "flux1": "Flux.1 D",
-    "Illustrious": "IL"
+    "illustrious": "IL",
+    "pony": "Pony"
 }
 
 def determine_base_model(version_string: Optional[str]) -> str: