src/api/models.py

from fastapi import APIRouter, HTTPException
from typing import List, Dict
from models.registry import GradioRegistry

router = APIRouter()
registry = GradioRegistry()

@router.get("/models", response_model=List[Dict])
async def list_models():
    """List all available models."""
    return registry.list_models()

@router.get("/models/{model_id}")
async def get_model_info(model_id: str):
    """Get information about a specific model."""
    model_info = registry.get_model_info(model_id)
    if not model_info:
        raise HTTPException(status_code=404, detail="Model not found")
    return model_info

@router.get("/models/{model_id}/status")
async def get_model_status(model_id: str):
    """Get the current status of a model."""
    model_info = registry.get_model_info(model_id)
    if not model_info:
        raise HTTPException(status_code=404, detail="Model not found")
    return {"status": model_info.status}

@router.post("/models/{model_id}/load")
async def load_model(model_id: str):
    """Load a model into memory."""
    model = registry.get_model(model_id)
    if not model:
        raise HTTPException(status_code=404, detail="Model not found")
    
    try:
        model.load_model()
        registry.update_model_status(model_id, "loaded")
        return {"status": "loaded"}
    except Exception as e:
        registry.update_model_status(model_id, "error")
        raise HTTPException(status_code=500, detail=str(e))