Update app.py

app.py

@@ -6,16 +6,15 @@ from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
 import torchaudio
 import torch
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
 
 # Use writable cache paths
 os.environ['HF_HOME'] = '/app/cache'
 os.environ['TORCH_HOME'] = '/app/cache'
 
-# FastAPI app setup
+# FastAPI setup
 app = FastAPI()
 
-# CORS: allow frontend from localhost
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["http://localhost:8080"],
@@ -24,11 +23,17 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load phoneme model + processor
+# Load models
 try:
-    processor = Wav2Vec2Processor.from_pretrained("vitouphy/wav2vec2-xls-r-300m-timit-phoneme")
-    model = Wav2Vec2ForCTC.from_pretrained("vitouphy/wav2vec2-xls-r-300m-timit-phoneme")
-    print("✅ Model loaded successfully.")
+    # Phoneme model
+    phoneme_processor = Wav2Vec2Processor.from_pretrained("vitouphy/wav2vec2-xls-r-300m-timit-phoneme")
+    phoneme_model = Wav2Vec2ForCTC.from_pretrained("vitouphy/wav2vec2-xls-r-300m-timit-phoneme")
+
+    # Speech-to-text model
+    stt_processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
+    stt_model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
+
+    print("✅ Models loaded successfully.")
 except Exception as e:
     print("❌ Model load error:", str(e))
     raise
@@ -57,18 +62,33 @@ async def transcribe(audio: UploadFile = File(...)):
             waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(waveform)
             sample_rate = 16000
 
-        input_values = processor(waveform.squeeze(), sampling_rate=sample_rate, return_tensors="pt").input_values
+        # Run phoneme model
+        phoneme_inputs = phoneme_processor(waveform.squeeze(), sampling_rate=sample_rate, return_tensors="pt").input_values
         with torch.no_grad():
-            logits = model(input_values).logits
-        predicted_ids = torch.argmax(logits, dim=-1)
-        phonemes = processor.decode(predicted_ids[0])
+            phoneme_logits = phoneme_model(phoneme_inputs).logits
+        phoneme_ids = torch.argmax(phoneme_logits, dim=-1)
+        phonemes = phoneme_processor.decode(phoneme_ids[0])
 
-        return {"phonemes": phonemes}
+        # Run speech-to-text model
+        stt_inputs = stt_processor(waveform.squeeze(), sampling_rate=sample_rate, return_tensors="pt").input_values
+        with torch.no_grad():
+            stt_logits = stt_model(stt_inputs).logits
+        stt_ids = torch.argmax(stt_logits, dim=-1)
+        transcript = stt_processor.decode(stt_ids[0])
+
+        return {
+            "phonemes": phonemes,
+            "transcript": transcript
+        }
 
     except Exception as e:
         print("❌ Transcription error:", str(e))
-        return {"phonemes": "[Error: " + str(e) + "]"}
+        return {
+            "phonemes": "[Error]",
+            "transcript": "[Error: " + str(e) + "]"
+        }
 
 @app.get("/")
 def root():
     return {"message": "Backend is running"}
+