Update handler.py

handler.py

@@ -7,10 +7,9 @@ from diffusers.image_processor import VaeImageProcessor
 
 class EndpointHandler:
     def __init__(self, path=""):
-        kwargs = {"torch_dtype": torch.float16}
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.dtype = kwargs["torch_dtype"]
-        self.vae = AutoencoderKL.from_pretrained(path, **kwargs).to(self.device).eval()
+        self.device = "cuda"
+        self.dtype = torch.float16
+        self.vae = AutoencoderKL.from_pretrained(path, torch_dtype=self.dtype).to(self.device, self.dtype).eval()
 
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
         self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)