Add MT with NLLB-200-distilled-600M and TTS with MMS-TTS

app.py

@@ -12,7 +12,7 @@ import soundfile as sf
 import librosa
 from fastapi import FastAPI, HTTPException, UploadFile, File, Form
 from fastapi.responses import JSONResponse
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, Tuple
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -31,11 +31,15 @@ model_status = {
 }
 error_message = None
 
-# STT model and processor (will be loaded in background)
+# Model instances
 stt_processor = None
 stt_model = None
+mt_model = None
+mt_tokenizer = None
+tts_model = None
+tts_tokenizer = None
 
-# Define the valid languages
+# Define the valid languages and mappings
 LANGUAGE_MAPPING = {
     "English": "eng",
     "Tagalog": "tgl",
@@ -45,9 +49,19 @@ LANGUAGE_MAPPING = {
     "Pangasinan": "pag"
 }
 
+NLLB_LANGUAGE_CODES = {
+    "eng": "eng_Latn",
+    "tgl": "tgl_Latn",
+    "ceb": "ceb_Latn",
+    "ilo": "ilo_Latn",
+    "war": "war_Latn",
+    "pag": "pag_Latn"
+}
+
 # Function to load models in background
 def load_models_task():
-    global models_loaded, loading_in_progress, model_status, error_message, stt_processor, stt_model
+    global models_loaded, loading_in_progress, model_status, error_message
+    global stt_processor, stt_model, mt_model, mt_tokenizer, tts_model, tts_tokenizer
     
     try:
         loading_in_progress = True
@@ -71,10 +85,51 @@ def load_models_task():
             error_message = f"STT model loading failed: {str(e)}"
             return
 
-        # Skip MT and TTS models for now to save memory
-        model_status["mt"] = "skipped"
-        model_status["tts"] = "skipped"
-        logger.info("MT and TTS models skipped to save memory")
+        # Load MT model
+        logger.info("Starting to load MT model...")
+        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+        
+        try:
+            logger.info("Loading NLLB-200-distilled-600M model...")
+            model_status["mt"] = "loading"
+            mt_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
+            mt_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M", clean_up_tokenization_spaces=True)
+            mt_model.to(device)
+            logger.info("MT model loaded successfully")
+            model_status["mt"] = "loaded"
+        except Exception as e:
+            logger.error(f"Failed to load MT model: {str(e)}")
+            model_status["mt"] = "failed"
+            error_message = f"MT model loading failed: {str(e)}"
+            return
+
+        # Load TTS model (default to Tagalog, will be updated by /update-languages)
+        logger.info("Starting to load TTS model...")
+        from transformers import VitsModel, AutoTokenizer
+        
+        try:
+            logger.info("Loading MMS-TTS model for Tagalog...")
+            model_status["tts"] = "loading"
+            tts_model = VitsModel.from_pretrained("facebook/mms-tts-tgl")
+            tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-tgl", clean_up_tokenization_spaces=True)
+            tts_model.to(device)
+            logger.info("TTS model loaded successfully")
+            model_status["tts"] = "loaded"
+        except Exception as e:
+            logger.error(f"Failed to load TTS model: {str(e)}")
+            # Fallback to English TTS if the target language fails
+            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", clean_up_tokenization_spaces=True)
+                tts_model.to(device)
+                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"
+                error_message = f"TTS model loading failed: {str(e)} (fallback also failed: {str(e2)})"
+                return
         
         models_loaded = True
         logger.info("Model loading completed successfully")
@@ -121,14 +176,46 @@ async def health_check():
 
 @app.post("/update-languages")
 async def update_languages(source_lang: str = Form(...), target_lang: str = Form(...)):
+    global tts_model, tts_tokenizer
+    
     if source_lang not in LANGUAGE_MAPPING or target_lang not in LANGUAGE_MAPPING:
         raise HTTPException(status_code=400, detail="Invalid language selected")
+    
+    source_code = LANGUAGE_MAPPING[source_lang]
+    target_code = LANGUAGE_MAPPING[target_lang]
+    
+    # Update the TTS model based on the target language
+    try:
+        logger.info(f"Loading MMS-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}", clean_up_tokenization_spaces=True)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        tts_model.to(device)
+        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", clean_up_tokenization_spaces=True)
+            tts_model.to(device)
+            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"
+            error_message = f"TTS model loading failed: {str(e)} (fallback also failed: {str(e2)})"
+    
     logger.info(f"Updating languages: {source_lang} → {target_lang}")
     return {"status": f"Languages updated to {source_lang} → {target_lang}"}
 
 @app.post("/translate-text")
 async def translate_text(text: str = Form(...), source_lang: str = Form(...), target_lang: str = Form(...)):
-    """Endpoint that creates a placeholder for text translation"""
+    """Endpoint to translate text and convert to speech"""
+    global mt_model, mt_tokenizer, tts_model, tts_tokenizer
+    
     if not text:
         raise HTTPException(status_code=400, detail="No text provided")
     if source_lang not in LANGUAGE_MAPPING or target_lang not in LANGUAGE_MAPPING:
@@ -136,19 +223,61 @@ async def translate_text(text: str = Form(...), source_lang: str = Form(...), ta
     
     logger.info(f"Translate-text requested: {text} from {source_lang} to {target_lang}")
     request_id = str(uuid.uuid4())
+    
+    # Translate the text
+    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:
+            source_nllb_code = NLLB_LANGUAGE_CODES[source_code]
+            target_nllb_code = NLLB_LANGUAGE_CODES[target_code]
+            mt_tokenizer.src_lang = source_nllb_code
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            inputs = mt_tokenizer(text, return_tensors="pt", clean_up_tokenization_spaces=True).to(device)
+            with torch.no_grad():
+                generated_tokens = mt_model.generate(
+                    **inputs,
+                    forced_bos_token_id=mt_tokenizer.convert_tokens_to_ids(target_nllb_code),
+                    max_length=448
+                )
+            translated_text = mt_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+            logger.info(f"Translation completed: {translated_text}")
+        except Exception as e:
+            logger.error(f"Error during translation: {str(e)}")
+            translated_text = f"Translation failed: {str(e)}"
+    else:
+        logger.warning("MT model not loaded, skipping translation")
+
+    # Convert translated text to speech
+    output_audio = None
+    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 = (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": "processing",
-        "message": "Translation not implemented yet (MT model not loaded).",
+        "status": "completed",
+        "message": "Translation and TTS completed (or partially completed).",
         "source_text": text,
-        "translated_text": "Translation not available",
-        "output_audio": None
+        "translated_text": translated_text,
+        "output_audio": output_audio
     }
 
 @app.post("/translate-audio")
 async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form(...), target_lang: str = Form(...)):
-    """Endpoint to transcribe audio using STT"""
-    global stt_processor, stt_model
+    """Endpoint to transcribe, translate, and convert audio to speech"""
+    global stt_processor, stt_model, mt_model, mt_tokenizer, tts_model, tts_tokenizer
     
     if not audio:
         raise HTTPException(status_code=400, detail="No audio file provided")
@@ -184,7 +313,7 @@ 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
+        # 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)
@@ -192,31 +321,38 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
         with torch.no_grad():
             generated_ids = stt_model.generate(**inputs)
             transcription = stt_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        
         logger.info(f"Transcription completed: {transcription}")
-        return {
-            "request_id": request_id,
-            "status": "completed",
-            "message": "Transcription completed successfully. Translation and TTS not implemented yet.",
-            "source_text": transcription,
-            "translated_text": "Translation not available",
-            "output_audio": None
-        }
-    except Exception as e:
-        logger.error(f"Error during transcription: {str(e)}")
-        return {
-            "request_id": request_id,
-            "status": "failed",
-            "message": f"Transcription failed: {str(e)}",
-            "source_text": "Transcription not available",
-            "translated_text": "Translation not available",
-            "output_audio": None
-        }
-    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)
+        # Translate the transcribed text
+        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:
+                source_nllb_code = NLLB_LANGUAGE_CODES[source_code]
+                target_nllb_code = NLLB_LANGUAGE_CODES[target_code]
+                mt_tokenizer.src_lang = source_nllb_code
+                inputs = mt_tokenizer(transcription, return_tensors="pt", clean_up_tokenization_spaces=True).to(device)
+                with torch.no_grad():
+                    generated_tokens = mt_model.generate(
+                        **inputs,
+                        forced_bos_token_id=mt_tokenizer.convert_tokens_to_ids(target_nllb_code),
+                        max_length=448
+                    )
+                translated_text = mt_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+                logger.info(f"Translation completed: {translated_text}")
+            except Exception as e:
+                logger.error(f"Error during translation: {str(e)}")
+                translated_text = f"Translation failed: {str(e)}"
+        else:
+            logger.warning("MT model not loaded, skipping translation")
+
+        # Convert translated text to speech
+        output_audio = None
+        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