Update handler.py

handler.py

@@ -2,6 +2,8 @@ from transformers import AutoModelForCausalLM, AutoProcessor
 from PIL import Image
 import requests
 import torch
+from urllib3.util.retry import Retry
+from requests.adapters import HTTPAdapter
 
 class EndpointHandler:
     def __init__(self, model_dir):
@@ -15,6 +17,18 @@ class EndpointHandler:
             trust_remote_code=True
         )
         self.device = device
+        self.session = self._create_session()
+
+    def _create_session(self):
+        session = requests.Session()
+        retries = Retry(
+            total=3,
+            backoff_factor=0.5,
+            status_forcelist=[429, 500, 502, 503, 504]
+        )
+        session.mount('http://', HTTPAdapter(max_retries=retries))
+        session.mount('https://', HTTPAdapter(max_retries=retries))
+        return session
 
     def __call__(self, data):
         try:
@@ -22,24 +36,28 @@ class EndpointHandler:
             params = data.get("parameters", {})
             
             image_url = inputs_data.get("url")
-            if not image_url or not image_url.startswith("http"):
-                raise ValueError("Invalid or missing 'url' field")
+            if not image_url:
+                return {"error": "Missing URL in inputs"}
                 
             image = self.load_image(image_url)
+            if not image:
+                return {"error": "Failed to load image"}
+
             model_inputs = self.processor(
                 text=inputs_data.get("task_prompt", "<MORE_DETAILED_CAPTION>"),
                 images=image,
                 return_tensors="pt"
             ).to(self.device)
 
-            generated_ids = self.model.generate(
-                input_ids=model_inputs["input_ids"],
-                pixel_values=model_inputs["pixel_values"],
-                max_new_tokens=params.get("max_new_tokens", 512),
-                num_beams=params.get("num_beams", 3),
-                early_stopping=params.get("early_stopping", True),
-                do_sample=params.get("do_sample", False)
-            )
+            with torch.inference_mode():
+                generated_ids = self.model.generate(
+                    input_ids=model_inputs["input_ids"],
+                    pixel_values=model_inputs["pixel_values"],
+                    max_new_tokens=params.get("max_new_tokens", 512),
+                    num_beams=params.get("num_beams", 3),
+                    early_stopping=params.get("early_stopping", True),
+                    do_sample=params.get("do_sample", False)
+                )
 
             generated_text = self.processor.batch_decode(
                 generated_ids,
@@ -49,12 +67,23 @@ class EndpointHandler:
             return {"caption": generated_text}
             
         except Exception as e:
-            return {"error": str(e)}
+            return {"error": f"Processing error: {str(e)}"}
 
     def load_image(self, image_url):
         try:
-            response = requests.get(image_url, stream=True)
+            headers = {
+                "User-Agent": "Mozilla/5.0",
+                "Accept": "image/jpeg,image/png,image/*",
+                "Referer": image_url
+            }
+            response = self.session.get(
+                image_url,
+                stream=True,
+                headers=headers,
+                timeout=15,
+                verify=False  # Added for SSL issues
+            )
             response.raise_for_status()
             return Image.open(response.raw).convert("RGB")
         except Exception as e:
-            raise ValueError(f"Failed to load image: {str(e)}")
+            return None