Update handler.py

handler.py

@@ -11,7 +11,7 @@ class EndpointHandler():
         model_id = path
         self.model = LlavaForConditionalGeneration.from_pretrained(
             model_id, 
-            torch_dtype=torch.float64, 
+            torch_dtype=torch.float16, 
             low_cpu_mem_usage=True,
             load_in_4bit=True
         )
@@ -30,7 +30,7 @@ class EndpointHandler():
                 response.raise_for_status()  # Raise an exception for 4xx or 5xx status codes
                 raw_image = Image.open(response.raw)
                 # Process image and generate output
-                inputs = self.processor(prompt, raw_image, return_tensors='pt').to(0, torch.float64)
+                inputs = self.processor(prompt, raw_image, return_tensors='pt').to(0, torch.float16)
                 output = self.model.generate(**inputs, max_new_tokens=200, do_sample=False)
                 readable = self.processor.decode(output[0][2:], skip_special_tokens=True)
                 outputs.append(readable)