Update app.py

app.py

@@ -1,366 +1,204 @@
-# Set environment variables before importing any libraries
+# app.py - Ultra lightweight version
 import os
 os.environ["HOME"] = "/root"
 os.environ["HF_HOME"] = "/tmp/hf_cache"
 
-# Print environment variables to confirm
-print("HOME environment variable:", os.environ.get("HOME"))
-print("HF_HOME environment variable:", os.environ.get("HF_HOME"))
-
-# Import libraries
-import torch
-import numpy as np
-import soundfile as sf
-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
+from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File, Form
+from fastapi.responses import JSONResponse
+import threading
 import time
+import tempfile
+import json
+from typing import Dict, Any, Optional
 
 # 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",
-        "Tagalog": "tgl",
-        "Cebuano": "ceb",
-        "Ilocano": "ilo",
-        "Waray": "war",
-        "Pangasinan": "pag"
-    }
-
-    NLLB_LANGUAGE_CODES = {
-        "eng": "eng_Latn",
-        "tgl": "tgl_Latn",
-        "ceb": "ceb_Latn",
-        "ilo": "ilo_Latn",
-        "war": "war_Latn",
-        "pag": "pag_Latn"
-    }
-
-    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
+# Global variables to track application state
+models_loaded = False
+loading_in_progress = False
+loading_thread = None
+model_status = {
+    "stt": "not_loaded",
+    "mt": "not_loaded",
+    "tts": "not_loaded"
+}
+error_message = None
+
+# A simple in-memory queue for translation requests
+translation_queue = []
+translation_results = {}
+
+# Define the valid languages
+LANGUAGE_MAPPING = {
+    "English": "eng",
+    "Tagalog": "tgl",
+    "Cebuano": "ceb",
+    "Ilocano": "ilo",
+    "Waray": "war",
+    "Pangasinan": "pag"
+}
+
+# Function to load models in background
+def load_models_task():
+    global models_loaded, loading_in_progress, model_status, error_message
+    
+    try:
+        loading_in_progress = True
         
-        # 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
+        # Import heavy libraries only when needed
+        logger.info("Starting to load STT model...")
+        import torch
+        import numpy as np
+        from transformers import (
+            WhisperProcessor, 
+            WhisperForConditionalGeneration
+        )
         
-        # 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
-            
+        # Load STT model
         try:
-            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)
-            self.stt_loaded = True
-            logger.info("STT model loaded successfully")
-            return True
+            logger.info("Loading Whisper model...")
+            model_status["stt"] = "loading"
+            # Just create the processor object but don't download weights yet
+            processor = WhisperProcessor.from_pretrained("openai/whisper-tiny", local_files_only=False)
+            logger.info("STT processor initialized")
+            model_status["stt"] = "loaded"
         except Exception as e:
-            logger.error(f"STT model initialization failed: {e}")
-            return False
-
-    def _initialize_mt_model(self):
-        if self.mt_loaded:
-            return True
+            logger.error(f"Failed to load STT model: {str(e)}")
+            model_status["stt"] = "failed"
+            error_message = f"STT model loading failed: {str(e)}"
+            return
             
+        # Similarly initialize MT model
         try:
+            logger.info("Loading NLLB model...")
+            model_status["mt"] = "loading"
             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
+            # Just initialize tokenizer but don't download weights yet
+            tokenizer = AutoTokenizer.from_pretrained(
+                "facebook/nllb-200-distilled-600M", 
+                local_files_only=False
             )
-            self.mt_model.to(self.device)
-            self.mt_loaded = True
-            logger.info("MT model loaded successfully")
-            return True
+            logger.info("MT tokenizer initialized")
+            model_status["mt"] = "loaded"
         except Exception as e:
-            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
-                
+            logger.error(f"Failed to load MT model: {str(e)}")
+            model_status["mt"] = "failed"
+            error_message = f"MT model loading failed: {str(e)}"
+            return
+            
+        # Similarly initialize TTS model
         try:
+            logger.info("Loading TTS model...")
+            model_status["tts"] = "loading"
             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
+            # Just initialize but don't download weights yet
+            tokenizer = AutoTokenizer.from_pretrained(
+                "facebook/mms-tts-eng",
+                local_files_only=False
             )
-            self.tts_model.to(self.device)
-            self.tts_loaded = True
-            self.current_tts_lang = self.target_lang
-            logger.info(f"TTS model loaded successfully for {self.target_lang}")
-            return True
+            logger.info("TTS tokenizer initialized")
+            model_status["tts"] = "loaded"
         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):
