Update app.py

app.py

@@ -1,6 +1,8 @@
 import os
+import json
+import uvicorn
 import torch
-from fastapi import FastAPI, HTTPException, UploadFile, File, Depends, BackgroundTasks, Request
+from fastapi import FastAPI, HTTPException, UploadFile, File, Depends, BackgroundTasks, Request, status
 from fastapi.responses import StreamingResponse, JSONResponse, FileResponse, HTMLResponse
 from pydantic import BaseModel, validator, Field, root_validator, EmailStr, constr
 from transformers import (
@@ -31,7 +33,7 @@ from PIL import Image
 from typing import Optional, List, Union, Dict, Any
 import uuid
 import logging
-import sqlite3
+from fastapi.exceptions import RequestValidationError
 from passlib.context import CryptContext
 from jose import JWTError, jwt
 from datetime import datetime, timedelta
@@ -41,7 +43,6 @@ from fastapi.middleware.gzip import GZipMiddleware
 from fastapi.security import APIKeyHeader, OAuth2PasswordBearer, OAuth2PasswordRequestForm
 from starlette.middleware.cors import CORSMiddleware
 
-
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(filename)s - %(lineno)d - %(message)s')
 logger = logging.getLogger(__name__)
 
@@ -53,7 +54,7 @@ ACCESS_TOKEN_EXPIRE_MINUTES = 30
 
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
-conn = sqlite3.connect('users.db')
+conn = sqlite3.connect('users.db', check_same_thread=False)
 cursor = conn.cursor()
 cursor.execute('''
     CREATE TABLE IF NOT EXISTS users (
@@ -80,7 +81,13 @@ TEMPLATES = Jinja2Templates(directory="templates")
 app = FastAPI()
 app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
 app.add_middleware(GZipMiddleware)
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 
 class User(BaseModel):
@@ -112,14 +119,34 @@ class GenerateRequest(BaseModel):
     mask_image: Optional[UploadFile] = None
     low_res_image: Optional[UploadFile] = None
 
-    @field_validator('task_type')
+    @validator('task_type')
     def validate_task_type(cls, value):
-        allowed_types = ["text", "image", "audio", "video", "classification", "translation", "question-answering", "speech-to-text", "text-to-speech", "image-segmentation", "feature-extraction", "token-classification", "fill-mask", "image-inpainting", "image-super-resolution", "object-detection", "image-captioning", "audio-transcription", "summarization"]
+        allowed_types = [
+            "text",
+            "image",
+            "audio",
+            "video",
+            "classification",
+            "translation",
+            "question-answering",
+            "speech-to-text",
+            "text-to-speech",
+            "image-segmentation",
+            "feature-extraction",
+            "token-classification",
+            "fill-mask",
+            "image-inpainting",
+            "image-super-resolution",
+            "object-detection",
+            "image-captioning",
+            "audio-transcription",
+            "summarization",
+        ]
         if value not in allowed_types:
             raise ValueError(f"Invalid task_type. Allowed types are: {allowed_types}")
         return value
 
-    @model_validator(mode='after')
+    @root_validator(pre=True)
     def check_input(cls, values):
         task_type = values.get("task_type")
         if task_type == "text" and values.get("input_text") is None:
@@ -287,7 +314,6 @@ async def verify_api_key(api_key: str = Depends(api_key_header)):
     if api_key != API_KEY:
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API Key")
 
-
 @app.post("/generate", dependencies=[Depends(verify_api_key)])
 async def generate(request: GenerateRequest, background_tasks: BackgroundTasks, model_data=Depends(get_model_data)):
     try:
@@ -373,7 +399,8 @@ async def generate(request: GenerateRequest, background_tasks: BackgroundTasks,
                 transcription = pipeline_func(contents, sampling_rate=16000)[0]["text"]
                 return JSONResponse({"transcription": transcription})
             except Exception as e:
-                raise HTTPException(status_code=500, detail=f"Error during speech-to-text: {str(e)}")
+                logger.exception(f"Error during speech-to-text: {e}")
+                raise HTTPException(status_code=500, detail=f"Error during speech-to-text: {str(e)}") from e
         elif request.task_type == "text-to-speech":
             if not request.input_text:
                 raise HTTPException(status_code=400, detail="Input text is required for text-to-speech.")
@@ -634,7 +661,11 @@ async def get_all_users_route():
 
 @app.exception_handler(RequestValidationError)
 async def validation_exception_handler(request: Request, exc: RequestValidationError):
-    return JSONResponse(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, content=json.dumps({"detail": exc.errors(), "body": exc.body}))
+    return JSONResponse(
+        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+        content=json.dumps({"detail": exc.errors(), "body": exc.body}),
+    )
+
 
 if __name__ == "__main__":
-    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)
+    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)