Add support for Apple's MPS backend

[daikiad]

This pull request introduces support for Apple's Metal Performance Shaders (MPS) and refines device compatibility. The main changes include:

• Replaced the CUDA kernel dependency for connected components calculation with scipy functions, enabling functionality on Apple's MPS backend.
• Modified setup.py to install the CUDA kernel only in CUDA environments. For other environments, scipy functions will be used.
• Replaced .cuda() calls with .to(device) to remove the strict CUDA dependency, allowing for flexible device usage.
• Added configuration for using MPS in Jupyter notebooks

These updates enhance the flexibility and portability of the project, ensuring it runs efficiently on both CUDA and MPS environments

[facebook-github-bot]

Hi @daikiad!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[mashb1t]

You may want to consider using torch.backends.cuda.is_available() directly here instead of checking for nvcc.
Docs: https://pytorch.org/docs/stable/backends.html

[daikiad]

Thank you for the advice! I have updated setup.py to use torch.cuda.is_available() instead of checking nvcc.

[facebook-github-bot]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[robechun]

When trying to work off this branch, I ran into an issue with matplotlib>=3.9.1. Doing matplotlib>=3.9.0 helped.

Separately, I had to comment out ext_modules in the main setup() function.

[rauberdaniel]

I ran into the same problem with matplotlib>=3.9.1 and >=3.9.0 fixed it. Thanks!
Didn’t have any issues with ext_modules though.

[albertjo]

@daikiad Thanks for creating this branch. I'm on a M3 pro device and I was able to get SAM2 up and running with device="cpu".

I'm having trouble setting it up with "mps" though, currently seeing RuntimeError: Invalid buffer size: 9.42 GB. Here are my repro steps:

1. predictor = build_sam2_video_predictor("sam2_hiera_s.yaml", "checkpoints/sam2_hiera_small.pt", device="mps")
2. os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
3. inference_state = predictor.init_state(video_path="path/to/frames")

[daikiad]

Closing this pull request as support for both CPU and MPS has been added in PR #192
Thank you all for your comments and feedback.

[djdoogle123]

@daikiad Thanks for creating this branch. I'm on a M3 pro device and I was able to get SAM2 up and running with device="cpu".

I'm having trouble setting it up with "mps" though, currently seeing RuntimeError: Invalid buffer size: 9.42 GB. Here are my repro steps:

1. predictor = build_sam2_video_predictor("sam2_hiera_s.yaml", "checkpoints/sam2_hiera_small.pt", device="mps")
2. os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
3. inference_state = predictor.init_state(video_path="path/to/frames")

@daikiad How were you able to get SAMv2.1 running on Mac M3 I’m running to error message:

File "/Applications/Nuke15.1v5/Nuke15.1v5.app/Contents/MacOS/plugins/ENV_DIR/venv/lib/python3.10/site-packages/torch/cuda/init.py", line 310, in _lazy_init
raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled