Update app.py

app.py

@@ -4,15 +4,14 @@ import numpy as np
 from PIL import Image
 from diffusers import StableDiffusionInpaintPipeline
 from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection, pipeline
-from transformers import SamProcessor, SamModel, pipeline as sam_pipeline
 import spaces
 
 @spaces.GPU
 def remove_object_with_text(input_image, prompt):
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    original_size = input_image.size
 
     # 1️⃣ Grounding DINO zero-shot detection
-    dino_id = "pengxian/grounding-dino"
     dino_processor = AutoProcessor.from_pretrained("IDEA-Research/grounding-dino-tiny")
     dino = AutoModelForZeroShotObjectDetection.from_pretrained("IDEA-Research/grounding-dino-tiny").to(device)
 
@@ -32,35 +31,27 @@ def remove_object_with_text(input_image, prompt):
         x1, y1, x2, y2 = [int(v) for v in box]
         mask.paste(255, (x1, y1, x2, y2))
 
-    # 2️⃣ SAM automatic mask refinement
-    sam_pipe = sam_pipeline("mask-generation", model="facebook/sam-vit-huge", device=0 if device=="cuda" else -1)
-    sam_out = sam_pipe(image=input_image, points_per_batch=256)
-    # Combine SAM masks that overlap the boxes
-    final_mask = Image.new("L", input_image.size, 0)
-    for m in sam_out["masks"]:
-        arr = np.array(m)
-        # apply only in box regions
-        for box in boxes:
-            x1,y1,x2,y2 = [int(v) for v in box]
-            sub = arr[y1:y2, x1:x2]
-            if sub.sum() > 1000:
-                final_mask.paste(Image.fromarray((arr*255).astype("uint8")), (0,0), Image.fromarray((arr*255).astype("uint8")))
-
-    # 3️⃣ Inpainting with Stable Diffusion
+    # 2️⃣ Inpainting with Stable Diffusion (high-res)
     pipe = StableDiffusionInpaintPipeline.from_pretrained(
         "stabilityai/stable-diffusion-2-inpainting",
         torch_dtype=torch.float16 if device == "cuda" else torch.float32
     ).to(device)
 
-    img_resized = input_image.resize((512,512))
-    mask_resized = final_mask.resize((512,512))
+    # Prepare for inpainting: resize to model's input size (512x512 or 768x768)
+    target_res = (768, 768)  # Higher resolution = better quality
+    img_resized = input_image.resize(target_res, Image.LANCZOS)
+    mask_resized = mask.resize(target_res, Image.LANCZOS)
 
     output = pipe(prompt="background", image=img_resized, mask_image=mask_resized).images[0]
-    return output, "Object removed."
+
+    # Resize back to original image size
+    final_output = output.resize(original_size, Image.LANCZOS)
+
+    return final_output, "Object removed and image size preserved."
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Text-Based Object Removal + Inpainting")
+    gr.Markdown("## Object Removal with Text + Original Size Preservation")
     inp = gr.Image(type="pil")
     txt = gr.Textbox(label="Describe object to remove", placeholder="e.g. a cat")
     btn = gr.Button("Remove")