quoccuongle/Grounded-SAM-2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

also catch errors during installation in case CUDAExtension cannot be loaded (#175)

Browse Source 
Previously we only catch build errors in `BuildExtension` in https://github.com/facebookresearch/segment-anything-2/pull/155. However, in some cases, the `CUDAExtension` instance might not load. So in this PR, we also catch such errors for `CUDAExtension`.
This commit is contained in:
Ronghang Hu
2024-08-07 12:26:11 -07:00
committed by GitHub
parent 6ecb5ff8d0
commit 6186d1529a
1 changed files with 32 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
setup.py
View File

@@ -44,11 +44,22 @@ BUILD_CUDA = os.getenv("SAM2_BUILD_CUDA", "1") == "1"
# You may force stopping on errors with `export SAM2_BUILD_ALLOW_ERRORS=0`. # You may force stopping on errors with `export SAM2_BUILD_ALLOW_ERRORS=0`.
BUILD_ALLOW_ERRORS = os.getenv("SAM2_BUILD_ALLOW_ERRORS", "1") == "1" BUILD_ALLOW_ERRORS = os.getenv("SAM2_BUILD_ALLOW_ERRORS", "1") == "1"
# 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")
CUDA_ERROR_MSG = (
"{}\n\n"
"Failed to build the SAM 2 CUDA extension due to the error above. "
"You can still use SAM 2, but some post-processing functionality may be limited "
"(see https://github.com/facebookresearch/segment-anything-2/blob/main/INSTALL.md).\n"
)
def get_extensions(): def get_extensions():
if not BUILD_CUDA: if not BUILD_CUDA:
return [] return []
try:
srcs = ["sam2/csrc/connected_components.cu"] srcs = ["sam2/csrc/connected_components.cu"]
compile_args = { compile_args = {
"cxx": [], "cxx": [],
@@ -60,39 +71,37 @@ def get_extensions():
], ],
} }
ext_modules = [CUDAExtension("sam2._C", srcs, extra_compile_args=compile_args)] ext_modules = [CUDAExtension("sam2._C", srcs, extra_compile_args=compile_args)]
except Exception as e:
if BUILD_ALLOW_ERRORS:
print(CUDA_ERROR_MSG.format(e))
ext_modules = []
else:
raise e
return ext_modules return ext_modules
class BuildExtensionIgnoreErrors(BuildExtension): 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"
"Failed to build the SAM 2 CUDA extension due to the error above. "
"You can still use SAM 2, but some post-processing functionality may be limited "
"(see https://github.com/facebookresearch/segment-anything-2/blob/main/INSTALL.md).\n"
)
def finalize_options(self): def finalize_options(self):
try: try:
super().finalize_options() super().finalize_options()
except Exception as e: except Exception as e:
print(self.ERROR_MSG.format(e)) print(CUDA_ERROR_MSG.format(e))
self.extensions = [] self.extensions = []
def build_extensions(self): def build_extensions(self):
try: try:
super().build_extensions() super().build_extensions()
except Exception as e: except Exception as e:
print(self.ERROR_MSG.format(e)) print(CUDA_ERROR_MSG.format(e))
self.extensions = [] self.extensions = []
def get_ext_filename(self, ext_name): def get_ext_filename(self, ext_name):
try: try:
return super().get_ext_filename(ext_name) return super().get_ext_filename(ext_name)
except Exception as e: except Exception as e:
print(self.ERROR_MSG.format(e)) print(CUDA_ERROR_MSG.format(e))
self.extensions = [] self.extensions = []
return "_C.so" return "_C.so"