mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
40 lines
1.0 KiB
Python
40 lines
1.0 KiB
Python
import os
|
|
from setuptools import setup
|
|
|
|
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
USE_CUDA = True
|
|
|
|
ext_modules = []
|
|
cmdclass = {}
|
|
setup_requires = []
|
|
if USE_CUDA:
|
|
try:
|
|
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
|
|
|
ext_modules = [
|
|
CUDAExtension(
|
|
name="pyrad_cuda",
|
|
sources=[
|
|
"pyrad/cuda/pyrad_cuda.cpp",
|
|
"pyrad/cuda/pyrad_cuda_kernel.cu",
|
|
],
|
|
include_dirs=[os.path.join("pyrad", "cuda", "include")],
|
|
extra_compile_args={
|
|
"cxx": [],
|
|
"nvcc": [],
|
|
},
|
|
verbose=True,
|
|
),
|
|
]
|
|
cmdclass = {"build_ext": BuildExtension}
|
|
setup_requires = ["pybind11>=2.5.0"]
|
|
except:
|
|
print("Failed to build the CUDA extension.")
|
|
|
|
if __name__ == "__main__":
|
|
setup(
|
|
ext_modules=ext_modules,
|
|
cmdclass=cmdclass,
|
|
setup_requires=setup_requires,
|
|
)
|