merged: zerogpu doesnt like strangers

app.py

@@ -7,12 +7,13 @@ from gradio_imageslider import ImageSlider
 from huggingface_hub import hf_hub_download
 from controlnet_union import ControlNetModel_Union
 from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
+from gradio_image_prompter import ImagePrompter
 from PIL import Image, ImageDraw
 import numpy as np
-from sam2_mask import create_sam2_mask_interface
-
-#from sam2.sam2_image_predictor import SAM2ImagePredictor
+from sam2.sam2_image_predictor import SAM2ImagePredictor
 
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+SAM_MODEL = "facebook/sam2.1-hiera-large"
 
 MODELS = {
     "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
@@ -61,6 +62,38 @@ def load_default_pipeline():
     ).to("cuda")
     return gr.update(value="Default pipeline loaded!")
 
+@spaces.GPU()
+def predict_masks(image, points):
+    """Predict a single mask from the image based on selected points."""
+    image_np = np.array(image)
+    points_list = [[point["x"], point["y"]] for point in points]
+    input_labels = [1] * len(points_list)
+
+    with torch.inference_mode():
+        PREDICTOR.set_image(image_np)
+        masks, _, _ = PREDICTOR.predict(
+            point_coords=points_list, point_labels=input_labels, multimask_output=False
+        )
+
+    # Prepare the overlay image
+    red_mask = np.zeros_like(image_np)
+    if masks and len(masks) > 0:
+        red_mask[:, :, 0] = masks[0].astype(np.uint8) * 255  # Apply the red channel
+        red_mask = PILImage.fromarray(red_mask)
+        original_image = PILImage.fromarray(image_np)
+        blended_image = PILImage.blend(original_image, red_mask, alpha=0.5)
+        return np.array(blended_image)
+    else:
+        return image_np
+
+
+def update_mask(prompts):
+    """Update the mask based on the prompts."""
+    image = prompts["image"]
+    points = prompts["points"]
+    return predict_masks(image, points)
+
+
 @spaces.GPU(duration=12)
 def fill_image(prompt, image, model_selection, paste_back):
     print(f"Received image: {image}")
@@ -489,7 +522,25 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                         history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
                         preview_image = gr.Image(label="Preview")
         with gr.TabItem("SAM2 Mask"):
-            sam2_mask_tab
+            gr.Markdown("# Object Segmentation with SAM2")
+            gr.Markdown(
+                """
+                This application utilizes **Segment Anything V2 (SAM2)** to allow you to upload an image and interactively generate a segmentation mask based on multiple points you select on the image.
+                """
+            )
+            with gr.Row():
+                with gr.Column():
+                    upload_image_input = ImagePrompter(show_label=False)
+                with gr.Column():
+                    image_output = gr.Image(label="Segmented Image", type="pil", height=400)
+            
+            # Define the action triggered by the upload_image_input change
+            upload_image_input.change(
+                fn=update_mask,
+                inputs=[upload_image_input],
+                outputs=[image_output],
+                show_progress=True,
+            )
         with gr.TabItem("Misc"):
             with gr.Column():
                 clear_cache_button = gr.Button("Clear CUDA Cache")