Update app.py

app.py

@@ -1,25 +1,26 @@
 import gradio as gr
 import torch
 import numpy as np
-from diffusers import StableDiffusionXLInpaintPipeline
+from diffusers import DiffusionPipeline
 from PIL import Image, ImageDraw
 from transformers import DetrImageProcessor, DetrForObjectDetection
 import spaces
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load RealVisXL Inpainting model
-pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
-    "SG161222/RealVisXL_V4.0",  # ✅ Use inpainting-specific model
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+# Load Inpainting Pipeline
+pipe = DiffusionPipeline.from_pretrained(
+    "SG161222/RealVisXL_V4.0",  # ✅ Realistic human generation model
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    use_safetensors=True
 ).to(device)
 
-# Load the DETR object detection model
+# Load DETR for human detection
 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):
+def detect_and_replace(input_image, prompt, negative_prompt=""):
     if input_image is None or prompt == "":
         return None
 
@@ -34,36 +35,53 @@ def detect_and_replace(input_image, prompt):
     mask = Image.new("L", input_image.size, 0)
     draw = ImageDraw.Draw(mask)
 
+    boxes = []
     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)
 
-    if np.array(mask).sum() == 0:
+    if not boxes:
         return "No human detected."
 
-    output = pipe(
-        prompt=prompt,
-        negative_prompt="low quality, blurry, extra limbs, bad anatomy, ugly, deformed, poorly drawn",
-        image=input_image,
-        mask_image=mask,
-        guidance_scale=7.5,
-        num_inference_steps=30
-    ).images[0]
+    output_image = input_image.copy()
 
-    return output
+    for box in boxes:
+        x1, y1, x2, y2 = box
+        width, height = x2 - x1, y2 - y1
+
+        # Generate imaginary person image
+        generated_image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            width=512,
+            height=768,
+            guidance_scale=7.5,
+            num_inference_steps=30,
+            output_type="pil"
+        ).images[0]
+
+        # Resize generated image to fit the detected box
+        resized_generated = generated_image.resize((width, height))
+
+        # Paste the generated image on the original image at the detected location
+        output_image.paste(resized_generated, (x1, y1))
+
+    return output_image
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Realistic Human Replacement with SDXL")
+    gr.Markdown("## Replace Bride and Groom with Imaginary Realistic Characters")
 
     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 fantasy characters")
+    prompt_text = gr.Textbox(label="Prompt", placeholder="Describe the imaginary bride/groom")
+    negative_prompt_text = gr.Textbox(label="Negative Prompt", placeholder="Optional negative prompt")
     submit = gr.Button("Submit")
 
-    submit.click(detect_and_replace, inputs=[input_image, prompt_text], outputs=output_image)
+    submit.click(detect_and_replace, inputs=[input_image, prompt_text, negative_prompt_text], outputs=output_image)
 
 demo.launch()