modify handler to receive a list of url instead of image bytes

handler.py

@@ -2,15 +2,16 @@ import requests
 from typing import Dict, Any
 from PIL import Image
 import torch
-import base64
 from io import BytesIO
 from transformers import BlipForConditionalGeneration, BlipProcessor
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+
 class EndpointHandler():
     def __init__(self, path=""):
-        self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+        self.processor = BlipProcessor.from_pretrained(
+            "Salesforce/blip-image-captioning-large")
         self.model = BlipForConditionalGeneration.from_pretrained(
             "Salesforce/blip-image-captioning-large"
         ).to(device)
@@ -18,29 +19,48 @@ class EndpointHandler():
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         input_data = data.get("inputs", {})
-        encoded_images = input_data.get("images")
-        
-        if not encoded_images:
-            return {"captions": [], "error": "No images provided"} 
+        image_urls = input_data.get("image_urls", [])
+
+        if not image_urls:
+            return {"captions": [], "error": "No images provided"}
+
+        texts = input_data.get(
+            "texts", [""] * len(image_urls))
+
+        if len(image_urls) != len(texts):
+            return {
+                "captions": [],
+                "error": "Texts and images should have the same length"
+            }
 
-        texts = input_data.get("texts", ["a photography of"] * len(encoded_images)) 
+        images_data = [requests.get(url).content for url in image_urls]
 
         try:
-            raw_images = [Image.open(BytesIO(base64.b64decode(img))).convert("RGB") for img in encoded_images]
+            raw_images = [
+                Image.open(BytesIO((img))).convert("RGB")
+                for img in images_data]
             processed_inputs = [
-                self.processor(image, text, return_tensors="pt") for image, text in zip(raw_images, texts)
+                self.processor(image, text, return_tensors="pt")
+                for image, text in zip(raw_images, texts)
             ]
             processed_inputs = {
-                "pixel_values": torch.cat([inp["pixel_values"] for inp in processed_inputs], dim=0).to(device),
-                "input_ids": torch.cat([inp["input_ids"] for inp in processed_inputs], dim=0).to(device),
-                "attention_mask": torch.cat([inp["attention_mask"] for inp in processed_inputs], dim=0).to(device)
+                "pixel_values": torch.cat(
+                    [inp["pixel_values"]
+                     for inp in processed_inputs], dim=0).to(device),
+                "input_ids": torch.cat(
+                    [inp["input_ids"]
+                     for inp in processed_inputs], dim=0).to(device),
+                "attention_mask": torch.cat(
+                    [inp["attention_mask"]
+                     for inp in processed_inputs], dim=0).to(device)
             }
 
             with torch.no_grad():
                 out = self.model.generate(**processed_inputs)
 
-            captions = self.processor.batch_decode(out, skip_special_tokens=True)
+            captions = self.processor.batch_decode(
+                out, skip_special_tokens=True)
             return {"captions": captions}
         except Exception as e:
             print(f"Error during processing: {str(e)}")
-            return {"captions": [], "error": str(e)}
+            return {"captions": [], "error": str(e)}