updated app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from diffusers import StableDiffusionInpaintPipeline
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 import gradio as gr
 from PIL import Image
 import numpy as np
@@ -9,10 +9,17 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def load_model():
     """Load ControlNet Inpainting model."""
-    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    controlnet = ControlNetModel.from_pretrained(
         "lllyasviel/control_v11p_sd15_inpaint",
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    )
+    
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        controlnet=controlnet,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
     ).to(device)
+    
     return pipe
 
 pipe = load_model()
@@ -36,7 +43,7 @@ def inpaint(image_editor_output, prompt):
         mask = mask.resize(image.size, Image.LANCZOS)
     
     # Perform inpainting
-    result = pipe(prompt=prompt, image=image, mask_image=mask).images[0]
+    result = pipe(prompt=prompt, image=image, control_image=mask).images[0]
     return result
 
 # Gradio UI with ImageEditor for inpainting