Spaces:

lightly-ai
/

SAMv2-Mask-Generator

Running

SauravMaheshkar commited on Jul 30, 2024

Commit

a96b320

unverified ·

1 Parent(s): f102315

feat: specify CUDA

Files changed (1) hide show

app.py CHANGED Viewed

@@ -23,7 +23,7 @@ choice_mapping = {
 @spaces.GPU
 def predict(model_choice: str, annotations, image):
     config_file, ckpt_path = choice_mapping[str(model_choice)]
-    sam2_model = build_sam2(config_file, ckpt_path, device="cpu")
     predictor = SAM2ImagePredictor(sam2_model)
     predictor.set_image(image)
     coordinates = np.array(

 @spaces.GPU
 def predict(model_choice: str, annotations, image):
     config_file, ckpt_path = choice_mapping[str(model_choice)]
+    sam2_model = build_sam2(config_file, ckpt_path, device="cuda")
     predictor = SAM2ImagePredictor(sam2_model)
     predictor.set_image(image)
     coordinates = np.array(