From 9817bac2fe325b7090d03dd8c6378b6c2f1b1478 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Wed, 27 Aug 2025 17:44:29 +0800 Subject: [PATCH] feat: Add metadata endpoint and implement model metadata retrieval functionality --- py/routes/base_model_routes.py | 27 +++++++++++++++++++++++++++ py/services/base_model_service.py | 11 +++++++++++ py/services/checkpoint_service.py | 3 +-- py/services/embedding_service.py | 3 +-- py/services/lora_service.py | 3 +-- py/utils/routes_common.py | 8 ++++---- static/js/api/apiConfig.js | 4 ++++ static/js/api/baseModelApi.js | 22 ++++++++++++++++++++++ 8 files changed, 71 insertions(+), 10 deletions(-) diff --git a/py/routes/base_model_routes.py b/py/routes/base_model_routes.py index 86f3ad43..5f770d51 100644 --- a/py/routes/base_model_routes.py +++ b/py/routes/base_model_routes.py @@ -69,6 +69,7 @@ class BaseModelRoutes(ABC): 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) + app.router.add_get(f'/api/{prefix}/metadata', self.get_model_metadata) # Autocomplete route app.router.add_get(f'/api/{prefix}/relative-paths', self.get_relative_paths) @@ -1138,6 +1139,32 @@ class BaseModelRoutes(ABC): 'error': str(e) }, status=500) + async def get_model_metadata(self, request: web.Request) -> web.Response: + """Get filtered CivitAI metadata for a model by file path""" + try: + file_path = request.query.get('file_path') + if not file_path: + return web.Response(text='File path is required', status=400) + + metadata = await self.service.get_model_metadata(file_path) + if metadata is not None: + return web.json_response({ + 'success': True, + 'metadata': metadata + }) + else: + return web.json_response({ + 'success': False, + 'error': f'{self.model_type.capitalize()} not found or no CivitAI metadata available' + }, status=404) + + except Exception as e: + logger.error(f"Error getting {self.model_type} metadata: {e}", exc_info=True) + return web.json_response({ + 'success': False, + 'error': str(e) + }, status=500) + async def get_relative_paths(self, request: web.Request) -> web.Response: """Get model relative file paths for autocomplete functionality""" try: diff --git a/py/services/base_model_service.py b/py/services/base_model_service.py index 93a475b1..bb52df50 100644 --- a/py/services/base_model_service.py +++ b/py/services/base_model_service.py @@ -4,6 +4,7 @@ import logging import os from ..utils.models import BaseModelMetadata +from ..utils.routes_common import ModelRouteUtils from ..utils.constants import NSFW_LEVELS from .settings_manager import settings from ..utils.utils import fuzzy_match @@ -379,6 +380,16 @@ class BaseModelService(ABC): return {'civitai_url': None, 'model_id': None, 'version_id': None} + async def get_model_metadata(self, file_path: str) -> Optional[Dict]: + """Get filtered CivitAI metadata for a model by file path""" + cache = await self.scanner.get_cached_data() + + for model in cache.raw_data: + if model.get('file_path') == file_path: + return ModelRouteUtils.filter_civitai_data(model.get("civitai", {})) + + return None + async def search_relative_paths(self, search_term: str, limit: int = 15) -> List[str]: """Search model relative file paths for autocomplete functionality""" cache = await self.scanner.get_cached_data() diff --git a/py/services/checkpoint_service.py b/py/services/checkpoint_service.py index cfc4b118..ef3dc4a8 100644 --- a/py/services/checkpoint_service.py +++ b/py/services/checkpoint_service.py @@ -34,12 +34,11 @@ class CheckpointService(BaseModelService): "file_size": checkpoint_data.get("size", 0), "modified": checkpoint_data.get("modified", ""), "tags": checkpoint_data.get("tags", []), - "modelDescription": checkpoint_data.get("modelDescription", ""), "from_civitai": checkpoint_data.get("from_civitai", True), "notes": checkpoint_data.get("notes", ""), "model_type": checkpoint_data.get("model_type", "checkpoint"), "favorite": checkpoint_data.get("favorite", False), - "civitai": ModelRouteUtils.filter_civitai_data(checkpoint_data.get("civitai", {})) + "civitai": ModelRouteUtils.filter_civitai_data(checkpoint_data.get("civitai", {}), minimal=True) } def find_duplicate_hashes(self) -> Dict: diff --git a/py/services/embedding_service.py b/py/services/embedding_service.py index b63559d2..bab067d9 100644 --- a/py/services/embedding_service.py +++ b/py/services/embedding_service.py @@ -34,12 +34,11 @@ class EmbeddingService(BaseModelService): "file_size": embedding_data.get("size", 0), "modified": embedding_data.get("modified", ""), "tags": embedding_data.get("tags", []), - "modelDescription": embedding_data.get("modelDescription", ""), "from_civitai": embedding_data.get("from_civitai", True), "notes": embedding_data.get("notes", ""), "model_type": embedding_data.get("model_type", "embedding"), "favorite": embedding_data.get("favorite", False), - "civitai": ModelRouteUtils.filter_civitai_data(embedding_data.get("civitai", {})) + "civitai": ModelRouteUtils.filter_civitai_data(embedding_data.get("civitai", {}), minimal=True) } def find_duplicate_hashes(self) -> Dict: diff --git a/py/services/lora_service.py b/py/services/lora_service.py index f4c71096..d1e522a3 100644 --- a/py/services/lora_service.py +++ b/py/services/lora_service.py @@ -34,12 +34,11 @@ class LoraService(BaseModelService): "file_size": lora_data.get("size", 0), "modified": lora_data.get("modified", ""), "tags": lora_data.get("tags", []), - "modelDescription": lora_data.get("modelDescription", ""), "from_civitai": lora_data.get("from_civitai", True), "usage_tips": lora_data.get("usage_tips", ""), "notes": lora_data.get("notes", ""), "favorite": lora_data.get("favorite", False), - "civitai": ModelRouteUtils.filter_civitai_data(lora_data.get("civitai", {})) + "civitai": ModelRouteUtils.filter_civitai_data(lora_data.get("civitai", {}), minimal=True) } async def _apply_specific_filters(self, data: List[Dict], **kwargs) -> List[Dict]: diff --git a/py/utils/routes_common.py b/py/utils/routes_common.py index 47554903..8898ed58 100644 --- a/py/utils/routes_common.py +++ b/py/utils/routes_common.py @@ -229,13 +229,13 @@ class ModelRouteUtils: await client.close() @staticmethod - def filter_civitai_data(data: Dict) -> Dict: + def filter_civitai_data(data: Dict, minimal: bool = False) -> Dict: """Filter relevant fields from CivitAI data""" if not data: return {} - - fields = [ - "id", "modelId", "name", "createdAt", "updatedAt", + + fields = ["name", "trainedWords"] if minimal else [ + "id", "modelId", "name", "createdAt", "updatedAt", "publishedAt", "trainedWords", "baseModel", "description", "model", "images", "customImages", "creator" ] diff --git a/static/js/api/apiConfig.js b/static/js/api/apiConfig.js index b0fcc30a..d2e5ab2b 100644 --- a/static/js/api/apiConfig.js +++ b/static/js/api/apiConfig.js @@ -88,6 +88,7 @@ export function getApiEndpoints(modelType) { duplicates: `/api/${modelType}/find-duplicates`, conflicts: `/api/${modelType}/find-filename-conflicts`, verify: `/api/${modelType}/verify-duplicates`, + metadata: `/api/${modelType}/metadata`, // Model-specific endpoints (will be merged with specific configs) specific: {} @@ -104,6 +105,7 @@ export const MODEL_SPECIFIC_ENDPOINTS = { triggerWords: `/api/${MODEL_TYPES.LORA}/get-trigger-words`, previewUrl: `/api/${MODEL_TYPES.LORA}/preview-url`, civitaiUrl: `/api/${MODEL_TYPES.LORA}/civitai-url`, + metadata: `/api/${MODEL_TYPES.LORA}/metadata`, modelDescription: `/api/${MODEL_TYPES.LORA}/model-description`, getTriggerWordsPost: `/api/${MODEL_TYPES.LORA}/get_trigger_words`, civitaiModelByVersion: `/api/${MODEL_TYPES.LORA}/civitai/model/version`, @@ -113,8 +115,10 @@ export const MODEL_SPECIFIC_ENDPOINTS = { info: `/api/${MODEL_TYPES.CHECKPOINT}/info`, checkpoints_roots: `/api/${MODEL_TYPES.CHECKPOINT}/checkpoints_roots`, unet_roots: `/api/${MODEL_TYPES.CHECKPOINT}/unet_roots`, + metadata: `/api/${MODEL_TYPES.CHECKPOINT}/metadata`, }, [MODEL_TYPES.EMBEDDING]: { + metadata: `/api/${MODEL_TYPES.EMBEDDING}/metadata`, } }; diff --git a/static/js/api/baseModelApi.js b/static/js/api/baseModelApi.js index 793fe536..98c82eea 100644 --- a/static/js/api/baseModelApi.js +++ b/static/js/api/baseModelApi.js @@ -948,4 +948,26 @@ export class BaseModelApiClient { completionMessage: 'Example images download complete' }); } + + async fetchModelMetadata(filePath) { + try { + const params = new URLSearchParams({ file_path: filePath }); + const response = await fetch(`${this.apiConfig.endpoints.metadata}?${params}`); + + if (!response.ok) { + throw new Error(`Failed to fetch ${this.apiConfig.config.singularName} metadata: ${response.statusText}`); + } + + const data = await response.json(); + + if (data.success) { + return data.metadata; + } else { + throw new Error(data.error || `No metadata found for ${this.apiConfig.config.singularName}`); + } + } catch (error) { + console.error(`Error fetching ${this.apiConfig.config.singularName} metadata:`, error); + throw error; + } + } } \ No newline at end of file