Update main.py

main.py

@@ -35,8 +35,6 @@ class ErrorResponse(BaseModel):
 # Initialize model from environment variable
 def initialize_model():
     global MODEL, PROCESSOR
-    
-    # Get HF token from environment variable
     hf_token = os.environ.get("HF_TOKEN", None)
     model_name = os.environ.get("MODEL_NAME", "jonatasgrosman/wav2vec2-large-xlsr-53-arabic")
     
@@ -44,7 +42,7 @@ def initialize_model():
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         print(f"Loading model on device: {device}")
         
-        # Load model and processor using the updated parameter `token`
+        # Load model and processor using updated parameter `token`
         if hf_token:
             PROCESSOR = Wav2Vec2Processor.from_pretrained(model_name, token=hf_token)
             MODEL = Wav2Vec2ForCTC.from_pretrained(model_name, token=hf_token)
@@ -59,6 +57,34 @@ def initialize_model():
         print(f"Error loading model: {e}")
         raise e
 
+# Lifespan event handler to initialize the model at startup
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    initialize_model()
+    yield
+
+# Create the FastAPI app with the lifespan handler and CORS middleware
+app = FastAPI(
+    title="Quran Recitation Comparison API",
+    description="API for comparing similarity between Quran recitations using Wav2Vec2 embeddings",
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],  # Allows all headers
+)
+
+# Root endpoint
+@app.get("/")
+async def root():
+    """Welcome endpoint."""
+    return {"message": "Welcome to the Quran Recitation Comparison API"}
+
 # Load audio file
 def load_audio(file_path, target_sr=16000, trim_silence=True, normalize=True):
     """Load and preprocess an audio file."""
@@ -118,14 +144,12 @@ def custom_dtw(X, Y, metric='euclidean'):
     wp : list
         The warping path
     """
-    # Initialize cost matrix
     n, m = len(X[0]), len(Y[0])
     D = np.zeros((n+1, m+1))
     D[0, :] = np.inf
     D[:, 0] = np.inf
     D[0, 0] = 0
     
-    # Fill cost matrix
     for i in range(1, n+1):
         for j in range(1, m+1):
             if metric == 'euclidean':
@@ -137,7 +161,6 @@ def custom_dtw(X, Y, metric='euclidean'):
                 
             D[i, j] = cost + min(D[i-1, j], D[i, j-1], D[i-1, j-1])
     
-    # Backtrack to find warping path
     i, j = n, m
     wp = [(i, j)]
     while i > 1 or j > 1:
@@ -153,7 +176,6 @@ def custom_dtw(X, Y, metric='euclidean'):
 def compute_dtw_distance(features1, features2):
     """Compute the DTW distance between two sequences of features."""
     try:
-        # Use custom DTW implementation instead of librosa's
         D, wp = custom_dtw(features1, features2, metric='euclidean')
         distance = D[-1, -1]
         normalized_distance = distance / len(wp)
@@ -195,7 +217,7 @@ def cleanup_temp_files(file_paths):
             except Exception as e:
                 print(f"Error removing temporary file {file_path}: {e}")
 
-# API endpoints
+# API endpoint for comparing recitations
 @app.post("/compare", response_model=SimilarityResponse)
 async def compare_recitations(
     background_tasks: BackgroundTasks,
@@ -212,37 +234,29 @@ async def compare_recitations(
     - **similarity_score**: Score between 0-100 indicating similarity
     - **interpretation**: Text interpretation of the similarity
     """
-    # Check if model is initialized
     if MODEL is None or PROCESSOR is None:
         raise HTTPException(status_code=500, detail="Model not initialized")
     
-    # Temporary file paths
     temp_file1 = os.path.join(UPLOAD_DIR, f"{uuid.uuid4()}.wav")
     temp_file2 = os.path.join(UPLOAD_DIR, f"{uuid.uuid4()}.wav")
     
     try:
-        # Save uploaded files
         with open(temp_file1, "wb") as f:
             shutil.copyfileobj(file1.file, f)
         
         with open(temp_file2, "wb") as f:
             shutil.copyfileobj(file2.file, f)
         
-        # Load audio files
         audio1 = load_audio(temp_file1)
         audio2 = load_audio(temp_file2)
         
-        # Extract embeddings
         embedding1 = get_deep_embedding(audio1)
         embedding2 = get_deep_embedding(audio2)
         
-        # Compute DTW distance
         norm_distance = compute_dtw_distance(embedding1.T, embedding2.T)
         
-        # Interpret results
         interpretation, similarity_score = interpret_similarity(norm_distance)
         
-        # Add cleanup task
         background_tasks.add_task(cleanup_temp_files, [temp_file1, temp_file2])
         
         return {
@@ -251,10 +265,10 @@ async def compare_recitations(
         }
     
     except Exception as e:
-        # Ensure files are cleaned up even in case of error
         background_tasks.add_task(cleanup_temp_files, [temp_file1, temp_file2])
         raise HTTPException(status_code=500, detail=str(e))
 
+# Health check endpoint
 @app.get("/health")
 async def health_check():
     """Health check endpoint."""
@@ -265,24 +279,8 @@ async def health_check():
         )
     return {"status": "ok", "model_loaded": True}
 
-# Use lifespan context manager instead of on_event decorators
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    initialize_model()
-    yield
-
-# Initialize FastAPI app with lifespan handler
-app = FastAPI(
-    title="Quran Recitation Comparison API",
-    description="API for comparing similarity between Quran recitations using Wav2Vec2 embeddings",
-    version="1.0.0",
-    lifespan=lifespan
-)
-
-# Note: Ensure that all route definitions are declared AFTER the app initialization above.
-
 # Run the FastAPI app
 if __name__ == "__main__":
     import uvicorn
-    port = int(os.environ.get("PORT", 7860))  # Default to port 7860 for Hugging Face Spaces
+    port = int(os.environ.get("PORT", 7860))  # Default to port 7860
     uvicorn.run("main:app", host="0.0.0.0", port=port, reload=False)