Implement lazy loading of ML models to fix startup timeout on HF Spaces

app.py

@@ -3,34 +3,36 @@ import os
 os.environ["HOME"] = "/root"
 os.environ["HF_HOME"] = "/tmp/hf_cache"
 
-# Debug: Print environment variables to confirm
+# Print environment variables to confirm
 print("HOME environment variable:", os.environ.get("HOME"))
 print("HF_HOME environment variable:", os.environ.get("HF_HOME"))
 
-# Now import other libraries
+# Import libraries
 import torch
 import numpy as np
 import soundfile as sf
-from transformers import (
-    AutoModelForSeq2SeqLM,
-    AutoTokenizer,
-    VitsModel,
-    AutoProcessor,
-    AutoModelForCTC,
-    WhisperProcessor,
-    WhisperForConditionalGeneration
-)
-from typing import Optional, Tuple, Dict
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from typing import Optional, Tuple, Dict, Any
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks
 from fastapi.responses import JSONResponse
 import tempfile
 import logging
+from threading import Thread
+import time
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("talklas-api")
 
 # Configure transformers logging to reduce verbosity
 logging.getLogger("transformers").setLevel(logging.ERROR)
 
 app = FastAPI(title="Talklas API")
 
+# Global variables to track model loading status
+is_loading = False
+loading_complete = False
+loading_error = None
+
 class TalklasTranslator:
     LANGUAGE_MAPPING = {
         "English": "eng",
@@ -52,72 +54,113 @@ class TalklasTranslator:
 
     def __init__(self, source_lang: str = "eng", target_lang: str = "tgl"):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Using device: {self.device}")
         self.source_lang = source_lang
         self.target_lang = target_lang
         self.sample_rate = 16000
-        self.mt_model = None  # Initialize as None
-        self.mt_tokenizer = None  # Initialize as None
-        self._initialize_stt_model()
-        self._initialize_mt_model()
-        self._initialize_tts_model()
-        print("All models loaded successfully, starting FastAPI app")
+        
+        # Initialize all models as None - will be lazy loaded
+        self.stt_processor = None
+        self.stt_model = None
+        self.mt_model = None
+        self.mt_tokenizer = None
+        self.tts_model = None
+        self.tts_tokenizer = None
+        
+        # Flags to track which models are loaded
+        self.stt_loaded = False
+        self.mt_loaded = False
+        self.tts_loaded = False
 
     def _initialize_stt_model(self):
+        if self.stt_loaded:
+            return True
+            
         try:
-            print("Trying to load openai/whisper-tiny...")
+            from transformers import WhisperProcessor, WhisperForConditionalGeneration
+            logger.info("Loading STT model: 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-tiny successfully")
+            self.stt_loaded = True
+            logger.info("STT model loaded successfully")
+            return True
         except Exception as e:
-            raise RuntimeError(f"STT model initialization failed: {e}")
+            logger.error(f"STT model initialization failed: {e}")
+            return False
 
     def _initialize_mt_model(self):
+        if self.mt_loaded:
+            return True
+            
         try:
-            print("Trying to load facebook/nllb-200-distilled-600M...")
+            from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+            logger.info("Loading MT model: facebook/nllb-200-distilled-600M...")
             self.mt_model = AutoModelForSeq2SeqLM.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")
+            self.mt_loaded = True
+            logger.info("MT model loaded successfully")
+            return True
         except Exception as e:
-            print(f"Failed to load facebook/nllb-200-distilled-600M: {e}")
-            print("Translation model not loaded, translation will return source text as a fallback")
-            self.mt_model = None
-            self.mt_tokenizer = None
+            logger.error(f"MT model initialization failed: {e}")
+            return False
 
     def _initialize_tts_model(self):
+        if self.tts_loaded:
+            # Check if we need to reload for a different language
+            if hasattr(self, 'current_tts_lang') and self.current_tts_lang == self.target_lang:
+                return True
+                
         try:
-            print(f"Trying to load facebook/mms-tts-{self.target_lang}...")
+            from transformers import VitsModel, AutoTokenizer
+            logger.info(f"Loading TTS model: facebook/mms-tts-{self.target_lang}...")
             self.tts_model = VitsModel.from_pretrained(f"facebook/mms-tts-{self.target_lang}")
             self.tts_tokenizer = AutoTokenizer.from_pretrained(
                 f"facebook/mms-tts-{self.target_lang}",
                 clean_up_tokenization_spaces=True
             )
             self.tts_model.to(self.device)
-            print(f"Loaded TTS model facebook/mms-tts-{self.target_lang} successfully")
-        except Exception:
-            print(f"Failed to load facebook/mms-tts-{self.target_lang}, falling back to English TTS")
-            self.tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
-            self.tts_tokenizer = AutoTokenizer.from_pretrained(
-                "facebook/mms-tts-eng",
-                clean_up_tokenization_spaces=True
-            )
-            self.tts_model.to(self.device)
-            print("Loaded fallback TTS model facebook/mms-tts-eng successfully")
+            self.tts_loaded = True
+            self.current_tts_lang = self.target_lang
+            logger.info(f"TTS model loaded successfully for {self.target_lang}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to load TTS model for {self.target_lang}: {e}")
+            try:
+                logger.info("Falling back to English TTS model...")
+                self.tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
+                self.tts_tokenizer = AutoTokenizer.from_pretrained(
+                    "facebook/mms-tts-eng",
+                    clean_up_tokenization_spaces=True
+                )
+                self.tts_model.to(self.device)
+                self.tts_loaded = True
+                self.current_tts_lang = "eng"
+                logger.info("Loaded fallback TTS model successfully")
+                return True
+            except Exception as fallback_error:
+                logger.error(f"Fallback TTS model initialization failed: {fallback_error}")
+                return False
 
     def update_languages(self, source_lang: str, target_lang: str):
-        print(f"Updating languages: source_lang={source_lang}, target_lang={target_lang}")
+        logger.info(f"Updating languages: source_lang={source_lang}, target_lang={target_lang}")
         self.source_lang = source_lang
         self.target_lang = target_lang
-        print("Calling _initialize_tts_model...")
-        self._initialize_tts_model()
-        print("Languages updated successfully")
+        
+        # Only reload TTS model if target language changed
+        if hasattr(self, 'current_tts_lang') and self.current_tts_lang != target_lang:
+            self._initialize_tts_model()
+        
         return f"Languages updated to {source_lang} → {target_lang}"
 
     def speech_to_text(self, audio_path: str) -> str:
+        if not self._initialize_stt_model():
+            raise Exception("STT model failed to initialize")
+            
         waveform, sample_rate = sf.read(audio_path)
         if sample_rate != 16000:
             import librosa
@@ -129,9 +172,10 @@ class TalklasTranslator:
         return transcription
 
     def translate_text(self, text: str) -> str:
-        if self.mt_model is None or self.mt_tokenizer is None:
-            print("Translation model not loaded, returning source text as fallback")
+        if not self._initialize_mt_model():
+            logger.warning("Translation model not loaded, returning source text as fallback")
             return text
+            
         source_code = self.NLLB_LANGUAGE_CODES[self.source_lang]
         target_code = self.NLLB_LANGUAGE_CODES[self.target_lang]
         self.mt_tokenizer.src_lang = source_code
@@ -145,6 +189,9 @@ class TalklasTranslator:
         return self.mt_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
     def text_to_speech(self, text: str) -> Tuple[int, np.ndarray]:
+        if not self._initialize_tts_model():
+            raise Exception("TTS model failed to initialize")
+            
         inputs = self.tts_tokenizer(text, return_tensors="pt", clean_up_tokenization_spaces=True).to(self.device)
         with torch.no_grad():
             output = self.tts_model(**inputs)
@@ -173,11 +220,75 @@ class TalklasTranslator:
             "performance": "Translation successful"
         }
 
