change pipeline to AutoPipelineForInpainting

app.py

@@ -4,7 +4,7 @@ import mediapipe as mp
 import torch
 
 from PIL import Image
-from diffusers import AutoPipelineForImage2Image, DPMSolverMultistepScheduler
+from diffusers import AutoPipelineForInpainting, DPMSolverMultistepScheduler
 from mediapipe.tasks import python
 from mediapipe.tasks.python import vision
 from scipy.ndimage import binary_dilation
@@ -21,14 +21,14 @@ MASK_CATEGORY = segmenter.labels
 base_model = "SG161222/RealVisXL_V4.0"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-pipeline = AutoPipelineForImage2Image.from_pretrained(
+pipeline = AutoPipelineForInpainting.from_pretrained(
     base_model, torch_dtype=torch.float16, use_safetensors=True
 )
 pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
 pipeline.to(device)
 generator = torch.Generator(device).manual_seed(0)
 
-def image_to_image(input_image, mask_image, prompt, negative_prompt, category, guidance_scale, num_inference_steps):
+def image_to_image(input_image, mask_image, prompt, negative_prompt, guidance_scale, num_inference_steps):
     # Generate the output image
     output_image = pipeline(
         generator=generator,
@@ -42,7 +42,7 @@ def image_to_image(input_image, mask_image, prompt, negative_prompt, category, g
     
     return output_image
 
-def segment_image(input_image, prompt, negative_prompt, category, guidance_scale, num_inference_steps):
+def segment_image(input_image, category):
     image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(input_image))
     segmentation_result = segmenter.segment(image)
     category_mask = segmentation_result.category_mask
@@ -80,11 +80,11 @@ with gr.Blocks() as grApp:
     
     generate_btn.click(
         fn=segment_image,
-        inputs=[input_image, prompt, category, guidance_scale, num_inference_steps],
+        inputs=[input_image, category],
         outputs=[mask_image],
     ).then(
         fn=image_to_image,
-        inputs=[input_image, mask_image, prompt, negative_prompt, category, guidance_scale, num_inference_steps],
+        inputs=[input_image, mask_image, prompt, negative_prompt, guidance_scale, num_inference_steps],
         outputs=[output_image],
     )