Spaces:

ybhavsar2009
/

woundview

Sleeping

ybhavsar2009 commited on Apr 13

Commit

1a0610b

verified ·

1 Parent(s): f034c52

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -72,7 +72,8 @@ from sam2.sam2_image_predictor import SAM2ImagePredictor
 sam2_checkpoint = "sam2_hiera_small.pt"
 model_cfg = "sam2_hiera_s.yaml"
-sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
 predictor = SAM2ImagePredictor(sam2_model)
 checkpoint_path = "sam2_lr0.0001_wd0.01_900.torch"

 sam2_checkpoint = "sam2_hiera_small.pt"
 model_cfg = "sam2_hiera_s.yaml"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device)
 predictor = SAM2ImagePredictor(sam2_model)
 checkpoint_path = "sam2_lr0.0001_wd0.01_900.torch"