test-

src/server/main.py

@@ -269,10 +269,9 @@ class TTSManager:
     async def load(self):
         if not self.model:
             logger.info("Loading TTS model IndicF5 asynchronously...")
-            local_path = "/app/models/tts_model"
             self.model = await asyncio.to_thread(
                 AutoModel.from_pretrained,
-                local_path,
+                self.repo_id,
                 trust_remote_code=True
             )
             self.model = self.model.to(self.device_type)
@@ -363,29 +362,29 @@ class TranslateManager:
     async def load(self):
         if not self.tokenizer or not self.model:
             if self.src_lang.startswith("eng") and not self.tgt_lang.startswith("eng"):
-                local_path = "/app/models/trans_en_indic"
+                model_name = "ai4bharat/indictrans2-en-indic-dist-200M" if self.use_distilled else "ai4bharat/indictrans2-en-indic-1B"
             elif not self.src_lang.startswith("eng") and self.tgt_lang.startswith("eng"):
-                local_path = "/app/models/trans_indic_en"
+                model_name = "ai4bharat/indictrans2-indic-en-dist-200M" if self.use_distilled else "ai4bharat/indictrans2-indic-en-1B"
             elif not self.src_lang.startswith("eng") and not self.tgt_lang.startswith("eng"):
-                local_path = "/app/models/trans_indic_indic"
+                model_name = "ai4bharat/indictrans2-indic-indic-dist-320M" if self.use_distilled else "ai4bharat/indictrans2-indic-indic-1B"
             else:
                 raise ValueError("Invalid language combination")
 
             self.tokenizer = await asyncio.to_thread(
                 AutoTokenizer.from_pretrained,
-                local_path,
+                model_name,
                 trust_remote_code=True
             )
             self.model = await asyncio.to_thread(
                 AutoModelForSeq2SeqLM.from_pretrained,
-                local_path,
+                model_name,
                 trust_remote_code=True,
                 torch_dtype=torch.float16,
                 attn_implementation="flash_attention_2"
             )
             self.model = self.model.to(self.device_type)
             self.model = torch.compile(self.model, mode="reduce-overhead")
-            logger.info(f"Translation model {local_path} loaded asynchronously")
+            logger.info(f"Translation model {model_name} loaded asynchronously")
 
 class ModelManager:
     def __init__(self, device_type=device, use_distilled=True, is_lazy_loading=False):