From 9fc2fb4d17c4c2569c594ff7a00e5da19bfb9eed Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Wed, 4 Jun 2025 18:38:45 +0800
Subject: [PATCH] Enhance model caching and exclusion functionality: update
 cache version, add excluded models to cache data, and ensure cache is saved
 to disk after model exclusion and deletion.

---
 py/services/model_scanner.py  | 9 +++++++--
 py/utils/routes_common.py     | 4 ++++
 static/js/api/baseModelApi.js | 8 ++++++++
 static/js/utils/modalUtils.js | 8 ++++++++
 4 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/py/services/model_scanner.py b/py/services/model_scanner.py
index e3d70f0c..4264bdc4 100644
--- a/py/services/model_scanner.py
+++ b/py/services/model_scanner.py
@@ -22,7 +22,8 @@ logger = logging.getLogger(__name__)
 # Version history:
 # 1 - Initial version
 # 2 - Added duplicate_filenames and duplicate_hashes tracking
-CACHE_VERSION = 2
+# 3 - Added _excluded_models list to cache
+CACHE_VERSION = 3
 
 class ModelScanner:
     """Base service for scanning and managing model files"""
@@ -115,7 +116,8 @@ class ModelScanner:
                     "duplicate_filenames": self._hash_index._duplicate_filenames
                 },
                 "tags_count": self._tags_count,
-                "dirs_last_modified": self._get_dirs_last_modified()
+                "dirs_last_modified": self._get_dirs_last_modified(),
+                "excluded_models": self._excluded_models  # Add excluded_models to cache data
             }
             
             # Preprocess data to handle large integers
@@ -217,6 +219,9 @@ class ModelScanner:
             # Load tags count
             self._tags_count = cache_data.get("tags_count", {})
             
+            # Load excluded models
+            self._excluded_models = cache_data.get("excluded_models", [])
+            
             # Resort the cache
             await self._cache.resort()
             
diff --git a/py/utils/routes_common.py b/py/utils/routes_common.py
index 2e96bf50..f5b7fe3d 100644
--- a/py/utils/routes_common.py
+++ b/py/utils/routes_common.py
@@ -303,6 +303,8 @@ class ModelRouteUtils:
             # Update hash index if available
             if hasattr(scanner, '_hash_index') and scanner._hash_index:
                 scanner._hash_index.remove_by_path(file_path)
+
+            await scanner._save_cache_to_disk()
             
             return web.json_response({
                 'success': True,
@@ -484,6 +486,8 @@ class ModelRouteUtils:
             
             # Add to excluded models list
             scanner._excluded_models.append(file_path)
+
+            await scanner._save_cache_to_disk()
             
             return web.json_response({
                 'success': True,
diff --git a/static/js/api/baseModelApi.js b/static/js/api/baseModelApi.js
index af9cfafb..230d642b 100644
--- a/static/js/api/baseModelApi.js
+++ b/static/js/api/baseModelApi.js
@@ -434,6 +434,8 @@ export function replaceModelPreview(filePath, modelType = 'lora') {
 // Delete a model (generic)
 export async function deleteModel(filePath, modelType = 'lora') {
     try {
+        state.loadingManager.showSimpleLoading(`Deleting ${modelType}...`);
+
         const endpoint = modelType === 'checkpoint' 
             ? '/api/checkpoints/delete' 
             : '/api/delete_model';
@@ -475,6 +477,8 @@ export async function deleteModel(filePath, modelType = 'lora') {
         console.error(`Error deleting ${modelType}:`, error);
         showToast(`Failed to delete ${modelType}: ${error.message}`, 'error');
         return false;
+    } finally {
+        state.loadingManager.hide();
     }
 }
 
@@ -662,6 +666,8 @@ export async function refreshSingleModelMetadata(filePath, modelType = 'lora') {
 // Generic function to exclude a model
 export async function excludeModel(filePath, modelType = 'lora') {
     try {
+        state.loadingManager.showSimpleLoading(`Excluding ${modelType}...`);
+
         const endpoint = modelType === 'checkpoint' 
             ? '/api/checkpoints/exclude' 
             : '/api/loras/exclude';
@@ -703,6 +709,8 @@ export async function excludeModel(filePath, modelType = 'lora') {
         console.error(`Error excluding ${modelType}:`, error);
         showToast(`Failed to exclude ${modelType}: ${error.message}`, 'error');
         return false;
+    } finally {
+        state.loadingManager.hide();
     }
 }
 
diff --git a/static/js/utils/modalUtils.js b/static/js/utils/modalUtils.js
index 3041be4d..041a8ee3 100644
--- a/static/js/utils/modalUtils.js
+++ b/static/js/utils/modalUtils.js
@@ -37,6 +37,10 @@ export async function confirmDelete() {
         }
         
         closeDeleteModal();
+
+        if (window.modelDuplicatesManager) {
+            window.modelDuplicatesManager.updateDuplicatesBadgeAfterRefresh();
+        }
     } catch (error) {
         console.error('Error deleting model:', error);
         alert(`Error deleting model: ${error}`);
@@ -86,6 +90,10 @@ export async function confirmExclude() {
         }
         
         closeExcludeModal();
+
+        if (window.modelDuplicatesManager) {
+            window.modelDuplicatesManager.updateDuplicatesBadgeAfterRefresh();
+        }
     } catch (error) {
         console.error('Error excluding model:', error);
     }