mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-21 21:22:11 -03:00
feat(civitai): enhance metadata fetching with error handling and cache validation
This commit is contained in:
@@ -255,20 +255,45 @@ class BaseModelRoutes(ABC):
|
|||||||
return await ModelRouteUtils.handle_exclude_model(request, self.service.scanner)
|
return await ModelRouteUtils.handle_exclude_model(request, self.service.scanner)
|
||||||
|
|
||||||
async def fetch_civitai(self, request: web.Request) -> web.Response:
|
async def fetch_civitai(self, request: web.Request) -> web.Response:
|
||||||
"""Handle CivitAI metadata fetch request"""
|
"""Handle CivitAI metadata fetch request - force refresh model metadata"""
|
||||||
response = await ModelRouteUtils.handle_fetch_civitai(request, self.service.scanner)
|
try:
|
||||||
|
data = await request.json()
|
||||||
# If successful, format the metadata before returning
|
file_path = data.get('file_path')
|
||||||
if response.status == 200:
|
if not file_path:
|
||||||
data = json.loads(response.body.decode('utf-8'))
|
return web.json_response({"success": False, "error": "File path is required"}, status=400)
|
||||||
if data.get("success") and data.get("metadata"):
|
|
||||||
formatted_metadata = await self.service.format_response(data["metadata"])
|
# Get model data from cache
|
||||||
return web.json_response({
|
cache = await self.service.scanner.get_cached_data()
|
||||||
"success": True,
|
model_data = next((item for item in cache.raw_data if item['file_path'] == file_path), None)
|
||||||
"metadata": formatted_metadata
|
|
||||||
})
|
if not model_data:
|
||||||
|
return web.json_response({"success": False, "error": "Model not found in cache"}, status=404)
|
||||||
return response
|
|
||||||
|
# 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:
|
async def relink_civitai(self, request: web.Request) -> web.Response:
|
||||||
"""Handle CivitAI metadata re-linking request"""
|
"""Handle CivitAI metadata re-linking request"""
|
||||||
@@ -652,12 +677,13 @@ class BaseModelRoutes(ABC):
|
|||||||
for model in to_process:
|
for model in to_process:
|
||||||
try:
|
try:
|
||||||
original_name = model.get('model_name')
|
original_name = model.get('model_name')
|
||||||
if await ModelRouteUtils.fetch_and_update_model(
|
result, error = await ModelRouteUtils.fetch_and_update_model(
|
||||||
sha256=model['sha256'],
|
sha256=model['sha256'],
|
||||||
file_path=model['file_path'],
|
file_path=model['file_path'],
|
||||||
model_data=model,
|
model_data=model,
|
||||||
update_cache_func=self.service.scanner.update_single_model_cache
|
update_cache_func=self.service.scanner.update_single_model_cache
|
||||||
):
|
)
|
||||||
|
if result:
|
||||||
success += 1
|
success += 1
|
||||||
if original_name != model.get('model_name'):
|
if original_name != model.get('model_name'):
|
||||||
needs_resort = True
|
needs_resort = True
|
||||||
|
|||||||
@@ -53,7 +53,7 @@ class MetadataUpdater:
|
|||||||
async def update_cache_func(old_path, new_path, metadata):
|
async def update_cache_func(old_path, new_path, metadata):
|
||||||
return await scanner.update_single_model_cache(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,
|
model_hash,
|
||||||
file_path,
|
file_path,
|
||||||
model_data,
|
model_data,
|
||||||
@@ -64,7 +64,7 @@ class MetadataUpdater:
|
|||||||
logger.info(f"Successfully refreshed metadata for {model_name}")
|
logger.info(f"Successfully refreshed metadata for {model_name}")
|
||||||
return True
|
return True
|
||||||
else:
|
else:
|
||||||
logger.warning(f"Failed to refresh metadata for {model_name}")
|
logger.warning(f"Failed to refresh metadata for {model_name}, {error}")
|
||||||
return False
|
return False
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -184,7 +184,7 @@ class ModelRouteUtils:
|
|||||||
file_path: str,
|
file_path: str,
|
||||||
model_data: dict,
|
model_data: dict,
|
||||||
update_cache_func: Callable[[str, str, Dict], Awaitable[bool]]
|
update_cache_func: Callable[[str, str, Dict], Awaitable[bool]]
|
||||||
) -> bool:
|
) -> tuple[bool, str]:
|
||||||
"""Fetch and update metadata for a single model
|
"""Fetch and update metadata for a single model
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
@@ -194,21 +194,22 @@ class ModelRouteUtils:
|
|||||||
update_cache_func: Function to update the cache with new metadata
|
update_cache_func: Function to update the cache with new metadata
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
bool: True if successful, False otherwise
|
tuple[bool, str]: (success, error_message). When success is True, error_message is None.
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
# Validate input parameters
|
# Validate input parameters
|
||||||
if not isinstance(model_data, dict):
|
if not isinstance(model_data, dict):
|
||||||
logger.error(f"Invalid model_data type: {type(model_data)}")
|
error_msg = f"Invalid model_data type: {type(model_data)}"
|
||||||
return False
|
logger.error(error_msg)
|
||||||
|
return False, error_msg
|
||||||
|
|
||||||
metadata_path = os.path.splitext(file_path)[0] + '.metadata.json'
|
metadata_path = os.path.splitext(file_path)[0] + '.metadata.json'
|
||||||
enable_metadata_archive_db = settings.get('enable_metadata_archive_db', False)
|
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 CivitAI deleted flag is set, skip CivitAI API provider
|
||||||
if not enable_metadata_archive_db:
|
if not enable_metadata_archive_db or model_data.get('db_checked') is True:
|
||||||
return False
|
return False, "CivitAI model is deleted and metadata archive DB is not enabled"
|
||||||
# Likely deleted from CivitAI, use archive_db if available
|
# Likely deleted from CivitAI, use archive_db if available
|
||||||
metadata_provider = await get_metadata_provider('sqlite')
|
metadata_provider = await get_metadata_provider('sqlite')
|
||||||
else:
|
else:
|
||||||
@@ -226,11 +227,16 @@ class ModelRouteUtils:
|
|||||||
data_to_save = model_data.copy()
|
data_to_save = model_data.copy()
|
||||||
data_to_save.pop('folder', None)
|
data_to_save.pop('folder', None)
|
||||||
await MetadataManager.save_metadata(file_path, data_to_save)
|
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['from_civitai'] = True
|
||||||
model_data['civitai_deleted'] = civitai_metadata.get('source') == 'archive_db'
|
model_data['civitai_deleted'] = civitai_metadata.get('source') == 'archive_db'
|
||||||
model_data['db_checked'] = enable_metadata_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 = model_data.copy()
|
||||||
local_metadata.pop('folder', None) # Remove 'folder' key if present
|
local_metadata.pop('folder', None) # Remove 'folder' key if present
|
||||||
@@ -254,14 +260,16 @@ class ModelRouteUtils:
|
|||||||
# Update cache using the provided function
|
# Update cache using the provided function
|
||||||
await update_cache_func(file_path, file_path, local_metadata)
|
await update_cache_func(file_path, file_path, local_metadata)
|
||||||
|
|
||||||
return True
|
return True, None
|
||||||
|
|
||||||
except KeyError as e:
|
except KeyError as e:
|
||||||
logger.error(f"Error fetching CivitAI data - Missing key: {e} in model_data={model_data}")
|
error_msg = f"Error fetching metadata - Missing key: {e} in model_data={model_data}"
|
||||||
return False
|
logger.error(error_msg)
|
||||||
|
return False, error_msg
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error fetching CivitAI data: {str(e)}", exc_info=True) # Include stack trace
|
error_msg = f"Error fetching metadata: {str(e)}"
|
||||||
return False
|
logger.error(error_msg, exc_info=True) # Include stack trace
|
||||||
|
return False, error_msg
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def filter_civitai_data(data: Dict, minimal: bool = False) -> Dict:
|
def filter_civitai_data(data: Dict, minimal: bool = False) -> Dict:
|
||||||
|
|||||||
Reference in New Issue
Block a user