Optimize app for Hugging Face Spaces: Switch to whisper-tiny, add health check, and improve model loading

- Switched STT model from openai/whisper-small to openai/whisper-tiny to reduce memory usage and prevent startup timeout
- Optimized update_languages method to only reload TTS model, not STT model
- Added debug logging in __init__ to confirm model loading
- Added /health endpoint for Hugging Face Spaces health check
- Fixed clean_up_tokenization_spaces warning in _initialize_mt_model by setting clean_up_tokenization_spaces=True

app.py

@@ -58,29 +58,25 @@ class TalklasTranslator:
         self._initialize_stt_model()
         self._initialize_mt_model()
         self._initialize_tts_model()
+        print("All models loaded successfully, starting FastAPI app")
 
     def _initialize_stt_model(self):
         try:
-            print("Trying to load openai/whisper-small...")
-            self.stt_processor = WhisperProcessor.from_pretrained("openai/whisper-small")
-            self.stt_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
+            print("Trying to load openai/whisper-tiny...")
+            self.stt_processor = WhisperProcessor.from_pretrained("openai/whisper-tiny")
+            self.stt_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny")
             self.stt_model.to(self.device)
-            print("Loaded openai/whisper-small successfully")
+            print("Loaded openai/whisper-tiny successfully")
         except Exception as e:
-            print(f"Failed to load openai/whisper-small: {e}")
-            print("Falling back to openai/whisper-base...")
-            try:
-                self.stt_processor = WhisperProcessor.from_pretrained("openai/whisper-base")
-                self.stt_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base")
-                self.stt_model.to(self.device)
-                print("Loaded openai/whisper-base successfully")
-            except Exception as e2:
-                raise RuntimeError(f"STT model initialization failed: {e2}")
+            raise RuntimeError(f"STT model initialization failed: {e}")
 
     def _initialize_mt_model(self):
         try:
             self.mt_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
-            self.mt_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
+            self.mt_tokenizer = AutoTokenizer.from_pretrained(
+                "facebook/nllb-200-distilled-600M",
+                clean_up_tokenization_spaces=True
+            )
             self.mt_model.to(self.device)
             print("Loaded NLLB translation model successfully")
         except Exception as e:
@@ -106,10 +102,12 @@ class TalklasTranslator:
             print("Loaded fallback TTS model facebook/mms-tts-eng successfully")
 
     def update_languages(self, source_lang: str, target_lang: str):
+        print(f"Updating languages: source_lang={source_lang}, target_lang={target_lang}")
         self.source_lang = source_lang
         self.target_lang = target_lang
-        self._initialize_stt_model()
+        print("Calling _initialize_tts_model...")
         self._initialize_tts_model()
+        print("Languages updated successfully")
         return f"Languages updated to {source_lang} → {target_lang}"
 
     def speech_to_text(self, audio_path: str) -> str:
@@ -167,6 +165,10 @@ class TalklasTranslator:
 
 translator = TalklasTranslator()
 
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
+
 @app.post("/update-languages")
 async def update_languages(source_lang: str = Form(...), target_lang: str = Form(...)):
     if source_lang not in TalklasTranslator.LANGUAGE_MAPPING or target_lang not in TalklasTranslator.LANGUAGE_MAPPING: