feat: Enhance CivitaiClient and ApiRoutes to handle model version errors and improve metadata fetching. Fixes https://github.com/willmiao/ComfyUI-Lora-Manager/issues/112

This commit is contained in:
Will Miao
2025-04-18 14:43:56 +08:00
parent 061660d47a
commit be8605d8c6
4 changed files with 70 additions and 43 deletions

View File

@@ -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', [])

View File

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

View File

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

View File

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