Update app.py

app.py

@@ -1,93 +1,87 @@
 import gradio as gr
-import torch
 import numpy as np
-from PIL import Image
+import torch
 from diffusers import StableDiffusionInpaintPipeline
-from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection, pipeline as hf_pipeline
+from PIL import Image
 from segment_anything import sam_model_registry, SamPredictor
-import spaces
-# Device setup
+from huggingface_hub import hf_hub_download
+
+# Device configuration
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load zero-shot detector
-dino_id = "IDEA-Research/grounding-dino-tiny"
-dino_processor = AutoProcessor.from_pretrained(dino_id)
-dino_model = AutoModelForZeroShotObjectDetection.from_pretrained(dino_id).to(device)
+# Step 1: Download SAM model checkpoint from Hugging Face
+checkpoint_path = hf_hub_download(
+    repo_id="Munaf1987/sam",  # ✅ Your model repo
+    filename="sam_vit_h_4b8939.pth",  # ✅ The exact filename in your repo
+)
 
-# Load SAM
-sam_checkpoint = "Munaf1987/sam"
-sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint)
+# Step 2: Load SAM model
+model_type = "vit_h"
+sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
 sam.to(device)
 predictor = SamPredictor(sam)
 
-# Load the inpainting pipeline
+# Step 3: Load Stable Diffusion Inpainting Pipeline
+sta_diff_model = "stabilityai/stable-diffusion-2-inpainting"
 pipe = StableDiffusionInpaintPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-inpainting",
-    torch_dtype=torch.float16 if device=="cuda" else torch.float32
-).to(device)
-@spaces.GPU
-def detect_and_mask(image, prompt="a person"):
-    inputs = dino_processor(images=image, text=prompt, return_tensors="pt").to(device)
-    with torch.no_grad():
-        outputs = dino_model(**inputs)
-    results = dino_processor.post_process_grounded_object_detection(
-        outputs, inputs.input_ids, box_threshold=0.3, text_threshold=0.25,
-        target_sizes=[image.size[::-1]]
+    sta_diff_model,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+)
+pipe = pipe.to(device)
+
+# Global variable to store selected pixels
+selected_pixels = []
+
+
+def generate_mask(input_image, evt: gr.SelectData):
+    """Generate mask based on user-selected points."""
+    selected_pixels.append((evt.index[0], evt.index[1]))  # evt.index is (x, y)
+
+    predictor.set_image(np.array(input_image))
+    input_points = np.array(selected_pixels)
+    input_labels = np.ones(input_points.shape[0])
+
+    masks, _, _ = predictor.predict(
+        point_coords=input_points,
+        point_labels=input_labels,
+        multimask_output=False,
     )
-    boxes = results[0]["boxes"]
-    height, width = image.size[1], image.size[0]
-
-    if len(boxes) == 0:
-        return None, None, "No humans detected."
-
-    # Build mask from boxes
-    mask_full = Image.new("L", image.size, 0)
-    for box in boxes:
-        x1, y1, x2, y2 = map(int, box)
-        mask_full.paste(255, (x1, y1, x2, y2))
-
-    predictor.set_image(np.array(image))
-    transformed = predictor.transform.apply_boxes(boxes.cpu().numpy(), image.size[::-1])
-    sam_masks, _, _ = predictor.predict(boxes=transformed, multimask_output=False)
-    combined = np.zeros_like(sam_masks[0], dtype=np.uint8)
-    for m in sam_masks:
-        combined = np.maximum(combined, m.astype(np.uint8))
-
-    mask_image = Image.fromarray(combined * 255).convert("L")
-    return image, mask_image, "Mask ready."
-@spaces.GPU
-def inpaint_background(image, mask, prompt="background"):
-    orig_size = image.size
-    # Resize to inpainting model's resolution
-    img512 = image.resize((512,512), Image.LANCZOS)
-    m512 = mask.resize((512,512), Image.LANCZOS)
-    result = pipe(prompt=prompt, image=img512, mask_image=m512).images[0]
-    return result.resize(orig_size, Image.LANCZOS), "Background inpainted."
-@spaces.GPU
-def replace_with_cartoon(image, mask, prompt="a cartoon human in place"):
-    orig_size = image.size
-    img512 = image.resize((512,512), Image.LANCZOS)
-    m512 = mask.resize((512,512), Image.LANCZOS)
-    result = pipe(prompt=prompt, image=img512, mask_image=m512).images[0]
-    return result.resize(orig_size, Image.LANCZOS), "Replaced with cartoon."
+
+    mask = masks[0] * 255
+    mask_image = Image.fromarray(mask.astype(np.uint8)).convert("L")
+    return mask_image
+
+
+def inpaint(input_image, mask_image, prompt):
+    """Run the inpainting model."""
+    if input_image is None or mask_image is None or prompt == "":
+        return None
+
+    # ✅ Resize mask but keep the input image original size
+    mask_image_resized = mask_image.resize(input_image.size)
+
+    output = pipe(
+        prompt=prompt,
+        image=input_image,
+        mask_image=mask_image_resized
+    ).images[0]
+
+    return output
+
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("### Remove or Replace Humans with a Cartoon Character")
-
-    img = gr.Image(type="pil")
-    detect_prompt = gr.Textbox(label="Detection text prompt", value="a person")
-    detect_btn = gr.Button("Detect Humans")
-    mask_out = gr.Image(type="pil", label="Detected Mask")
-    status1 = gr.Textbox(interactive=False)
-
-    bg_btn = gr.Button("Remove Humans (Background Fill)")
-    cartoon_btn = gr.Button("Replace with Cartoon")
-    out_img = gr.Image(type="pil", label="Final Output")
-    status2 = gr.Textbox(interactive=False)
-
-    detect_btn.click(detect_and_mask, inputs=[img, detect_prompt], outputs=[img, mask_out, status1])
-    bg_btn.click(inpaint_background, inputs=[img, mask_out], outputs=[out_img, status2])
-    cartoon_btn.click(replace_with_cartoon, inputs=[img, mask_out], outputs=[out_img, status2])
+    gr.Markdown("## Stable Diffusion Inpainting with SAM Mask Selection")
+
+    with gr.Row():
+        input_image = gr.Image(type="pil", label="Input Image", interactive=True)
+        mask_display = gr.Image(type="pil", label="Generated Mask")
+        output_image = gr.Image(type="pil", label="Output Image")
+
+    prompt_text = gr.Textbox(label="Prompt", placeholder="Enter a prompt for inpainting")
+    submit = gr.Button("Submit")
+
+    input_image.select(generate_mask, inputs=input_image, outputs=mask_display)
+    submit.click(inpaint, inputs=[input_image, mask_display, prompt_text], outputs=output_image)
 
 demo.launch()