Add tag filtering checkpoint

py/services/lora_scanner.py

@@ -34,6 +34,7 @@ class LoraScanner:
             self._initialization_task: Optional[asyncio.Task] = None
             self._initialized = True
             self.file_monitor = None  # Add this line
+            self._tags_count = {}  # Add a dictionary to store tag counts
 
     def set_file_monitor(self, monitor):
         """Set file monitor instance"""
@@ -90,13 +91,21 @@ class LoraScanner:
             # Clear existing hash index
             self._hash_index.clear()
             
+            # Clear existing tags count
+            self._tags_count = {}
+            
             # Scan for new data
             raw_data = await self.scan_all_loras()
             
-            # Build hash index
+            # Build hash index and tags count
             for lora_data in raw_data:
                 if 'sha256' in lora_data and 'file_path' in lora_data:
                     self._hash_index.add_entry(lora_data['sha256'], lora_data['file_path'])
+                
+                # Count tags
+                if 'tags' in lora_data and lora_data['tags']:
+                    for tag in lora_data['tags']:
+                        self._tags_count[tag] = self._tags_count.get(tag, 0) + 1
             
             # Update cache
             self._cache = LoraCache(
@@ -158,7 +167,7 @@ class LoraScanner:
 
     async def get_paginated_data(self, page: int, page_size: int, sort_by: str = 'name', 
                                folder: str = None, search: str = None, fuzzy: bool = False,
-                               recursive: bool = False, base_models: list = None):
+                               recursive: bool = False, base_models: list = None, tags: list = None) -> Dict:
         """Get paginated and filtered lora data
         
         Args:
@@ -170,6 +179,7 @@ class LoraScanner:
             fuzzy: Use fuzzy matching for search
             recursive: Include subfolders when folder filter is applied
             base_models: List of base models to filter by
+            tags: List of tags to filter by
         """
         cache = await self.get_cached_data()
 
@@ -198,6 +208,13 @@ class LoraScanner:
                 if item.get('base_model') in base_models
             ]
         
+        # Apply tag filtering
+        if tags and len(tags) > 0:
+            filtered_data = [
+                item for item in filtered_data
+                if any(tag in item.get('tags', []) for tag in tags)
+            ]
+        
         # 应用搜索过滤
         if search:
             if fuzzy:
@@ -311,12 +328,67 @@ class LoraScanner:
         
         # Convert to dict and add folder info
         lora_data = metadata.to_dict()
+        # Try to fetch missing metadata from Civitai if needed
+        await self._fetch_missing_metadata(file_path, lora_data)
         rel_path = os.path.relpath(file_path, root_path)
         folder = os.path.dirname(rel_path)
         lora_data['folder'] = folder.replace(os.path.sep, '/')
         
         return lora_data
 
+    async def _fetch_missing_metadata(self, file_path: str, lora_data: Dict) -> None:
+        """Fetch missing description and tags from Civitai if needed
+        
+        Args:
+            file_path: Path to the lora file
+            lora_data: Lora metadata dictionary to update
+        """
+        try:
+            # Check if we need to fetch additional metadata from Civitai
+            needs_metadata_update = False
+            model_id = None
+            
+            # Check if we have Civitai model ID but missing metadata
+            if lora_data.get('civitai'):
+                # Try to get model ID directly from the correct location
+                model_id = lora_data['civitai'].get('modelId')
+                
+                if model_id:
+                    model_id = str(model_id)
+                    # Check if tags are missing or empty
+                    tags_missing = not lora_data.get('tags') or len(lora_data.get('tags', [])) == 0
+                    
+                    # Check if description is missing or empty
+                    desc_missing = not lora_data.get('modelDescription') or lora_data.get('modelDescription') in (None, "")
+                    
+                    needs_metadata_update = tags_missing or desc_missing
+            
+            # Fetch missing metadata if needed
+            if needs_metadata_update and model_id:
+                logger.info(f"Fetching missing metadata for {file_path} with model ID {model_id}")
+                from ..services.civitai_client import CivitaiClient
+                client = CivitaiClient()
+                model_metadata = await client.get_model_metadata(model_id)
+                await client.close()
+                
+                if model_metadata:
+                    logger.info(f"Updating metadata for {file_path} with model ID {model_id}")
+                    
+                    # Update tags if they were missing
+                    if model_metadata.get('tags') and (not lora_data.get('tags') or len(lora_data.get('tags', [])) == 0):
+                        lora_data['tags'] = model_metadata['tags']
+                    
+                    # Update description if it was missing
+                    if model_metadata.get('description') and (not lora_data.get('modelDescription') or lora_data.get('modelDescription') in (None, "")):
+                        lora_data['modelDescription'] = model_metadata['description']
+                    
+                    # Save the updated metadata back to file
+                    metadata_path = os.path.splitext(file_path)[0] + '.metadata.json'
+                    with open(metadata_path, 'w', encoding='utf-8') as f:
+                        json.dump(lora_data, f, indent=2, ensure_ascii=False)
+        except Exception as e:
+            logger.error(f"Failed to update metadata from Civitai for {file_path}: {e}")
+
     async def update_preview_in_cache(self, file_path: str, preview_url: str) -> bool:
         """Update preview URL in cache for a specific lora
         
@@ -427,6 +499,15 @@ class LoraScanner:
     async def update_single_lora_cache(self, original_path: str, new_path: str, metadata: Dict) -> bool:
         cache = await self.get_cached_data()
         
+        # Find the existing item to remove its tags from count
+        existing_item = next((item for item in cache.raw_data if item['file_path'] == original_path), None)
+        if existing_item and 'tags' in existing_item:
+            for tag in existing_item.get('tags', []):
+                if tag in self._tags_count:
+                    self._tags_count[tag] = max(0, self._tags_count[tag] - 1)
+                    if self._tags_count[tag] == 0:
+                        del self._tags_count[tag]
+        
         # Remove old path from hash index if exists
         self._hash_index.remove_by_path(original_path)
         
@@ -460,6 +541,11 @@ class LoraScanner:
             # Update folders list
             all_folders = set(item['folder'] for item in cache.raw_data)
             cache.folders = sorted(list(all_folders), key=lambda x: x.lower())
+            
+            # Update tags count with the new/updated tags
+            if 'tags' in metadata:
+                for tag in metadata.get('tags', []):
+                    self._tags_count[tag] = self._tags_count.get(tag, 0) + 1
         
         # Resort cache
         await cache.resort()
@@ -505,3 +591,26 @@ class LoraScanner:
         """Get hash for a LoRA by its file path"""
         return self._hash_index.get_hash(file_path)
 
+    # Add new method to get top tags
+    async def get_top_tags(self, limit: int = 20) -> List[Dict[str, any]]:
+        """Get top tags sorted by count
+        
+        Args:
+            limit: Maximum number of tags to return
+            
+        Returns:
+            List of dictionaries with tag name and count, sorted by count
+        """
+        # Make sure cache is initialized
+        await self.get_cached_data()
+        
+        # Sort tags by count in descending order
+        sorted_tags = sorted(
+            [{"tag": tag, "count": count} for tag, count in self._tags_count.items()],
+            key=lambda x: x['count'],
+            reverse=True
+        )
+        
+        # Return limited number
+        return sorted_tags[:limit]
+