File size: 4,415 Bytes
2398fdd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
import glob
import os
import os.path as osp
import platform
import sys
from setuptools import find_packages, setup
__version__ = "1.4.0"
# exec(open("gsplat/version.py", "r").read())
URL = "https://github.com/nerfstudio-project/gsplat" # TODO
BUILD_NO_CUDA = os.getenv("BUILD_NO_CUDA", "0") == "1"
WITH_SYMBOLS = os.getenv("WITH_SYMBOLS", "0") == "1"
LINE_INFO = os.getenv("LINE_INFO", "0") == "1"
def get_ext():
from torch.utils.cpp_extension import BuildExtension
return BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False)
def get_extensions():
import torch
from torch.__config__ import parallel_info
from torch.utils.cpp_extension import CUDAExtension
# Set default CUDA architectures if not specified
if "TORCH_CUDA_ARCH_LIST" not in os.environ:
# T4 GPU (7.5), A100 (8.0), RTX 30xx (8.6)
os.environ["TORCH_CUDA_ARCH_LIST"] = "7.5;8.0;8.6"
extensions_dir = osp.join("src", "gsplat", "cuda", "csrc")
sources = glob.glob(osp.join(extensions_dir, "*.cu")) + glob.glob(
osp.join(extensions_dir, "*.cpp")
)
# remove generated 'hip' files, in case of rebuilds
sources = [path for path in sources if "hip" not in path]
undef_macros = []
define_macros = []
if sys.platform == "win32":
define_macros += [("gsplat_EXPORTS", None)]
extra_compile_args = {"cxx": ["-O3"]}
if not os.name == "nt": # Not on Windows:
extra_compile_args["cxx"] += ["-Wno-sign-compare"]
extra_link_args = [] if WITH_SYMBOLS else ["-s"]
info = parallel_info()
if (
"backend: OpenMP" in info
and "OpenMP not found" not in info
and sys.platform != "darwin"
):
extra_compile_args["cxx"] += ["-DAT_PARALLEL_OPENMP"]
if sys.platform == "win32":
extra_compile_args["cxx"] += ["/openmp"]
else:
extra_compile_args["cxx"] += ["-fopenmp"]
else:
print("Compiling without OpenMP...")
# Compile for mac arm64
if sys.platform == "darwin" and platform.machine() == "arm64":
extra_compile_args["cxx"] += ["-arch", "arm64"]
extra_link_args += ["-arch", "arm64"]
nvcc_flags = os.getenv("NVCC_FLAGS", "")
nvcc_flags = [] if nvcc_flags == "" else nvcc_flags.split(" ")
nvcc_flags += ["-O3", "--use_fast_math"]
if LINE_INFO:
nvcc_flags += ["-lineinfo"]
if torch.version.hip:
# USE_ROCM was added to later versions of PyTorch.
# Define here to support older PyTorch versions as well:
define_macros += [("USE_ROCM", None)]
undef_macros += ["__HIP_NO_HALF_CONVERSIONS__"]
else:
nvcc_flags += ["--expt-relaxed-constexpr"]
extra_compile_args["nvcc"] = nvcc_flags
if sys.platform == "win32":
extra_compile_args["nvcc"] += ["-DWIN32_LEAN_AND_MEAN"]
extension = CUDAExtension(
"csrc",
sources,
include_dirs=[
extensions_dir, # Include the csrc directory itself for headers
osp.join(extensions_dir, "third_party", "glm"),
],
define_macros=define_macros,
undef_macros=undef_macros,
extra_compile_args=extra_compile_args,
extra_link_args=extra_link_args,
)
return [extension]
setup(
name="gsplat",
version=__version__,
description=" Python package for differentiable rasterization of gaussians",
url=URL,
download_url=f"{URL}/archive/gsplat-{__version__}.tar.gz",
python_requires=">=3.7",
# install_requires=[
# "jaxtyping",
# "rich>=12",
# "torch",
# "typing_extensions; python_version<'3.8'",
# ],
# extras_require={
# # dev dependencies. Install them by `pip install gsplat[dev]`
# "dev": [
# "black[jupyter]==22.3.0",
# "isort==5.10.1",
# "pylint==2.13.4",
# "pytest==7.1.2",
# "pytest-xdist==2.5.0",
# "typeguard>=2.13.3",
# "pyyaml==6.0",
# "build",
# "twine",
# "ninja",
# ],
# },
ext_modules=get_extensions() if not BUILD_NO_CUDA else [],
cmdclass={"build_ext": get_ext()} if not BUILD_NO_CUDA else {},
packages=find_packages(where="src"),
package_dir={"": "src"},
# https://github.com/pypa/setuptools/issues/1461#issuecomment-954725244
include_package_data=True,
)
|