Update handler.py

handler.py

@@ -1,9 +1,10 @@
 from typing import  Dict, List, Any
 import torch
-from diffusers import StableDiffusionXLImg2ImgPipeline
+from diffusers import StableDiffusionXLImg2ImgPipeline, DiffusionPipeline, AutoencoderKL
 from PIL import Image
 import base64
 from io import BytesIO
+from diffusers.image_processor import VaeImageProcessor
 
 
 # set device
@@ -31,12 +32,22 @@ class EndpointHandler():
         
         prompt = data.pop("prompt", "")
 
-
         if encoded_image is not None:
             image = self.decode_base64_image(encoded_image)
-            print(image.shape)
+
+            image_processor = VaeImageProcessor();
+            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
+            
             self.smooth_pipe.enable_xformers_memory_efficient_attention()
-            out = self.smooth_pipe(prompt, image=image).images[0]
+            out = self.smooth_pipe(prompt, image=latents_dist).images[0]
 
             return out