Update handler.py

handler.py

@@ -27,19 +27,23 @@ class EndpointHandler():
         
         # load StableDiffusionInpaintPipeline pipeline
         self.pipe = AutoPipelineForInpainting.from_pretrained(
-            "runwayml/stable-diffusion-inpainting",
-            revision="fp16",
+            "kandinsky-community/kandinsky-2-2-decoder-inpaint",
             torch_dtype=torch.float16,
         )
         # use DPMSolverMultistepScheduler
-        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
-        # move to device        
+        # self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
+
+        self.pipe.enable_model_cpu_offload()
+
+        self.pipe.enable_xformers_memory_efficient_attention()
+        
+        # move to device
         self.pipe = self.pipe.to(device)
 
-        self.pipe2 = AutoPipelineForInpainting.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-        self.pipe2.to("cuda")
+        # self.pipe2 = AutoPipelineForInpainting.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
+        # self.pipe2.to("cuda")
 
-        self.pipe3 = AutoPipelineForImage2Image.from_pipe(self.pipe2)
+        # self.pipe3 = AutoPipelineForImage2Image.from_pipe(self.pipe2)
         
 
 
@@ -92,8 +96,6 @@ class EndpointHandler():
         """
 
         #pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to("cuda")
-
-        self.pipe.enable_xformers_memory_efficient_attention()
         
         # run inference pipeline
         out = self.pipe(prompt=prompt, negative_prompt=negative_prompt, image=image, mask_image=mask_image)
@@ -103,7 +105,7 @@ class EndpointHandler():
         image = out.images[0].resize((1024, 1024))
 
         print("image resizing successful!")
-
+"""
         self.pipe2.enable_xformers_memory_efficient_attention()
 
         image = self.pipe2(
@@ -130,10 +132,10 @@ class EndpointHandler():
         ).images[0]
 
         print("3rd pipeline part successful!")
-        
+        """
             
         # return first generate PIL image
-        return image2
+        return image
     
     # helper to decode input image
     def decode_base64_image(self, image_string):