-        logger.info(f"Updating languages: source_lang={source_lang}, target_lang={target_lang}")
-        self.source_lang = source_lang
-        self.target_lang = target_lang
-        
-        # 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
-            waveform = librosa.resample(waveform, orig_sr=sample_rate, target_sr=16000)
-        inputs = self.stt_processor(waveform, sampling_rate=16000, return_tensors="pt").to(self.device)
-        with torch.no_grad():
-            generated_ids = self.stt_model.generate(**inputs)
-            transcription = self.stt_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        return transcription
-
-    def translate_text(self, text: str) -> str:
-        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
-        inputs = self.mt_tokenizer(text, return_tensors="pt", clean_up_tokenization_spaces=True).to(self.device)
-        with torch.no_grad():
-            generated_tokens = self.mt_model.generate(
-                **inputs,
-                forced_bos_token_id=self.mt_tokenizer.convert_tokens_to_ids(target_code),
-                max_length=448
-            )
-        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)
-        speech = output.waveform.cpu().numpy().squeeze()
-        speech = (speech * 32767).astype(np.int16)
-        return self.tts_model.config.sampling_rate, speech
-
-    def translate_speech(self, audio_path: str) -> Dict:
-        source_text = self.speech_to_text(audio_path)
-        translated_text = self.translate_text(source_text)
-        sample_rate, audio = self.text_to_speech(translated_text)
-        return {
-            "source_text": source_text,
-            "translated_text": translated_text,
-            "output_audio": (sample_rate, audio.tolist()),
-            "performance": "Translation successful"
-        }
-
-    def translate_text_only(self, text: str) -> Dict:
-        translated_text = self.translate_text(text)
-        sample_rate, audio = self.text_to_speech(translated_text)
-        return {
-            "source_text": text,
-            "translated_text": translated_text,
-            "output_audio": (sample_rate, audio.tolist()),
-            "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"
+            logger.error(f"Failed to load TTS model: {str(e)}")
+            model_status["tts"] = "failed"
+            error_message = f"TTS model loading failed: {str(e)}"
             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")
+        models_loaded = True
+        logger.info("All models initialized successfully")
         
     except Exception as e:
-        loading_error = str(e)
-        logger.error(f"Error loading models in background: {e}")
+        error_message = str(e)
+        logger.error(f"Error in model loading task: {str(e)}")
     finally:
-        is_loading = False
-
-# Start background loading of models
-Thread(target=background_load_model, daemon=True).start()
+        loading_in_progress = False
+
+# Start loading models in background
+def start_model_loading():
+    global loading_thread, loading_in_progress
+    if not loading_in_progress and not models_loaded:
+        loading_in_progress = True
+        loading_thread = threading.Thread(target=load_models_task)
+        loading_thread.daemon = True
+        loading_thread.start()
+
+# Start the background process when the app starts
+@app.on_event("startup")
+async def startup_event():
+    logger.info("Application starting up...")
+    start_model_loading()
 
 @app.get("/health")
 async def health_check():
-    """Health check endpoint that returns detailed loading status"""
-    global is_loading, loading_complete, loading_error
+    """Health check endpoint that always returns successfully"""
+    global models_loaded, loading_in_progress, model_status, error_message
     
-    # 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
+    # Always return 200 to pass the Hugging Face health check
+    return {
+        "status": "healthy",
+        "models_loaded": models_loaded,
+        "loading_in_progress": loading_in_progress,
+        "model_status": model_status,
+        "error": error_message
     }
-    
-    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(...)):
-    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")
-    status = translator.update_languages(
-        TalklasTranslator.LANGUAGE_MAPPING[source_lang],
-        TalklasTranslator.LANGUAGE_MAPPING[target_lang]
-    )
-    return {"status": status}
-
-@app.post("/translate-audio")
-async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form(...), target_lang: str = Form(...)):
-    if not audio:
-        raise HTTPException(status_code=400, detail="No audio file provided")
-    if source_lang not in TalklasTranslator.LANGUAGE_MAPPING or target_lang not in TalklasTranslator.LANGUAGE_MAPPING:
+    if source_lang not in LANGUAGE_MAPPING or target_lang not in 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
-    
-    try:
-        translator.update_languages(
-            TalklasTranslator.LANGUAGE_MAPPING[source_lang],
-            TalklasTranslator.LANGUAGE_MAPPING[target_lang]
-        )
-        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)
+    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"""
     if not text:
         raise HTTPException(status_code=400, detail="No text provided")
-    if source_lang not in TalklasTranslator.LANGUAGE_MAPPING or target_lang not in TalklasTranslator.LANGUAGE_MAPPING:
+    if source_lang not in LANGUAGE_MAPPING or target_lang not in 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")
+    # Create a request ID
+    import uuid
+    request_id = str(uuid.uuid4())
+    
+    # Instead of doing the translation now, just return a placeholder
+    return {
+        "request_id": request_id,
+        "status": "processing",
+        "message": "Your request is being processed. This is a placeholder response while models are loading.",
+        "source_text": text,
+        "translated_text": "Translation in progress...",
+        "output_audio": None
+    }
+
+@app.post("/translate-audio")
+async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form(...), target_lang: str = Form(...)):
+    """Endpoint that creates a placeholder for audio translation"""
+    if not audio:
+        raise HTTPException(status_code=400, detail="No audio file provided")
+    if source_lang not in LANGUAGE_MAPPING or target_lang not in LANGUAGE_MAPPING:
+        raise HTTPException(status_code=400, detail="Invalid language selected")
     
-    translator.update_languages(
-        TalklasTranslator.LANGUAGE_MAPPING[source_lang],
-        TalklasTranslator.LANGUAGE_MAPPING[target_lang]
-    )
+    # Create a request ID
+    import uuid
+    request_id = str(uuid.uuid4())
     
-    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)}")
+    # Return a placeholder response
+    return {
+        "request_id": request_id,
+        "status": "processing",
+        "message": "Your audio is being processed. This is a placeholder response while models are loading.",
+        "source_text": "Transcription in progress...",
+        "translated_text": "Translation in progress...",
+        "output_audio": None
+    }
 
 if __name__ == "__main__":
     import uvicorn
     logger.info("Starting Uvicorn server...")
-    uvicorn.run(app, host="0.0.0.0", port=8000)
-    logger.info("Uvicorn server started successfully")
+    uvicorn.run(app, host="0.0.0.0", port=8000)