diff --git a/README.md b/README.md index 9aad5341b..b53b32f2e 100644 --- a/README.md +++ b/README.md @@ -68,7 +68,10 @@ This breakthrough is driven by an innovative data engine that has automatically - Python 3.12 or higher - PyTorch 2.7 or higher -- CUDA-compatible GPU with CUDA 12.6 or higher +- A CUDA-compatible GPU (CUDA 12.6 or higher) is recommended for best + performance. CPU-only environments are also supported — the device is + selected automatically based on `torch.cuda.is_available()` (inference will + be slower without a GPU). 1. **Create a new Conda environment:** @@ -84,6 +87,13 @@ conda activate sam3 pip install torch==2.10.0 torchvision --index-url https://download.pytorch.org/whl/cu128 ``` +> **CPU-only:** if you don't have a CUDA-compatible GPU, install the CPU build +> instead. SAM 3 will automatically fall back to running on CPU. +> +> ```bash +> pip install torch==2.10.0 torchvision --index-url https://download.pytorch.org/whl/cpu +> ``` + 3. **Clone the repository and install the package:** ```bash diff --git a/sam3/model/decoder.py b/sam3/model/decoder.py index b90e55f5c..0c7dfd4c5 100644 --- a/sam3/model/decoder.py +++ b/sam3/model/decoder.py @@ -280,7 +280,9 @@ def __init__( if resolution is not None and stride is not None: feat_size = resolution // stride coords_h, coords_w = self._get_coords( - feat_size, feat_size, device="cuda" + feat_size, + feat_size, + device="cuda" if torch.cuda.is_available() else "cpu", ) self.compilable_cord_cache = (coords_h, coords_w) self.compilable_stored_size = (feat_size, feat_size) diff --git a/sam3/model/position_encoding.py b/sam3/model/position_encoding.py index f38c62b32..ad3cb881e 100644 --- a/sam3/model/position_encoding.py +++ b/sam3/model/position_encoding.py @@ -51,8 +51,9 @@ def __init__( (int(precompute_resolution // 28), int(precompute_resolution // 28)), (precompute_resolution // 32, precompute_resolution // 32), ] + device = "cuda" if torch.cuda.is_available() else "cpu" for size in precompute_sizes: - tensors = torch.zeros((1, 1) + size, device="cuda") + tensors = torch.zeros((1, 1) + size, device=device) self.forward(tensors) # further clone and detach it in the cache (just to be safe) self.cache[size] = self.cache[size].clone().detach() diff --git a/sam3/model/sam3_image_processor.py b/sam3/model/sam3_image_processor.py index 83bbb0d85..57bffa7bb 100644 --- a/sam3/model/sam3_image_processor.py +++ b/sam3/model/sam3_image_processor.py @@ -14,7 +14,13 @@ class Sam3Processor: """ """ - def __init__(self, model, resolution=1008, device="cuda", confidence_threshold=0.5): + def __init__( + self, + model, + resolution=1008, + device="cuda" if torch.cuda.is_available() else "cpu", + confidence_threshold=0.5, + ): self.model = model self.resolution = resolution self.device = device diff --git a/test/test_cpu_position_encoding.py b/test/test_cpu_position_encoding.py new file mode 100644 index 000000000..1956da8ab --- /dev/null +++ b/test/test_cpu_position_encoding.py @@ -0,0 +1,35 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +"""Regression test for CPU-only construction of PositionEmbeddingSine. + +Previously, PositionEmbeddingSine precomputed its cache with a hardcoded +``device="cuda"``, which raised a RuntimeError on CPU-only machines when a +``precompute_resolution`` was provided (see GitHub issue #587). +""" + +import unittest +from unittest.mock import patch + +import torch +from sam3.model.position_encoding import PositionEmbeddingSine + + +class TestPositionEmbeddingSineCpu(unittest.TestCase): + """PositionEmbeddingSine should build without CUDA when none is available.""" + + def test_precompute_builds_on_cpu_when_cuda_unavailable(self) -> None: + """Constructing with precompute_resolution must not require CUDA.""" + with patch("torch.cuda.is_available", return_value=False): + pos_enc = PositionEmbeddingSine( + num_pos_feats=256, + precompute_resolution=112, + ) + + # The cache is populated and lives on CPU. + self.assertGreater(len(pos_enc.cache), 0) + for cached in pos_enc.cache.values(): + self.assertEqual(cached.device.type, "cpu") + + +if __name__ == "__main__": + unittest.main()