Initial Real-ESRGAN API project setup

2026-02-16 19:56:25 +01:00
commit 0e59652575
34 changed files with 3668 additions and 0 deletions
--- a/app/routers/models.py
+++ b/app/routers/models.py
@@ -0,0 +1,109 @@
+"""Model management endpoints."""
+import logging
+
+from fastapi import APIRouter, HTTPException
+
+from app.schemas.models import ModelDownloadRequest, ModelDownloadResponse, ModelListResponse
+from app.services import model_manager
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(prefix='/api/v1', tags=['models'])
+
+
+@router.get('/models')
+async def list_models() -> ModelListResponse:
+    """List all available models."""
+    try:
+        available = model_manager.get_available_models()
+        local_count = sum(1 for m in available if m['available'])
+        
+        return ModelListResponse(
+            available_models=available,
+            total_models=len(available),
+            local_models=local_count,
+        )
+    except Exception as e:
+        logger.error(f'Failed to list models: {e}', exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.post('/models/download')
+async def download_models(request: ModelDownloadRequest) -> ModelDownloadResponse:
+    """Download one or more models."""
+    if not request.models:
+        raise HTTPException(status_code=400, detail='No models specified')
+
+    try:
+        logger.info(f'Downloading models: {request.models}')
+        results = await model_manager.download_models(request.models)
+        
+        downloaded = []
+        failed = []
+        errors = {}
+        
+        for model_name, (success, message) in results.items():
+            if success:
+                downloaded.append(model_name)
+            else:
+                failed.append(model_name)
+                errors[model_name] = message
+        
+        return ModelDownloadResponse(
+            success=len(failed) == 0,
+            message=f'Downloaded {len(downloaded)} model(s)',
+            downloaded=downloaded,
+            failed=failed,
+            errors=errors,
+        )
+    except Exception as e:
+        logger.error(f'Model download failed: {e}', exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.get('/models/{model_name}')
+async def get_model_info(model_name: str):
+    """Get information about a specific model."""
+    models = model_manager.get_available_models()
+    
+    for model in models:
+        if model['name'] == model_name:
+            return model
+    
+    raise HTTPException(status_code=404, detail=f'Model not found: {model_name}')
+
+
+@router.post('/models/{model_name}/download')
+async def download_model(model_name: str):
+    """Download a specific model."""
+    try:
+        success, message = await model_manager.download_model(model_name)
+        
+        if not success:
+            raise HTTPException(status_code=500, detail=message)
+        
+        return {
+            'success': True,
+            'message': message,
+            'model': model_name,
+        }
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f'Failed to download model {model_name}: {e}', exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.get('/models-info')
+async def get_models_directory_info():
+    """Get information about the models directory."""
+    try:
+        info = model_manager.get_models_directory_info()
+        return {
+            'models_directory': info['path'],
+            'total_size_mb': round(info['size_mb'], 2),
+            'model_count': info['model_count'],
+        }
+    except Exception as e:
+        logger.error(f'Failed to get models directory info: {e}', exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))