Update main.py

main.py

@@ -1,15 +1,17 @@
 # main.py
-import io
 import json
 import torch
 import torchaudio
 import requests
 import numpy as np
+import tempfile
+import os
 
-from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi import FastAPI, UploadFile, File
 from transformers import AutoProcessor, AutoModelForCTC
 from pydantic import BaseModel
 
+
 with open("config.json", "r") as f:
     CONFIG = json.load(f)
 
@@ -26,7 +28,7 @@ OVERLAP_SECONDS = 2
 MAX_SAMPLES = SAMPLE_RATE * MAX_AUDIO_SECONDS
 
 
-app = FastAPI(title="Universal Audio STT", version="1.3.0")
+app = FastAPI(title="Universal Audio STT", version="1.4.0")
 
 
 processor = AutoProcessor.from_pretrained(MODEL_NAME)
@@ -37,18 +39,28 @@ model.eval()
 class STTResponse(BaseModel):
     transcript: str
     downstream_response: dict | None = None
+    error: str | None = None
 
 
+def load_audio_safe(file_bytes: bytes) -> tuple[np.ndarray | None, str | None]:
+    if not file_bytes:
+        return None, "Empty audio file"
 
-def load_audio(file_bytes: bytes) -> np.ndarray:
-    buffer = io.BytesIO(file_bytes)
+    with tempfile.NamedTemporaryFile(delete=False) as tmp:
+        tmp.write(file_bytes)
+        tmp_path = tmp.name
 
     try:
-        waveform, sr = torchaudio.load(buffer)
+        waveform, sr = torchaudio.load(tmp_path)
     except Exception:
-        raise ValueError("Unsupported or corrupted audio file")
+        return None, "Unsupported or corrupted audio format"
+    finally:
+        os.unlink(tmp_path)
+
+    if waveform.numel() == 0:
+        return None, "Audio contains no samples"
 
-    waveform = waveform.mean(dim=0)  # why: force mono
+    waveform = waveform.mean(dim=0)
 
     if sr != SAMPLE_RATE:
         waveform = torchaudio.functional.resample(
@@ -56,9 +68,9 @@ def load_audio(file_bytes: bytes) -> np.ndarray:
         )
 
     if waveform.numel() > MAX_SAMPLES:
-        raise ValueError("Audio exceeds 5 minute limit")
+        return None, "Audio exceeds 5 minute limit"
 
-    return waveform.numpy()
+    return waveform.numpy(), None
 
 
 def chunk_audio(audio: np.ndarray):
@@ -67,7 +79,7 @@ def chunk_audio(audio: np.ndarray):
     step = chunk_size - overlap
 
     for start in range(0, len(audio), step):
-        chunk = audio[start : start + chunk_size]
+        chunk = audio[start:start + chunk_size]
         if len(chunk) < SAMPLE_RATE:
             break
         yield chunk
@@ -82,9 +94,7 @@ def transcribe_chunk(chunk: np.ndarray) -> str:
     )
 
     with torch.no_grad():
-        logits = model(
-            inputs.input_values.to(DEVICE)
-        ).logits
+        logits = model(inputs.input_values.to(DEVICE)).logits
 
     predicted_ids = torch.argmax(logits, dim=-1)
     return processor.batch_decode(predicted_ids)[0].strip()
@@ -112,12 +122,18 @@ def forward_to_text_model(text: str):
 
 @app.post("/stt", response_model=STTResponse)
 async def stt(audio: UploadFile = File(...)):
-    try:
-        audio_bytes = await audio.read()
-        audio_data = load_audio(audio_bytes)
-        transcript = transcribe_long(audio_data)
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=str(e))
+    audio_bytes = await audio.read()
+
+    audio_data, error = load_audio_safe(audio_bytes)
+
+    if error:
+        return STTResponse(
+            transcript="",
+            downstream_response=None,
+            error=error,
+        )
+
+    transcript = transcribe_long(audio_data)
 
     downstream = None
     try:
@@ -128,5 +144,7 @@ async def stt(audio: UploadFile = File(...)):
     return STTResponse(
         transcript=transcript,
         downstream_response=downstream,
+        error=None,
     )
 
+