Update app.py

app.py

@@ -1,26 +1,26 @@
 import gradio as gr
 import torch
 import numpy as np
-from diffusers import DiffusionPipeline
+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 realistic human generation model
-pipe = DiffusionPipeline.from_pretrained(
-    "SG161222/RealVisXL_V4.0",
+# Load the Stable Diffusion XL Inpainting model
+pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
     torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    use_safetensors=True
+    variant="fp16",
 ).to(device)
 
-# Load DETR for human detection
+# Load the DETR object detection model
 processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 detector = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50").to(device)
 
 @spaces.GPU
-def detect_and_replace(input_image, prompt, negative_prompt=""):
+def detect_and_replace_humans(input_image, prompt):
     if input_image is None or prompt == "":
         return None
 
@@ -35,56 +35,51 @@ def detect_and_replace(input_image, prompt, negative_prompt=""):
     mask = Image.new("L", input_image.size, 0)
     draw = ImageDraw.Draw(mask)
 
-    boxes = []
+    found = False
     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()]
-            boxes.append(box)
             draw.rectangle(box, fill=255)
+            found = True
 
-    if not boxes:
+    if not found:
         return "No human detected."
 
-    output_image = input_image.copy()
-
-    for idx, box in enumerate(boxes):
-        x1, y1, x2, y2 = box
-        width, height = x2 - x1, y2 - y1
-
-        # Generate one realistic human image per person
-        generated_image = pipe(
-            prompt=f"{prompt}, full body, plain background, isolated subject",
-            negative_prompt=negative_prompt,
-            width=512,
-            height=768,
-            guidance_scale=7.5,
-            num_inference_steps=30,
-            output_type="pil"
-        ).images[0]
-
-        # Crop the subject to avoid white borders
-        cropped_generated = generated_image.crop(generated_image.getbbox())
-
-        # Resize generated image to fit detected box
-        resized_generated = cropped_generated.resize((width, height))
-
-        # Paste the resized image at the correct location
-        output_image.paste(resized_generated, (x1, y1), mask=None)  # You can add mask for soft edges later
-
-    return output_image
+    # Pre-defined positive and negative prompts
+    positive_prompt = (
+        "Replace the masked humans with imaginary Indian bride and groom wearing traditional Indian wedding attire, "
+        "with detailed embroidery, colorful saree and sherwani, realistic faces, natural skin texture, matching pose, "
+        "perfect lighting, and the same camera perspective. Keep the background unchanged."
+    )
+
+    negative_prompt = (
+        "blurry, distorted, deformed, double face, extra limbs, low quality, bad proportions, low resolution, "
+        "changed background, multiple faces, duplicate body parts, cartoon, watermark, text"
+    )
+
+    # Inpainting process
+    output = pipe(
+        prompt=positive_prompt,
+        negative_prompt=negative_prompt,
+        image=input_image,
+        mask_image=mask,
+        num_inference_steps=40,
+        guidance_scale=8.5
+    ).images[0]
+
+    return output
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Replace Bride and Groom with Imaginary Realistic Characters")
+    gr.Markdown("## Replace Humans with Imaginary Indian Bride and Groom (Background Preserved)")
 
     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="Describe the imaginary bride/groom")
-    negative_prompt_text = gr.Textbox(label="Negative Prompt", placeholder="Optional negative prompt")
+    prompt_text = gr.Textbox(label="Prompt (Optional, Predefined Prompt Used)", placeholder="You can leave this blank")
     submit = gr.Button("Submit")
 
-    submit.click(detect_and_replace, inputs=[input_image, prompt_text, negative_prompt_text], outputs=output_image)
+    submit.click(detect_and_replace_humans, inputs=[input_image, prompt_text], outputs=output_image)
 
 demo.launch()