fix

src/server/main.py

@@ -68,7 +68,7 @@ class Settings(BaseSettings):
 
 settings = Settings()
 
-# Quantization config for LLM (unchanged from gemma_llm.py)
+# Quantization config for LLM
 quantization_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
@@ -76,7 +76,7 @@ quantization_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-# LLM Manager (from gemma_llm.py with async load)
+# LLM Manager
 class LLMManager:
     def __init__(self, model_name: str, device: str = "cuda" if torch.cuda.is_available() else "cpu"):
         self.model_name = model_name
@@ -259,7 +259,7 @@ class LLMManager:
         logger.info(f"Chat_v2 response: {decoded}")
         return decoded
 
-# TTS Manager (async load)
+# TTS Manager
 class TTSManager:
     def __init__(self, device_type=device):
         self.device_type = device_type
@@ -348,7 +348,7 @@ SUPPORTED_LANGUAGES = {
     "por_Latn", "rus_Cyrl", "pol_Latn"
 }
 
-# Translation Manager (async load)
+# Translation Manager
 class TranslateManager:
     def __init__(self, src_lang, tgt_lang, device_type=device, use_distilled=True):
         self.device_type = device_type
@@ -413,7 +413,7 @@ class ModelManager:
             return 'indic_indic'
         raise ValueError("Invalid language combination")
 
-# ASR Manager (async load)
+# ASR Manager
 class ASRModelManager:
     def __init__(self, device_type="cuda"):
         self.device_type = device_type
@@ -498,7 +498,7 @@ async def lifespan(app: FastAPI):
     logger.info("Starting model loading in background...")
     asyncio.create_task(load_all_models())
     yield
-    llm_manager.unload()  # Synchronous unload as per original gemma_llm.py
+    llm_manager.unload()
     logger.info("Server shutdown complete")
 
 # FastAPI App
@@ -582,9 +582,17 @@ async def translate(request: TranslationRequest, translate_manager: TranslateMan
     return TranslationResponse(translations=translations)
 
 async def perform_internal_translation(sentences: List[str], src_lang: str, tgt_lang: str) -> List[str]:
-    translate_manager = model_manager.get_model(src_lang, tgt_lang)
+    try:
+        translate_manager = model_manager.get_model(src_lang, tgt_lang)
+    except ValueError as e:
+        logger.info(f"Model not preloaded: {str(e)}, loading now...")
+        key = model_manager._get_model_key(src_lang, tgt_lang)
+        await model_manager.load_model(src_lang, tgt_lang, key)
+        translate_manager = model_manager.get_model(src_lang, tgt_lang)
+    
     if not translate_manager.model:  # Ensure model is loaded
         await translate_manager.load()
+    
     request = TranslationRequest(sentences=sentences, src_lang=src_lang, tgt_lang=tgt_lang)
     response = await translate(request, translate_manager)
     return response.translations
@@ -601,7 +609,7 @@ async def home():
 async def unload_all_models():
     try:
         logger.info("Starting to unload all models...")
-        llm_manager.unload()  # Synchronous as per original
+        llm_manager.unload()
         logger.info("All models unloaded successfully")
         return {"status": "success", "message": "All models unloaded"}
     except Exception as e:
@@ -612,7 +620,7 @@ async def unload_all_models():
 async def load_all_models():
     try:
         logger.info("Starting to load all models...")
-        await llm_manager.load()  # Async load
+        await llm_manager.load()
         logger.info("All models loaded successfully")
         return {"status": "success", "message": "All models loaded"}
     except Exception as e: