Update app.py

app.py

@@ -28,72 +28,29 @@ SARVAM_LANGUAGES = INDIC_LANGUAGES
 TORCH_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 DEVICE_MAP = "auto" if torch.cuda.is_available() else None
 
-class ModelManager:
-    def __init__(self):
-        self.indictrans_model = None
-        self.indictrans_tokenizer = None
-        self.sarvam_model = None
-        self.sarvam_tokenizer = None
-        self.current_model = None
-        
-    def load_indictrans_model(self):
-        if self.indictrans_model is None:
-            try:
-                self.indictrans_model = AutoModelForCausalLM.from_pretrained(
-                    "ai4bharat/IndicTrans3-beta",
-                    torch_dtype=TORCH_DTYPE,
-                    device_map=DEVICE_MAP,
-                    token=HF_TOKEN,
-                    low_cpu_mem_usage=True,
-                    trust_remote_code=True
-                )
-                self.indictrans_tokenizer = AutoTokenizer.from_pretrained(
-                    "ai4bharat/IndicTrans3-beta",
-                    trust_remote_code=True
-                )
-                # Enable optimizations
-                if hasattr(self.indictrans_model, 'eval'):
-                    self.indictrans_model.eval()
-                if torch.cuda.is_available():
-                    torch.cuda.empty_cache()
-            except Exception as e:
-                print(f"Error loading IndicTrans model: {e}")
-                
-    def load_sarvam_model(self):
-        if self.sarvam_model is None:
-            try:
-                self.sarvam_model = AutoModelForCausalLM.from_pretrained(
-                    "sarvamai/sarvam-translate",
-                    torch_dtype=TORCH_DTYPE,
-                    device_map=DEVICE_MAP,
-                    token=HF_TOKEN,
-                    low_cpu_mem_usage=True,
-                    trust_remote_code=True
-                )
-                self.sarvam_tokenizer = AutoTokenizer.from_pretrained(
-                    "sarvamai/sarvam-translate",
-                    trust_remote_code=True
-                )
-                # Enable optimizations
-                if hasattr(self.sarvam_model, 'eval'):
-                    self.sarvam_model.eval()
-                if torch.cuda.is_available():
-                    torch.cuda.empty_cache()
-            except Exception as e:
-                print(f"Error loading Sarvam model: {e}")
-                
-    def get_model_and_tokenizer(self, model_type):
-        if model_type == "indictrans":
-            if self.indictrans_model is None:
-                self.load_indictrans_model()
-            return self.indictrans_model, self.indictrans_tokenizer
-        else:  # sarvam
-            if self.sarvam_model is None:
-                self.load_sarvam_model()
-            return self.sarvam_model, self.sarvam_tokenizer
+indictrans_model = AutoModelForCausalLM.from_pretrained(
+    "ai4bharat/IndicTrans3-beta",
+    torch_dtype=TORCH_DTYPE,
+    device_map=DEVICE_MAP,
+    token=HF_TOKEN,
+    low_cpu_mem_usage=True,
+    trust_remote_code=True
+)
+
+sarvam_model = AutoModelForCausalLM.from_pretrained(
+    "sarvamai/sarvam-translate",
+    torch_dtype=TORCH_DTYPE,
+    device_map=DEVICE_MAP,
+    token=HF_TOKEN,
+    low_cpu_mem_usage=True,
+    trust_remote_code=True
+)
+
+tokenizer = AutoTokenizer.from_pretrained(
+    "ai4bharat/IndicTrans3-beta",
+    trust_remote_code=True
+)
 
-# Global model manager
-model_manager = ModelManager()
 
 def format_message_for_translation(message, target_lang):
     return f"Translate the following text to {target_lang}: {message}"
@@ -175,7 +132,10 @@ def translate_message(
     model_type: str = "indictrans"
 ) -> Iterator[str]:
     
-    model, tokenizer = model_manager.get_model_and_tokenizer(model_type)
+    if model_type == "indictrans":
+        model = indictrans_model
+    elif model_type == "sarvam":
+        model = sarvam_model
     
     if model is None or tokenizer is None:
         yield "Error: Model failed to load. Please try again."