fixup3

INSTALL.md

@@ -11,7 +11,7 @@ Then, install SAM 2 from the root of this repository via
 pip install -e ".[demo]"
 ```
 
-If your environment doesn't support CUDA, you may turn off the CUDA extension building during installation (via `SAM2_BUILD_CUDA` environment variable) as follows:
+If your environment doesn't support CUDA, you may turn off the CUDA extension building during installation via environment variable `SAM2_BUILD_CUDA=0`, as follows:
 ```bash
 # skip CUDA extension
 SAM2_BUILD_CUDA=0 pip install -e ".[demo]"
@@ -21,9 +21,11 @@ SAM2_BUILD_CUDA=0 pip install -e ".[demo]"
 
 If you see a message like `Failed to build SAM 2 CUDA extensions due to the error above` during installation or `Skipping the post-processing step due to the error above` at runtime, it indicates that the SAM 2 CUDA extension failed to build in your environment. In this case, you can still use SAM 2 for both image and video applications, but the post-processing step (removing small holes and sprinkles in the output masks) will be skipped. This shouldn't affect the results in most cases.
 
-If you would like to enable this post-processing step, you can reinstall SAM 2 with CUDA extension as follows
+If you would like to enable this post-processing step, you can reinstall SAM 2 with environment variable `SAM2_BUILD_ALLOW_ERRORS=0` to force building CUDA extension, as follows
 ```bash
-pip uninstall SAM-2; pip install -e ".[demo]"
+pip uninstall SAM-2;
+
+SAM2_BUILD_ALLOW_ERRORS=0 pip install -e ".[demo]"
 ```
 Note that PyTorch needs to be installed first before building the SAM 2 CUDA extension. It also requires installing [CUDA toolkits](https://developer.nvidia.com/cuda-toolkit-archive) that match the CUDA version for your PyTorch installation. (This should typically be CUDA 12.1 if you follow the default installation command.) Please check the section below on common installation issues if the CUDA extension fails to build during installation or load at runtime.
 

setup.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 import os
-import sys
 
 from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
@@ -38,11 +37,13 @@
     "dev": ["black==24.2.0", "usort==1.0.2", "ufmt==2.0.0b2"],
 }
 
+# allow turning off CUDA build with `export SAM2_BUILD_CUDA=0`
+BUILD_CUDA = os.getenv("SAM2_BUILD_CUDA", "1") == "1"
+BUILD_ALLOW_ERRORS = os.getenv("SAM2_BUILD_ALLOW_ERRORS", "1") == "1"
+
 
 def get_extensions():
-    # allow turning off CUDA build with `export SAM2_BUILD_CUDA=0`
-    build_cuda = os.getenv("SAM2_BUILD_CUDA", "1") == "1"
-    if not build_cuda:
+    if not BUILD_CUDA:
         return []
 
     srcs = ["sam2/csrc/connected_components.cu"]
@@ -61,11 +62,13 @@ def get_extensions():
 
 class BuildExtensionIgnoreErrors(BuildExtension):
     # Catch and skip errors during extension building and print a warning message
+    # (note that this message only shows up under verbose build mode
+    # "pip install -v -e ." or "python setup.py build_ext -v")
     ERROR_MSG = (
-        "{}\n"
+        "{}\n\n"
         "Failed to build SAM 2 CUDA extensions due to the error above. "
         "You can still use SAM 2, but some functionality may be limited (see "
-        "https://github.com/facebookresearch/segment-anything-2/blob/main/INSTALL.md)."
+        "https://github.com/facebookresearch/segment-anything-2/blob/main/INSTALL.md).\n"
     )
 
     def finalize_options(self):
@@ -108,6 +111,10 @@ def get_ext_filename(self, ext_name):
     python_requires=">=3.10.0",
     ext_modules=get_extensions(),
     cmdclass={
-        "build_ext": BuildExtensionIgnoreErrors.with_options(no_python_abi_suffix=True)
+        "build_ext": (
+            BuildExtensionIgnoreErrors.with_options(no_python_abi_suffix=True)
+            if BUILD_ALLOW_ERRORS
+            else BuildExtension.with_options(no_python_abi_suffix=True)
+        ),
     },
 )