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

This commit is contained in:
Will Miao
2025-06-19 15:19:24 +08:00
parent a7304ccf47
commit 605a06317b
10 changed files with 238 additions and 136 deletions

View File

@@ -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

View File

@@ -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()