Add support for NVIDIA Blackwell GPUs (sm_120, RTX 50-series)

.github/workflows/doc.yml

@@ -15,7 +15,7 @@ jobs:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v3
         with:
-          python-version: '3.9'
+          python-version: '3.10'
       - name: Install dependencies
         run: |
           pip install uv

.github/workflows/publish.yml

@@ -16,7 +16,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v1
         with:
-          python-version: '3.8'
+          python-version: '3.10'
       - name: Install dependencies
         run: |
           python -m pip install build twine

Dockerfile

@@ -85,7 +85,7 @@ RUN pip install --no-cache-dir --upgrade pip 'setuptools<70.0.0' && \
     pip install --no-cache-dir torch==2.1.2+cu118 torchvision==0.16.2+cu118 'numpy<2.0.0' --extra-index-url https://download.pytorch.org/whl/cu118 && \
     git clone --branch master --recursive https://github.com/cvg/Hierarchical-Localization.git /opt/hloc && \
     cd /opt/hloc && git checkout v1.4 && python3.10 -m pip install --no-cache-dir . && cd ~ && \
-    TCNN_CUDA_ARCHITECTURES="${CUDA_ARCHITECTURES}" pip install --no-cache-dir "git+https://github.com/NVlabs/tiny-cuda-nn.git@b3473c81396fe927293bdfd5a6be32df8769927c#subdirectory=bindings/torch" && \
+    TCNN_CUDA_ARCHITECTURES="${CUDA_ARCHITECTURES}" pip install --no-cache-dir --no-build-isolation "git+https://github.com/NVlabs/tiny-cuda-nn.git@b3473c81396fe927293bdfd5a6be32df8769927c#subdirectory=bindings/torch" && \
     pip install --no-cache-dir pycolmap==0.6.1 pyceres==2.1 omegaconf==2.3.0
 
 # Install gsplat and nerfstudio.
@@ -98,7 +98,7 @@ COPY --from=source /tmp/nerfstudio/ /tmp/nerfstudio
 RUN export TORCH_CUDA_ARCH_LIST="$(echo "$CUDA_ARCHITECTURES" | tr ';' '\n' | awk '$0 > 70 {print substr($0,1,1)"."substr($0,2)}' | tr '\n' ' ' | sed 's/ $//')" && \
     export MAX_JOBS=4 && \
     GSPLAT_VERSION="$(sed -n 's/.*gsplat==\s*\([^," '"'"']*\).*/\1/p' /tmp/nerfstudio/pyproject.toml)" && \
-    pip install --no-cache-dir git+https://github.com/nerfstudio-project/gsplat.git@v${GSPLAT_VERSION} && \
+    pip install --no-cache-dir --no-build-isolation git+https://github.com/nerfstudio-project/gsplat.git@v${GSPLAT_VERSION} && \
     pip install --no-cache-dir /tmp/nerfstudio 'numpy<2.0.0' && \
     rm -rf /tmp/nerfstudio
 

Dockerfile.blackwell

