fix(routes): enhance model processing to include checks for missing tags, description, and creator

This commit is contained in:
Will Miao
2025-09-10 23:30:08 +08:00
parent a7df8ae15c
commit 1b5e608a27
2 changed files with 29 additions and 31 deletions

View File

@@ -611,10 +611,19 @@ class BaseModelRoutes(ABC):
success = 0 success = 0
needs_resort = False needs_resort = False
# Prepare models to process, only those without CivitAI data # Prepare models to process, only those without CivitAI data or missing tags, description, or creator
to_process = [ to_process = [
model for model in cache.raw_data model for model in cache.raw_data
if model.get('sha256') and (not model.get('civitai') or 'id' not in model.get('civitai')) if (
model.get('sha256')
and (
not model.get('civitai')
or not model['civitai'].get('id')
or not model.get('tags')
or not model.get('modelDescription')
or not (model.get('civitai') and model['civitai'].get('creator'))
)
)
] ]
total_to_process = len(to_process) total_to_process = len(to_process)

View File

@@ -83,38 +83,27 @@ class ModelRouteUtils:
local_metadata['civitai'] = merged_civitai local_metadata['civitai'] = merged_civitai
local_metadata['from_civitai'] = True local_metadata['from_civitai'] = True
# Update model name if available # Update model-related metadata from civitai_metadata.model
if 'model' in civitai_metadata: if 'model' in civitai_metadata and civitai_metadata['model']:
if civitai_metadata.get('model', {}).get('name'): model_data = civitai_metadata['model']
local_metadata['model_name'] = civitai_metadata['model']['name']
# Extract model metadata directly from civitai_metadata if available
model_metadata = None
if 'model' in civitai_metadata and civitai_metadata.get('model'): # Update model name if available and not already set
# Data is already available in the response from get_model_version if model_data.get('name'):
model_metadata = { local_metadata['model_name'] = model_data['name']
'description': civitai_metadata.get('model', {}).get('description', ''),
'tags': civitai_metadata.get('model', {}).get('tags', []),
'creator': civitai_metadata.get('creator', {})
}
# If we have modelId and don't have enough metadata, fetch additional data # Update modelDescription if missing or empty in local_metadata
if not model_metadata or not model_metadata.get('description'): if not local_metadata.get('modelDescription') and model_data.get('description'):
model_id = civitai_metadata.get('modelId') local_metadata['modelDescription'] = model_data['description']
if model_id and metadata_provider:
fetched_metadata, _ = await metadata_provider.get_model_metadata(str(model_id))
if fetched_metadata:
model_metadata = fetched_metadata
# Update local metadata with the model information # Update tags if missing or empty in local_metadata
if model_metadata: if not local_metadata.get('tags') and model_data.get('tags'):
local_metadata['modelDescription'] = model_metadata.get('description', '') local_metadata['tags'] = model_data['tags']
# Only set tags if local_metadata['tags'] is empty
if not local_metadata.get('tags'): # Update creator in civitai metadata if missing
local_metadata['tags'] = model_metadata.get('tags', []) if model_data.get('creator') and not local_metadata.get('civitai', {}).get('creator'):
if 'creator' in model_metadata and model_metadata['creator']: if 'civitai' not in local_metadata:
local_metadata['civitai']['creator'] = model_metadata['creator'] local_metadata['civitai'] = {}
local_metadata['civitai']['creator'] = model_data['creator']
# Update base model # Update base model
local_metadata['base_model'] = determine_base_model(civitai_metadata.get('baseModel')) local_metadata['base_model'] = determine_base_model(civitai_metadata.get('baseModel'))