Update app.py

app.py

@@ -1,16 +1,16 @@
 import gradio as gr
 import torch
 import numpy as np
-from diffusers import StableDiffusionInpaintPipeline
+from diffusers import StableDiffusionXLInpaintPipeline
 from PIL import Image, ImageDraw
 from transformers import DetrImageProcessor, DetrForObjectDetection
 import spaces
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the Stable Diffusion Inpainting model
-pipe = StableDiffusionInpaintPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-inpainting",
+# Load RealVisXL Inpainting model
+pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "SG161222/RealVisXL_V4.0_Inpainting",  # ✅ Use inpainting-specific model
     torch_dtype=torch.float16 if device == "cuda" else torch.float32
 ).to(device)
 
@@ -19,7 +19,7 @@ processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 detector = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50").to(device)
 
 @spaces.GPU
-def detect_and_remove(input_image, prompt):
+def detect_and_replace(input_image, prompt):
     if input_image is None or prompt == "":
         return None
 
@@ -34,7 +34,6 @@ def detect_and_remove(input_image, prompt):
     mask = Image.new("L", input_image.size, 0)
     draw = ImageDraw.Draw(mask)
 
-    # Draw boxes for "person" class only
     for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
         if detector.config.id2label[label.item()] == "person":
             box = [int(i) for i in box.tolist()]
@@ -43,26 +42,28 @@ def detect_and_remove(input_image, prompt):
     if np.array(mask).sum() == 0:
         return "No human detected."
 
-    # Inpainting
     output = pipe(
         prompt=prompt,
+        negative_prompt="low quality, blurry, extra limbs, bad anatomy, ugly, deformed, poorly drawn",
         image=input_image,
-        mask_image=mask
+        mask_image=mask,
+        guidance_scale=7.5,
+        num_inference_steps=30
     ).images[0]
 
     return output
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Automatic Human Removal and Inpainting")
+    gr.Markdown("## Realistic Human Replacement with SDXL")
 
     with gr.Row():
         input_image = gr.Image(type="pil", label="Input Image")
         output_image = gr.Image(type="pil", label="Output Image")
 
-    prompt_text = gr.Textbox(label="Prompt", placeholder="Example: Replace humans with cartoon background")
+    prompt_text = gr.Textbox(label="Prompt", placeholder="Example: Replace humans with fantasy characters")
     submit = gr.Button("Submit")
 
-    submit.click(detect_and_remove, inputs=[input_image, prompt_text], outputs=output_image)
+    submit.click(detect_and_replace, inputs=[input_image, prompt_text], outputs=output_image)
 
-demo.launch()
+demo.launch()