Updated handler.py

handler.py

@@ -32,15 +32,20 @@ class EndpointHandler:
         parameters = data.pop("parameters", {})
  
         raw_images = inputs
+
+        inputs = self.processor(inputs, return_tensors="pt").to("cuda")
+
+        out = self.model.generate(**inputs)
                                      
         processed_image = self.processor(images=raw_images, return_tensors="pt").to(device)
-        processed_image["pixel_values"] = processed_image["pixel_values"].to(device)
-        processed_image = {**processed_image, **parameters}
+        # processed_image["pixel_values"] = processed_image["pixel_values"].to(device)
+        # processed_image = {**processed_image, **parameters}
+        
+        # with torch.no_grad():
+        #     out = self.model.generate(
+        #         **processed_image
+        #     )
+        captions = self.processor.decode(out[0], skip_special_tokens=True)
         
-        with torch.no_grad():
-            out = self.model.generate(
-                **processed_image
-            )
-        captions = self.processor.batch_decode(out, skip_special_tokens=True)
         # postprocess the prediction
         return {"captions": captions}