feat: add model notes, preview URL, and Civitai URL endpoints to BaseModelRoutes and BaseModelService

This commit is contained in:
Will Miao
2025-08-15 18:58:49 +08:00
parent dda19b3920
commit ed1cd39a6c
4 changed files with 128 additions and 49 deletions

View File

@@ -65,6 +65,9 @@ class BaseModelRoutes(ABC):
app.router.add_get(f'/api/{prefix}/unified-folder-tree', self.get_unified_folder_tree) app.router.add_get(f'/api/{prefix}/unified-folder-tree', self.get_unified_folder_tree)
app.router.add_get(f'/api/{prefix}/find-duplicates', self.find_duplicate_models) app.router.add_get(f'/api/{prefix}/find-duplicates', self.find_duplicate_models)
app.router.add_get(f'/api/{prefix}/find-filename-conflicts', self.find_filename_conflicts) app.router.add_get(f'/api/{prefix}/find-filename-conflicts', self.find_filename_conflicts)
app.router.add_get(f'/api/{prefix}/get-notes', self.get_model_notes)
app.router.add_get(f'/api/{prefix}/preview-url', self.get_model_preview_url)
app.router.add_get(f'/api/{prefix}/civitai-url', self.get_model_civitai_url)
# Common Download management # Common Download management
app.router.add_post(f'/api/download-model', self.download_model) app.router.add_post(f'/api/download-model', self.download_model)
@@ -977,6 +980,84 @@ class BaseModelRoutes(ABC):
'error': str(e) 'error': str(e)
}) })
return web.json_response({
'success': False,
'error': str(e)
}, status=500)
async def get_model_notes(self, request: web.Request) -> web.Response:
"""Get notes for a specific model file"""
try:
model_name = request.query.get('name')
if not model_name:
return web.Response(text=f'{self.model_type.capitalize()} file name is required', status=400)
notes = await self.service.get_model_notes(model_name)
if notes is not None:
return web.json_response({
'success': True,
'notes': notes
})
else:
return web.json_response({
'success': False,
'error': f'{self.model_type.capitalize()} not found in cache'
}, status=404)
except Exception as e:
logger.error(f"Error getting {self.model_type} notes: {e}", exc_info=True)
return web.json_response({
'success': False,
'error': str(e)
}, status=500)
async def get_model_preview_url(self, request: web.Request) -> web.Response:
"""Get the static preview URL for a model file"""
try:
model_name = request.query.get('name')
if not model_name:
return web.Response(text=f'{self.model_type.capitalize()} file name is required', status=400)
preview_url = await self.service.get_model_preview_url(model_name)
if preview_url:
return web.json_response({
'success': True,
'preview_url': preview_url
})
else:
return web.json_response({
'success': False,
'error': f'No preview URL found for the specified {self.model_type}'
}, status=404)
except Exception as e:
logger.error(f"Error getting {self.model_type} preview URL: {e}", exc_info=True)
return web.json_response({
'success': False,
'error': str(e)
}, status=500)
async def get_model_civitai_url(self, request: web.Request) -> web.Response:
"""Get the Civitai URL for a model file"""
try:
model_name = request.query.get('name')
if not model_name:
return web.Response(text=f'{self.model_type.capitalize()} file name is required', status=400)
result = await self.service.get_model_civitai_url(model_name)
if result['civitai_url']:
return web.json_response({
'success': True,
**result
})
else:
return web.json_response({
'success': False,
'error': f'No Civitai data found for the specified {self.model_type}'
}, status=404)
except Exception as e:
logger.error(f"Error getting {self.model_type} Civitai URL: {e}", exc_info=True)
return web.json_response({ return web.json_response({
'success': False, 'success': False,
'error': str(e) 'error': str(e)

View File

@@ -43,10 +43,7 @@ class LoraRoutes(BaseModelRoutes):
"""Setup LoRA-specific routes""" """Setup LoRA-specific routes"""
# LoRA-specific query routes # LoRA-specific query routes
app.router.add_get(f'/api/{prefix}/letter-counts', self.get_letter_counts) app.router.add_get(f'/api/{prefix}/letter-counts', self.get_letter_counts)
app.router.add_get(f'/api/{prefix}/get-notes', self.get_lora_notes)
app.router.add_get(f'/api/{prefix}/get-trigger-words', self.get_lora_trigger_words) app.router.add_get(f'/api/{prefix}/get-trigger-words', self.get_lora_trigger_words)
app.router.add_get(f'/api/{prefix}/preview-url', self.get_lora_preview_url)
app.router.add_get(f'/api/{prefix}/civitai-url', self.get_lora_civitai_url)
app.router.add_get(f'/api/{prefix}/model-description', self.get_lora_model_description) app.router.add_get(f'/api/{prefix}/model-description', self.get_lora_model_description)
# CivitAI integration with LoRA-specific validation # CivitAI integration with LoRA-specific validation

View File

@@ -330,4 +330,50 @@ class BaseModelService(ABC):
current_level[part] = {} current_level[part] = {}
current_level = current_level[part] current_level = current_level[part]
return unified_tree return unified_tree
async def get_model_notes(self, model_name: str) -> Optional[str]:
"""Get notes for a specific model file"""
cache = await self.scanner.get_cached_data()
for model in cache.raw_data:
if model['file_name'] == model_name:
return model.get('notes', '')
return None
async def get_model_preview_url(self, model_name: str) -> Optional[str]:
"""Get the static preview URL for a model file"""
cache = await self.scanner.get_cached_data()
for model in cache.raw_data:
if model['file_name'] == model_name:
preview_url = model.get('preview_url')
if preview_url:
from ..config import config
return config.get_preview_static_url(preview_url)
return None
async def get_model_civitai_url(self, model_name: str) -> Dict[str, Optional[str]]:
"""Get the Civitai URL for a model file"""
cache = await self.scanner.get_cached_data()
for model in cache.raw_data:
if model['file_name'] == model_name:
civitai_data = model.get('civitai', {})
model_id = civitai_data.get('modelId')
version_id = civitai_data.get('id')
if model_id:
civitai_url = f"https://civitai.com/models/{model_id}"
if version_id:
civitai_url += f"?modelVersionId={version_id}"
return {
'civitai_url': civitai_url,
'model_id': str(model_id),
'version_id': str(version_id) if version_id else None
}
return {'civitai_url': None, 'model_id': None, 'version_id': None}

View File

@@ -147,16 +147,6 @@ class LoraService(BaseModelService):
return letters return letters
async def get_lora_notes(self, lora_name: str) -> Optional[str]:
"""Get notes for a specific LoRA file"""
cache = await self.scanner.get_cached_data()
for lora in cache.raw_data:
if lora['file_name'] == lora_name:
return lora.get('notes', '')
return None
async def get_lora_trigger_words(self, lora_name: str) -> List[str]: async def get_lora_trigger_words(self, lora_name: str) -> List[str]:
"""Get trigger words for a specific LoRA file""" """Get trigger words for a specific LoRA file"""
cache = await self.scanner.get_cached_data() cache = await self.scanner.get_cached_data()
@@ -168,41 +158,6 @@ class LoraService(BaseModelService):
return [] return []
async def get_lora_preview_url(self, lora_name: str) -> Optional[str]:
"""Get the static preview URL for a LoRA file"""
cache = await self.scanner.get_cached_data()
for lora in cache.raw_data:
if lora['file_name'] == lora_name:
preview_url = lora.get('preview_url')
if preview_url:
return config.get_preview_static_url(preview_url)
return None
async def get_lora_civitai_url(self, lora_name: str) -> Dict[str, Optional[str]]:
"""Get the Civitai URL for a LoRA file"""
cache = await self.scanner.get_cached_data()
for lora in cache.raw_data:
if lora['file_name'] == lora_name:
civitai_data = lora.get('civitai', {})
model_id = civitai_data.get('modelId')
version_id = civitai_data.get('id')
if model_id:
civitai_url = f"https://civitai.com/models/{model_id}"
if version_id:
civitai_url += f"?modelVersionId={version_id}"
return {
'civitai_url': civitai_url,
'model_id': str(model_id),
'version_id': str(version_id) if version_id else None
}
return {'civitai_url': None, 'model_id': None, 'version_id': None}
def find_duplicate_hashes(self) -> Dict: def find_duplicate_hashes(self) -> Dict:
"""Find LoRAs with duplicate SHA256 hashes""" """Find LoRAs with duplicate SHA256 hashes"""
return self.scanner._hash_index.get_duplicate_hashes() return self.scanner._hash_index.get_duplicate_hashes()