update handler.py

handler.py

@@ -1,54 +1,88 @@
-from typing import Dict, Any, List
 from PIL import Image
 import torch
+import base64
 from io import BytesIO
-from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
+from transformers import BlipForConditionalGeneration, BlipProcessor
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Source: https://www.philschmid.de/custom-inference-handler
+class EndpointHandler():
+    def __init__(self):
+        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()
 
+    def __call__(self, image_data: str) -> dict:
+        try:
+            raw_image = Image.open(BytesIO(base64.b64decode(image_data))).convert("RGB")
 
-class EndpointHandler:
-    def __init__(self, path="nlpconnect/vit-gpt2-image-captioning"):
-        self.model = VisionEncoderDecoderModel.from_pretrained(path)
+            processed_input = self.processor(raw_image, return_tensors="pt").to(device)
 
-        # Using ViTImageProcessor instead of ViTFeatureExtractor
-        self.feature_extractor = ViTImageProcessor.from_pretrained(path)
+            with torch.no_grad():
+                out = self.model.generate(**processed_input)
 
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
+            caption = self.processor.batch_decode(out, skip_special_tokens=True)[0]
+            return {"caption": caption}
+        except Exception as e:
+            print(f"Error during processing: {str(e)}")
+            return {"caption": "", "error": str(e)}
 
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.to(self.device)
 
-        self.max_length = 16
-        self.num_beams = 4
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        """
-        Args:
-            data (:obj:):
-                includes the input image data.
-        Return:
-            A :obj:`dict` with the caption.
-        """
-        image_bytes = data.get("inputs", None)
 
-        # Convert image bytes to PIL Image
-        image = Image.open(BytesIO(image_bytes))
-        if image.mode != "RGB":
-            image = image.convert(mode="RGB")
+# from typing import Dict, Any, List
+# from PIL import Image
+# import torch
+# from io import BytesIO
+# from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
+
+
+# # Source: https://www.philschmid.de/custom-inference-handler
+
+
+# class EndpointHandler:
+#     def __init__(self, path="nlpconnect/vit-gpt2-image-captioning"):
+#         self.model = VisionEncoderDecoderModel.from_pretrained(path)
+
+#         # Using ViTImageProcessor instead of ViTFeatureExtractor
+#         self.feature_extractor = ViTImageProcessor.from_pretrained(path)
+
+#         self.tokenizer = AutoTokenizer.from_pretrained(path)
 
-        pixel_values = self.feature_extractor(
-            images=image, return_tensors="pt"
-        ).pixel_values
-        pixel_values = pixel_values.to(self.device)
+#         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#         self.model.to(self.device)
+
+#         self.max_length = 16
+#         self.num_beams = 4
 
-        gen_kwargs = {"max_length": self.max_length, "num_beams": self.num_beams}
-        output_ids = self.model.generate(pixel_values, **gen_kwargs)
+#     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+#         """
+#         Args:
+#             data (:obj:):
+#                 includes the input image data.
+#         Return:
+#             A :obj:`dict` with the caption.
+#         """
+#         image_bytes = data.get("inputs", None)
 
-        caption = self.tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
+#         # Convert image bytes to PIL Image
+#         image = Image.open(BytesIO(image_bytes))
+#         if image.mode != "RGB":
+#             image = image.convert(mode="RGB")
 
-        return {"caption": caption}
+#         pixel_values = self.feature_extractor(
+#             images=image, return_tensors="pt"
+#         ).pixel_values
+#         pixel_values = pixel_values.to(self.device)
+
+#         gen_kwargs = {"max_length": self.max_length, "num_beams": self.num_beams}
+#         output_ids = self.model.generate(pixel_values, **gen_kwargs)
+
+#         caption = self.tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
+
+#         return {"caption": caption}
 
 
 # from typing import Dict, Any, List