cuda

app.py

@@ -12,8 +12,21 @@ from PIL import Image, ImageDraw
 import numpy as np
 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"
+# class SAM2PredictorSingleton:
+#     _instance = None
+
+#     def __new__(cls):
+#         if cls._instance is None:
+#             cls._instance = super(SAM2PredictorSingleton, cls).__new__(cls)
+#             cls._instance._initialize_predictor()
+#         return cls._instance
+
+#     def _initialize_predictor(self):
+#         MODEL = "facebook/sam2-hiera-large"
+#         DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#         self.predictor = SAM2ImagePredictor.from_pretrained(MODEL, device=DEVICE)
+
+
 
 MODELS = {
     "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
@@ -51,6 +64,8 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 pipe.to("cuda")
 print(pipe)
+DEVICE = torch.device("cuda")
+SAM_MODEL = "facebook/sam2.1-hiera-large"
 PREDICTOR = SAM2ImagePredictor.from_pretrained(SAM_MODEL, device=DEVICE)
 
 def load_default_pipeline():
@@ -64,8 +79,11 @@ def load_default_pipeline():
     return gr.update(value="Default pipeline loaded!")
 
 @spaces.GPU()
-def predict_masks(image, points):
+def predict_masks(prompts):
     """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
 
@@ -74,29 +92,28 @@ def predict_masks(image, 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):
-        points_list = [[point[0], point[1]] 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(points_list)
+    input_labels = [1] * len(input_points)
 
     with torch.inference_mode():
-        PREDICTOR.set_image(np.array(image))
+        PREDICTOR.set_image(image)
         masks, _, _ = PREDICTOR.predict(
-            point_coords=points_list, point_labels=input_labels, multimask_output=False
+            point_coords=input_points, point_labels=input_labels, multimask_output=False
         )
 
     # Prepare the overlay image
-    image_np = np.array(image)
-    red_mask = np.zeros_like(image_np)
+    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_np)
+        original_image = PILImage.fromarray(image)
         blended_image = PILImage.blend(original_image, red_mask, alpha=0.5)
         return np.array(blended_image)
     else:
-        return image_np
+        return image
 
 def update_mask(prompts):
     """Update the mask based on the prompts."""