+# Create translator instance but don't load models yet
 translator = TalklasTranslator()
 
+def background_load_model():
+    """Background task to load models"""
+    global is_loading, loading_complete, loading_error
+    
+    try:
+        is_loading = True
+        # Load STT model first to make health check pass quickly
+        success = translator._initialize_stt_model()
+        if not success:
+            loading_error = "Failed to load STT model"
+            return
+            
+        # Then load MT model
+        success = translator._initialize_mt_model()
+        if not success:
+            logger.warning("MT model failed to load, will use fallback")
+            
+        # Finally load TTS model
+        success = translator._initialize_tts_model()
+        if not success:
+            loading_error = "Failed to load TTS model"
+            return
+            
+        loading_complete = True
+        logger.info("All models loaded successfully in background")
+        
+    except Exception as e:
+        loading_error = str(e)
+        logger.error(f"Error loading models in background: {e}")
+    finally:
+        is_loading = False
+
+# Start background loading of models
+Thread(target=background_load_model, daemon=True).start()
+
 @app.get("/health")
 async def health_check():
-    return {"status": "healthy"}
+    """Health check endpoint that returns detailed loading status"""
+    global is_loading, loading_complete, loading_error
+    
+    # Check if at least the STT model is loaded (minimum requirement)
+    if translator.stt_loaded:
+        status = "healthy"
+    elif loading_error:
+        status = "error"
+    elif is_loading:
+        status = "loading"
+    else:
+        status = "not_initialized"
+        
+    response = {
+        "status": status,
+        "models": {
+            "stt": "loaded" if translator.stt_loaded else "not_loaded",
+            "mt": "loaded" if translator.mt_loaded else "not_loaded",
+            "tts": "loaded" if translator.tts_loaded else "not_loaded",
+        },
+        "loading": is_loading,
+        "complete": loading_complete
+    }
+    
+    if loading_error:
+        response["error"] = loading_error
+    
+    # Hugging Face Spaces considers a service healthy if the health endpoint returns a 200 status
+    return response
 
 @app.post("/update-languages")
 async def update_languages(source_lang: str = Form(...), target_lang: str = Form(...)):
