Upload handler.py

handler.py

@@ -10,15 +10,10 @@ class EndpointHandler():
         """
         path: 
         """
-        # Preload all the elements you are going to need at inference.
-        # pseudo:
-        # self.model= load_model(path)
-        #self.processor = Blip2Processor.from_pretrained(path)
-        #self.pipeline = pipeline(model = path)
-        self.path = path
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        #self.processor = Blip2Processor.from_pretrained(path)
-        #self.model = Blip2Model.from_pretrained(path, torch_dtype=torch.float16)
+        self.processor = Blip2Processor.from_pretrained(path)
+        self.model = Blip2Model.from_pretrained(path, torch_dtype=torch.float16)
+        self.model.to(self.device)
         
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -30,10 +25,8 @@ class EndpointHandler():
         """
         inputs = data.pop("inputs", data)
         image_url = inputs['image_url']
-        #image = Image.open(requests.get(image_url, stream=True).raw)
-        #processed_image = self.processor(images=image, return_tensors="pt").to(self.device, torch.float16)
-
-        #generated_ids = self.pipeline(**inputs)
-        #generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
-
-        return image_url, self.path, self.device
+        image = Image.open(requests.get(image_url, stream=True).raw)
+        processed_image = self.processor(images=image, return_tensors="pt").to(self.device, torch.float16)
+        generated_ids = self.model.generate(**processed_image)
+        generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+        return image_url, generated_text