Update carebridge_client.py

carebridge_client.py

@@ -8,63 +8,66 @@ from gtts import gTTS
 import tempfile
 
 class CareBridgeTranslator:
-    def __init__(self, model_id="google/translategemma-4b-it", device=None):
-        self.model_id = model_id
-        if device is None:
-            self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        else:
-            self.device = device
-        print(f"[SIMBOTI] Loading model {model_id}...")
-        self.processor = AutoProcessor.from_pretrained(model_id)
-        self.model = AutoModelForImageTextToText.from_pretrained(model_id, device_map=self.device, torch_dtype=torch.float16 if self.device == "cuda" else torch.float32)
-        print("[SIMBOTI] Model loaded successfully.")
-        self.LANG_MAP = {"English": "en", "Polish": "pl", "Romanian": "ro", "Punjabi": "pa", "Urdu": "ur", "Portuguese": "pt", "Spanish": "es", "Arabic": "ar", "Bengali": "bn", "Gujarati": "gu", "Italian": "it"}
+        def __init__(self, model_id="google/translategemma-4b-it"):
+                    self.model_id = model_id
+                    self.model = None
+                    self.processor = None
+                    self.LANG_MAP = {"English": "en", "Polish": "pl", "Romanian": "ro", "Punjabi": "pa", "Urdu": "ur", "Portuguese": "pt", "Spanish": "es", "Arabic": "ar", "Bengali": "bn", "Gujarati": "gu", "Italian": "it"}
+                    print("[SIMBOTI] Translator initialized. Model will load on first use.")
 
-    def translate_text(self, text, source_lang_name, target_lang_name):
-        src_code = self.LANG_MAP.get(source_lang_name)
-        tgt_code = self.LANG_MAP.get(target_lang_name)
-        if not src_code or not tgt_code:
-            return f"Error: Language not supported."
-        message = {"role": "user", "content": [{"type": "text", "source_lang_code": src_code, "target_lang_code": tgt_code, "text": text}]}
+        def _load_model(self):
+                    if self.model is None:
+                                    print(f"[SIMBOTI] Loading model {self.model_id}...")
+                                    self.processor = AutoProcessor.from_pretrained(self.model_id)
+                                    self.model = AutoModelForImageTextToText.from_pretrained(self.model_id, device_map="cuda", torch_dtype=torch.float16)
+                                    print("[SIMBOTI] Model loaded successfully.")
+
+                def translate_text(self, text, source_lang_name, target_lang_name):
+                            src_code = self.LANG_MAP.get(source_lang_name)
+                            tgt_code = self.LANG_MAP.get(target_lang_name)
+                            if not src_code or not tgt_code:
+                                            return "Error: Language not supported."
+                                        message = {"role": "user", "content": [{"type": "text", "source_lang_code": src_code, "target_lang_code": tgt_code, "text": text}]}
         return self._run_inference([message])
 
     def translate_image(self, image_path, source_lang_name, target_lang_name):
-        src_code = self.LANG_MAP.get(source_lang_name)
+                src_code = self.LANG_MAP.get(source_lang_name)
         tgt_code = self.LANG_MAP.get(target_lang_name)
         if not src_code or not tgt_code:
-            return "Error: Language not supported."
+                        return "Error: Language not supported."
         if isinstance(image_path, str):
-            image = Image.open(image_path)
-        else:
+                        image = Image.open(image_path)
+else:
             image = image_path
         message = {"role": "user", "content": [{"type": "image", "source_lang_code": src_code, "target_lang_code": tgt_code, "image": image}]}
         return self._run_inference([message])
 
     def translate_audio(self, audio_path, source_lang_name, target_lang_name):
-        src_code = self.LANG_MAP.get(source_lang_name)
+                src_code = self.LANG_MAP.get(source_lang_name)
         tgt_code = self.LANG_MAP.get(target_lang_name)
         if not src_code or not tgt_code:
-            return "Error: Language not supported."
+                        return "Error: Language not supported."
         audio, sr = librosa.load(audio_path, sr=16000)
         message = {"role": "user", "content": [{"type": "audio", "source_lang_code": src_code, "target_lang_code": tgt_code, "audio": audio}]}
         return self._run_inference([message])
 
     def speak_text(self, text, lang_name):
-        lang_code = self.LANG_MAP.get(lang_name, "en")
+                lang_code = self.LANG_MAP.get(lang_name, "en")
         try:
-            tts = gTTS(text=text, lang=lang_code)
+                        tts = gTTS(text=text, lang=lang_code)
             temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
             tts.save(temp_file.name)
             return temp_file.name
-        except Exception as e:
+except Exception as e:
             print(f"TTS Error: {e}")
             return None
 
     @spaces.GPU()
     def _run_inference(self, messages):
-        inputs = self.processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt").to(self.device)
+                self._load_model()
+        inputs = self.processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt").to("cuda")
         with torch.no_grad():
-            outputs = self.model.generate(**inputs, max_new_tokens=512, do_sample=False)
+                        outputs = self.model.generate(**inputs, max_new_tokens=512, do_sample=False)
         input_len = inputs["input_ids"].shape[-1]
         decoded = self.processor.decode(outputs[0][input_len:], skip_special_tokens=True)
         return decoded.strip()