Update app.py

app.py

@@ -116,7 +116,7 @@ def load_models_task():
             logger.info("TTS model loaded successfully")
             model_status["tts"] = "loaded"
         except Exception as e:
-            logger.error(f"Failed to load TTS model: {str(e)}")
+            logger.error(f"Failed to load TTS model for Tagalog: {str(e)}")
             # Fallback to English TTS if the target language fails
             try:
                 logger.info("Falling back to MMS-TTS English model...")
@@ -304,8 +304,12 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
         temp_file.write(await audio.read())
         temp_path = temp_file.name
     
+    transcription = "Transcription not available"
+    translated_text = "Translation not available"
+    output_audio = None
+    
     try:
-        # Read and preprocess the audio
+        # Step 1: Transcribe the audio (STT)
         logger.info(f"Reading audio file: {temp_path}")
         waveform, sample_rate = sf.read(temp_path)
         logger.info(f"Audio loaded: sample_rate={sample_rate}, waveform_shape={waveform.shape}")
@@ -313,7 +317,6 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
             logger.info(f"Resampling audio from {sample_rate} Hz to 16000 Hz")
             waveform = librosa.resample(waveform, orig_sr=sample_rate, target_sr=16000)
         
-        # Process the audio with Whisper (STT)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         logger.info(f"Using device: {device}")
         inputs = stt_processor(waveform, sampling_rate=16000, return_tensors="pt").to(device)
@@ -323,10 +326,9 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
             transcription = stt_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         logger.info(f"Transcription completed: {transcription}")
 
-        # Translate the transcribed text
+        # Step 2: Translate the transcribed text (MT)
         source_code = LANGUAGE_MAPPING[source_lang]
         target_code = LANGUAGE_MAPPING[target_lang]
-        translated_text = "Translation not available"
         
         if model_status["mt"] == "loaded" and mt_model is not None and mt_tokenizer is not None:
             try:
@@ -348,11 +350,43 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
         else:
             logger.warning("MT model not loaded, skipping translation")
 
-        # Convert translated text to speech
-        output_audio = None
+        # Step 3: Convert translated text to speech (TTS)
         if model_status["tts"].startswith("loaded") and tts_model is not None and tts_tokenizer is not None:
             try:
                 inputs = tts_tokenizer(translated_text, return_tensors="pt", clean_up_tokenization_spaces=True).to(device)
                 with torch.no_grad():
                     output = tts_model(**inputs)
-                speech = output.waveform.cpu().numpy().squeeze
+                speech = output.waveform.cpu().numpy().squeeze()
+                speech = (speech * 32767).astype(np.int16)
+                output_audio = (tts_model.config.sampling_rate, speech.tolist())
+                logger.info("TTS conversion completed")
+            except Exception as e:
+                logger.error(f"Error during TTS conversion: {str(e)}")
+                output_audio = None
+        
+        return {
+            "request_id": request_id,
+            "status": "completed",
+            "message": "Transcription, translation, and TTS completed (or partially completed).",
+            "source_text": transcription,
+            "translated_text": translated_text,
+            "output_audio": output_audio
+        }
+    except Exception as e:
+        logger.error(f"Error during processing: {str(e)}")
+        return {
+            "request_id": request_id,
+            "status": "failed",
+            "message": f"Processing failed: {str(e)}",
+            "source_text": transcription,
+            "translated_text": translated_text,
+            "output_audio": output_audio
+        }
+    finally:
+        logger.info(f"Cleaning up temporary file: {temp_path}")
+        os.unlink(temp_path)
+
+if __name__ == "__main__":
+    import uvicorn
+    logger.info("Starting Uvicorn server...")
+    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)