| 24 | requirements = ["torch", "torchvision"] |
| 25 | |
| 26 | def get_extensions(): |
| 27 | this_dir = os.path.dirname(os.path.abspath(__file__)) |
| 28 | extensions_dir = os.path.join(this_dir, "src") |
| 29 | |
| 30 | main_file = glob.glob(os.path.join(extensions_dir, "*.cpp")) |
| 31 | source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp")) |
| 32 | source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu")) |
| 33 | |
| 34 | sources = main_file + source_cpu |
| 35 | extension = CppExtension |
| 36 | extra_compile_args = {"cxx": []} |
| 37 | define_macros = [] |
| 38 | |
| 39 | # Force cuda since torch ask for a device, not if cuda is in fact available. |
| 40 | if (os.environ.get('FORCE_CUDA') or torch.cuda.is_available()) and CUDA_HOME is not None: |
| 41 | extension = CUDAExtension |
| 42 | sources += source_cuda |
| 43 | define_macros += [("WITH_CUDA", None)] |
| 44 | extra_compile_args["nvcc"] = [ |
| 45 | "-DCUDA_HAS_FP16=1", |
| 46 | "-D__CUDA_NO_HALF_OPERATORS__", |
| 47 | "-D__CUDA_NO_HALF_CONVERSIONS__", |
| 48 | "-D__CUDA_NO_HALF2_OPERATORS__", |
| 49 | ] |
| 50 | else: |
| 51 | if CUDA_HOME is None: |
| 52 | raise NotImplementedError('CUDA_HOME is None. Please set environment variable CUDA_HOME.') |
| 53 | else: |
| 54 | raise NotImplementedError('No CUDA runtime is found. Please set FORCE_CUDA=1 or test it by running torch.cuda.is_available().') |
| 55 | |
| 56 | sources = [os.path.join(extensions_dir, s) for s in sources] |
| 57 | include_dirs = [extensions_dir] |
| 58 | ext_modules = [ |
| 59 | extension( |
| 60 | "MultiScaleDeformableAttention", |
| 61 | sources, |
| 62 | include_dirs=include_dirs, |
| 63 | define_macros=define_macros, |
| 64 | extra_compile_args=extra_compile_args, |
| 65 | ) |
| 66 | ] |
| 67 | return ext_modules |
| 68 | |
| 69 | setup( |
| 70 | name="MultiScaleDeformableAttention", |