Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

better support for non-CUDA devices (CPU, MPS) #192

Merged
merged 4 commits into from
Aug 12, 2024

Conversation

ronghanghu
Copy link
Contributor

@ronghanghu ronghanghu commented Aug 10, 2024

In this PR, we provide better support for non-CUDA devices (CPU, MPS) in the SAM 2 image/video predictors and the notebook examples:

  • we add compute_device parameter to image loading, and cast to the model's device (based on its parameters) instead of hard-coding to CUDA
  • the 3 notebooks are updated to automatic select the available devices
    Screenshot 2024-08-09 at 8 06 08 PM
  • post-processing is only enabled on CUDA and skipped on CPU or MPS
    • also, use fixed numpy random seed in image_predictor_example.ipynb and automatic_mask_generator_example.ipynb to get deterministic mask coloring (so that we only have small text changes on notebook files that could be easier to review in the future)
  • reduce video predictor notebook size with higher frame stride in visualization and smaller figure size (to make sure it renders on GitHub)
  • tested on CUDA, CPU and MPS (Apple M1 Pro)

@ronghanghu ronghanghu marked this pull request as ready for review August 10, 2024 03:23
@ronghanghu ronghanghu force-pushed the ronghanghu/non_cuda branch 3 times, most recently from 0f67b96 to 7213380 Compare August 10, 2024 03:57
@ronghanghu ronghanghu force-pushed the ronghanghu/non_cuda branch from 7213380 to 3d18489 Compare August 10, 2024 06:43
Copy link
Contributor

@haithamkhedr haithamkhedr left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good! thanks for adding support to non CUDA devices

sam2/sam2_video_predictor.py Outdated Show resolved Hide resolved
@ronghanghu ronghanghu merged commit 1034ee2 into facebookresearch:main Aug 12, 2024
2 checks passed
fbcotter added a commit to wayveai/segment-anything-2 that referenced this pull request Sep 6, 2024
… into facebookresearch-main

* 'main' of github.com:facebookresearch/segment-anything-2: (40 commits)
  open `README.md` with unicode (to support Hugging Face emoji); fix various typos (facebookresearch#218)
  accept kwargs in auto_mask_generator
  Fix HF image predictor
  improving warning message and adding further tips for installation (facebookresearch#204)
  better support for non-CUDA devices (CPU, MPS) (facebookresearch#192)
  Update hieradet.py
  add Colab support to the notebooks; pack config files in `sam2_configs` package during installation (facebookresearch#176)
  also catch errors during installation in case `CUDAExtension` cannot be loaded (facebookresearch#175)
  Add interface for box prompt in SAM 2 video predictor (facebookresearch#174)
  Address comment
  Update hieradet.py
  Update docstrings
  Revert code snippet
  Updated INSTALL.md with CUDA_HOME-related troubleshooting (facebookresearch#140)
  Format using ufmt
  Update INSTALL.md (facebookresearch#156)
  Update README
  Make it optional to build CUDA extension for SAM 2; also fallback to all available kernels if Flash Attention fails (facebookresearch#155)
  Clean up
  Address comment
  ...
xydy666 pushed a commit to xydy666/segment-anything-2 that referenced this pull request Sep 17, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants