Update handler.py

handler.py

@@ -1,20 +1,25 @@
 from typing import Dict, List, Any
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import torch
 
-
-class PreTrainedPipeline():
+class EndpointHandler():
     def __init__(self, path=""):
       self.processor = TrOCRProcessor.from_pretrained(path)
       self.model = VisionEncoderDecoderModel.from_pretrained(path)
+
+      device = "cuda" if torch.cuda.is_available() else "cpu"
+      self.model.to(device)
     
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        image = data.pop("inputs", data)
+        inputs = data.pop("inputs", data)
+        image_input = inputs.get('image')
+
 
         # process image
-        pixel_values = self.processor(images=image, return_tensors="pt").pixel_values
+        pixel_values = self.processor(images=image_input, return_tensors="pt").pixel_values
 
         # run prediction
-        generated_ids = self.model.generate(pixel_values)
+        generated_ids = self.model.generate(pixel_values.to(device))
         
         # decode output
         prediction = generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)