feat(scanner): index cached models by version id

py/services/model_cache.py

@@ -1,6 +1,6 @@
 import asyncio
-from typing import List, Dict, Tuple
-from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+from dataclasses import dataclass, field
 from operator import itemgetter
 from natsort import natsorted
 
@@ -17,10 +17,12 @@ SUPPORTED_SORT_MODES = [
 
 @dataclass
 class ModelCache:
-    """Cache structure for model data with extensible sorting"""
+    """Cache structure for model data with extensible sorting."""
+
     raw_data: List[Dict]
     folders: List[str]
-    
+    version_index: Dict[int, Dict] = field(default_factory=dict)
+
     def __post_init__(self):
         self._lock = asyncio.Lock()
         # Cache for last sort: (sort_key, order) -> sorted list
@@ -28,6 +30,58 @@ class ModelCache:
         self._last_sorted_data: List[Dict] = []
         # Default sort on init
         asyncio.create_task(self.resort())
+        self.rebuild_version_index()
+
+    @staticmethod
+    def _normalize_version_id(value: Any) -> Optional[int]:
+        """Normalize a potential version identifier into an integer."""
+
+        if isinstance(value, int):
+            return value
+        if isinstance(value, str):
+            try:
+                return int(value)
+            except ValueError:
+                return None
+        return None
+
+    def rebuild_version_index(self) -> None:
+        """Rebuild the version index from the current raw data."""
+
+        self.version_index = {}
+        for item in self.raw_data:
+            self.add_to_version_index(item)
+
+    def add_to_version_index(self, item: Dict) -> None:
+        """Register a cache item in the version index if possible."""
+
+        civitai_data = item.get('civitai') if isinstance(item, dict) else None
+        if not isinstance(civitai_data, dict):
+            return
+
+        version_id = self._normalize_version_id(civitai_data.get('id'))
+        if version_id is None:
+            return
+
+        self.version_index[version_id] = item
+
+    def remove_from_version_index(self, item: Dict) -> None:
+        """Remove a cache item from the version index if present."""
+
+        civitai_data = item.get('civitai') if isinstance(item, dict) else None
+        if not isinstance(civitai_data, dict):
+            return
+
+        version_id = self._normalize_version_id(civitai_data.get('id'))
+        if version_id is None:
+            return
+
+        existing = self.version_index.get(version_id)
+        if existing is item or (
+            isinstance(existing, dict)
+            and existing.get('file_path') == item.get('file_path')
+        ):
+            self.version_index.pop(version_id, None)
 
     async def resort(self):
         """Resort cached data according to last sort mode if set"""
@@ -41,6 +95,7 @@ class ModelCache:
 
             all_folders = set(l['folder'] for l in self.raw_data)
             self.folders = sorted(list(all_folders), key=lambda x: x.lower())
+            self.rebuild_version_index()
 
     def _sort_data(self, data: List[Dict], sort_key: str, order: str) -> List[Dict]:
         """Sort data by sort_key and order"""

py/services/model_scanner.py

@@ -634,7 +634,8 @@ class ModelScanner:
                                 if model_data:
                                     # Add to cache
                                     self._cache.raw_data.append(model_data)
-                                    
+                                    self._cache.add_to_version_index(model_data)
+
                                     # Update hash index if available
                                     if 'sha256' in model_data and 'file_path' in model_data:
                                         self._hash_index.add_entry(model_data['sha256'].lower(), model_data['file_path'])
@@ -661,7 +662,9 @@ class ModelScanner:
                 for path in missing_files:
                     try:
                         model_to_remove = path_to_item[path]
-                        
+
+                        self._cache.remove_from_version_index(model_to_remove)
+
                         # Update tags count
                         for tag in model_to_remove.get('tags', []):
                             if tag in self._tags_count:
@@ -684,6 +687,8 @@ class ModelScanner:
                 all_folders = set(item.get('folder', '') for item in self._cache.raw_data)
                 self._cache.folders = sorted(list(all_folders), key=lambda x: x.lower())
 
+                self._cache.rebuild_version_index()
+
                 # Resort cache
                 await self._cache.resort()
 
@@ -829,6 +834,8 @@ class ModelScanner:
         else:
             self._cache.raw_data = list(scan_result.raw_data)
 
+        self._cache.rebuild_version_index()
+
         await self._cache.resort()
 
     async def _gather_model_data(
@@ -934,7 +941,8 @@ class ModelScanner:
             
             # Add to cache
             self._cache.raw_data.append(metadata_dict)
-            
+            self._cache.add_to_version_index(metadata_dict)
+
             # Resort cache data
             await self._cache.resort()
             
@@ -1076,6 +1084,9 @@ class ModelScanner:
         cache = await self.get_cached_data()
 
         existing_item = next((item for item in cache.raw_data if item['file_path'] == original_path), None)
+        if existing_item:
+            cache.remove_from_version_index(existing_item)
+
         if existing_item and 'tags' in existing_item:
             for tag in existing_item.get('tags', []):
                 if tag in self._tags_count:
@@ -1106,6 +1117,7 @@ class ModelScanner:
             )
 
             cache.raw_data.append(cache_entry)
+            cache.add_to_version_index(cache_entry)
 
             sha_value = cache_entry.get('sha256')
             if sha_value:
@@ -1117,6 +1129,8 @@ class ModelScanner:
             for tag in cache_entry.get('tags', []):
                 self._tags_count[tag] = self._tags_count.get(tag, 0) + 1
 
+        cache.rebuild_version_index()
+
         await cache.resort()
 
         if cache_modified:
@@ -1339,11 +1353,12 @@ class ModelScanner:
             # Update hash index
             for model in models_to_remove:
                 file_path = model['file_path']
+                self._cache.remove_from_version_index(model)
                 if hasattr(self, '_hash_index') and self._hash_index:
                     # Get the hash and filename before removal for duplicate checking
                     file_name = os.path.splitext(os.path.basename(file_path))[0]
                     hash_val = model.get('sha256', '').lower()
-                    
+
                     # Remove from hash index
                     self._hash_index.remove_by_path(file_path, hash_val)
                     
@@ -1352,8 +1367,9 @@ class ModelScanner:
             
             # Update cache data
             self._cache.raw_data = [item for item in self._cache.raw_data if item['file_path'] not in file_paths]
-            
+
             # Resort cache
+            self._cache.rebuild_version_index()
             await self._cache.resort()
 
             await self._persist_current_cache()
@@ -1393,16 +1409,17 @@ class ModelScanner:
         Returns:
             bool: True if the model version exists, False otherwise
         """
+        try:
+            normalized_id = int(model_version_id)
+        except (TypeError, ValueError):
+            return False
+
         try:
             cache = await self.get_cached_data()
-            if not cache or not cache.raw_data:
+            if not cache:
                 return False
 
-            for item in cache.raw_data:
-                if item.get('civitai') and item['civitai'].get('id') == model_version_id:
-                    return True
-
-            return False
+            return normalized_id in cache.version_index
         except Exception as e:
             logger.error(f"Error checking model version existence: {e}")
             return False