temp seg solution

app.py

@@ -10,10 +10,12 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from gradio_image_prompter import ImagePrompter
 from PIL import Image, ImageDraw
 import numpy as np
-from sam2.sam2_image_predictor import SAM2ImagePredictor
-from sam2_mask import create_sam2_tab
+# from sam2.sam2_image_predictor import SAM2ImagePredictor
+# from sam2_mask import create_sam2_tab
 import subprocess
+import os
 
+HF_TOKEN = os.getenv("HF_TOKEN")
 subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
 
 # class SAM2PredictorSingleton:
@@ -79,57 +81,57 @@ def load_default_pipeline():
     ).to("cuda")
     return gr.update(value="Default pipeline loaded!")
 
-@spaces.GPU()
-def predict_masks(prompts):
-
-    DEVICE = torch.device("cuda")
-    SAM_MODEL = "facebook/sam2.1-hiera-large"
-    # if PREDICTOR is None:
-    #     PREDICTOR = SAM2ImagePredictor.from_pretrained(SAM_MODEL, device=DEVICE)
-    # else:
-    #     PREDICTOR = PREDICTOR
-    PREDICTOR = SAM2ImagePredictor.from_pretrained(SAM_MODEL, device=DEVICE)
-
-    """Predict a single mask from the image based on selected points."""
-    image = np.array(prompts["image"])  # Convert the image to a numpy array
-    points = prompts["points"]  # Get the points from prompts
-
-    if not points:
-        return image  # Return the original image if no points are selected
-
-    # Debugging: Print the structure of points
-    print(f"Points structure: {points}")
-
-    # Ensure points is a list of lists with at least two elements
-    if isinstance(points, list) and all(isinstance(point, list) and len(point) >= 2 for point in points):
-        input_points = [[point[0], point[1]] for point in points]
-    else:
-        return image  # Return the original image if points structure is unexpected
-
-    input_labels = [1] * len(input_points)
-
-    with torch.inference_mode():
-        PREDICTOR.set_image(image)
-        masks, _, _ = PREDICTOR.predict(
-            point_coords=input_points, point_labels=input_labels, multimask_output=False
-        )
-
-    # Prepare the overlay image
-    red_mask = np.zeros_like(image)
-    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)
-        blended_image = PILImage.blend(original_image, red_mask, alpha=0.5)
-        return np.array(blended_image)
-    else:
-        return image
-
-def update_mask(prompts):
-    """Update the mask based on the prompts."""
-    image = prompts["image"]
-    points = prompts["points"]
-    return predict_masks(image, points)
+# @spaces.GPU()
+# def predict_masks(prompts):
+
+#     DEVICE = torch.device("cuda")
+#     SAM_MODEL = "facebook/sam2.1-hiera-large"
+#     # if PREDICTOR is None:
+#     #     PREDICTOR = SAM2ImagePredictor.from_pretrained(SAM_MODEL, device=DEVICE)
+#     # else:
+#     #     PREDICTOR = PREDICTOR
+#     PREDICTOR = SAM2ImagePredictor.from_pretrained(SAM_MODEL, device=DEVICE)
+
+#     """Predict a single mask from the image based on selected points."""
+#     image = np.array(prompts["image"])  # Convert the image to a numpy array
+#     points = prompts["points"]  # Get the points from prompts
+
+#     if not points:
+#         return image  # Return the original image if no points are selected
+
+#     # Debugging: Print the structure of points
+#     print(f"Points structure: {points}")
+
+#     # Ensure points is a list of lists with at least two elements
+#     if isinstance(points, list) and all(isinstance(point, list) and len(point) >= 2 for point in points):
+#         input_points = [[point[0], point[1]] for point in points]
+#     else:
+#         return image  # Return the original image if points structure is unexpected
+
+#     input_labels = [1] * len(input_points)
+
+#     with torch.inference_mode():
+#         PREDICTOR.set_image(image)
+#         masks, _, _ = PREDICTOR.predict(
+#             point_coords=input_points, point_labels=input_labels, multimask_output=False
+#         )
+
+#     # Prepare the overlay image
+#     red_mask = np.zeros_like(image)
+#     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)
+#         blended_image = PILImage.blend(original_image, red_mask, alpha=0.5)
+#         return np.array(blended_image)
+#     else:
+#         return image
+
+# 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)
@@ -558,33 +560,35 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                         use_as_input_button_outpaint = gr.Button("Use as Input Image", visible=False)
                         history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
                         preview_image = gr.Image(label="Preview")
-        with gr.TabItem("SAM2 Masking"):
-            input_image, points_map, output_result_mask = create_sam2_tab()
-        with gr.TabItem("SAM2 Mask"):
-            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():
-                    image_input = gr.State()
-                    # Input: ImagePrompter for uploaded image
-                    upload_image_input = ImagePrompter(show_label=False)
-                with gr.Column():
-                    image_output = gr.Image(label="Segmented Image", type="pil", height=400)
-                with gr.Row():
-                # Button to trigger the prediction
-                    predict_button = gr.Button("Predict Mask")
+        # with gr.TabItem("SAM2 Masking"):
+        #     input_image, points_map, output_result_mask = create_sam2_tab()
+        # with gr.TabItem("SAM2 Mask"):
+        #     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():
+        #             image_input = gr.State()
+        #             # Input: ImagePrompter for uploaded image
+        #             upload_image_input = ImagePrompter(show_label=False)
+        #         with gr.Column():
+        #             image_output = gr.Image(label="Segmented Image", type="pil", height=400)
+        #         with gr.Row():
+        #         # Button to trigger the prediction
+        #             predict_button = gr.Button("Predict Mask")
                 
-                # Define the action triggered by the predict button
-                predict_button.click(
-                    fn=predict_masks,
-                    inputs=[upload_image_input],
-                    outputs=[image_output],
-                    show_progress=True,
-                )
+        #         # Define the action triggered by the predict button
+        #         predict_button.click(
+        #             fn=predict_masks,
+        #             inputs=[upload_image_input],
+        #             outputs=[image_output],
+        #             show_progress=True,
+        #         )
+        with gr.Tab("SAM2.1 Segmented Mask"):
+            temp_space = gr.load("LPX55/SAM2-Image-Predictor-CPU", src="spaces", token=HF_TOKEN)
             # Define the action triggered by the upload_image_input change
             # upload_image_input.change(
             #     fn=update_mask,