Update app.py

app.py

@@ -1,5 +1,5 @@
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import FileResponse, JSONResponse
+from fastapi import FastAPI, HTTPException, UploadFile, File
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
 import torch
 from transformers import (
@@ -7,11 +7,15 @@ from transformers import (
     AutoProcessor,
     BarkModel,
     pipeline,
-    AutoModelForSequenceClassification
+    AutoModelForSequenceClassification,
+    Wav2Vec2Processor,
+    Wav2Vec2ForCTC
 )
 import scipy.io.wavfile as wavfile
 import uuid
 import os
+from io import BytesIO
+import soundfile as sf
 from typing import Optional
 
 # FastAPI instance
@@ -21,6 +25,7 @@ app = FastAPI(title="Kinyarwanda Engine", version="1.0")
 MODEL_PATH = "/app/models/suno-bark"
 SENTIMENT_MODEL_PATH = "/app/models/sentiment"
 SAMPLE_RATE = 24000
+ASR_MODEL_PATH = "jonatasgrosman/wav2vec2-large-xlsr-53-Kinyarwanda"
 
 # Ensure working directory for audio
 AUDIO_DIR = "/tmp/audio"
@@ -44,6 +49,11 @@ try:
         max_length=512
     )
 
+    # STT
+    asr_processor = Wav2Vec2Processor.from_pretrained(ASR_MODEL_PATH)
+    asr_model = Wav2Vec2ForCTC.from_pretrained(ASR_MODEL_PATH)
+    asr_model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+
     # Device config
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
@@ -70,27 +80,44 @@ def root():
 # Text-to-Speech endpoint
 @app.post("/tts/")
 def text_to_speech(request: TTSRequest):
-    output_file = os.path.join(AUDIO_DIR, f"tts_{uuid.uuid4().hex}.wav")
-
     try:
         inputs = processor(request.text, return_tensors="pt").to(device)
+
         with torch.no_grad():
             audio_array = model.generate(**inputs)
 
-        wavfile.write(output_file, rate=SAMPLE_RATE, data=audio_array.cpu().numpy().squeeze())
+        audio_data = audio_array.cpu().numpy().squeeze()
+        buffer = BytesIO()
+        wavfile.write(buffer, rate=SAMPLE_RATE, data=audio_data)
+        buffer.seek(0)
 
-        return FileResponse(
-            output_file,
+        return StreamingResponse(
+            buffer,
             media_type="audio/wav",
-            filename=os.path.basename(output_file)
+            headers={"Content-Disposition": f"attachment; filename=tts_{uuid.uuid4().hex}.wav"}
         )
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"TTS generation failed: {str(e)}")
 
-    finally:
-        if os.path.exists(output_file):
-            os.remove(output_file)
+# Speech-to-Text endpoint
+@app.post("/stt/")
+def speech_to_text(audio_file: UploadFile = File(...)):
+    try:
+        audio_bytes = audio_file.file.read()
+        audio, sample_rate = sf.read(BytesIO(audio_bytes))
+
+        inputs = asr_processor(audio, sampling_rate=sample_rate, return_tensors="pt", padding=True).input_values.to(device)
+
+        with torch.no_grad():
+            logits = asr_model(inputs).logits
+            predicted_ids = torch.argmax(logits, dim=-1)
+
+        transcription = asr_processor.batch_decode(predicted_ids)[0]
+        return {"transcription": transcription}
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"STT failed: {str(e)}")
 
 # Sentiment Analysis endpoint
 @app.post("/sentiment/")