updated handler

handler.py

@@ -1,29 +1,34 @@
 import io
 import torch
+import requests
 from PIL import Image
 from transformers import BlipProcessor, BlipForConditionalGeneration
 
 class EndpointHandler:
     def __init__(self, model_dir):
-        # Load the processor and model from the specified directory
         self.processor = BlipProcessor.from_pretrained(model_dir)
         self.model = BlipForConditionalGeneration.from_pretrained(model_dir)
         self.model.eval()
-        # Move model to GPU if available
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model.to(self.device)
 
     def __call__(self, data):
-        # Read the image from the incoming request
-        image_bytes = data
-        image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
+        # Expecting: { "image_url": "https://example.com/image.jpg" }
+        if isinstance(data, dict) and "image_url" in data:
+            image_url = data["image_url"]
+            try:
+                response = requests.get(image_url)
+                response.raise_for_status()
+                image = Image.open(io.BytesIO(response.content)).convert("RGB")
+            except Exception as e:
+                return {"error": f"Failed to load image from URL: {str(e)}"}
+        else:
+            return {"error": "Please provide an 'image_url' in the JSON payload."}
 
-        # Preprocess the image
         inputs = self.processor(images=image, return_tensors="pt").to(self.device)
 
-        # Generate caption
         with torch.no_grad():
             out = self.model.generate(**inputs)
-        caption = self.processor.decode(out[0], skip_special_tokens=True)
 
+        caption = self.processor.decode(out[0], skip_special_tokens=True)
         return {"caption": caption}