Add tag filtering checkpoint

This commit is contained in:
Will Miao
2025-03-10 13:18:56 +08:00
parent 0069f84630
commit 721bef3ff8
15 changed files with 482 additions and 50 deletions

View File

@@ -45,6 +45,7 @@ class ApiRoutes:
app.router.add_post('/loras/api/save-metadata', routes.save_metadata)
app.router.add_get('/api/lora-preview-url', routes.get_lora_preview_url) # Add new route
app.router.add_post('/api/move_models_bulk', routes.move_models_bulk)
app.router.add_get('/api/top-tags', routes.get_top_tags) # Add new route for top tags
# Add update check routes
UpdateRoutes.setup_routes(app)
@@ -142,6 +143,10 @@ class ApiRoutes:
'error': 'Invalid sort parameter'
}, status=400)
# Parse tags filter parameter
tags = request.query.get('tags', '').split(',')
tags = [tag.strip() for tag in tags if tag.strip()]
# Get paginated data with search and filters
result = await self.scanner.get_paginated_data(
page=page,
@@ -151,7 +156,8 @@ class ApiRoutes:
search=search,
fuzzy=fuzzy,
recursive=recursive,
base_models=base_models # Pass base models filter
base_models=base_models, # Pass base models filter
tags=tags # Add tags parameter
)
# Format the response data
@@ -190,6 +196,8 @@ class ApiRoutes:
"file_path": lora["file_path"].replace(os.sep, "/"),
"file_size": lora["size"],
"modified": lora["modified"],
"tags": lora["tags"],
"modelDescription": lora["modelDescription"],
"from_civitai": lora.get("from_civitai", True),
"usage_tips": lora.get("usage_tips", ""),
"notes": lora.get("notes", ""),
@@ -335,6 +343,14 @@ class ApiRoutes:
local_metadata['model_name'] = civitai_metadata['model'].get('name',
local_metadata.get('model_name'))
# Fetch additional model metadata (description and tags) if we have model ID
model_id = civitai_metadata['modelId']
if model_id:
model_metadata = await client.get_model_metadata(str(model_id))
if model_metadata:
local_metadata['modelDescription'] = model_metadata.get('description', '')
local_metadata['tags'] = model_metadata.get('tags', [])
# Update base model
local_metadata['base_model'] = civitai_metadata.get('baseModel')
@@ -708,6 +724,7 @@ class ApiRoutes:
# Check if we already have the description stored in metadata
description = None
tags = []
if file_path:
metadata_path = os.path.splitext(file_path)[0] + '.metadata.json'
if os.path.exists(metadata_path):
@@ -715,38 +732,70 @@ class ApiRoutes:
with open(metadata_path, 'r', encoding='utf-8') as f:
metadata = json.load(f)
description = metadata.get('modelDescription')
tags = metadata.get('tags', [])
except Exception as e:
logger.error(f"Error loading metadata from {metadata_path}: {e}")
# If description is not in metadata, fetch from CivitAI
if not description:
logger.info(f"Fetching model description for model ID: {model_id}")
description = await self.civitai_client.get_model_description(model_id)
logger.info(f"Fetching model metadata for model ID: {model_id}")
model_metadata = await self.civitai_client.get_model_metadata(model_id)
# Save the description to metadata if we have a file path and got a description
if file_path and description:
try:
metadata_path = os.path.splitext(file_path)[0] + '.metadata.json'
if os.path.exists(metadata_path):
with open(metadata_path, 'r', encoding='utf-8') as f:
metadata = json.load(f)
metadata['modelDescription'] = description
with open(metadata_path, 'w', encoding='utf-8') as f:
json.dump(metadata, f, indent=2, ensure_ascii=False)
logger.info(f"Saved model description to metadata for {file_path}")
except Exception as e:
logger.error(f"Error saving model description to metadata: {e}")
if model_metadata:
description = model_metadata.get('description')
tags = model_metadata.get('tags', [])
# Save the metadata to file if we have a file path and got metadata
if file_path:
try:
metadata_path = os.path.splitext(file_path)[0] + '.metadata.json'
if os.path.exists(metadata_path):
with open(metadata_path, 'r', encoding='utf-8') as f:
metadata = json.load(f)
metadata['modelDescription'] = description
metadata['tags'] = tags
with open(metadata_path, 'w', encoding='utf-8') as f:
json.dump(metadata, f, indent=2, ensure_ascii=False)
logger.info(f"Saved model metadata to file for {file_path}")
except Exception as e:
logger.error(f"Error saving model metadata: {e}")
return web.json_response({
'success': True,
'description': description or "<p>No model description available.</p>"
'description': description or "<p>No model description available.</p>",
'tags': tags
})
except Exception as e:
logger.error(f"Error getting model description: {e}", exc_info=True)
logger.error(f"Error getting model metadata: {e}", exc_info=True)
return web.json_response({
'success': False,
'error': str(e)
}, status=500)
async def get_top_tags(self, request: web.Request) -> web.Response:
"""Handle request for top tags sorted by frequency"""
try:
# Parse query parameters
limit = int(request.query.get('limit', '20'))
# Validate limit
if limit < 1 or limit > 100:
limit = 20 # Default to a reasonable limit
# Get top tags
top_tags = await self.scanner.get_top_tags(limit)
return web.json_response({
'success': True,
'tags': top_tags
})
except Exception as e:
logger.error(f"Error getting top tags: {str(e)}", exc_info=True)
return web.json_response({
'success': False,
'error': 'Internal server error'
}, status=500)

View File

@@ -30,6 +30,11 @@ class LoraRoutes:
"folder": lora["folder"],
"sha256": lora["sha256"],
"file_path": lora["file_path"].replace(os.sep, "/"),
"size": lora["size"],
"tags": lora["tags"],
"modelDescription": lora["modelDescription"],
"usage_tips": lora["usage_tips"],
"notes": lora["notes"],
"modified": lora["modified"],
"from_civitai": lora.get("from_civitai", True),
"civitai": self._filter_civitai_data(lora.get("civitai", {}))