Update handler.py

handler.py

@@ -13,7 +13,7 @@ from peft import PeftModel, PeftConfig
 
 # -
 
-# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 class EndpointHandler():
     def __init__(self, path=""):
@@ -50,12 +50,19 @@ class EndpointHandler():
         image_io = io.BytesIO(image_bytes)
         image = Image.open(image_io)
 
-        inputs = self.processor(images=image, return_tensors="pt")
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        inputs = self.processor(images=image, return_tensors="pt").to(device, torch.float16)
         pixel_values = inputs.pixel_values
         
-        generated_ids = self.model.generate(pixel_values=pixel_values, max_length=25)
+        generated_ids = self.model.generate(pixel_values=pixel_values, max_length=100)
         generated_caption = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        print(generated_caption)
+
+        # device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        # inputs = processor(images=image, return_tensors="pt").to(device, torch.float16)
+        # pixel_values = inputs.pixel_values
+        
+        # generated_ids = model.generate(pixel_values=pixel_values, max_length=100)
+        # generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
  
         # raw_images = [Image.open(BytesIO(_img)) for _img in inputs]