@@ -196,6 +307,17 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
     if source_lang not in TalklasTranslator.LANGUAGE_MAPPING or target_lang not in TalklasTranslator.LANGUAGE_MAPPING:
         raise HTTPException(status_code=400, detail="Invalid language selected")
     
+    # Check if models are loaded
+    if not translator.stt_loaded:
+        if loading_error:
+            raise HTTPException(status_code=500, detail=f"Model loading failed: {loading_error}")
+        elif is_loading:
+            raise HTTPException(status_code=503, detail="Models are still loading, please try again later")
+        else:
+            # Try to load models now
+            if not translator._initialize_stt_model():
+                raise HTTPException(status_code=500, detail="Failed to initialize STT model")
+    
     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
         temp_file.write(await audio.read())
         temp_path = temp_file.name
@@ -207,6 +329,8 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
         )
         result = translator.translate_speech(temp_path)
         return JSONResponse(content=result)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Translation failed: {str(e)}")
     finally:
         os.unlink(temp_path)
 
@@ -217,15 +341,26 @@ async def translate_text(text: str = Form(...), source_lang: str = Form(...), ta
     if source_lang not in TalklasTranslator.LANGUAGE_MAPPING or target_lang not in TalklasTranslator.LANGUAGE_MAPPING:
         raise HTTPException(status_code=400, detail="Invalid language selected")
     
+    # Check if models are loaded
+    if not translator.mt_loaded or not translator.tts_loaded:
+        if loading_error:
+            raise HTTPException(status_code=500, detail=f"Model loading failed: {loading_error}")
+        elif is_loading:
+            raise HTTPException(status_code=503, detail="Models are still loading, please try again later")
+    
     translator.update_languages(
         TalklasTranslator.LANGUAGE_MAPPING[source_lang],
         TalklasTranslator.LANGUAGE_MAPPING[target_lang]
     )
-    result = translator.translate_text_only(text)
-    return JSONResponse(content=result)
+    
+    try:
+        result = translator.translate_text_only(text)
+        return JSONResponse(content=result)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Translation failed: {str(e)}")
 
 if __name__ == "__main__":
     import uvicorn
-    print("Starting Uvicorn server...")
+    logger.info("Starting Uvicorn server...")
     uvicorn.run(app, host="0.0.0.0", port=8000)
-    print("Uvicorn server started successfully")
+    logger.info("Uvicorn server started successfully")