Fix TTS language issue by dynamically updating model in /translate-audio endpoint

app.py

@@ -39,6 +39,7 @@ model_status = {
     "tts": "not_loaded"
 }
 error_message = None
+current_tts_language = "tgl"  # Track the current TTS language
 
 # Model instances
 stt_processor = None
@@ -179,7 +180,7 @@ def load_models_task():
             error_message = f"MT model loading failed: {str(e)}"
             return
 
-        # Load TTS model (default to Tagalog, will be updated by /update-languages)
+        # Load TTS model (default to Tagalog, will be updated dynamically)
         logger.info("Starting to load TTS model...")
         from transformers import VitsModel, AutoTokenizer
         
@@ -201,6 +202,7 @@ def load_models_task():
                 tts_model.to(device)
                 logger.info("Fallback TTS model loaded successfully")
                 model_status["tts"] = "loaded (fallback)"
+                current_tts_language = "eng"
             except Exception as e2:
                 logger.error(f"Failed to load fallback TTS model: {str(e2)}")
                 model_status["tts"] = "failed"
@@ -259,7 +261,7 @@ async def health_check():
 
 @app.post("/update-languages")
 async def update_languages(source_lang: str = Form(...), target_lang: str = Form(...)):
-    global stt_processor, stt_model, tts_model, tts_tokenizer
+    global stt_processor, stt_model, tts_model, tts_tokenizer, current_tts_language
     
     if source_lang not in LANGUAGE_MAPPING or target_lang not in LANGUAGE_MAPPING:
         raise HTTPException(status_code=400, detail="Invalid language selected")
@@ -314,6 +316,7 @@ async def update_languages(source_lang: str = Form(...), target_lang: str = Form
         tts_model = VitsModel.from_pretrained(f"facebook/mms-tts-{target_code}")
         tts_tokenizer = AutoTokenizer.from_pretrained(f"facebook/mms-tts-{target_code}")
         tts_model.to(device)
+        current_tts_language = target_code
         logger.info(f"TTS model updated to {target_code}")
         model_status["tts"] = "loaded"
     except Exception as e:
@@ -323,6 +326,7 @@ async def update_languages(source_lang: str = Form(...), target_lang: str = Form
             tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
             tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
             tts_model.to(device)
+            current_tts_language = "eng"
             logger.info("Fallback TTS model loaded successfully")
             model_status["tts"] = "loaded (fallback)"
         except Exception as e2:
@@ -337,7 +341,7 @@ async def update_languages(source_lang: str = Form(...), target_lang: str = Form
 @app.post("/translate-text")
 async def translate_text(text: str = Form(...), source_lang: str = Form(...), target_lang: str = Form(...)):
     """Endpoint to translate text and convert to speech"""
-    global mt_model, mt_tokenizer, tts_model, tts_tokenizer
+    global mt_model, mt_tokenizer, tts_model, tts_tokenizer, current_tts_language
     
     if not text:
         raise HTTPException(status_code=400, detail="No text provided")
@@ -373,6 +377,31 @@ async def translate_text(text: str = Form(...), source_lang: str = Form(...), ta
     else:
         logger.warning("MT model not loaded, skipping translation")
 
+    # Update TTS model if the target language doesn't match the current TTS language
+    if current_tts_language != target_code:
+        try:
+            logger.info(f"Updating TTS model for {target_code}...")
+            from transformers import VitsModel, AutoTokenizer
+            tts_model = VitsModel.from_pretrained(f"facebook/mms-tts-{target_code}")
+            tts_tokenizer = AutoTokenizer.from_pretrained(f"facebook/mms-tts-{target_code}")
+            tts_model.to(device)
+            current_tts_language = target_code
+            logger.info(f"TTS model updated to {target_code}")
+            model_status["tts"] = "loaded"
+        except Exception as e:
+            logger.error(f"Failed to load TTS model for {target_code}: {str(e)}")
+            try:
+                logger.info("Falling back to MMS-TTS English model...")
+                tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
+                tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
+                tts_model.to(device)
+                current_tts_language = "eng"
+                logger.info("Fallback TTS model loaded successfully")
+                model_status["tts"] = "loaded (fallback)"
+            except Exception as e2:
+                logger.error(f"Failed to load fallback TTS model: {str(e2)}")
+                model_status["tts"] = "failed"
+
     # Convert translated text to speech
     output_audio_url = None
     if model_status["tts"].startswith("loaded") and tts_model is not None and tts_tokenizer is not None:
@@ -409,7 +438,7 @@ async def translate_text(text: str = Form(...), source_lang: str = Form(...), ta
 @app.post("/translate-audio")
 async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form(...), target_lang: str = Form(...)):
     """Endpoint to transcribe, translate, and convert audio to speech"""
-    global stt_processor, stt_model, mt_model, mt_tokenizer, tts_model, tts_tokenizer
+    global stt_processor, stt_model, mt_model, mt_tokenizer, tts_model, tts_tokenizer, current_tts_language
     
     if not audio:
         raise HTTPException(status_code=400, detail="No audio file provided")
@@ -506,7 +535,32 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
         else:
             logger.warning("MT model not loaded, skipping translation")
 
-        # Step 5: Convert translated text to speech (TTS)
+        # Step 5: Update TTS model if the target language doesn't match the current TTS language
+        if current_tts_language != target_code:
+            try:
+                logger.info(f"Updating TTS model for {target_code}...")
+                from transformers import VitsModel, AutoTokenizer
+                tts_model = VitsModel.from_pretrained(f"facebook/mms-tts-{target_code}")
+                tts_tokenizer = AutoTokenizer.from_pretrained(f"facebook/mms-tts-{target_code}")
+                tts_model.to(device)
+                current_tts_language = target_code
+                logger.info(f"TTS model updated to {target_code}")
+                model_status["tts"] = "loaded"
+            except Exception as e:
+                logger.error(f"Failed to load TTS model for {target_code}: {str(e)}")
+                try:
+                    logger.info("Falling back to MMS-TTS English model...")
+                    tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
+                    tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
+                    tts_model.to(device)
+                    current_tts_language = "eng"
+                    logger.info("Fallback TTS model loaded successfully")
+                    model_status["tts"] = "loaded (fallback)"
+                except Exception as e2:
+                    logger.error(f"Failed to load fallback TTS model: {str(e2)}")
+                    model_status["tts"] = "failed"
+
+        # Step 6: 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").to(device)