4
0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-10 18:17:31 +08:00
This commit is contained in:
yesMohanHere 2026-07-18 13:26:01 +00:00 committed by GitHub
commit 28b1281195
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
14 changed files with 249 additions and 39 deletions

View File

@ -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

View File

@ -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

View File

@ -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

181
Dockerfile.blackwell Normal file
View File

@ -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

View File

@ -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,7 +51,7 @@ 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.
:::::
::::::
@ -59,7 +59,7 @@ Nerfstudio requires `python >= 3.8`. We recommend using conda to manage dependen
## 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:

View File

@ -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)

View File

@ -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:

View File

@ -429,7 +429,7 @@ class Trainer:
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 @@ class Trainer:
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"])

View File

@ -22,7 +22,18 @@ import torch
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):

View File

@ -514,7 +514,7 @@ class Mill19Download(DatasetDownload):
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}")

View File

@ -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
)

View File

@ -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

View File

@ -160,9 +160,11 @@ class RenderStateMachine(threading.Thread):
[color[0] / 255.0, color[1] / 255.0, color[2] / 255.0],
device=self.viewer.get_model().device,
)
with background_color_override_context(
background_color
), torch.no_grad(), viewer_utils.SetTrace(self.check_interrupt):
with (
background_color_override_context(background_color),
torch.no_grad(),
viewer_utils.SetTrace(self.check_interrupt),
):
outputs = self.viewer.get_model().get_outputs_for_camera(camera, obb_box=obb)
else:
with torch.no_grad(), viewer_utils.SetTrace(self.check_interrupt):

View File

@ -8,7 +8,7 @@ version = "1.1.5"
description = "All-in-one repository for state-of-the-art NeRFs"
readme = "README.md"
license = { text="Apache 2.0"}
requires-python = ">=3.8.0"
requires-python = ">=3.10"
classifiers = [
"Development Status :: 3 - Alpha",
"Programming Language :: Python",
@ -31,7 +31,7 @@ dependencies = [
"mediapy>=1.1.0",
"msgpack>=1.0.4",
"msgpack_numpy>=0.4.8",
"nerfacc==0.5.2",
"nerfacc>=0.5.3",
"open3d>=0.16.0",
"opencv-python-headless==4.10.0.84",
"Pillow>=10.3.0",
@ -64,7 +64,7 @@ dependencies = [
"xatlas",
"trimesh>=3.20.2",
"timm==0.6.7",
"gsplat==1.4.0",
"gsplat==1.5.3",
"pytorch-msssim",
"pathos",
"packaging",