ChaoningZhang / MobileSAM

This is the official code for MobileSAM project that makes SAM lightweight for mobile applications and beyond!
Apache License 2.0
4.86k stars 505 forks source link

Apple MPS support? #74

Open shersoni610 opened 1 year ago

shersoni610 commented 1 year ago

Hello,

Does anyone knows how to use apple "MPS" instead of cpu/cuda device?

JunkyByte commented 1 year ago

Install torch with mps support then it should be enough to choose mps as backend

from mobile_sam import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

model_type = "vit_t"
sam_checkpoint = "./weights/mobile_sam.pt"

device = "mps"

mobile_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
mobile_sam.to(device=device)
mobile_sam.eval()

predictor = SamPredictor(mobile_sam)
predictor.set_image(<your_image>)
masks, _, _ = predictor.predict(<input_prompts>)