@@ -0,0 +1,181 @@
+# syntax=docker/dockerfile:1
+# Blackwell (sm_120, RTX 50-series) variant of the nerfstudio Docker image.
+#
+# Differences from the stock Dockerfile:
+# - CUDA 12.8 base images (sm_120 requires CUDA >= 12.8)
+# - PyTorch 2.8 with cu128 wheels (ships sm_120 kernels)
+# - COLMAP 3.12.6 / GLOMAP 1.2.0 (CUDA 12.8-compatible releases)
+# - tiny-cuda-nn pinned to a Blackwell-capable master commit
+# - CUDA_ARCHITECTURES defaults to "120"
+# - nerfacc's JIT CUDA kernels are pre-compiled during the build (the
+#   runtime stage has no nvcc, so they could not be built on first use)
+ARG UBUNTU_VERSION=22.04
+ARG NVIDIA_CUDA_VERSION=12.8.1
+# CUDA architectures, required by Colmap and tiny-cuda-nn. Use >= 8.0 for faster TCNN.
+ARG CUDA_ARCHITECTURES="120"
+ARG NERFSTUDIO_VERSION=""
+
+# Pull source either provided or from git.
+FROM scratch as source_copy
+ONBUILD COPY . /tmp/nerfstudio
+FROM alpine/git as source_no_copy
+ARG NERFSTUDIO_VERSION
+ONBUILD RUN git clone --branch ${NERFSTUDIO_VERSION} --recursive https://github.com/nerfstudio-project/nerfstudio.git /tmp/nerfstudio
+ARG NERFSTUDIO_VERSION
+FROM source_${NERFSTUDIO_VERSION:+no_}copy as source
+
+FROM nvidia/cuda:${NVIDIA_CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION} as builder
+ARG CUDA_ARCHITECTURES
+ARG NVIDIA_CUDA_VERSION
+ARG UBUNTU_VERSION
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV QT_XCB_GL_INTEGRATION=xcb_egl
+# Work around apt _apt sandbox verification failing under docker's seccomp
+# profile on newer kernels (spurious "invalid signature" GPG errors).
+RUN echo 'APT::Sandbox::User "root";' > /etc/apt/apt.conf.d/99docker-sandbox && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends --no-install-suggests \
+        git \
+        wget \
+        ninja-build \
+        build-essential \
+        libboost-program-options-dev \
+        libboost-filesystem-dev \
+        libboost-graph-dev \
+        libboost-system-dev \
+        libeigen3-dev \
+        libflann-dev \
+        libfreeimage-dev \
+        libmetis-dev \
+        libgoogle-glog-dev \
+        libgtest-dev \
+        libsqlite3-dev \
+        libglew-dev \
+        qtbase5-dev \
+        libqt5opengl5-dev \
+        libcgal-dev \
+        libceres-dev \
+        python3.10-dev \
+        python3-pip
+
+# Build and install CMake
+RUN wget https://github.com/Kitware/CMake/releases/download/v3.31.3/cmake-3.31.3-linux-x86_64.sh \
+    -q -O /tmp/cmake-install.sh \
+    && chmod u+x /tmp/cmake-install.sh \
+    && mkdir /opt/cmake-3.31.3 \
+    && /tmp/cmake-install.sh --skip-license --prefix=/opt/cmake-3.31.3 \
+    && rm /tmp/cmake-install.sh \
+    && ln -s /opt/cmake-3.31.3/bin/* /usr/local/bin
+
+# Build and install GLOMAP.
+RUN git clone https://github.com/colmap/glomap.git && \
+    cd glomap && \
+    git checkout "1.2.0" && \
+    mkdir build && \
+    cd build && \
+    mkdir -p /build && \
+    cmake .. -GNinja "-DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES}" \
+        -DCMAKE_INSTALL_PREFIX=/build/glomap && \
+    ninja install -j4 && \
+    cd ~
+
+# Build and install COLMAP.
+RUN git clone https://github.com/colmap/colmap.git && \
+    cd colmap && \
+    git checkout "3.12.6" && \
+    mkdir build && \
+    cd build && \
+    mkdir -p /build && \
+    cmake .. -GNinja "-DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES}" \
+        -DCMAKE_INSTALL_PREFIX=/build/colmap && \
+    ninja install -j4 && \
+    cd ~
+
+# Upgrade pip and install dependencies (torch with sm_120 kernels via cu128 wheels).
+RUN pip install --no-cache-dir --upgrade pip 'setuptools<70.0.0' && \
+    pip install --no-cache-dir torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu128 && \
+    git clone --branch master --recursive https://github.com/cvg/Hierarchical-Localization.git /opt/hloc && \
+    cd /opt/hloc && git checkout v1.4 && python3.10 -m pip install --no-cache-dir . && cd ~ && \
+    TCNN_CUDA_ARCHITECTURES="${CUDA_ARCHITECTURES}" pip install --no-cache-dir --no-build-isolation "git+https://github.com/NVlabs/tiny-cuda-nn.git@749dd70c5afc5a9dadb85e5652ed65d55e0ba187#subdirectory=bindings/torch" && \
+    pip install --no-cache-dir pycolmap==0.6.1 pyceres==2.1 omegaconf==2.3.0
+
+# Install gsplat and nerfstudio.
+# NOTE: both are installed jointly in order to prevent docker cache with latest
+# gsplat version (we do not expliticly specify the commit hash).
+#
+# We set MAX_JOBS to reduce resource usage for GH actions:
+# - https://github.com/nerfstudio-project/gsplat/blob/db444b904976d6e01e79b736dd89a1070b0ee1d0/setup.py#L13-L23
+COPY --from=source /tmp/nerfstudio/ /tmp/nerfstudio
+RUN export TORCH_CUDA_ARCH_LIST="$(echo "$CUDA_ARCHITECTURES" | tr ';' '\n' | awk '{if ($0 > 99) print substr($0,1,2)"."substr($0,3); else if ($0 > 70) print substr($0,1,1)"."substr($0,2)}' | tr '\n' ' ' | sed 's/ $//')" && \
+    export MAX_JOBS=4 && \
+    GSPLAT_VERSION="$(sed -n 's/.*gsplat==\s*\([^," '"'"']*\).*/\1/p' /tmp/nerfstudio/pyproject.toml)" && \
+    pip install --no-cache-dir --no-build-isolation git+https://github.com/nerfstudio-project/gsplat.git@v${GSPLAT_VERSION} && \
+    pip install --no-cache-dir /tmp/nerfstudio 'numpy<2.0.0' && \
+    rm -rf /tmp/nerfstudio
+
+# Pre-compile nerfacc's JIT CUDA kernels. nerfacc compiles its kernels on first
+# use, but the runtime stage has no nvcc, so build them now and copy the torch
+# extension cache into the runtime image. The build container has no GPU, so
+# use the CUDA driver stub to satisfy the import after compilation.
+RUN export TORCH_CUDA_ARCH_LIST="$(echo "$CUDA_ARCHITECTURES" | tr ';' '\n' | awk '{if ($0 > 99) print substr($0,1,2)"."substr($0,3); else if ($0 > 70) print substr($0,1,1)"."substr($0,2)}' | tr '\n' ' ' | sed 's/ $//')" && \
+    export MAX_JOBS=4 && \
+    ln -sf /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 && \
+    LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:${LD_LIBRARY_PATH} \
+        python3.10 -c "from nerfacc.cuda._backend import _C; print('nerfacc CUDA kernels pre-compiled')" && \
+    find /root/.cache/torch_extensions -name "nerfacc_cuda.so" | grep .
+
+# Fix permissions
+RUN chmod -R go=u /usr/local/lib/python3.10 && \
+    chmod -R go=u /build
+
+#
+# Docker runtime stage.
+#
+FROM nvidia/cuda:${NVIDIA_CUDA_VERSION}-runtime-ubuntu${UBUNTU_VERSION} as runtime
+ARG CUDA_ARCHITECTURES
+ARG NVIDIA_CUDA_VERSION
+ARG UBUNTU_VERSION
+
+LABEL org.opencontainers.image.source = "https://github.com/nerfstudio-project/nerfstudio"
+LABEL org.opencontainers.image.licenses = "Apache License 2.0"
+LABEL org.opencontainers.image.base.name="docker.io/library/nvidia/cuda:${NVIDIA_CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION}"
+LABEL org.opencontainers.image.documentation = "https://docs.nerf.studio/"
+
+# Minimal dependencies to run COLMAP binary compiled in the builder stage.
+# Note: this reduces the size of the final image considerably, since all the
+# build dependencies are not needed.
+RUN echo 'APT::Sandbox::User "root";' > /etc/apt/apt.conf.d/99docker-sandbox && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends --no-install-suggests \
+        libboost-filesystem1.74.0 \
+        libboost-program-options1.74.0 \
+        libc6 \
+        libceres2 \
+        libfreeimage3 \
+        libgcc-s1 \
+        libgl1 \
+        libglew2.2 \
+        libgoogle-glog0v5 \
+        libqt5core5a \
+        libqt5gui5 \
+        libqt5widgets5 \
+        python3.10 \
+        python3.10-dev \
+        build-essential \
+        python-is-python3 \
+        ffmpeg
+
+# Copy packages from builder stage.
+COPY --from=builder /build/colmap/ /usr/local/
+COPY --from=builder /build/glomap/ /usr/local/
+COPY --from=builder /usr/local/lib/python3.10/dist-packages/ /usr/local/lib/python3.10/dist-packages/
+COPY --from=builder /usr/local/bin/ns* /usr/local/bin/
+# Pre-compiled nerfacc JIT kernels (see builder stage).
+COPY --from=builder /root/.cache/torch_extensions /root/.cache/torch_extensions
+
+# Install nerfstudio cli auto completion
+RUN /bin/bash -c 'ns-install-cli --mode install'
+
+# Bash as default entrypoint.
+CMD /bin/bash -l

docs/quickstart/installation.md

@@ -5,7 +5,7 @@
 ::::::{tab-set}
 :::::{tab-item} Linux
 
-Nerfstudio requires `python >= 3.8`. We recommend using conda to manage dependencies. Make sure to install [Conda](https://docs.conda.io/en/latest/miniconda.html) before proceeding.
+Nerfstudio requires `python >= 3.10`. We recommend using conda to manage dependencies. Make sure to install [Conda](https://docs.conda.io/en/latest/miniconda.html) before proceeding.
 
 :::::
 :::::{tab-item} Windows
@@ -51,15 +51,15 @@ For example:
 When updating, or if you close your terminal before you finish the installation and run your first `splatfacto`, you have to re-do this environment activation step.
 :::
 
-Nerfstudio requires `python >= 3.8`. We recommend using conda to manage dependencies. Make sure to install [Conda](https://docs.conda.io/en/latest/miniconda.html) before proceeding.
+Nerfstudio requires `python >= 3.10`. We recommend using conda to manage dependencies. Make sure to install [Conda](https://docs.conda.io/en/latest/miniconda.html) before proceeding.
 
 :::::
 ::::::
 
 ## Create environment
 
 ```bash
-conda create --name nerfstudio -y python=3.8
+conda create --name nerfstudio -y python=3.10
 conda activate nerfstudio
 python -m pip install --upgrade pip
 
@@ -79,6 +79,28 @@ pip uninstall torch torchvision functorch tinycudann
 ```
 
 ::::{tab-set}
+:::{tab-item} Torch 2.8+ with CUDA 12.8 (Blackwell / RTX 50-series)
+
+Required for NVIDIA Blackwell GPUs (RTX 50-series, compute capability sm_120), which need CUDA 12.8 or newer:
+
+```bash
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
+```
+
+To build the necessary CUDA extensions, `cuda-toolkit` 12.8 is also required. We
+recommend installing with conda:
+
+```bash
+conda install -c "nvidia/label/cuda-12.8.0" cuda-toolkit
+```
+
+When building CUDA extensions (gsplat, tiny-cuda-nn, ...) for a Blackwell GPU, make sure the architecture list includes sm_120:
+
+```bash
+export TORCH_CUDA_ARCH_LIST="12.0+PTX"
+```
+
+:::
 :::{tab-item} Torch 2.1.2 with CUDA 11.8 (recommended)
 
 Install PyTorch 2.1.2 with CUDA 11.8:
@@ -245,6 +267,14 @@ docker build \
     --file Dockerfile .
 ```
 
+For NVIDIA Blackwell GPUs (RTX 50-series, compute capability sm_120), use the CUDA 12.8-based
+`Dockerfile.blackwell` variant instead — the stock image targets CUDA 11.8, which does not
+support sm_120:
+
+```bash
+docker build --tag nerfstudio-blackwell -f Dockerfile.blackwell .
+```
+
 ### Using an interactive container
 
 The docker container can be launched with an interactive terminal where nerfstudio commands can be entered as usual. Some parameters are required and some are strongly recommended for usage as following:

nerfstudio/data/utils/colmap_parsing_utils.py

@@ -95,6 +95,8 @@ def write_next_bytes(fid, data, format_char_sequence, endian_character="<"):
     """
     if isinstance(data, (list, tuple)):
         bytes = struct.pack(endian_character + format_char_sequence, *data)
+    elif isinstance(data, np.ndarray):
+        bytes = struct.pack(endian_character + format_char_sequence, *data.tolist())
     else:
         bytes = struct.pack(endian_character + format_char_sequence, data)
     fid.write(bytes)

nerfstudio/data/utils/data_utils.py

@@ -36,24 +36,7 @@ def pil_to_numpy(im: PILImage) -> np.ndarray:
     # Load in image completely (PIL defaults to lazy loading)
     im.load()
 
-    # Unpack data
-    e = Image._getencoder(im.mode, "raw", im.mode)
-    e.setimage(im.im)
-
-    # NumPy buffer for the result
-    shape, typestr = Image._conv_type_shape(im)
-    data = np.empty(shape, dtype=np.dtype(typestr))
-    mem = data.data.cast("B", (data.data.nbytes,))
-
-    bufsize, s, offset = 65536, 0, 0
-    while not s:
-        _, s, d = e.encode(bufsize)
-        mem[offset : offset + len(d)] = d
-        offset += len(d)
-    if s < 0:
-        raise RuntimeError("encoder error %d in tobytes" % s)
-
-    return data
+    return np.asarray(im)
 
 
 def get_image_mask_tensor_from_path(filepath: Union[Path, IO[bytes]], scale_factor: float = 1.0) -> torch.Tensor:

nerfstudio/engine/trainer.py

@@ -429,7 +429,7 @@ def _load_checkpoint(self) -> None:
                 load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(load_dir))[-1]
             load_path: Path = load_dir / f"step-{load_step:09d}.ckpt"
             assert load_path.exists(), f"Checkpoint {load_path} does not exist"
-            loaded_state = torch.load(load_path, map_location="cpu")
+            loaded_state = torch.load(load_path, map_location="cpu", weights_only=False)
             self._start_step = loaded_state["step"] + 1
             # load the checkpoints for pipeline, optimizers, and gradient scalar
             self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"])
@@ -440,7 +440,7 @@ def _load_checkpoint(self) -> None:
             CONSOLE.print(f"Done loading Nerfstudio checkpoint from {load_path}")
         elif load_checkpoint is not None:
             assert load_checkpoint.exists(), f"Checkpoint {load_checkpoint} does not exist"
-            loaded_state = torch.load(load_checkpoint, map_location="cpu")
+            loaded_state = torch.load(load_checkpoint, map_location="cpu", weights_only=False)
             self._start_step = loaded_state["step"] + 1
             # load the checkpoints for pipeline, optimizers, and gradient scalar
             self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"])

nerfstudio/field_components/activations.py

@@ -22,7 +22,18 @@
 from jaxtyping import Float
 from torch import Tensor
 from torch.autograd import Function
-from torch.cuda.amp import custom_bwd, custom_fwd
+
+try:
+    # torch >= 2.4: custom_fwd/custom_bwd live in torch.amp and require device_type
+    from torch.amp import custom_bwd as _custom_bwd, custom_fwd as _custom_fwd
+
+    def custom_fwd(*args, **kwargs):
+        return _custom_fwd(*args, device_type="cuda", **kwargs)
+
+    custom_bwd = _custom_bwd(device_type="cuda")
+except ImportError:
+    # torch < 2.4: custom_fwd/custom_bwd live in torch.cuda.amp (no device_type argument)
+    from torch.cuda.amp import custom_bwd, custom_fwd  # type: ignore
 
 
 class _TruncExp(Function):

nerfstudio/scripts/downloads/download_data.py

@@ -514,7 +514,7 @@ def download(self, save_dir: Path) -> None:
             split_filepaths = []
             for image_path, new_image_path in copied_images.items():
                 metadata_path = image_path.parent.parent / "metadata" / f"{image_path.stem}.pt"
-                metadata = torch.load(metadata_path, map_location="cpu")
+                metadata = torch.load(metadata_path, map_location="cpu", weights_only=False)
                 c2w = torch.eye(4)
                 c2w[:3] = metadata["c2w"]
                 file_path = str(Path("images") / f"{new_image_path.name}")

nerfstudio/scripts/render.py

@@ -186,9 +186,10 @@ def _render_trajectory_video(
                         max_idx = true_max_idx
 
                 if crop_data is not None:
-                    with renderers.background_color_override_context(
-                        crop_data.background_color.to(pipeline.device)
-                    ), torch.no_grad():
+                    with (
+                        renderers.background_color_override_context(crop_data.background_color.to(pipeline.device)),
+                        torch.no_grad(),
+                    ):
                         outputs = pipeline.model.get_outputs_for_camera(
                             cameras[camera_idx : camera_idx + 1], obb_box=obb_box
                         )

nerfstudio/utils/eval_utils.py

@@ -59,7 +59,7 @@ def eval_load_checkpoint(config: TrainerConfig, pipeline: Pipeline) -> Tuple[Pat
         load_step = config.load_step
     load_path = config.load_dir / f"step-{load_step:09d}.ckpt"
     assert load_path.exists(), f"Checkpoint {load_path} does not exist"
-    loaded_state = torch.load(load_path, map_location="cpu")
+    loaded_state = torch.load(load_path, map_location="cpu", weights_only=False)
     pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"])
     CONSOLE.print(f":white_check_mark: Done loading checkpoint from {load_path}")
     return load_path, load_step