Update handler.py

handler.py

@@ -64,6 +64,8 @@ class EndpointHandler:
         self.padding_token = 128263
         self.start_audio_token = 128257  # Start of Audio token
         self.end_audio_token = 128258  # End of Audio token
+
+        self._warmed_up = False
         
         logger.info("Handler initialization complete")
         
@@ -268,6 +270,10 @@ class EndpointHandler:
         """
         Main entry point for the handler
         """
+        # Run warmup only once, the first time __call__ is triggered
+        if not self._warmed_up:
+            self._warmup()
+        
         try:
             logger.info(f"Received request: {type(data)}")
             
@@ -286,3 +292,11 @@ class EndpointHandler:
             logger.error(traceback.format_exc())
             return {"error": str(e)}
 
+    def _warmup(self):
+        try:
+            dummy_prompt = "tara: Hello"
+            input_ids = self.tokenizer(dummy_prompt, return_tensors="pt").input_ids.to(self.device)
+            _ = self.model.generate(input_ids=input_ids, max_new_tokens=5)
+            self._warmed_up = True
+        except Exception as e:
+            print(f"[WARMUP ERROR] {str(e)}")