Update handler.py

handler.py

@@ -20,6 +20,10 @@ class EndpointHandler():
           "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16
         )
         self.smooth_pipe.to("cuda")
+
+        self.vae = AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
+                                                 subfolder="vae",  use_safetensors=True,
+                                                 ).to("cuda")
         
 
 
@@ -39,12 +43,8 @@ class EndpointHandler():
             latents = image_processor.preprocess(image)
             latents = latents.to(device="cuda")
             
-            vae =  AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
-                                                 subfolder="vae",  use_safetensors=True,
-                                                 ).to("cuda")
-            
             with torch.no_grad():
-                latents_dist = vae.encode(latents).latent_dist.sample() * vae.config.scaling_factor
+                latents_dist = self.vae.encode(latents).latent_dist.sample() * vae.config.scaling_factor
             
             self.smooth_pipe.enable_xformers_memory_efficient_attention()
             out = self.smooth_pipe(prompt, image=latents_dist).images