From a8d7070832d9eb5496aaeeef4c5ee949c9a7867e Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Sat, 20 Sep 2025 21:35:34 +0800 Subject: [PATCH] feat(civitai): enhance metadata fetching with error handling and cache validation --- py/routes/base_model_routes.py | 58 +++++++++++++++++++++-------- py/utils/example_images_metadata.py | 4 +- py/utils/routes_common.py | 34 ++++++++++------- 3 files changed, 65 insertions(+), 31 deletions(-) diff --git a/py/routes/base_model_routes.py b/py/routes/base_model_routes.py index db914960..6c3f11b4 100644 --- a/py/routes/base_model_routes.py +++ b/py/routes/base_model_routes.py @@ -255,20 +255,45 @@ class BaseModelRoutes(ABC): return await ModelRouteUtils.handle_exclude_model(request, self.service.scanner) async def fetch_civitai(self, request: web.Request) -> web.Response: - """Handle CivitAI metadata fetch request""" - response = await ModelRouteUtils.handle_fetch_civitai(request, self.service.scanner) - - # If successful, format the metadata before returning - if response.status == 200: - data = json.loads(response.body.decode('utf-8')) - if data.get("success") and data.get("metadata"): - formatted_metadata = await self.service.format_response(data["metadata"]) - return web.json_response({ - "success": True, - "metadata": formatted_metadata - }) - - return response + """Handle CivitAI metadata fetch request - force refresh model metadata""" + try: + data = await request.json() + file_path = data.get('file_path') + if not file_path: + return web.json_response({"success": False, "error": "File path is required"}, status=400) + + # Get model data from cache + cache = await self.service.scanner.get_cached_data() + model_data = next((item for item in cache.raw_data if item['file_path'] == file_path), None) + + if not model_data: + return web.json_response({"success": False, "error": "Model not found in cache"}, status=404) + + # Check if model has SHA256 hash + if not model_data.get('sha256'): + return web.json_response({"success": False, "error": "No SHA256 hash found"}, status=400) + + # Use fetch_and_update_model to get and update metadata + success, error = await ModelRouteUtils.fetch_and_update_model( + sha256=model_data['sha256'], + file_path=file_path, + model_data=model_data, + update_cache_func=self.service.scanner.update_single_model_cache + ) + + if not success: + return web.json_response({"success": False, "error": error}) + + # Format the updated metadata for response + formatted_metadata = await self.service.format_response(model_data) + return web.json_response({ + "success": True, + "metadata": formatted_metadata + }) + + except Exception as e: + logger.error(f"Error fetching from CivitAI: {e}", exc_info=True) + return web.json_response({"success": False, "error": str(e)}, status=500) async def relink_civitai(self, request: web.Request) -> web.Response: """Handle CivitAI metadata re-linking request""" @@ -652,12 +677,13 @@ class BaseModelRoutes(ABC): for model in to_process: try: original_name = model.get('model_name') - if await ModelRouteUtils.fetch_and_update_model( + result, error = await ModelRouteUtils.fetch_and_update_model( sha256=model['sha256'], file_path=model['file_path'], model_data=model, update_cache_func=self.service.scanner.update_single_model_cache - ): + ) + if result: success += 1 if original_name != model.get('model_name'): needs_resort = True diff --git a/py/utils/example_images_metadata.py b/py/utils/example_images_metadata.py index 496d5ad0..71566bff 100644 --- a/py/utils/example_images_metadata.py +++ b/py/utils/example_images_metadata.py @@ -53,7 +53,7 @@ class MetadataUpdater: async def update_cache_func(old_path, new_path, metadata): return await scanner.update_single_model_cache(old_path, new_path, metadata) - success = await ModelRouteUtils.fetch_and_update_model( + success, error = await ModelRouteUtils.fetch_and_update_model( model_hash, file_path, model_data, @@ -64,7 +64,7 @@ class MetadataUpdater: logger.info(f"Successfully refreshed metadata for {model_name}") return True else: - logger.warning(f"Failed to refresh metadata for {model_name}") + logger.warning(f"Failed to refresh metadata for {model_name}, {error}") return False except Exception as e: diff --git a/py/utils/routes_common.py b/py/utils/routes_common.py index d6635118..b5f6af30 100644 --- a/py/utils/routes_common.py +++ b/py/utils/routes_common.py @@ -184,7 +184,7 @@ class ModelRouteUtils: file_path: str, model_data: dict, update_cache_func: Callable[[str, str, Dict], Awaitable[bool]] - ) -> bool: + ) -> tuple[bool, str]: """Fetch and update metadata for a single model Args: @@ -194,21 +194,22 @@ class ModelRouteUtils: update_cache_func: Function to update the cache with new metadata Returns: - bool: True if successful, False otherwise + tuple[bool, str]: (success, error_message). When success is True, error_message is None. """ try: # Validate input parameters if not isinstance(model_data, dict): - logger.error(f"Invalid model_data type: {type(model_data)}") - return False + error_msg = f"Invalid model_data type: {type(model_data)}" + logger.error(error_msg) + return False, error_msg metadata_path = os.path.splitext(file_path)[0] + '.metadata.json' enable_metadata_archive_db = settings.get('enable_metadata_archive_db', False) - if model_data.get('civitai_deleted') is True and model_data.get('db_checked') is False: + if model_data.get('civitai_deleted') is True: # If CivitAI deleted flag is set, skip CivitAI API provider - if not enable_metadata_archive_db: - return False + if not enable_metadata_archive_db or model_data.get('db_checked') is True: + return False, "CivitAI model is deleted and metadata archive DB is not enabled" # Likely deleted from CivitAI, use archive_db if available metadata_provider = await get_metadata_provider('sqlite') else: @@ -226,11 +227,16 @@ class ModelRouteUtils: data_to_save = model_data.copy() data_to_save.pop('folder', None) await MetadataManager.save_metadata(file_path, data_to_save) - return False + + # For other errors, log and return False with error message + error_msg = f"Error fetching metadata: {error} (model_name={model_data.get('model_name', '')})" + logger.error(error_msg) + return False, error_msg model_data['from_civitai'] = True model_data['civitai_deleted'] = civitai_metadata.get('source') == 'archive_db' model_data['db_checked'] = enable_metadata_archive_db + model_data['last_checked_at'] = datetime.now().timestamp() local_metadata = model_data.copy() local_metadata.pop('folder', None) # Remove 'folder' key if present @@ -254,14 +260,16 @@ class ModelRouteUtils: # Update cache using the provided function await update_cache_func(file_path, file_path, local_metadata) - return True + return True, None except KeyError as e: - logger.error(f"Error fetching CivitAI data - Missing key: {e} in model_data={model_data}") - return False + error_msg = f"Error fetching metadata - Missing key: {e} in model_data={model_data}" + logger.error(error_msg) + return False, error_msg except Exception as e: - logger.error(f"Error fetching CivitAI data: {str(e)}", exc_info=True) # Include stack trace - return False + error_msg = f"Error fetching metadata: {str(e)}" + logger.error(error_msg, exc_info=True) # Include stack trace + return False, error_msg @staticmethod def filter_civitai_data(data: Dict, minimal: bool = False) -> Dict: