Update app.py

app.py

@@ -35,16 +35,20 @@ models_loaded = False
 loading_in_progress = False
 loading_thread = None
 model_status = {
-    "stt": "not_loaded",
+    "stt_mms": "not_loaded",
+    "stt_whisper": "not_loaded",
     "mt": "not_loaded",
     "tts": "not_loaded"
 }
 error_message = None
 current_tts_language = "tgl"  # Track the current TTS language
+current_stt_model = None  # Track which STT model is active ("mms" or "whisper")
 
 # Model instances
-stt_processor = None
-stt_model = None
+stt_mms_processor = None
+stt_mms_model = None
+stt_whisper_processor = None
+stt_whisper_model = None
 mt_model = None
 mt_tokenizer = None
 tts_model = None
@@ -152,38 +156,44 @@ def schedule_cleanup():
 # Function to load models in background
 def load_models_task():
     global models_loaded, loading_in_progress, model_status, error_message
-    global stt_processor, stt_model, mt_model, mt_tokenizer, tts_model, tts_tokenizer
+    global stt_mms_processor, stt_mms_model, stt_whisper_processor, stt_whisper_model
+    global mt_model, mt_tokenizer, tts_model, tts_tokenizer
     
     try:
         loading_in_progress = True
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         
-        # Load STT model (MMS with fallback to Whisper)
-        logger.info("Starting to load STT model...")
+        # Load STT models (MMS and Whisper Small)
+        logger.info("Starting to load STT models...")
         from transformers import AutoProcessor, AutoModelForCTC, WhisperProcessor, WhisperForConditionalGeneration
         
+        # Load MMS STT model
         try:
             logger.info("Loading MMS STT model...")
-            model_status["stt"] = "loading"
-            stt_processor = AutoProcessor.from_pretrained("facebook/mms-1b-all")
-            stt_model = AutoModelForCTC.from_pretrained("facebook/mms-1b-all")
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            stt_model.to(device)
+            model_status["stt_mms"] = "loading"
+            stt_mms_processor = AutoProcessor.from_pretrained("facebook/mms-1b-all")
+            stt_mms_model = AutoModelForCTC.from_pretrained("facebook/mms-1b-all")
+            stt_mms_model.to(device)
             logger.info("MMS STT model loaded successfully")
-            model_status["stt"] = "loaded_mms"
+            model_status["stt_mms"] = "loaded"
         except Exception as mms_error:
             logger.error(f"Failed to load MMS STT model: {str(mms_error)}")
-            logger.info("Falling back to Whisper STT model...")
-            try:
-                stt_processor = WhisperProcessor.from_pretrained("openai/whisper-tiny")
-                stt_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny")
-                stt_model.to(device)
-                logger.info("Whisper STT model loaded successfully as fallback")
-                model_status["stt"] = "loaded_whisper"
-            except Exception as whisper_error:
-                logger.error(f"Failed to load Whisper STT model: {str(whisper_error)}")
-                model_status["stt"] = "failed"
-                error_message = f"STT model loading failed: MMS error: {str(mms_error)}, Whisper error: {str(whisper_error)}"
-                return
+            model_status["stt_mms"] = "failed"
+            error_message = f"MMS STT model loading failed: {str(mms_error)}"
+        
+        # Load Whisper Small STT model
+        try:
+            logger.info("Loading Whisper Small STT model...")
+            model_status["stt_whisper"] = "loading"
+            stt_whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-small")
+            stt_whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
+            stt_whisper_model.to(device)
+            logger.info("Whisper Small STT model loaded successfully")
+            model_status["stt_whisper"] = "loaded"
+        except Exception as whisper_error:
+            logger.error(f"Failed to load Whisper Small STT model: {str(whisper_error)}")
+            model_status["stt_whisper"] = "failed"
+            error_message = f"Whisper Small STT model loading failed: {str(whisper_error)}"
 
         # Load MT model
         logger.info("Starting to load MT model...")
@@ -203,7 +213,7 @@ def load_models_task():
             error_message = f"MT model loading failed: {str(e)}"
             return
 
-        # Load TTS model (default to Tagalog, will be updated dynamically)
+        # Load TTS model (default to Tagalog)
         logger.info("Starting to load TTS model...")
         from transformers import VitsModel, AutoTokenizer
         
@@ -217,22 +227,25 @@ def load_models_task():
             model_status["tts"] = "loaded"
         except Exception as 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...")
                 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)"
-                current_tts_language = "eng"
             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
+        # Check if critical models are loaded
+        stt_loaded = model_status["stt_mms"] == "loaded" or model_status["stt_whisper"] == "loaded"
+        mt_loaded = model_status["mt"] == "loaded"
+        tts_loaded = model_status["tts"].startswith("loaded")
+        models_loaded = stt_loaded and mt_loaded and tts_loaded
         logger.info("Model loading completed successfully")
         
     except Exception as e:
@@ -284,7 +297,8 @@ 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, current_tts_language
+    global stt_mms_processor, stt_mms_model, stt_whisper_processor, stt_whisper_model
+    global tts_model, tts_tokenizer, current_tts_language, current_stt_model
     
     if source_lang not in LANGUAGE_MAPPING or target_lang not in LANGUAGE_MAPPING:
         raise HTTPException(status_code=400, detail="Invalid language selected")
@@ -292,43 +306,78 @@ async def update_languages(source_lang: str = Form(...), target_lang: str = Form
     source_code = LANGUAGE_MAPPING[source_lang]
     target_code = LANGUAGE_MAPPING[target_lang]
     
-    # Update the STT model based on the source language (MMS or Whisper)
+    # Update the STT model based on the source language
     try:
-        logger.info("Updating STT model for source language...")
+        logger.info(f"Updating STT model for source language {source_code}...")
         from transformers import AutoProcessor, AutoModelForCTC, WhisperProcessor, WhisperForConditionalGeneration
         device = "cuda" if torch.cuda.is_available() else "cpu"
         
-        try:
-            logger.info(f"Loading MMS STT model for {source_code}...")
-            stt_processor = AutoProcessor.from_pretrained("facebook/mms-1b-all")
-            stt_model = AutoModelForCTC.from_pretrained("facebook/mms-1b-all")
-            stt_model.to(device)
-            # Set the target language for MMS
-            if source_code in stt_processor.tokenizer.vocab.keys():
-                stt_processor.tokenizer.set_target_lang(source_code)
-                stt_model.load_adapter(source_code)
-                logger.info(f"MMS STT model updated to {source_code}")
-                model_status["stt"] = "loaded_mms"
-            else:
-                logger.warning(f"Language {source_code} not supported by MMS, using default")
-                model_status["stt"] = "loaded_mms_default"
-        except Exception as mms_error:
-            logger.error(f"Failed to load MMS STT model for {source_code}: {str(mms_error)}")
-            logger.info("Falling back to Whisper STT model...")
+        # Use Whisper Small for English or Tagalog, MMS for others
+        if source_code in ["eng", "tgl"]:
             try:
-                stt_processor = WhisperProcessor.from_pretrained("openai/whisper-tiny")
-                stt_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny")
-                stt_model.to(device)
-                logger.info("Whisper STT model loaded successfully as fallback")
-                model_status["stt"] = "loaded_whisper"
+                logger.info(f"Loading Whisper Small STT model for {source_code}...")
+                if model_status["stt_whisper"] != "loaded":
+                    stt_whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-small")
+                    stt_whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
+                    stt_whisper_model.to(device)
+                    model_status["stt_whisper"] = "loaded"
+                current_stt_model = "whisper"
+                logger.info("Whisper Small STT model selected")
             except Exception as whisper_error:
-                logger.error(f"Failed to load Whisper STT model: {str(whisper_error)}")
-                model_status["stt"] = "failed"
-                error_message = f"STT model update failed: MMS error: {str(mms_error)}, Whisper error: {str(whisper_error)}"
-                return {"status": "failed", "error": error_message}
+                logger.error(f"Failed to load Whisper Small STT model: {str(whisper_error)}")
+                try:
+                    logger.info(f"Falling back to MMS STT model for {source_code}...")
+                    if model_status["stt_mms"] != "loaded":
+                        stt_mms_processor = AutoProcessor.from_pretrained("facebook/mms-1b-all")
+                        stt_mms_model = AutoModelForCTC.from_pretrained("facebook/mms-1b-all")
+                        stt_mms_model.to(device)
+                        model_status["stt_mms"] = "loaded"
+                    if source_code in stt_mms_processor.tokenizer.vocab.keys():
+                        stt_mms_processor.tokenizer.set_target_lang(source_code)
+                        stt_mms_model.load_adapter(source_code)
+                    current_stt_model = "mms"
+                    logger.info("MMS STT model selected as fallback")
+                except Exception as mms_error:
+                    logger.error(f"Failed to load MMS STT model: {str(mms_error)}")
+                    model_status["stt_mms"] = "failed"
+                    model_status["stt_whisper"] = "failed"
+                    error_message = f"STT model update failed: Whisper error: {str(whisper_error)}, MMS error: {str(mms_error)}"
+                    return {"status": "failed", "error": error_message}
+        else:
+            try:
+                logger.info(f"Loading MMS STT model for {source_code}...")
+                if model_status["stt_mms"] != "loaded":
+                    stt_mms_processor = AutoProcessor.from_pretrained("facebook/mms-1b-all")
+                    stt_mms_model = AutoModelForCTC.from_pretrained("facebook/mms-1b-all")
+                    stt_mms_model.to(device)
+                    model_status["stt_mms"] = "loaded"
+                if source_code in stt_mms_processor.tokenizer.vocab.keys():
+                    stt_mms_processor.tokenizer.set_target_lang(source_code)
+                    stt_mms_model.load_adapter(source_code)
+                current_stt_model = "mms"
+                logger.info(f"MMS STT model selected for {source_code}")
+            except Exception as mms_error:
+                logger.error(f"Failed to load MMS STT model: {str(mms_error)}")
+                try:
+                    logger.info(f"Falling back to Whisper Small STT model for {source_code}...")
+                    if model_status["stt_whisper"] != "loaded":
+                        stt_whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-small")
+                        stt_whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
+                        stt_whisper_model.to(device)
+                        model_status["stt_whisper"] = "loaded"
+                    current_stt_model = "whisper"
+                    logger.info("Whisper Small STT model selected as fallback")
+                except Exception as whisper_error:
+                    logger.error(f"Failed to load Whisper Small STT model: {str(whisper_error)}")
+                    model_status["stt_mms"] = "failed"
+                    model_status["stt_whisper"] = "failed"
+                    error_message = f"STT model update failed: MMS error: {str(mms_error)}, Whisper error: {str(whisper_error)}"
+                    return {"status": "failed", "error": error_message}
+    
     except Exception as e:
         logger.error(f"Error updating STT model: {str(e)}")
-        model_status["stt"] = "failed"
+        model_status["stt_mms"] = "failed"
+        model_status["stt_whisper"] = "failed"
         error_message = f"STT model update failed: {str(e)}"
         return {"status": "failed", "error": error_message}
     
@@ -466,7 +515,8 @@ 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, current_tts_language
+    global stt_mms_processor, stt_mms_model, stt_whisper_processor, stt_whisper_model
+    global mt_model, mt_tokenizer, tts_model, tts_tokenizer, current_tts_language, current_stt_model
     
     if not audio:
         raise HTTPException(status_code=400, detail="No audio file provided")
@@ -477,17 +527,37 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
     request_id = str(uuid.uuid4())
     
     # Check if STT model is loaded
-    if model_status["stt"] not in ["loaded_mms", "loaded_mms_default", "loaded_whisper"] or stt_processor is None or stt_model is None:
-        logger.warning("STT model not loaded, returning placeholder response")
-        return {
-            "request_id": request_id,
-            "status": "processing",
-            "message": "STT model not loaded yet. Please try again later.",
-            "source_text": "Transcription not available",
-            "translated_text": "Translation not available",
-            "is_inappropriate": False,
-            "output_audio": None
-        }
+    source_code = LANGUAGE_MAPPING[source_lang]
+    use_whisper = source_code in ["eng", "tgl"]
+    
+    if use_whisper and (model_status["stt_whisper"] != "loaded" or stt_whisper_processor is None or stt_whisper_model is None):
+        logger.warning("Whisper Small STT model not loaded, falling back to MMS")
+        if model_status["stt_mms"] != "loaded" or stt_mms_processor is None or stt_mms_model is None:
+            logger.warning("MMS STT model not loaded either, returning placeholder response")
+            return {
+                "request_id": request_id,
+                "status": "processing",
+                "message": "STT models not loaded yet. Please try again later.",
+                "source_text": "Transcription not available",
+                "translated_text": "Translation not available",
+                "is_inappropriate": False,
+                "output_audio": None
+            }
+        use_whisper = False
+    elif not use_whisper and (model_status["stt_mms"] != "loaded" or stt_mms_processor is None or stt_mms_model is None):
+        logger.warning("MMS STT model not loaded, falling back to Whisper Small")
+        if model_status["stt_whisper"] != "loaded" or stt_whisper_processor is None or stt_whisper_model is None:
+            logger.warning("Whisper Small STT model not loaded either, returning placeholder response")
+            return {
+                "request_id": request_id,
+                "status": "processing",
+                "message": "STT models not loaded yet. Please try again later.",
+                "source_text": "Transcription not available",
+                "translated_text": "Translation not available",
+                "is_inappropriate": False,
+                "output_audio": None
+            }
+        use_whisper = True
     
     # Save the uploaded audio to a temporary file
     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
@@ -526,24 +596,30 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
 
         # Step 3: Transcribe the audio (STT)
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info(f"Using device: {device}")
-        inputs = stt_processor(waveform.numpy(), sampling_rate=16000, return_tensors="pt").to(device)
-        logger.info("Audio processed, generating transcription...")
+        logger.info(f"Using device: {device} with {'Whisper Small' if use_whisper else 'MMS'} model")
         
-        with torch.no_grad():
-            if model_status["stt"] == "loaded_whisper":
-                # Whisper model
-                generated_ids = stt_model.generate(**inputs, language="en")
-                transcription = stt_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-            else:
-                # MMS model
-                logits = stt_model(**inputs).logits
+        if use_whisper:
+            processor = stt_whisper_processor
+            model = stt_whisper_model
+            inputs = processor(waveform.numpy()[0], sampling_rate=16000, return_tensors="pt").to(device)
+            with torch.no_grad():
+                language = "en" if source_code == "eng" else "tl" if source_code == "tgl" else None
+                generated_ids = model.generate(**inputs, language=language, task="transcribe")
+                transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        else:
+            processor = stt_mms_processor
+            model = stt_mms_model
+            if source_code in processor.tokenizer.vocab.keys():
+                processor.tokenizer.set_target_lang(source_code)
+                model.load_adapter(source_code)
+            inputs = processor(waveform.numpy(), sampling_rate=16000, return_tensors="pt").to(device)
+            with torch.no_grad():
+                logits = model(**inputs).logits
                 predicted_ids = torch.argmax(logits, dim=-1)
-                transcription = stt_processor.batch_decode(predicted_ids)[0]
+                transcription = processor.batch_decode(predicted_ids)[0]
         logger.info(f"Transcription completed: {transcription}")
 
         # Step 4: Translate the transcribed text (MT)
-        source_code = LANGUAGE_MAPPING[source_lang]
         target_code = LANGUAGE_MAPPING[target_lang]
         
         if model_status["mt"] == "loaded" and mt_model is not None and mt_tokenizer is not None:
@@ -618,7 +694,8 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
             except Exception as e:
                 logger.error(f"Error during TTS conversion: {str(e)}")
                 output_audio_url = None
-                return {
+        
+        return {
             "request_id": request_id,
             "status": "completed",
             "message": "Transcription, translation, and TTS completed (or partially completed).",