mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-21 21:22:11 -03:00
679 lines
29 KiB
Python
679 lines
29 KiB
Python
import os
|
|
import json
|
|
import jinja2
|
|
from aiohttp import web
|
|
import logging
|
|
import asyncio
|
|
|
|
from ..utils.routes_common import ModelRouteUtils
|
|
from ..utils.constants import NSFW_LEVELS
|
|
from ..services.websocket_manager import ws_manager
|
|
from ..services.service_registry import ServiceRegistry
|
|
from ..config import config
|
|
from ..services.settings_manager import settings
|
|
from ..utils.utils import fuzzy_match
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class CheckpointsRoutes:
|
|
"""API routes for checkpoint management"""
|
|
|
|
def __init__(self):
|
|
self.scanner = None # Will be initialized in setup_routes
|
|
self.template_env = jinja2.Environment(
|
|
loader=jinja2.FileSystemLoader(config.templates_path),
|
|
autoescape=True
|
|
)
|
|
self.download_manager = None # Will be initialized in setup_routes
|
|
self._download_lock = asyncio.Lock()
|
|
|
|
async def initialize_services(self):
|
|
"""Initialize services from ServiceRegistry"""
|
|
self.scanner = await ServiceRegistry.get_checkpoint_scanner()
|
|
self.download_manager = await ServiceRegistry.get_download_manager()
|
|
|
|
def setup_routes(self, app):
|
|
"""Register routes with the aiohttp app"""
|
|
# Schedule service initialization on app startup
|
|
app.on_startup.append(lambda _: self.initialize_services())
|
|
|
|
app.router.add_get('/checkpoints', self.handle_checkpoints_page)
|
|
app.router.add_get('/api/checkpoints', self.get_checkpoints)
|
|
app.router.add_post('/api/checkpoints/fetch-all-civitai', self.fetch_all_civitai)
|
|
app.router.add_get('/api/checkpoints/base-models', self.get_base_models)
|
|
app.router.add_get('/api/checkpoints/top-tags', self.get_top_tags)
|
|
app.router.add_get('/api/checkpoints/scan', self.scan_checkpoints)
|
|
app.router.add_get('/api/checkpoints/info/{name}', self.get_checkpoint_info)
|
|
app.router.add_get('/api/checkpoints/roots', self.get_checkpoint_roots)
|
|
app.router.add_get('/api/checkpoints/civitai/versions/{model_id}', self.get_civitai_versions) # Add new route
|
|
|
|
# Add new routes for model management similar to LoRA routes
|
|
app.router.add_post('/api/checkpoints/delete', self.delete_model)
|
|
app.router.add_post('/api/checkpoints/fetch-civitai', self.fetch_civitai)
|
|
app.router.add_post('/api/checkpoints/replace-preview', self.replace_preview)
|
|
app.router.add_post('/api/checkpoints/download', self.download_checkpoint)
|
|
app.router.add_post('/api/checkpoints/save-metadata', self.save_metadata) # Add new route
|
|
|
|
# Add new WebSocket endpoint for checkpoint progress
|
|
app.router.add_get('/ws/checkpoint-progress', ws_manager.handle_checkpoint_connection)
|
|
|
|
async def get_checkpoints(self, request):
|
|
"""Get paginated checkpoint data"""
|
|
try:
|
|
# Parse query parameters
|
|
page = int(request.query.get('page', '1'))
|
|
page_size = min(int(request.query.get('page_size', '20')), 100)
|
|
sort_by = request.query.get('sort_by', 'name')
|
|
folder = request.query.get('folder', None)
|
|
search = request.query.get('search', None)
|
|
fuzzy_search = request.query.get('fuzzy_search', 'false').lower() == 'true'
|
|
base_models = request.query.getall('base_model', [])
|
|
tags = request.query.getall('tag', [])
|
|
|
|
# Process search options
|
|
search_options = {
|
|
'filename': request.query.get('search_filename', 'true').lower() == 'true',
|
|
'modelname': request.query.get('search_modelname', 'true').lower() == 'true',
|
|
'tags': request.query.get('search_tags', 'false').lower() == 'true',
|
|
'recursive': request.query.get('recursive', 'false').lower() == 'true',
|
|
}
|
|
|
|
# Process hash filters if provided
|
|
hash_filters = {}
|
|
if 'hash' in request.query:
|
|
hash_filters['single_hash'] = request.query['hash']
|
|
elif 'hashes' in request.query:
|
|
try:
|
|
hash_list = json.loads(request.query['hashes'])
|
|
if isinstance(hash_list, list):
|
|
hash_filters['multiple_hashes'] = hash_list
|
|
except (json.JSONDecodeError, TypeError):
|
|
pass
|
|
|
|
# Get data from scanner
|
|
result = await self.get_paginated_data(
|
|
page=page,
|
|
page_size=page_size,
|
|
sort_by=sort_by,
|
|
folder=folder,
|
|
search=search,
|
|
fuzzy_search=fuzzy_search,
|
|
base_models=base_models,
|
|
tags=tags,
|
|
search_options=search_options,
|
|
hash_filters=hash_filters
|
|
)
|
|
|
|
# Format response items
|
|
formatted_result = {
|
|
'items': [self._format_checkpoint_response(cp) for cp in result['items']],
|
|
'total': result['total'],
|
|
'page': result['page'],
|
|
'page_size': result['page_size'],
|
|
'total_pages': result['total_pages']
|
|
}
|
|
|
|
# Return as JSON
|
|
return web.json_response(formatted_result)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in get_checkpoints: {e}", exc_info=True)
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
async def get_paginated_data(self, page, page_size, sort_by='name',
|
|
folder=None, search=None, fuzzy_search=False,
|
|
base_models=None, tags=None,
|
|
search_options=None, hash_filters=None):
|
|
"""Get paginated and filtered checkpoint data"""
|
|
cache = await self.scanner.get_cached_data()
|
|
|
|
# Get default search options if not provided
|
|
if search_options is None:
|
|
search_options = {
|
|
'filename': True,
|
|
'modelname': True,
|
|
'tags': False,
|
|
'recursive': False,
|
|
}
|
|
|
|
# Get the base data set
|
|
filtered_data = cache.sorted_by_date if sort_by == 'date' else cache.sorted_by_name
|
|
|
|
# Apply hash filtering if provided (highest priority)
|
|
if hash_filters:
|
|
single_hash = hash_filters.get('single_hash')
|
|
multiple_hashes = hash_filters.get('multiple_hashes')
|
|
|
|
if single_hash:
|
|
# Filter by single hash
|
|
single_hash = single_hash.lower() # Ensure lowercase for matching
|
|
filtered_data = [
|
|
cp for cp in filtered_data
|
|
if cp.get('sha256', '').lower() == single_hash
|
|
]
|
|
elif multiple_hashes:
|
|
# Filter by multiple hashes
|
|
hash_set = set(hash.lower() for hash in multiple_hashes) # Convert to set for faster lookup
|
|
filtered_data = [
|
|
cp for cp in filtered_data
|
|
if cp.get('sha256', '').lower() in hash_set
|
|
]
|
|
|
|
# Jump to pagination
|
|
total_items = len(filtered_data)
|
|
start_idx = (page - 1) * page_size
|
|
end_idx = min(start_idx + page_size, total_items)
|
|
|
|
result = {
|
|
'items': filtered_data[start_idx:end_idx],
|
|
'total': total_items,
|
|
'page': page,
|
|
'page_size': page_size,
|
|
'total_pages': (total_items + page_size - 1) // page_size
|
|
}
|
|
|
|
return result
|
|
|
|
# Apply SFW filtering if enabled in settings
|
|
if settings.get('show_only_sfw', False):
|
|
filtered_data = [
|
|
cp for cp in filtered_data
|
|
if not cp.get('preview_nsfw_level') or cp.get('preview_nsfw_level') < NSFW_LEVELS['R']
|
|
]
|
|
|
|
# Apply folder filtering
|
|
if folder is not None:
|
|
if search_options.get('recursive', False):
|
|
# Recursive folder filtering - include all subfolders
|
|
filtered_data = [
|
|
cp for cp in filtered_data
|
|
if cp['folder'].startswith(folder)
|
|
]
|
|
else:
|
|
# Exact folder filtering
|
|
filtered_data = [
|
|
cp for cp in filtered_data
|
|
if cp['folder'] == folder
|
|
]
|
|
|
|
# Apply base model filtering
|
|
if base_models and len(base_models) > 0:
|
|
filtered_data = [
|
|
cp for cp in filtered_data
|
|
if cp.get('base_model') in base_models
|
|
]
|
|
|
|
# Apply tag filtering
|
|
if tags and len(tags) > 0:
|
|
filtered_data = [
|
|
cp for cp in filtered_data
|
|
if any(tag in cp.get('tags', []) for tag in tags)
|
|
]
|
|
|
|
# Apply search filtering
|
|
if search:
|
|
search_results = []
|
|
|
|
for cp in filtered_data:
|
|
# Search by file name
|
|
if search_options.get('filename', True):
|
|
if fuzzy_search:
|
|
if fuzzy_match(cp.get('file_name', ''), search):
|
|
search_results.append(cp)
|
|
continue
|
|
elif search.lower() in cp.get('file_name', '').lower():
|
|
search_results.append(cp)
|
|
continue
|
|
|
|
# Search by model name
|
|
if search_options.get('modelname', True):
|
|
if fuzzy_search:
|
|
if fuzzy_match(cp.get('model_name', ''), search):
|
|
search_results.append(cp)
|
|
continue
|
|
elif search.lower() in cp.get('model_name', '').lower():
|
|
search_results.append(cp)
|
|
continue
|
|
|
|
# Search by tags
|
|
if search_options.get('tags', False) and 'tags' in cp:
|
|
if any((fuzzy_match(tag, search) if fuzzy_search else search.lower() in tag.lower()) for tag in cp['tags']):
|
|
search_results.append(cp)
|
|
continue
|
|
|
|
filtered_data = search_results
|
|
|
|
# Calculate pagination
|
|
total_items = len(filtered_data)
|
|
start_idx = (page - 1) * page_size
|
|
end_idx = min(start_idx + page_size, total_items)
|
|
|
|
result = {
|
|
'items': filtered_data[start_idx:end_idx],
|
|
'total': total_items,
|
|
'page': page,
|
|
'page_size': page_size,
|
|
'total_pages': (total_items + page_size - 1) // page_size
|
|
}
|
|
|
|
return result
|
|
|
|
def _format_checkpoint_response(self, checkpoint):
|
|
"""Format checkpoint data for API response"""
|
|
return {
|
|
"model_name": checkpoint["model_name"],
|
|
"file_name": checkpoint["file_name"],
|
|
"preview_url": config.get_preview_static_url(checkpoint.get("preview_url", "")),
|
|
"preview_nsfw_level": checkpoint.get("preview_nsfw_level", 0),
|
|
"base_model": checkpoint.get("base_model", ""),
|
|
"folder": checkpoint["folder"],
|
|
"sha256": checkpoint.get("sha256", ""),
|
|
"file_path": checkpoint["file_path"].replace(os.sep, "/"),
|
|
"file_size": checkpoint.get("size", 0),
|
|
"modified": checkpoint.get("modified", ""),
|
|
"tags": checkpoint.get("tags", []),
|
|
"modelDescription": checkpoint.get("modelDescription", ""),
|
|
"from_civitai": checkpoint.get("from_civitai", True),
|
|
"notes": checkpoint.get("notes", ""),
|
|
"model_type": checkpoint.get("model_type", "checkpoint"),
|
|
"civitai": ModelRouteUtils.filter_civitai_data(checkpoint.get("civitai", {}))
|
|
}
|
|
|
|
async def fetch_all_civitai(self, request: web.Request) -> web.Response:
|
|
"""Fetch CivitAI metadata for all checkpoints in the background"""
|
|
try:
|
|
cache = await self.scanner.get_cached_data()
|
|
total = len(cache.raw_data)
|
|
processed = 0
|
|
success = 0
|
|
needs_resort = False
|
|
|
|
# Prepare checkpoints to process
|
|
to_process = [
|
|
cp for cp in cache.raw_data
|
|
if cp.get('sha256') and (not cp.get('civitai') or 'id' not in cp.get('civitai')) and cp.get('from_civitai', True)
|
|
]
|
|
total_to_process = len(to_process)
|
|
|
|
# Send initial progress
|
|
await ws_manager.broadcast({
|
|
'status': 'started',
|
|
'total': total_to_process,
|
|
'processed': 0,
|
|
'success': 0
|
|
})
|
|
|
|
# Process each checkpoint
|
|
for cp in to_process:
|
|
try:
|
|
original_name = cp.get('model_name')
|
|
if await ModelRouteUtils.fetch_and_update_model(
|
|
sha256=cp['sha256'],
|
|
file_path=cp['file_path'],
|
|
model_data=cp,
|
|
update_cache_func=self.scanner.update_single_model_cache
|
|
):
|
|
success += 1
|
|
if original_name != cp.get('model_name'):
|
|
needs_resort = True
|
|
|
|
processed += 1
|
|
|
|
# Send progress update
|
|
await ws_manager.broadcast({
|
|
'status': 'processing',
|
|
'total': total_to_process,
|
|
'processed': processed,
|
|
'success': success,
|
|
'current_name': cp.get('model_name', 'Unknown')
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error fetching CivitAI data for {cp['file_path']}: {e}")
|
|
|
|
if needs_resort:
|
|
await cache.resort(name_only=True)
|
|
|
|
# Send completion message
|
|
await ws_manager.broadcast({
|
|
'status': 'completed',
|
|
'total': total_to_process,
|
|
'processed': processed,
|
|
'success': success
|
|
})
|
|
|
|
return web.json_response({
|
|
"success": True,
|
|
"message": f"Successfully updated {success} of {processed} processed checkpoints (total: {total})"
|
|
})
|
|
|
|
except Exception as e:
|
|
# Send error message
|
|
await ws_manager.broadcast({
|
|
'status': 'error',
|
|
'error': str(e)
|
|
})
|
|
logger.error(f"Error in fetch_all_civitai for checkpoints: {e}")
|
|
return web.Response(text=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)
|
|
|
|
async def get_base_models(self, request: web.Request) -> web.Response:
|
|
"""Get base models used in loras"""
|
|
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 base models
|
|
base_models = await self.scanner.get_base_models(limit)
|
|
|
|
return web.json_response({
|
|
'success': True,
|
|
'base_models': base_models
|
|
})
|
|
except Exception as e:
|
|
logger.error(f"Error retrieving base models: {e}")
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': str(e)
|
|
}, status=500)
|
|
|
|
async def scan_checkpoints(self, request):
|
|
"""Force a rescan of checkpoint files"""
|
|
try:
|
|
await self.scanner.get_cached_data(force_refresh=True)
|
|
return web.json_response({"status": "success", "message": "Checkpoint scan completed"})
|
|
except Exception as e:
|
|
logger.error(f"Error in scan_checkpoints: {e}", exc_info=True)
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
async def get_checkpoint_info(self, request):
|
|
"""Get detailed information for a specific checkpoint by name"""
|
|
try:
|
|
name = request.match_info.get('name', '')
|
|
checkpoint_info = await self.scanner.get_checkpoint_info_by_name(name)
|
|
|
|
if checkpoint_info:
|
|
return web.json_response(checkpoint_info)
|
|
else:
|
|
return web.json_response({"error": "Checkpoint not found"}, status=404)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in get_checkpoint_info: {e}", exc_info=True)
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
async def handle_checkpoints_page(self, request: web.Request) -> web.Response:
|
|
"""Handle GET /checkpoints request"""
|
|
try:
|
|
# Check if the CheckpointScanner is initializing
|
|
# It's initializing if the cache object doesn't exist yet,
|
|
# OR if the scanner explicitly says it's initializing (background task running).
|
|
is_initializing = (
|
|
self.scanner._cache is None or
|
|
(hasattr(self.scanner, '_is_initializing') and self.scanner._is_initializing)
|
|
)
|
|
|
|
if is_initializing:
|
|
# If still initializing, return loading page
|
|
template = self.template_env.get_template('checkpoints.html')
|
|
rendered = template.render(
|
|
folders=[], # 空文件夹列表
|
|
is_initializing=True, # 新增标志
|
|
settings=settings, # Pass settings to template
|
|
request=request # Pass the request object to the template
|
|
)
|
|
|
|
logger.info("Checkpoints page is initializing, returning loading page")
|
|
else:
|
|
# 正常流程 - 获取已经初始化好的缓存数据
|
|
try:
|
|
cache = await self.scanner.get_cached_data(force_refresh=False)
|
|
template = self.template_env.get_template('checkpoints.html')
|
|
rendered = template.render(
|
|
folders=cache.folders,
|
|
is_initializing=False,
|
|
settings=settings, # Pass settings to template
|
|
request=request # Pass the request object to the template
|
|
)
|
|
except Exception as cache_error:
|
|
logger.error(f"Error loading checkpoints cache data: {cache_error}")
|
|
# 如果获取缓存失败，也显示初始化页面
|
|
template = self.template_env.get_template('checkpoints.html')
|
|
rendered = template.render(
|
|
folders=[],
|
|
is_initializing=True,
|
|
settings=settings,
|
|
request=request
|
|
)
|
|
logger.info("Checkpoints cache error, returning initialization page")
|
|
|
|
return web.Response(
|
|
text=rendered,
|
|
content_type='text/html'
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error handling checkpoints request: {e}", exc_info=True)
|
|
return web.Response(
|
|
text="Error loading checkpoints page",
|
|
status=500
|
|
)
|
|
|
|
async def delete_model(self, request: web.Request) -> web.Response:
|
|
"""Handle checkpoint model deletion request"""
|
|
return await ModelRouteUtils.handle_delete_model(request, self.scanner)
|
|
|
|
async def fetch_civitai(self, request: web.Request) -> web.Response:
|
|
"""Handle CivitAI metadata fetch request for checkpoints"""
|
|
return await ModelRouteUtils.handle_fetch_civitai(request, self.scanner)
|
|
|
|
async def replace_preview(self, request: web.Request) -> web.Response:
|
|
"""Handle preview image replacement for checkpoints"""
|
|
return await ModelRouteUtils.handle_replace_preview(request, self.scanner)
|
|
|
|
async def download_checkpoint(self, request: web.Request) -> web.Response:
|
|
"""Handle checkpoint download request"""
|
|
async with self._download_lock:
|
|
# Get the download manager from service registry if not already initialized
|
|
if self.download_manager is None:
|
|
self.download_manager = await ServiceRegistry.get_download_manager()
|
|
|
|
try:
|
|
data = await request.json()
|
|
|
|
# Create progress callback that uses checkpoint-specific WebSocket
|
|
async def progress_callback(progress):
|
|
await ws_manager.broadcast_checkpoint_progress({
|
|
'status': 'progress',
|
|
'progress': progress
|
|
})
|
|
|
|
# Check which identifier is provided
|
|
download_url = data.get('download_url')
|
|
model_hash = data.get('model_hash')
|
|
model_version_id = data.get('model_version_id')
|
|
|
|
# Validate that at least one identifier is provided
|
|
if not any([download_url, model_hash, model_version_id]):
|
|
return web.Response(
|
|
status=400,
|
|
text="Missing required parameter: Please provide either 'download_url', 'hash', or 'modelVersionId'"
|
|
)
|
|
|
|
result = await self.download_manager.download_from_civitai(
|
|
download_url=download_url,
|
|
model_hash=model_hash,
|
|
model_version_id=model_version_id,
|
|
save_dir=data.get('checkpoint_root'),
|
|
relative_path=data.get('relative_path', ''),
|
|
progress_callback=progress_callback,
|
|
model_type="checkpoint"
|
|
)
|
|
|
|
if not result.get('success', False):
|
|
error_message = result.get('error', 'Unknown error')
|
|
|
|
# Return 401 for early access errors
|
|
if 'early access' in error_message.lower():
|
|
logger.warning(f"Early access download failed: {error_message}")
|
|
return web.Response(
|
|
status=401,
|
|
text=f"Early Access Restriction: {error_message}"
|
|
)
|
|
|
|
return web.Response(status=500, text=error_message)
|
|
|
|
return web.json_response(result)
|
|
|
|
except Exception as e:
|
|
error_message = str(e)
|
|
|
|
# Check if this might be an early access error
|
|
if '401' in error_message:
|
|
logger.warning(f"Early access error (401): {error_message}")
|
|
return web.Response(
|
|
status=401,
|
|
text="Early Access Restriction: This model requires purchase. Please ensure you have purchased early access and are logged in to Civitai."
|
|
)
|
|
|
|
logger.error(f"Error downloading checkpoint: {error_message}")
|
|
return web.Response(status=500, text=error_message)
|
|
|
|
async def get_checkpoint_roots(self, request):
|
|
"""Return the checkpoint root directories"""
|
|
try:
|
|
if self.scanner is None:
|
|
self.scanner = await ServiceRegistry.get_checkpoint_scanner()
|
|
|
|
roots = self.scanner.get_model_roots()
|
|
return web.json_response({
|
|
"success": True,
|
|
"roots": roots
|
|
})
|
|
except Exception as e:
|
|
logger.error(f"Error getting checkpoint roots: {e}", exc_info=True)
|
|
return web.json_response({
|
|
"success": False,
|
|
"error": str(e)
|
|
}, status=500)
|
|
|
|
async def save_metadata(self, request: web.Request) -> web.Response:
|
|
"""Handle saving metadata updates for checkpoints"""
|
|
try:
|
|
if self.scanner is None:
|
|
self.scanner = await ServiceRegistry.get_checkpoint_scanner()
|
|
|
|
data = await request.json()
|
|
file_path = data.get('file_path')
|
|
if not file_path:
|
|
return web.Response(text='File path is required', status=400)
|
|
|
|
# Remove file path from data to avoid saving it
|
|
metadata_updates = {k: v for k, v in data.items() if k != 'file_path'}
|
|
|
|
# Get metadata file path
|
|
metadata_path = os.path.splitext(file_path)[0] + '.metadata.json'
|
|
|
|
# Load existing metadata
|
|
metadata = await ModelRouteUtils.load_local_metadata(metadata_path)
|
|
|
|
# Update metadata
|
|
metadata.update(metadata_updates)
|
|
|
|
# Save updated metadata
|
|
with open(metadata_path, 'w', encoding='utf-8') as f:
|
|
json.dump(metadata, f, indent=2, ensure_ascii=False)
|
|
|
|
# Update cache
|
|
await self.scanner.update_single_model_cache(file_path, file_path, metadata)
|
|
|
|
# If model_name was updated, resort the cache
|
|
if 'model_name' in metadata_updates:
|
|
cache = await self.scanner.get_cached_data()
|
|
await cache.resort(name_only=True)
|
|
|
|
return web.json_response({'success': True})
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error saving checkpoint metadata: {e}", exc_info=True)
|
|
return web.Response(text=str(e), status=500)
|
|
|
|
async def get_civitai_versions(self, request: web.Request) -> web.Response:
|
|
"""Get available versions for a Civitai checkpoint model with local availability info"""
|
|
try:
|
|
if self.scanner is None:
|
|
self.scanner = await ServiceRegistry.get_checkpoint_scanner()
|
|
|
|
# Get the civitai client from service registry
|
|
civitai_client = await ServiceRegistry.get_civitai_client()
|
|
|
|
model_id = request.match_info['model_id']
|
|
response = await civitai_client.get_model_versions(model_id)
|
|
if not response or not response.get('modelVersions'):
|
|
return web.Response(status=404, text="Model not found")
|
|
|
|
versions = response.get('modelVersions', [])
|
|
model_type = response.get('type', '')
|
|
|
|
# Check model type - should be Checkpoint
|
|
if model_type.lower() != 'checkpoint':
|
|
return web.json_response({
|
|
'error': f"Model type mismatch. Expected Checkpoint, got {model_type}"
|
|
}, status=400)
|
|
|
|
# Check local availability for each version
|
|
for version in versions:
|
|
# Find the primary model file (type="Model" and primary=true) in the files list
|
|
model_file = next((file for file in version.get('files', [])
|
|
if file.get('type') == 'Model' and file.get('primary') == True), None)
|
|
|
|
# If no primary file found, try to find any model file
|
|
if not model_file:
|
|
model_file = next((file for file in version.get('files', [])
|
|
if file.get('type') == 'Model'), None)
|
|
|
|
if model_file:
|
|
sha256 = model_file.get('hashes', {}).get('SHA256')
|
|
if sha256:
|
|
# Set existsLocally and localPath at the version level
|
|
version['existsLocally'] = self.scanner.has_hash(sha256)
|
|
if version['existsLocally']:
|
|
version['localPath'] = self.scanner.get_path_by_hash(sha256)
|
|
|
|
# Also set the model file size at the version level for easier access
|
|
version['modelSizeKB'] = model_file.get('sizeKB')
|
|
else:
|
|
# No model file found in this version
|
|
version['existsLocally'] = False
|
|
|
|
return web.json_response(versions)
|
|
except Exception as e:
|
|
logger.error(f"Error fetching checkpoint model versions: {e}")
|
|
return web.Response(status=500, text=str(e))
|