Update app.py

app.py

@@ -2,12 +2,11 @@ import logging
 import math
 import time
 import base64
-import io
 import os
 from typing import Dict, Any
 from functools import wraps
 
-from fastapi import FastAPI, Depends, HTTPException, File, UploadFile
+from fastapi import FastAPI, Depends, HTTPException, File, UploadFile, Form, Header
 from fastapi.encoders import jsonable_encoder
 from pydantic import BaseModel
 import jax.numpy as jnp
@@ -38,7 +37,7 @@ chunk_len = round(CHUNK_LENGTH_S * pipeline.feature_extractor.sampling_rate)
 stride_left = stride_right = round(stride_length_s * pipeline.feature_extractor.sampling_rate)
 step = chunk_len - stride_left - stride_right
 
-# do a pre-compile step so that the first user to use the demo isn't hit with a long transcription time
+# Pre-compile step
 logger.debug("Compiling forward call...")
 start = time.time()
 random_inputs = {
@@ -51,7 +50,7 @@ compile_time = time.time() - start
 logger.debug(f"Compiled in {compile_time}s")
 
 class TranscribeAudioRequest(BaseModel):
-    audio_base64: str = None
+    audio_base64: str
     task: str = "transcribe"
     return_timestamps: bool = False
 
@@ -69,40 +68,55 @@ def timeit(func):
         return result
     return wrapper
 
-def check_api_key():
+def check_api_key(x_api_key: str = Header(...)):
     api_key = os.environ.get("WHISPER_API_KEY")
-    if not api_key:
-        raise HTTPException(status_code=401, detail="API key not found in environment variables")
-    return api_key
+    if not api_key or x_api_key != api_key:
+        raise HTTPException(status_code=401, detail="Invalid or missing API key")
+    return x_api_key
 
-@app.post("/transcribe_audio")
+@app.post("/transcribe_audio_file")
 @timeit
-async def transcribe_chunked_audio(
-    request: TranscribeAudioRequest = None,
-    file: UploadFile = File(None),
+async def transcribe_audio_file(
+    file: UploadFile = File(...),
+    task: str = Form("transcribe"),
+    return_timestamps: bool = Form(False),
     api_key: str = Depends(check_api_key)
 ) -> Dict[str, Any]:
-    logger.debug("Starting transcribe_chunked_audio function")
-    logger.debug(f"Received parameters - task: {request.task if request else 'transcribe'}, return_timestamps: {request.return_timestamps if request else False}")
+    logger.debug("Starting transcribe_audio_file function")
+    logger.debug(f"Received parameters - task: {task}, return_timestamps: {return_timestamps}")
     
     try:
-        if file:
-            logger.debug("Processing uploaded file")
-            audio_data = await file.read()
-            file_size = len(audio_data)
-        elif request and request.audio_base64:
-            logger.debug("Processing base64 encoded audio")
-            audio_data = base64.b64decode(request.audio_base64)
-            file_size = len(audio_data)
-        else:
-            raise HTTPException(status_code=400, detail="No audio data provided")
-        
+        audio_data = await file.read()
+        file_size = len(audio_data)
+        file_size_mb = file_size / (1024 * 1024)
+        logger.debug(f"Audio file size: {file_size} bytes ({file_size_mb:.2f}MB)")
+    except Exception as e:
+        logger.error(f"Error reading audio file: {str(e)}", exc_info=True)
+        raise HTTPException(status_code=400, detail=f"Error reading audio file: {str(e)}")
+
+    return await process_audio(audio_data, file_size_mb, task, return_timestamps)
+
+@app.post("/transcribe_audio_base64")
+@timeit
+async def transcribe_audio_base64(
+    request: TranscribeAudioRequest,
+    api_key: str = Depends(check_api_key)
+) -> Dict[str, Any]:
+    logger.debug("Starting transcribe_audio_base64 function")
+    logger.debug(f"Received parameters - task: {request.task}, return_timestamps: {request.return_timestamps}")
+    
+    try:
+        audio_data = base64.b64decode(request.audio_base64)
+        file_size = len(audio_data)
         file_size_mb = file_size / (1024 * 1024)
-        logger.debug(f"Audio data size: {file_size} bytes ({file_size_mb:.2f}MB)")
+        logger.debug(f"Decoded audio data size: {file_size} bytes ({file_size_mb:.2f}MB)")
     except Exception as e:
-        logger.error(f"Error processing audio data: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=400, detail=f"Error processing audio data: {str(e)}")
+        logger.error(f"Error decoding base64 audio data: {str(e)}", exc_info=True)
+        raise HTTPException(status_code=400, detail=f"Error decoding base64 audio data: {str(e)}")
+
+    return await process_audio(audio_data, file_size_mb, request.task, request.return_timestamps)
 
+async def process_audio(audio_data: bytes, file_size_mb: float, task: str, return_timestamps: bool) -> Dict[str, Any]:
     if file_size_mb > FILE_LIMIT_MB:
         logger.warning(f"Max file size exceeded: {file_size_mb:.2f}MB > {FILE_LIMIT_MB}MB")
         raise HTTPException(status_code=400, detail=f"File size exceeds file size limit. Got file of size {file_size_mb:.2f}MB for a limit of {FILE_LIMIT_MB}MB.")
@@ -118,15 +132,13 @@ async def transcribe_chunked_audio(
 
     logger.debug("Calling tqdm_generate to transcribe audio")
     try:
-        task = request.task if request else "transcribe"
-        return_timestamps = request.return_timestamps if request else False
         text, runtime, timing_info = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
         logger.debug(f"Transcription completed. Runtime: {runtime:.2f}s")
     except Exception as e:
         logger.error(f"Error in tqdm_generate: {str(e)}", exc_info=True)
         raise HTTPException(status_code=500, detail=f"Error transcribing audio: {str(e)}")
 
-    logger.debug("Transcribe_chunked_audio function completed successfully")
+    logger.debug("Audio processing completed successfully")
     return jsonable_encoder({
         "text": text,
         "runtime": runtime,
@@ -211,4 +223,5 @@ def format_timestamp(seconds: float, always_include_hours: bool = False, decimal
         return f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}"
     else:
         # we have a malformed timestamp so just return it as is
-        return seconds
+        return seconds
+