Update app.py

app.py

@@ -2,17 +2,23 @@ import torch
 from diffusers import StableDiffusionInpaintPipeline
 import gradio as gr
 from PIL import Image
+import numpy as np
 
 # Ensure we use GPU if available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the inpainting model
-pipe = StableDiffusionInpaintPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-inpainting",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-).to(device)
+def load_model():
+    """Load Stable Diffusion 2 Inpainting model."""
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-2-inpainting",
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    ).to(device)
+    return pipe
+
+pipe = load_model()
 
 def inpaint(image, mask, prompt):
+    """Perform inpainting based on user input."""
     image = Image.fromarray(image)
     mask = Image.fromarray(mask).convert("L")  # Convert mask to grayscale
     
@@ -20,17 +26,17 @@ def inpaint(image, mask, prompt):
     result = pipe(prompt=prompt, image=image, mask_image=mask).images[0]
     return result
 
-# Create Gradio UI
+# Create Gradio UI with an editable brush tool
 demo = gr.Interface(
     fn=inpaint,
     inputs=[
-        gr.Image(type="numpy", label="Upload Image"),
-        gr.Image(type="numpy", label="Upload Mask (White for areas to be inpainted)"),
+        gr.Image(type="numpy", tool="editor", label="Upload Image & Draw Mask"),
+        gr.Image(type="numpy", tool="sketch", label="Paint Mask"),
         gr.Textbox(label="Prompt for Inpainting"),
     ],
     outputs=gr.Image(type="pil"),
-    title="Stable Diffusion 2 Inpainting",
-    description="Upload an image and a mask, then provide a text prompt to guide the inpainting process."
+    title="Stable Diffusion 2 Inpainting with Brush Tool",
+    description="Upload an image, paint over areas to be inpainted, and provide a text prompt.",
 )
 
 demo.launch()