From c0c329a2f111ea829862ceba9c598e4cadba202a Mon Sep 17 00:00:00 2001 From: yesMohanHere Date: Fri, 17 Jul 2026 15:53:36 +0530 Subject: [PATCH] Add support for NVIDIA Blackwell GPUs (sm_120, RTX 50-series) - Bump gsplat to >=1.5.3 (sm_120-capable release line) and nerfacc to >=0.5.3 - Drop EOL Python 3.8/3.9: requires-python >=3.10, CI publish/doc workflows - torch.load(..., weights_only=False) for torch>=2.6 checkpoint compatibility - Use torch.amp custom_fwd/custom_bwd with device_type="cuda" (torch>=2.4 API) - Replace deprecated PIL private encoder API in pil_to_numpy with np.asarray (fixes Pillow>=11 compatibility) - Docs: add Torch 2.8+ / CUDA 12.8 install tab for Blackwell GPUs Tested on RTX 5060 (sm_120) with torch 2.11.0+cu128: splatfacto and nerfacto train end-to-end and ns-render produces correct output. --- .github/workflows/doc.yml | 2 +- .github/workflows/publish.yml | 2 +- docs/quickstart/installation.md | 28 +++++++++++++++++-- nerfstudio/data/utils/data_utils.py | 19 +------------ nerfstudio/engine/trainer.py | 4 +-- nerfstudio/field_components/activations.py | 6 ++-- nerfstudio/scripts/downloads/download_data.py | 2 +- nerfstudio/utils/eval_utils.py | 2 +- pyproject.toml | 6 ++-- 9 files changed, 38 insertions(+), 33 deletions(-) diff --git a/.github/workflows/doc.yml b/.github/workflows/doc.yml index d2637429c..3974bd4ca 100644 --- a/.github/workflows/doc.yml +++ b/.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 diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index b75984535..e9039a2aa 100644 --- a/.github/workflows/publish.yml +++ b/.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 diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md index de7906421..c4e79a19f 100644 --- a/docs/quickstart/installation.md +++ b/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,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: diff --git a/nerfstudio/data/utils/data_utils.py b/nerfstudio/data/utils/data_utils.py index d1a672184..109c848ca 100644 --- a/nerfstudio/data/utils/data_utils.py +++ b/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: diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index a653e1de8..7f2f0d34d 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -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"]) diff --git a/nerfstudio/field_components/activations.py b/nerfstudio/field_components/activations.py index e413965e2..371999094 100644 --- a/nerfstudio/field_components/activations.py +++ b/nerfstudio/field_components/activations.py @@ -21,21 +21,21 @@ from typing import TYPE_CHECKING import torch from jaxtyping import Float from torch import Tensor +from torch.amp import custom_bwd, custom_fwd from torch.autograd import Function -from torch.cuda.amp import custom_bwd, custom_fwd class _TruncExp(Function): # Implementation from torch-ngp: # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py @staticmethod - @custom_fwd(cast_inputs=torch.float32) + @custom_fwd(cast_inputs=torch.float32, device_type="cuda") def forward(ctx, x): ctx.save_for_backward(x) return torch.exp(x) @staticmethod - @custom_bwd + @custom_bwd(device_type="cuda") def backward(ctx, g): x = ctx.saved_tensors[0] return g * torch.exp(x.clamp(-15, 15)) diff --git a/nerfstudio/scripts/downloads/download_data.py b/nerfstudio/scripts/downloads/download_data.py index a2c9d2758..653fbf94d 100644 --- a/nerfstudio/scripts/downloads/download_data.py +++ b/nerfstudio/scripts/downloads/download_data.py @@ -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}") diff --git a/nerfstudio/utils/eval_utils.py b/nerfstudio/utils/eval_utils.py index 11a8b2341..8d07fe842 100644 --- a/nerfstudio/utils/eval_utils.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 585eee4f9..fd1b60bb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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",