Update handler.py

handler.py

@@ -1,36 +1,38 @@
-import base64
-import requests
-from PIL import Image
-import warnings
-from typing import Dict, List, Any, Union
-import torch
-from io import BytesIO
-from transformers import BlipProcessor, BlipForConditionalGeneration, BitsAndBytesConfig
+    def __init__(self, path=""):
+        # load the optimized model
+        self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+        self.model = BlipForConditionalGeneration.from_pretrained(
+            "Salesforce/blip-image-captioning-large"
+        ).to(device)
+        self.model.eval()
+        self.model = self.model.to(device)
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Args:
+            data (dict):
+                Should contain:
+                - 'images': List[bytes] of images.
+                - 'texts': List[str] of associated texts. (Optional for unconditional captioning)
+        Return:
+            A dict with key "captions" and associated list of generated captions.
+        """
+        images = data.get("images")
+        texts = data.get("texts", ["a photography of"] * len(images))  # Default to "a photography of" if not provided
 
-class EndpointHandler():
-    def __init__(self, model_dir="Salesforce/blip-image-captioning-large"):
-        self.model = BlipForConditionalGeneration.from_pretrained(model_dir).to(device).eval()
-        self.processor = BlipProcessor.from_pretrained(model_dir)
+        raw_images = [Image.open(BytesIO(_img)).convert("RGB") for _img in images]
+
+        # Here, process both image and text
+        processed_inputs = [self.processor(img, txt, return_tensors="pt") for img, txt 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)
+        }
 
-    def __call__(self, data):
-        input_data = data['inputs'][0]
-        img_url = input_data.get('img_url')
-        text_prompt = input_data.get('text', None)
-        
-        raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
-        
-        if text_prompt:
-            inputs = self.processor(raw_image, text_prompt, return_tensors="pt").to(device)
-        else:
-            inputs = self.processor(raw_image, return_tensors="pt").to(device)
-        
         with torch.no_grad():
-            generated_ids = self.model.generate(
-                **inputs, 
-                max_new_tokens=150
-            )
-        captions = self.processor.decode(generated_ids[0], skip_special_tokens=True)
+            out = self.model.generate(**processed_inputs)
+
+        captions = self.processor.batch_decode(out, skip_special_tokens=True)
 
         return {"captions": captions}