From be8605d8c631377cfdd5d045a4155aa698e9e5cb Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Fri, 18 Apr 2025 14:43:56 +0800 Subject: [PATCH] feat: Enhance CivitaiClient and ApiRoutes to handle model version errors and improve metadata fetching. Fixes https://github.com/willmiao/ComfyUI-Lora-Manager/issues/112 --- py/routes/api_routes.py | 9 +++-- py/services/civitai_client.py | 35 ++++++++++++++++---- py/services/download_manager.py | 11 ++++--- py/services/recipe_scanner.py | 58 +++++++++++++++------------------ 4 files changed, 70 insertions(+), 43 deletions(-) diff --git a/py/routes/api_routes.py b/py/routes/api_routes.py index 834f8ad9..6c67e129 100644 --- a/py/routes/api_routes.py +++ b/py/routes/api_routes.py @@ -409,7 +409,12 @@ class ApiRoutes: return web.json_response(model) # Get model details from Civitai API - model = await self.civitai_client.get_model_version_info(model_version_id) + model, error_msg = await self.civitai_client.get_model_version_info(model_version_id) + + if not model: + status_code = 404 if error_msg and "model not found" in error_msg.lower() else 500 + return web.Response(status=status_code, text=error_msg or "Failed to fetch model information") + return web.json_response(model) except Exception as e: logger.error(f"Error fetching model details: {e}") @@ -773,7 +778,7 @@ class ApiRoutes: logger.info(f"Fetching model metadata for model ID: {model_id}") model_metadata, _ = await self.civitai_client.get_model_metadata(model_id) - if model_metadata: + if (model_metadata): description = model_metadata.get('description') tags = model_metadata.get('tags', []) diff --git a/py/services/civitai_client.py b/py/services/civitai_client.py index b6e5a5f6..f96e415a 100644 --- a/py/services/civitai_client.py +++ b/py/services/civitai_client.py @@ -210,8 +210,17 @@ class CivitaiClient: logger.error(f"Error fetching model versions: {e}") return None - async def get_model_version_info(self, version_id: str) -> Optional[Dict]: - """Fetch model version metadata from Civitai""" + async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]: + """Fetch model version metadata from Civitai + + Args: + version_id: The Civitai model version ID + + Returns: + Tuple[Optional[Dict], Optional[str]]: A tuple containing: + - The model version data or None if not found + - An error message if there was an error, or None on success + """ try: session = await self.session url = f"{self.base_url}/model-versions/{version_id}" @@ -219,11 +228,25 @@ class CivitaiClient: async with session.get(url, headers=headers) as response: if response.status == 200: - return await response.json() - return None + return await response.json(), None + + # Handle specific error cases + if response.status == 404: + # Try to parse the error message + try: + error_data = await response.json() + error_msg = error_data.get('error', f"Model not found (status 404)") + logger.warning(f"Model version not found: {version_id} - {error_msg}") + return None, error_msg + except: + return None, "Model not found (status 404)" + + # Other error cases + return None, f"Failed to fetch model info (status {response.status})" except Exception as e: - logger.error(f"Error fetching model version info: {e}") - return None + error_msg = f"Error fetching model version info: {e}" + logger.error(error_msg) + return None, error_msg async def get_model_metadata(self, model_id: str) -> Tuple[Optional[Dict], int]: """Fetch model metadata (description and tags) from Civitai API diff --git a/py/services/download_manager.py b/py/services/download_manager.py index 39794a01..d30426ca 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -86,21 +86,24 @@ class DownloadManager: # Get version info based on the provided identifier version_info = None + error_msg = None if download_url: # Extract version ID from download URL version_id = download_url.split('/')[-1] - version_info = await civitai_client.get_model_version_info(version_id) + version_info, error_msg = await civitai_client.get_model_version_info(version_id) elif model_version_id: # Use model version ID directly - version_info = await civitai_client.get_model_version_info(model_version_id) + version_info, error_msg = await civitai_client.get_model_version_info(model_version_id) elif model_hash: # Get model by hash version_info = await civitai_client.get_model_by_hash(model_hash) if not version_info: - return {'success': False, 'error': 'Failed to fetch model metadata'} + if error_msg and "model not found" in error_msg.lower(): + return {'success': False, 'error': f'Model not found on Civitai: {error_msg}'} + return {'success': False, 'error': error_msg or 'Failed to fetch model metadata'} # Check if this is an early access model if version_info.get('earlyAccessEndsAt'): @@ -202,7 +205,7 @@ class DownloadManager: # Check if it's a video or an image is_video = images[0].get('type') == 'video' - if is_video: + if (is_video): # For videos, use .mp4 extension preview_ext = '.mp4' preview_path = os.path.splitext(save_path)[0] + preview_ext diff --git a/py/services/recipe_scanner.py b/py/services/recipe_scanner.py index a4b2f59c..2a2e6063 100644 --- a/py/services/recipe_scanner.py +++ b/py/services/recipe_scanner.py @@ -341,6 +341,10 @@ class RecipeScanner: metadata_updated = False for lora in recipe_data['loras']: + # Skip deleted loras that were already marked + if lora.get('isDeleted', False): + continue + # Skip if already has complete information if 'hash' in lora and 'file_name' in lora and lora['file_name']: continue @@ -356,10 +360,17 @@ class RecipeScanner: metadata_updated = True else: # If not in cache, fetch from Civitai - hash_from_civitai = await self._get_hash_from_civitai(model_version_id) - if hash_from_civitai: - lora['hash'] = hash_from_civitai - metadata_updated = True + result = await self._get_hash_from_civitai(model_version_id) + if isinstance(result, tuple): + hash_from_civitai, is_deleted = result + if hash_from_civitai: + lora['hash'] = hash_from_civitai + metadata_updated = True + elif is_deleted: + # Mark the lora as deleted if it was not found on Civitai + lora['isDeleted'] = True + logger.warning(f"Marked lora with modelVersionId {model_version_id} as deleted") + metadata_updated = True else: logger.debug(f"Could not get hash for modelVersionId {model_version_id}") @@ -411,41 +422,26 @@ class RecipeScanner: logger.error("Failed to get CivitaiClient from ServiceRegistry") return None - version_info = await civitai_client.get_model_version_info(model_version_id) + version_info, error_msg = await civitai_client.get_model_version_info(model_version_id) - if not version_info or not version_info.get('files'): - logger.debug(f"No files found in version info for ID: {model_version_id}") - return None - + if not version_info: + if error_msg and "model not found" in error_msg.lower(): + logger.warning(f"Model with version ID {model_version_id} was not found on Civitai - marking as deleted") + return None, True # Return None hash and True for isDeleted flag + else: + logger.debug(f"Could not get hash for modelVersionId {model_version_id}: {error_msg}") + return None, False # Return None hash but not marked as deleted + # Get hash from the first file for file_info in version_info.get('files', []): if file_info.get('hashes', {}).get('SHA256'): - return file_info['hashes']['SHA256'] + return file_info['hashes']['SHA256'], False # Return hash with False for isDeleted flag logger.debug(f"No SHA256 hash found in version info for ID: {model_version_id}") - return None + return None, False except Exception as e: logger.error(f"Error getting hash from Civitai: {e}") - return None - - async def _get_model_version_name(self, model_version_id: str) -> Optional[str]: - """Get model version name from Civitai API""" - try: - # Get CivitaiClient from ServiceRegistry - civitai_client = await self._get_civitai_client() - if not civitai_client: - return None - - version_info = await civitai_client.get_model_version_info(model_version_id) - - if version_info and 'name' in version_info: - return version_info['name'] - - logger.debug(f"No version name found for modelVersionId {model_version_id}") - return None - except Exception as e: - logger.error(f"Error getting model version name from Civitai: {e}") - return None + return None, False async def _determine_base_model(self, loras: List[Dict]) -> Optional[str]: """Determine the most common base model among LoRAs"""