Update handler.py

handler.py

@@ -27,11 +27,12 @@ class EndpointHandler():
         
         # load StableDiffusionInpaintPipeline pipeline
         self.pipe = AutoPipelineForInpainting.from_pretrained(
-            "kandinsky-community/kandinsky-2-2-decoder-inpaint",
+            "runwayml/stable-diffusion-inpainting",
+            revision="fp16",
             torch_dtype=torch.float16,
         )
         # use DPMSolverMultistepScheduler
-        # self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
+        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
 
         self.pipe.enable_model_cpu_offload()
 
@@ -40,10 +41,15 @@ class EndpointHandler():
         # 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.enable_model_cpu_offload()
+        self.pipe2.enable_xformers_memory_efficient_attention()
+        
+        self.pipe2.to("cuda")
 
-        # self.pipe3 = AutoPipelineForImage2Image.from_pipe(self.pipe2)
+        self.pipe3 = AutoPipelineForImage2Image.from_pipe(self.pipe2)
+        self.pipe3.enable_model_cpu_offload()
+        self.pipe3.enable_xformers_memory_efficient_attention()
         
 
 
@@ -98,15 +104,13 @@ class EndpointHandler():
         #pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to("cuda")
         
         # run inference pipeline
-        out = self.pipe(prompt=prompt, negative_prompt=negative_prompt, image=image, mask_image=mask_image, strength=strength, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps)
+        out = self.pipe(prompt=prompt, negative_prompt=negative_prompt, image=image, mask_image=mask_image)
 
         print("1st pipeline part successful!")
 
         image = out.images[0].resize((1024, 1024))
 
         print("image resizing successful!")
-"""
-        self.pipe2.enable_xformers_memory_efficient_attention()
 
         image = self.pipe2(
             prompt=prompt,
@@ -121,8 +125,6 @@ class EndpointHandler():
 
         print("2nd pipeline part successful!")
         
-        self.pipe3.enable_xformers_memory_efficient_attention()
-        
         image2 = self.pipe3(
             prompt=prompt,
             image=image,
@@ -132,10 +134,9 @@ class EndpointHandler():
         ).images[0]
 
         print("3rd pipeline part successful!")
-        """
             
         # return first generate PIL image
-        return image
+        return image2
     
     # helper to decode input image
     def decode_base64_image(self, image_string):