Update main.py

main.py

@@ -1,8 +1,9 @@
 import os
 import torch
 import torchaudio
-from flask import Flask, request, jsonify
-from flask_cors import CORS
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
 from transformers import (
     WhisperProcessor,
     WhisperForConditionalGeneration,
@@ -30,8 +31,14 @@ logging.getLogger("transformers").setLevel(logging.ERROR)
 logging.getLogger("urllib3").setLevel(logging.ERROR)
 logging.getLogger("huggingface_hub").setLevel(logging.ERROR)
 
-app = Flask(__name__)
-CORS(app)
+app = FastAPI()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 # ========== Load Whisper Model (quantized) ==========
 def load_whisper_model(model_size="small", save_dir="/tmp/models_cache/whisper"):
@@ -131,32 +138,41 @@ def warm_up_models():
 warm_up_models()
 
 # ========== Flask Route ==========
-@app.route('/transcribe', methods=['POST'])
-def transcribe():
-    if 'audio' not in request.files:
-        return jsonify({"error": "No audio file provided."}), 400
+@app.post('/transcribe')
+async def transcribe(audio: UploadFile = File(...)):
+    if not audio:
+        raise HTTPException(status_code=400, detail="No audio file provided.")
 
-    audio_file = request.files['audio']
     os.makedirs("/tmp/temp_audio", exist_ok=True)
-    audio_path = f"/tmp/temp_audio/{audio_file.filename}"
-    audio_file.save(audio_path)
+    audio_path = f"/tmp/temp_audio/{audio.filename}"
 
+    # Save uploaded file
     try:
+        with open(audio_path, "wb") as f:
+            content = await audio.read()
+            f.write(content)
+
         transcription = transcribe_long_audio(audio_path, processor, whisper_model)
         corrected_text = correct_grammar(transcription, grammar_pipeline)
 
-        return jsonify({
+        return JSONResponse({
             "raw_transcription": transcription,
             "corrected_transcription": corrected_text
         })
 
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        raise HTTPException(status_code=500, detail=str(e))
 
     finally:
-        if os.path.exists(audio_path):
-            os.remove(audio_path)
+        try:
+            if os.path.exists(audio_path):
+                os.remove(audio_path)
+        except Exception:
+            pass
 
 # ========== Run App ==========
 if __name__ == '__main__':
-    app.run(host="0.0.0.0", debug=False, port=7860)
+    # Run with Uvicorn for FastAPI
+    import uvicorn
+
+    uvicorn.run("main:app", host="0.0.0.0", port=7860, log_level="info")