feat: Enhance media handling by adding NSFW level support and improving preview image management

py/services/model_cache.py

@@ -32,12 +32,13 @@ class ModelCache:
             all_folders = set(l['folder'] for l in self.raw_data)
             self.folders = sorted(list(all_folders), key=lambda x: x.lower())
 
-    async def update_preview_url(self, file_path: str, preview_url: str) -> bool:
+    async def update_preview_url(self, file_path: str, preview_url: str, preview_nsfw_level: int) -> bool:
         """Update preview_url for a specific model in all cached data
         
         Args:
             file_path: The file path of the model to update
             preview_url: The new preview URL
+            preview_nsfw_level: The NSFW level of the preview
             
         Returns:
             bool: True if the update was successful, False if the model wasn't found
@@ -47,19 +48,9 @@ class ModelCache:
             for item in self.raw_data:
                 if item['file_path'] == file_path:
                     item['preview_url'] = preview_url
+                    item['preview_nsfw_level'] = preview_nsfw_level
                     break
             else:
                 return False  # Model not found
-                
-            # Update in sorted lists (references to the same dict objects)
-            for item in self.sorted_by_name:
-                if item['file_path'] == file_path:
-                    item['preview_url'] = preview_url
-                    break
-                    
-            for item in self.sorted_by_date:
-                if item['file_path'] == file_path:
-                    item['preview_url'] = preview_url
-                    break
                     
             return True

py/services/model_scanner.py

@@ -1184,12 +1184,13 @@ class ModelScanner:
         """Get list of excluded model file paths"""
         return self._excluded_models.copy()
 
-    async def update_preview_in_cache(self, file_path: str, preview_url: str) -> bool:
+    async def update_preview_in_cache(self, file_path: str, preview_url: str, preview_nsfw_level: int) -> bool:
         """Update preview URL in cache for a specific lora
         
         Args:
             file_path: The file path of the lora to update
             preview_url: The new preview URL
+            preview_nsfw_level: The NSFW level of the preview
             
         Returns:
             bool: True if the update was successful, False if cache doesn't exist or lora wasn't found
@@ -1197,7 +1198,7 @@ class ModelScanner:
         if self._cache is None:
             return False
 
-        updated = await self._cache.update_preview_url(file_path, preview_url)
+        updated = await self._cache.update_preview_url(file_path, preview_url, preview_nsfw_level)
         if updated:
             # Save updated cache to disk
             await self._save_cache_to_disk()

py/utils/routes_common.py

@@ -409,6 +409,15 @@ class ModelRouteUtils:
                 raise ValueError("Expected 'model_path' field")
             model_path = (await field.read()).decode()
             
+            # Read NSFW level (new parameter)
+            nsfw_level = 0  # Default to 0 (unknown)
+            field = await reader.next()
+            if field and field.name == 'nsfw_level':
+                try:
+                    nsfw_level = int((await field.read()).decode())
+                except (ValueError, TypeError):
+                    logger.warning("Invalid NSFW level format, using default 0")
+            
             # Save preview file
             base_name = os.path.splitext(os.path.basename(model_path))[0]
             folder = os.path.dirname(model_path)
@@ -435,7 +444,7 @@ class ModelRouteUtils:
                 if os.path.exists(existing_preview):
                     try:
                         os.remove(existing_preview)
-                        logger.info(f"Deleted existing preview: {existing_preview}")
+                        logger.debug(f"Deleted existing preview: {existing_preview}")
                     except Exception as e:
                         logger.warning(f"Failed to delete existing preview {existing_preview}: {e}")
             
@@ -444,26 +453,28 @@ class ModelRouteUtils:
             with open(preview_path, 'wb') as f:
                 f.write(optimized_data)
             
-            # Update preview path in metadata
+            # Update preview path and NSFW level in metadata
             metadata_path = os.path.splitext(model_path)[0] + '.metadata.json'
             if os.path.exists(metadata_path):
                 try:
                     with open(metadata_path, 'r', encoding='utf-8') as f:
                         metadata = json.load(f)
                     
-                    # Update preview_url directly in the metadata dict
+                    # Update preview_url and preview_nsfw_level in the metadata dict
                     metadata['preview_url'] = preview_path
+                    metadata['preview_nsfw_level'] = nsfw_level
                     
                     await MetadataManager.save_metadata(model_path, metadata)
                 except Exception as e:
                     logger.error(f"Error updating metadata: {e}")
             
             # Update preview URL in scanner cache
-            await scanner.update_preview_in_cache(model_path, preview_path)
+            await scanner.update_preview_in_cache(model_path, preview_path, nsfw_level)
             
             return web.json_response({
                 "success": True,
-                "preview_url": config.get_preview_static_url(preview_path)
+                "preview_url": config.get_preview_static_url(preview_path),
+                "preview_nsfw_level": nsfw_level
             })
             
         except Exception as e: