Update app.py

app.py

@@ -124,15 +124,16 @@ def rmbg(image=None, url=None):
 
 def mask_generation(image=None, d=None):
     d = eval(d)  # convert this to dictionary
-    predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2.1-hiera-large")
-    predictor.set_image(image)
-    input_point = np.array(d["input_points"])
-    input_label = np.array(d["input_labels"])
-    masks, scores, logits = predictor.predict(
-        point_coords=input_point,
-        point_labels=input_label,
-        multimask_output=True,
-    )
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2.1-hiera-large")
+        predictor.set_image(image)
+        input_point = np.array(d["input_points"])
+        input_label = np.array(d["input_labels"])
+        masks, scores, logits = predictor.predict(
+            point_coords=input_point,
+            point_labels=input_label,
+            multimask_output=True,
+        )
     sorted_ind = np.argsort(scores)[::-1]
     masks = masks[sorted_ind]
     scores = scores[sorted_ind]