From 36514797fda7c2d9e875e162351149588fbe3d9d Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Mon, 20 Feb 2023 10:12:41 -0800 Subject: [PATCH 1/9] Support multiple versions of stable diffusion (#1461) Co-authored-by: David McAllister --- nerfstudio/generative/stable_diffusion.py | 18 ++-- scripts/generative/trace_stable_diffusion.py | 105 ++++++++++++------- 2 files changed, 76 insertions(+), 47 deletions(-) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index fec3aa588..a3955712b 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -48,9 +48,11 @@ except ImportError: logging.set_verbosity_error() IMG_DIM = 512 CONST_SCALE = 0.18215 - -SD_SOURCE = "runwayml/stable-diffusion-v1-5" -CLIP_SOURCE = "openai/clip-vit-large-patch14" +SD_IDENTIFIERS = { + "1-5": "runwayml/stable-diffusion-v1-5", + "2-0": "stabilityai/stable-diffusion-2-base", + "2-1": "stabilityai/stable-diffusion-2-1-base", +} @dataclass @@ -87,7 +89,7 @@ class StableDiffusion(nn.Module): num_train_timesteps: number of training timesteps """ - def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000) -> None: + def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000, version="1-5") -> None: super().__init__() self.device = device @@ -104,14 +106,16 @@ class StableDiffusion(nn.Module): ) self.alphas = self.scheduler.alphas_cumprod.to(self.device) # type: ignore - pipe = StableDiffusionPipeline.from_pretrained(SD_SOURCE, torch_dtype=torch.float16) + sd_id = SD_IDENTIFIERS[version] + pipe = StableDiffusionPipeline.from_pretrained(sd_id, torch_dtype=torch.float16) assert pipe is not None pipe = pipe.to(self.device) pipe.enable_attention_slicing() # use jitted unet - unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / "sd_unet_traced.pt" + filename_sd_id = sd_id.split("/")[-1] + unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / f"{filename_sd_id}_unet_traced.pt" if unet_traced_filename.exists(): CONSOLE.print("Loading traced UNet.") unet_traced = torch.jit.load(unet_traced_filename) @@ -133,7 +137,7 @@ class StableDiffusion(nn.Module): del pipe.unet else: CONSOLE.print("[bold yellow] Warning: Loading UNet without JIT acceleration.") - CONSOLE.print(r"Run [yellow]python scripts/generative/trace_stable_diffusion.py[/yellow] for a speedup!") + CONSOLE.print(f"Run [yellow]ns-trace-sd --sd-version {version} [/yellow] for a speedup!") self.unet = pipe.unet self.unet.to(memory_format=torch.channels_last) diff --git a/scripts/generative/trace_stable_diffusion.py b/scripts/generative/trace_stable_diffusion.py index a501bd424..eee069eb3 100644 --- a/scripts/generative/trace_stable_diffusion.py +++ b/scripts/generative/trace_stable_diffusion.py @@ -2,6 +2,7 @@ import functools import sys +from dataclasses import dataclass from pathlib import Path import appdirs @@ -9,6 +10,8 @@ import torch import tyro from rich.console import Console +from nerfstudio.generative.stable_diffusion import SD_IDENTIFIERS + CONSOLE = Console(width=120) try: @@ -20,55 +23,77 @@ except ImportError: sys.exit(1) -def jit_unet(save_dir: Path = Path(appdirs.user_data_dir("nerfstudio"))): - """Trace Stable Diffusion UNet for speed improvement +@dataclass +class TraceSD: + """Traces the Stable Diffusion U-Net for better performance.""" - Args: - save_dir: directory to save the traced model - """ + sd_version: str + """Stable diffusion version.""" + save_dir: Path = Path(appdirs.user_data_dir("nerfstudio")) + """Directory to save the traced model.""" - # torch disable grad - torch.set_grad_enabled(False) + def main(self): + """Trace the Stable Diffusion U-Net and save it to disk.""" + # torch disable grad + torch.set_grad_enabled(False) - # load inputs - def generate_inputs(): - sample = torch.randn(2, 4, 64, 64).half().cuda() - timestep = torch.rand(1).half().cuda() * 999 - encoder_hidden_states = torch.randn(2, 77, 768).half().cuda() - return sample, timestep, encoder_hidden_states + if self.sd_version not in SD_IDENTIFIERS: + CONSOLE.print(f"[bold red]Invalid Stable Diffusion version; choose from {set(SD_IDENTIFIERS.keys())}") + sys.exit(1) - pipe = StableDiffusionPipeline.from_pretrained( - "runwayml/stable-diffusion-v1-5", - torch_dtype=torch.float16, - ).to("cuda") + sd_id = SD_IDENTIFIERS[self.sd_version] + pipe = StableDiffusionPipeline.from_pretrained( + sd_id, + torch_dtype=torch.float16, + ) + assert pipe is not None + pipe = pipe.to("cuda") + pipe.enable_attention_slicing() - pipe.enable_attention_slicing() + # load inputs + def generate_inputs(): + sample = torch.randn(2, 4, 64, 64).half().cuda() + timestep = torch.rand(1).half().cuda() * 999 - unet = pipe.unet - unet.eval() - unet.to(memory_format=torch.channels_last) # use channels_last memory format - unet.forward = functools.partial(unet.forward, return_dict=False) # set return_dict=False as default + # get shape of encoder hidden state (varies by SD version) + text_input = pipe.tokenizer("", return_tensors="pt") + _, w, d = pipe.text_encoder(text_input.input_ids.cuda())[0].shape - # warmup - inputs = None - for _ in range(3): - with torch.inference_mode(): - inputs = generate_inputs() - _ = unet(*inputs) + encoder_hidden_states = torch.randn(2, w, d).half().cuda() + return sample, timestep, encoder_hidden_states - # trace - print("tracing..") - assert inputs is not None - unet_traced = torch.jit.trace(unet, inputs) - unet_traced.eval() - print("done tracing") + unet = pipe.unet + unet.eval() + unet.to(memory_format=torch.channels_last) # use channels_last memory format + unet.forward = functools.partial(unet.forward, return_dict=False) # set return_dict=False as default - # save the traced model - unet_traced_filename = save_dir / "sd_unet_traced.pt" - if not unet_traced_filename.exists(): - unet_traced_filename.parent.mkdir(parents=True, exist_ok=True) - unet_traced.save(unet_traced_filename) + # warmup + inputs = None + for _ in range(3): + with torch.inference_mode(): + inputs = generate_inputs() + _ = unet(*inputs) + + # trace + print("tracing..") + assert inputs is not None + unet_traced = torch.jit.trace(unet, inputs) + unet_traced.eval() # type: ignore + print("done tracing") + + # save the traced model + filename_sd_id = sd_id.split("/")[-1] + unet_traced_filename = self.save_dir / f"{filename_sd_id}_unet_traced.pt" + if not unet_traced_filename.exists(): + unet_traced_filename.parent.mkdir(parents=True, exist_ok=True) + unet_traced.save(unet_traced_filename) # type: ignore + + +def entrypoint(): + """Entrypoint for use with pyproject scripts.""" + tyro.extras.set_accent_color("bright_yellow") + tyro.cli(TraceSD).main() if __name__ == "__main__": - tyro.cli(jit_unet) + entrypoint() From 796fad8bcbca62c4baf681081251690f93b82de3 Mon Sep 17 00:00:00 2001 From: Rohan Mathur Date: Mon, 20 Feb 2023 10:58:42 -0800 Subject: [PATCH 2/9] Equirectangular input (#1446) * started adding equirectangular to process_data * was multiplying frame by 255 twice * added equi2planar configs for the 3 presets * fixed presets * setup process data * changed variable name and comment * fixed middle preset * removed third preset * added progress bar and fixed video for equirectangular input * fixed whitespace stuff * added pyequilib dependencey * Updated for JPG support for equirectangular image input Updated process_data_utils for JPG support for equirectangular image input, since even though .jpg images worked, .JPG images (output from 360 camera software) were not working. * changed file type matching * added equirect utils * added license header * added docstring on equirect utils * added heuristic for image size * added heuristic for videos * fixed linty stuff * shortened code * raised error if no ims in directory * Updated ProcessVideo for fixing directories for equirectangular input Updated the ProcessVideo for equirectangular input to make the image directory have the correct perspective projections. The original equirectangular images are copied into a temp_images folder which also holds the perspective projections. This folder is deleted once the perspective projections are copied into the process data images folder. --------- Co-authored-by: cvachha Co-authored-by: cvachha --- nerfstudio/process_data/equirect_utils.py | 120 ++++++++++++++++++ nerfstudio/process_data/process_data_utils.py | 1 + pyproject.toml | 1 + scripts/process_data.py | 56 +++++++- 4 files changed, 173 insertions(+), 5 deletions(-) create mode 100644 nerfstudio/process_data/equirect_utils.py diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py new file mode 100644 index 000000000..596af58fb --- /dev/null +++ b/nerfstudio/process_data/equirect_utils.py @@ -0,0 +1,120 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Helper utils for processing equirectangular data.""" + +import os +from pathlib import Path +from typing import Tuple + +import cv2 +import numpy as np +import torch +from equilib import Equi2Pers +from rich.progress import ( + BarColumn, + Progress, + TaskProgressColumn, + TextColumn, + TimeRemainingColumn, +) + +from nerfstudio.utils.rich_utils import ItersPerSecColumn + + +def generate_planar_projections_from_equirectangular( + image_dir: Path, + planar_image_size: Tuple[int, int], + samples_per_im: int, +) -> Path: + """Generate planar projections from an equirectangular image. + + Args: + image_dir: The directory containing the equirectangular image. + planar_image_size: The size of the planar projections [width, height]. + samples_per_im: The number of samples to take per image. + returns: + The path to the planar projections directory. + """ + + device = torch.device("cuda") + + fov = 120 + yaw_pitch_pairs = [] + if samples_per_im == 8: + fov = 120 + for i in np.arange(-180, 180, 90): + yaw_pitch_pairs.append((i, 0)) + for i in np.arange(-180, 180, 180): + yaw_pitch_pairs.append((i, 45)) + for i in np.arange(-180, 180, 180): + yaw_pitch_pairs.append((i, -45)) + elif samples_per_im == 14: + fov = 110 + for i in np.arange(-180, 180, 60): + yaw_pitch_pairs.append((i, 0)) + for i in np.arange(-180, 180, 90): + yaw_pitch_pairs.append((i, 45)) + for i in np.arange(-180, 180, 90): + yaw_pitch_pairs.append((i, -45)) + + equi2pers = Equi2Pers(height=planar_image_size[1], width=planar_image_size[0], fov_x=fov, mode="bilinear") + frame_dir = image_dir + output_dir = image_dir / "planar_projections" + output_dir.mkdir(exist_ok=True) + num_ims = len(os.listdir(frame_dir)) + progress = Progress( + TextColumn("[bold blue]Generating Planar Images", justify="right"), + BarColumn(), + TaskProgressColumn(show_speed=True), + ItersPerSecColumn(suffix="equirect frames/s"), + TimeRemainingColumn(elapsed_when_finished=True, compact=True), + ) + + curr_im = 0 + with progress: + for i in progress.track(os.listdir(frame_dir), description="", total=num_ims): + if i.lower().endswith((".jpg", ".png", ".jpeg")): + im = np.array(cv2.imread(os.path.join(frame_dir, i))) + im = torch.tensor(im, dtype=torch.float32, device=device) + im = torch.permute(im, (2, 0, 1)) / 255.0 + count = 0 + for u_deg, v_deg in yaw_pitch_pairs: + v_rad = torch.pi * v_deg / 180.0 + u_rad = torch.pi * u_deg / 180.0 + pers_image = equi2pers(im, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) * 255.0 + pers_image = (pers_image.permute(1, 2, 0)).type(torch.uint8).to("cpu").numpy() + cv2.imwrite(f"{output_dir}/{i[:-4]}_{count}.jpg", pers_image) + count += 1 + curr_im += 1 + + return output_dir + + +def compute_resolution_from_equirect(image_dir: Path, num_images: int) -> Tuple[int, int]: + """Compute the resolution of the persepctive projections of equirectangular images + from the heuristic: num_image * res**2 = orig_height * orig_width. + + Args: + image_dir: The directory containing the equirectangular images. + returns: + The target resolution of the perspective projections. + """ + + for i in os.listdir(image_dir): + if i.lower().endswith((".jpg", ".png", ".jpeg")): + im = np.array(cv2.imread(os.path.join(image_dir, i))) + res_squared = (im.shape[0] * im.shape[1]) / num_images + return (int(np.sqrt(res_squared)), int(np.sqrt(res_squared))) + raise ValueError("No images found in the directory.") diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index 4919edd93..2d78469eb 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -43,6 +43,7 @@ class CameraModel(Enum): CAMERA_MODELS = { "perspective": CameraModel.OPENCV, "fisheye": CameraModel.OPENCV_FISHEYE, + "equirectangular": CameraModel.OPENCV, } diff --git a/pyproject.toml b/pyproject.toml index 73bdc21a6..2dae98036 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ dependencies = [ "Pillow>=9.3.0", "plotly>=5.7.0", "protobuf<=3.20.3,!=3.20.0", + "pyequilib>=0.5.6", "pymeshlab>=2022.2.post2", "pyngrok>=5.1.0", "python-socketio>=5.7.1", diff --git a/scripts/process_data.py b/scripts/process_data.py index 752e1d15c..7b6bb7922 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -3,6 +3,7 @@ import json +import shutil import sys import zipfile from dataclasses import dataclass @@ -16,6 +17,7 @@ from typing_extensions import Annotated, Literal from nerfstudio.process_data import ( colmap_utils, + equirect_utils, hloc_utils, insta360_utils, metashape_utils, @@ -45,7 +47,7 @@ class ProcessImages: """Path the data, either a video file or a directory of images.""" output_dir: Path """Path to the output directory.""" - camera_type: Literal["perspective", "fisheye"] = "perspective" + camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective" """Camera model to use.""" matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree" """Feature matching method to use. Vocab tree is recommended for a balance of speed and @@ -83,6 +85,10 @@ class ProcessImages: """ colmap_cmd: str = "colmap" """How to call the COLMAP executable.""" + images_per_equirect: Literal[8, 14] = 8 + """Number of samples per image to take from each equirectangular image. + Used only when camera-type is equirectangular. + """ gpu: bool = True """If True, use GPU.""" verbose: bool = False @@ -110,6 +116,14 @@ class ProcessImages: image_dir = self.output_dir / "images" image_dir.mkdir(parents=True, exist_ok=True) + # Generate planar projections if equirectangular + if self.camera_type == "equirectangular": + pers_size = equirect_utils.compute_resolution_from_equirect(self.data, self.images_per_equirect) + CONSOLE.log(f"Generating {self.images_per_equirect} {pers_size} sized images per equirectangular image") + self.data = equirect_utils.generate_planar_projections_from_equirectangular( + self.data, pers_size, self.images_per_equirect + ) + summary_log = [] # Copy and downscale images @@ -208,7 +222,7 @@ class ProcessVideo: """Path to the output directory.""" num_frames_target: int = 300 """Target number of frames to use for the dataset, results may not be exact.""" - camera_type: Literal["perspective", "fisheye"] = "perspective" + camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective" """Camera model to use.""" matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree" """Feature matching method to use. Vocab tree is recommended for a balance of speed and @@ -240,6 +254,10 @@ class ProcessVideo: """If True, skips COLMAP and generates transforms.json if possible.""" colmap_cmd: str = "colmap" """How to call the COLMAP executable.""" + images_per_equirect: Literal[8, 14] = 8 + """Number of samples per image to take from each equirectangular image. + Used only when camera-type is equirectangular. + """ percent_radius_crop: float = 1.0 """Create circle crop mask. The radius is the percent of the image diagonal.""" percent_crop: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) @@ -260,9 +278,37 @@ class ProcessVideo: summary_log = [] # Convert video to images - summary_log, num_extracted_frames = process_data_utils.convert_video_to_images( - self.data, image_dir=image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose - ) + if self.camera_type == "equirectangular": + # create temp images folder to store the equirect and perspective images + temp_image_dir = self.output_dir / "temp_images" + temp_image_dir.mkdir(parents=True, exist_ok=True) + summary_log, num_extracted_frames = process_data_utils.convert_video_to_images( + self.data, image_dir=temp_image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose + ) + else: + summary_log, num_extracted_frames = process_data_utils.convert_video_to_images( + self.data, image_dir=image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose + ) + + # Generate planar projections if equirectangular + if self.camera_type == "equirectangular": + perspective_image_size = equirect_utils.compute_resolution_from_equirect( + self.output_dir / "temp_images", self.images_per_equirect + ) + image_dir = equirect_utils.generate_planar_projections_from_equirectangular( + self.output_dir / "temp_images", perspective_image_size, self.images_per_equirect + ) + + # copy the perspective images to the image directory + process_data_utils.copy_images( + self.output_dir / "temp_images" / "planar_projections", + image_dir=self.output_dir / "images", + verbose=False, + ) + image_dir = self.output_dir / "images" + + # remove the temp_images folder + shutil.rmtree(self.output_dir / "temp_images", ignore_errors=True) # Create mask mask_path = process_data_utils.save_mask( From 151b1060fe8a205f6b31b22f7e1945faae6c3bd6 Mon Sep 17 00:00:00 2001 From: Sacha Lewin Date: Tue, 21 Feb 2023 18:17:57 +0100 Subject: [PATCH 3/9] Uniform Sampler in Nerfacto (#1464) * Allow black/white background color in Nerfacto * Use white background in Nerfacto Blender Benchmark * Fixed unused import * Updated index for license headers script * Removed debug code in NGP models * Uniform initial sampler config in Nerfacto * Undo changes in settings.json * Removed unwanted settings file * Removed unneeded config parameter * Use uniform sampler on Blender benchmark --- nerfstudio/model_components/ray_samplers.py | 1 - nerfstudio/models/nerfacto.py | 15 +++++++++++++-- scripts/benchmarking/launch_train_blender.sh | 2 +- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index ef0079ffd..866218ba1 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -393,7 +393,6 @@ class VolumetricSampler(Sampler): self.occupancy_grid = occupancy_grid if self.scene_aabb is not None: self.scene_aabb = self.scene_aabb.to("cuda").flatten() - print(self.scene_aabb) def get_sigma_fn(self, origins, directions) -> Optional[Callable]: """Returns a function that returns the density of a point. diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 4972dd5be..bf92968da 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -46,7 +46,10 @@ from nerfstudio.model_components.losses import ( orientation_loss, pred_normal_loss, ) -from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler +from nerfstudio.model_components.ray_samplers import ( + ProposalNetworkSampler, + UniformSampler, +) from nerfstudio.model_components.renderers import ( AccumulationRenderer, DepthRenderer, @@ -95,6 +98,8 @@ class NerfactoModelConfig(ModelConfig): ] ) """Arguments for the proposal density fields.""" + proposal_initial_sampler: Literal["piecewise", "uniform"] = "piecewise" + """Initial sampler for the proposal network. Piecewise is preferred for unbounded scenes.""" interlevel_loss_mult: float = 1.0 """Proposal loss multiplier.""" distortion_loss_mult: float = 0.002 @@ -171,12 +176,19 @@ class NerfactoModel(Model): 1, self.config.proposal_update_every, ) + + # Change proposal network initial sampler if uniform + initial_sampler = None # None is for piecewise as default (see ProposalNetworkSampler) + if self.config.proposal_initial_sampler == "uniform": + initial_sampler = UniformSampler(single_jitter=self.config.use_single_jitter) + self.proposal_sampler = ProposalNetworkSampler( num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray, num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray, num_proposal_network_iterations=self.config.num_proposal_iterations, single_jitter=self.config.use_single_jitter, update_sched=update_schedule, + initial_sampler=initial_sampler, ) # Collider @@ -312,7 +324,6 @@ class NerfactoModel(Model): def get_image_metrics_and_images( self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor] ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: - image = batch["image"].to(self.device) rgb = outputs["rgb"] acc = colormaps.apply_colormap(outputs["accumulation"]) diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh index 298e21749..8f7148c03 100644 --- a/scripts/benchmarking/launch_train_blender.sh +++ b/scripts/benchmarking/launch_train_blender.sh @@ -28,7 +28,7 @@ fi method_opts=() if [ "$method_name" = "nerfacto" ]; then # https://github.com/nerfstudio-project/nerfstudio/issues/806#issuecomment-1284327844 - method_opts=(--pipeline.model.background-color white --pipeline.model.near-plane 2. --pipeline.model.far-plane 6. --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-average-appearance-embedding False) + method_opts=(--pipeline.model.background-color white --pipeline.model.proposal-initial-sampler uniform --pipeline.model.near-plane 2. --pipeline.model.far-plane 6. --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-average-appearance-embedding False) fi shift $((OPTIND-1)) From 53e99c396b213c59bd0bf1945460738fdc2e9530 Mon Sep 17 00:00:00 2001 From: Dreek Fire Date: Tue, 21 Feb 2023 09:24:09 -0800 Subject: [PATCH 4/9] Add Windows installation tips to quickstart. (#1427) * add hashmap size and num_levels to instant_ngp_field, issue 1158 * Add Windows installation tips to quickstart * Include notes for CUDA 11.7 * separate duplicate targets in installation.md * change installation page to use tabs --- docs/quickstart/installation.md | 109 ++++++++++++++++++++++++++++++-- 1 file changed, 103 insertions(+), 6 deletions(-) diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md index 5d53a4d0c..81a64b024 100644 --- a/docs/quickstart/installation.md +++ b/docs/quickstart/installation.md @@ -1,8 +1,35 @@ # Installation ## Prerequisites +::::::{tab-set} +:::::{tab-item} Linux -CUDA must be installed on the system. This library has been tested with version 11.3. You can find more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html) +Install CUDA. This library has been tested with version 11.3. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html). + +::::: +:::::{tab-item} Windows +Install [Git](https://git-scm.com/downloads). + +Install CUDA. This library and the Windows installation process has been tested with version 11.3. + +::::{tab-set} +:::{tab-item} 11.3 with Visual Studio 2019 v16.9 + +Install Visual Studio 2019 ver 16.9. This must be done before installing CUDA. The necessary components are included in the `Desktop Development with C++` workflow (also called `C++ Build Tools` in the BuildTools edition). You can find older versions of Visual Studio 2019 [here](https://learn.microsoft.com/en-us/visualstudio/releases/2019/history#release-dates-and-build-numbers). The latest version, 16.11, may cause errors when installing tiny-cuda-nn. + +Install CUDA 11.3. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html). + +::: +:::{tab-item} 11.7 with Visual Studio 2022 + +Install Visual Studio 2022. This must be done before installing CUDA. The necessary components are included in the `Desktop Development with C++` workflow (also called `C++ Build Tools` in the BuildTools edition). + +Install CUDA 11.7. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html). + +::: +:::: +::::: +:::::: ## Create environment @@ -20,7 +47,8 @@ python -m pip install --upgrade pip (pytorch)= ### pytorch -Either pytorch 1.12.1 (with CUDA up to 11.3) or 1.13.1 (with CUDA 11.6/7/8) can be used. +::::{tab-set} +:::{tab-item} Torch 1.12.1 with CUDA 11.3 - To install 1.12.1 with CUDA 11.3: @@ -28,6 +56,9 @@ Either pytorch 1.12.1 (with CUDA up to 11.3) or 1.13.1 (with CUDA 11.6/7/8) can pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html ``` +::: +:::{tab-item} Torch 1.13.1 with CUDA 11.7 + - To install 1.13.1 with CUDA 11.7: Note that if a pytorch version prior to 1.13 is installed, @@ -37,12 +68,15 @@ it should be uninstalled first to avoid upgrade issues (e.g. with functorch) pip uninstall torch torchvision functorch ``` -Install pytorch 1.13.1 with CUDA (this repo has been tested with CUDA 11.7.1) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn) +Install pytorch 1.13.1 with CUDA and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn) ```bash pip install torch torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117 ``` +::: +:::: + ### tinycudann After pytorch, install the torch bindings for tinycudann: @@ -102,7 +136,7 @@ or be built from the repository using docker build --tag nerfstudio -f Dockerfile . ``` ### 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 reocmmended for usage as following: +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: ```bash docker run --gpus all \ # Give the container access to nvidia GPU (required). -v /folder/of/your/data:/workspace/ \ # Mount a folder from the local machine into the container to be able to process them (required). @@ -148,13 +182,15 @@ K80 | 37 ## Installation FAQ -- [TinyCUDA installation errors out with cuda mismatch](tiny-cuda-error) +- [TinyCUDA installation errors out with cuda mismatch](tiny-cuda-mismatch-error) +- [TinyCUDA installation errors out with no CUDA toolset found](tiny-cuda-integration-error) +- [TinyCUDA installation errors out with syntax errors](tiny-cuda-syntax-error) - [Installation errors, File "setup.py" not found](pip-install-error) - [Runtime errors, "len(sources) > 0".](cuda-sources-error)
-(tiny-cuda-error)= +(tiny-cuda-mismatch-error)= **TinyCUDA installation errors out with cuda mismatch** @@ -167,6 +203,67 @@ See [pytorch](pytorch) under Dependencies, above.
+(tiny-cuda-integration-error)= + +**(Windows) TinyCUDA installation errors out with no CUDA toolset found** + +While installing tiny-cuda on Windows, you run into: `No CUDA toolset found.` + +**Solution**: + +Confirm that you have Visual Studio installed (CUDA 11.3 is only compatible up to 2019 ver 16.9, not 16.11 or 2022). + +Make sure CUDA Visual Studio integration is enabled. This should be done automatically by the CUDA installer if it is run after Visual Studio is installed. You can also manually enable integration. + +::::{tab-set} +:::{tab-item} Visual Studio 2019 + +To manually enable integration for Visual Studio 2019, copy all 4 files from + +``` +C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.3\extras\visual_studio_integration\MSBuildExtensions +``` + +to + +``` +C:\Program Files (x86)\Microsoft Visual Studio\2019\BuildTools\MSBuild\Microsoft\VC\v160\BuildCustomizations +``` + +::: +:::{tab-item} Visual Studio 2022 + +To manually enable integration for Visual Studio 2022, copy all 4 files from + +``` +C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\extras\visual_studio_integration\MSBuildExtensions +``` + +to + +``` +C:\Program Files\Microsoft Visual Studio\2022\[Community, Professional, Enterprise, or BuildTools]\MSBuild\Microsoft\VC\v160\BuildCustomizations +``` + +::: +:::: + +
+ +(tiny-cuda-syntax-error)= + +**(Windows) TinyCUDA installation errors out with syntax errors** + +While installing tiny-cuda on Windows, you run into: `Expected a "("` + +**Solution**: + +If you are using CUDA 11.3 with Visual Studio 2019, confirm that Visual Studio is on version 16.9 and not 16.11. + +If you are using a different version of CUDA and Visual Studio, try switching to CUDA 11.3 with Visual Studio 2019 ver 16.9. + +
+ (pip-install-error)= **Installation errors, File "setup.py" not found** From f3a34bbf95e65d3c7be6066437716cfa16286152 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Tue, 21 Feb 2023 12:51:51 -0800 Subject: [PATCH 5/9] Refactor schedulers (#1466) * Refactor schedulers * lint * lint --- nerfstudio/cameras/camera_optimizers.py | 7 +- nerfstudio/configs/method_configs.py | 6 +- nerfstudio/engine/optimizers.py | 6 +- nerfstudio/engine/schedulers.py | 118 ++++++++++-------------- nerfstudio/models/tensorf.py | 10 +- 5 files changed, 68 insertions(+), 79 deletions(-) diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 20aaef3bc..afd002a71 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -31,7 +31,10 @@ from typing_extensions import Literal, assert_never from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3 from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.engine.optimizers import AdamOptimizerConfig -from nerfstudio.engine.schedulers import SchedulerConfig +from nerfstudio.engine.schedulers import ( + ExponentialDecaySchedulerConfig, + SchedulerConfig, +) from nerfstudio.utils import poses as pose_utils @@ -53,7 +56,7 @@ class CameraOptimizerConfig(InstantiateConfig): optimizer: AdamOptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15) """ADAM parameters for camera optimization.""" - scheduler: SchedulerConfig = SchedulerConfig(max_steps=10000) + scheduler: SchedulerConfig = ExponentialDecaySchedulerConfig(max_steps=10000) """Learning rate scheduler for camera optimizer..""" param_group: tyro.conf.Suppress[str] = "camera_opt" diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 4f41ba4af..6b1a71180 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -43,7 +43,7 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import ( PhototourismDataParserConfig, ) from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig -from nerfstudio.engine.schedulers import SchedulerConfig +from nerfstudio.engine.schedulers import ExponentialDecaySchedulerConfig from nerfstudio.engine.trainer import TrainerConfig from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig @@ -267,11 +267,11 @@ method_configs["tensorf"] = TrainerConfig( optimizers={ "fields": { "optimizer": AdamOptimizerConfig(lr=0.001), - "scheduler": SchedulerConfig(lr_final=0.0001, max_steps=30000), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000), }, "encodings": { "optimizer": AdamOptimizerConfig(lr=0.02), - "scheduler": SchedulerConfig(lr_final=0.002, max_steps=30000), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=30000), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py index 6878eccc2..786c51115 100644 --- a/nerfstudio/engine/optimizers.py +++ b/nerfstudio/engine/optimizers.py @@ -86,8 +86,10 @@ class Optimizers: self.optimizers[param_group_name] = config[param_group_name]["optimizer"].setup(params=params) self.parameters[param_group_name] = params if config[param_group_name]["scheduler"]: - self.schedulers[param_group_name] = config[param_group_name]["scheduler"].setup( - optimizer=self.optimizers[param_group_name], lr_init=lr_init + self.schedulers[param_group_name] = ( + config[param_group_name]["scheduler"] + .setup() + .get_scheduler(optimizer=self.optimizers[param_group_name], lr_init=lr_init) ) def optimizer_step(self, param_group_name: str) -> None: diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py index c545b9124..decdcd101 100644 --- a/nerfstudio/engine/schedulers.py +++ b/nerfstudio/engine/schedulers.py @@ -14,8 +14,9 @@ """Scheduler Classes""" +from abc import abstractmethod from dataclasses import dataclass, field -from typing import Any, Optional, Type +from typing import Type import numpy as np from torch.optim import Optimizer, lr_scheduler @@ -25,92 +26,71 @@ from nerfstudio.configs.base_config import InstantiateConfig @dataclass class SchedulerConfig(InstantiateConfig): - """Basic scheduler config with self-defined exponential decay schedule""" + """Basic scheduler config""" - _target: Type = field(default_factory=lambda: ExponentialDecaySchedule) + _target: Type = field(default_factory=lambda: Scheduler) + """_target: target class to instantiate""" + + +class Scheduler: # pylint: disable=too-few-public-methods + """Exponential learning rate decay function.""" + + config: SchedulerConfig + + def __init__(self, config: SchedulerConfig): + super().__init__() + self.config = config + + @abstractmethod + def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler: + """Abstract method that returns a scheduler object. + + Args: + optimizer: The optimizer to use. + lr_init: The initial learning rate. + Returns: + The scheduler object. + """ + + +@dataclass +class ExponentialDecaySchedulerConfig(SchedulerConfig): + """Exponential learning rate decay config""" + + _target: Type = field(default_factory=lambda: ExponentialDecayScheduler) lr_final: float = 0.000005 + """The final learning rate.""" max_steps: int = 1000000 - - # TODO: somehow make this more generic. i dont like the idea of overriding the setup function - # but also not sure how to go about passing things into predefined torch objects. - def setup(self, optimizer=None, lr_init=None, **kwargs) -> Any: - """Returns the instantiated object using the config.""" - return self._target(optimizer, lr_init, self.lr_final, self.max_steps) + """The maximum number of steps.""" + lr_delay_steps: int = 0 + """The number of steps to delay the learning rate.""" + lr_delay_mult: float = 1.0 + """The multiplier for the learning rate after the delay.""" -class ExponentialDecaySchedule(lr_scheduler.LambdaLR): +class ExponentialDecayScheduler(Scheduler): # pylint: disable=too-few-public-methods """Exponential learning rate decay function. See https://github.com/google-research/google-research/blob/ fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360 for details. - - Args: - optimizer: The optimizer to update. - lr_init: The initial learning rate. - lr_final: The final learning rate. - max_steps: The maximum number of steps. - lr_delay_steps: The number of steps to delay the learning rate. - lr_delay_mult: The multiplier for the learning rate after the delay. """ - config: SchedulerConfig + config: ExponentialDecaySchedulerConfig - def __init__(self, optimizer, lr_init, lr_final, max_steps, lr_delay_steps=0, lr_delay_mult=1.0) -> None: + def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler: def func(step): - if lr_delay_steps > 0: - delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin( - 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1) + if self.config.lr_delay_steps > 0: + delay_rate = self.config.lr_delay_mult + (1 - self.config.lr_delay_mult) * np.sin( + 0.5 * np.pi * np.clip(step / self.config.lr_delay_steps, 0, 1) ) else: delay_rate = 1.0 - t = np.clip(step / max_steps, 0, 1) - log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) + t = np.clip(step / self.config.max_steps, 0, 1) + log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(self.config.lr_final) * t) multiplier = ( log_lerp / lr_init ) # divided by lr_init because the multiplier is with the initial learning rate return delay_rate * multiplier - super().__init__(optimizer, lr_lambda=func) - - -class DelayerScheduler(lr_scheduler.LambdaLR): - """Starts with a flat lr schedule until it reaches N epochs then applies a given scheduler""" - - def __init__( - self, - optimizer: Optimizer, - lr_init, # pylint: disable=unused-argument - lr_final, # pylint: disable=unused-argument - max_steps, # pylint: disable=unused-argument - delay_epochs: int = 500, - after_scheduler: Optional[lr_scheduler.LambdaLR] = None, - ) -> None: - def func(step): - if step > delay_epochs: - if after_scheduler is not None: - multiplier = after_scheduler.lr_lambdas[0](step - delay_epochs) # type: ignore - return multiplier - return 1.0 - return 0.0 - - super().__init__(optimizer, lr_lambda=func) - - -class DelayedExponentialScheduler(DelayerScheduler): - """Delayer Scheduler with an Exponential Scheduler initialized afterwards.""" - - def __init__( - self, - optimizer: Optimizer, - lr_init, - lr_final, - max_steps, - delay_epochs: int = 200, - ): - after_scheduler = ExponentialDecaySchedule( - optimizer, - lr_init, - lr_final, - max_steps, - ) - super().__init__(optimizer, lr_init, lr_final, max_steps, delay_epochs, after_scheduler=after_scheduler) + scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=func) + return scheduler diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index d5af1579d..48310fde2 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -130,9 +130,13 @@ class TensoRFModel(Model): "optimizer" ].setup(params=enc) if optimizers_config["encodings"]["scheduler"]: - training_callback_attributes.optimizers.schedulers["encodings"] = optimizers_config["encodings"][ - "scheduler" - ].setup(optimizer=training_callback_attributes.optimizers.optimizers["encodings"], lr_init=lr_init) + training_callback_attributes.optimizers.schedulers["encodings"] = ( + optimizers_config["encodings"]["scheduler"] + .setup() + .get_scheduler( + optimizer=training_callback_attributes.optimizers.optimizers["encodings"], lr_init=lr_init + ) + ) callbacks = [ TrainingCallback( From 0d336a181e76588f3b3cb6ba6419d7efc2bcf297 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Tue, 21 Feb 2023 16:09:38 -0800 Subject: [PATCH 6/9] Update exponential scheduler (#1468) --- nerfstudio/engine/schedulers.py | 69 ++++++++++++++++++++------------- 1 file changed, 41 insertions(+), 28 deletions(-) diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py index decdcd101..d4554c7d8 100644 --- a/nerfstudio/engine/schedulers.py +++ b/nerfstudio/engine/schedulers.py @@ -12,14 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +# pylint: disable=too-few-public-methods + """Scheduler Classes""" from abc import abstractmethod from dataclasses import dataclass, field -from typing import Type +from typing import Optional, Type import numpy as np from torch.optim import Optimizer, lr_scheduler +from typing_extensions import Literal from nerfstudio.configs.base_config import InstantiateConfig @@ -29,11 +32,11 @@ class SchedulerConfig(InstantiateConfig): """Basic scheduler config""" _target: Type = field(default_factory=lambda: Scheduler) - """_target: target class to instantiate""" + """target class to instantiate""" -class Scheduler: # pylint: disable=too-few-public-methods - """Exponential learning rate decay function.""" +class Scheduler: + """Base scheduler""" config: SchedulerConfig @@ -55,42 +58,52 @@ class Scheduler: # pylint: disable=too-few-public-methods @dataclass class ExponentialDecaySchedulerConfig(SchedulerConfig): - """Exponential learning rate decay config""" + """Config for exponential decay scheduler with warmup""" _target: Type = field(default_factory=lambda: ExponentialDecayScheduler) - lr_final: float = 0.000005 - """The final learning rate.""" - max_steps: int = 1000000 + """target class to instantiate""" + lr_pre_warmup: float = 1e-8 + """Learning rate before warmup.""" + lr_final: Optional[float] = None + """Final learning rate. If not provided, it will be set to the optimizers learning rate.""" + warmup_steps: int = 0 + """Number of warmup steps.""" + max_steps: int = 100000 """The maximum number of steps.""" - lr_delay_steps: int = 0 - """The number of steps to delay the learning rate.""" - lr_delay_mult: float = 1.0 - """The multiplier for the learning rate after the delay.""" + ramp: Literal["linear", "cosine"] = "cosine" + """The ramp function to use during the warmup.""" -class ExponentialDecayScheduler(Scheduler): # pylint: disable=too-few-public-methods - """Exponential learning rate decay function. - See https://github.com/google-research/google-research/blob/ - fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360 - for details. +class ExponentialDecayScheduler(Scheduler): + """Exponential decay scheduler with linear warmup. Scheduler first ramps up to `lr_init` in `warmup_steps` + steps, then exponentially decays to `lr_final` in `max_steps` steps. """ config: ExponentialDecaySchedulerConfig def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler: + if self.config.lr_final is None: + lr_final = lr_init + else: + lr_final = self.config.lr_final + def func(step): - if self.config.lr_delay_steps > 0: - delay_rate = self.config.lr_delay_mult + (1 - self.config.lr_delay_mult) * np.sin( - 0.5 * np.pi * np.clip(step / self.config.lr_delay_steps, 0, 1) - ) + if step < self.config.warmup_steps: + if self.config.ramp == "cosine": + lr = self.config.lr_pre_warmup + (1 - self.config.lr_pre_warmup) * np.sin( + 0.5 * np.pi * np.clip(step / self.config.warmup_steps, 0, 1) + ) + else: + lr = ( + self.config.lr_pre_warmup + + (lr_init - self.config.lr_pre_warmup) * step / self.config.warmup_steps + ) else: - delay_rate = 1.0 - t = np.clip(step / self.config.max_steps, 0, 1) - log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(self.config.lr_final) * t) - multiplier = ( - log_lerp / lr_init - ) # divided by lr_init because the multiplier is with the initial learning rate - return delay_rate * multiplier + t = np.clip( + (step - self.config.warmup_steps) / (self.config.max_steps - self.config.warmup_steps), 0, 1 + ) + lr = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) + return lr / lr_init # divided by lr_init because the multiplier is with the initial learning rate scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=func) return scheduler From 207959da16f820efabe41d6d950569c341a33d79 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Tue, 21 Feb 2023 17:12:23 -0800 Subject: [PATCH 7/9] Warmups --- nerfstudio/configs/method_configs.py | 13 ++++++------ .../datamanagers/dreamfusion_datamanager.py | 11 +++++++--- .../generative/stable_diffusion_utils.py | 2 +- nerfstudio/models/dreamfusion.py | 21 ++++++++++--------- 4 files changed, 27 insertions(+), 20 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 99fbc1c75..e16f2c302 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -349,21 +349,22 @@ method_configs["dreamfusion"] = TrainerConfig( dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, + horizontal_rotation_warmup=4000, ), model=DreamFusionModelConfig( eval_num_rays_per_chunk=1 << 15, distortion_loss_mult=0.02, - interlevel_loss_mult=1.0, + interlevel_loss_mult=100.0, orientation_loss_mult=0.1, max_res=256, sphere_collider=True, initialize_density=False, random_background=True, - proposal_warmup=1000, + proposal_warmup=500, proposal_update_every=2, proposal_weights_anneal_max_num_iters=100, - start_lambertian_training=1000, - start_normals_training=500, + start_lambertian_training=500, + start_normals_training=1000, opacity_loss_mult=0.001, positional_prompting="discrete", guidance_scale=100, @@ -375,8 +376,8 @@ method_configs["dreamfusion"] = TrainerConfig( "scheduler": None, }, "fields": { - "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15), - "scheduler": None, + "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig(warmup_steps=2000, lr_final=1e-6, max_steps=20000), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py index 153913d42..9d7af9d3f 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py @@ -66,7 +66,7 @@ def random_train_pose( radius_mean: float = 1.0, radius_std: float = 0.1, central_rotation_range: Tuple[float, float] = (0, 360), - vertical_rotation_range: Tuple[float, float] = (-80, 20), + vertical_rotation_range: Tuple[float, float] = (-90, 0), focal_range: Tuple[float, float] = (0.75, 1.35), jitter_std: float = 0.01, center: Tuple[float, float, float] = (0, 0, 0), @@ -169,11 +169,14 @@ class DreamFusionDataManagerConfig(VanillaDataManagerConfig): """Std of radius of camera orbit""" focal_range: Tuple[float, float] = (0.6, 1.2) """Range of focal length""" - vertical_rotation_range: Tuple[float, float] = (-80, 30) + vertical_rotation_range: Tuple[float, float] = (-90, 0) """Range of vertical rotation""" jitter_std: float = 0.05 """Std of camera direction jitter, so we don't just point the cameras towards the center every time""" center: Tuple[float, float, float] = (0, 0, 0) + """Center coordinate of the camera sphere""" + horizontal_rotation_warmup: int = 0 + """How many steps until the full horizontal rotation range is used""" class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-method @@ -244,7 +247,8 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me # ).flatten() # return ray_bundle, {"initialization": False} - # TODO below + horizontal_range = min((step / self.config.horizontal_rotation_warmup), 1) * 180 + cameras, vertical_rotation, central_rotation = random_train_pose( self.config.train_images_per_batch, self.config.train_resolution, @@ -255,6 +259,7 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me vertical_rotation_range=self.config.vertical_rotation_range, jitter_std=self.config.jitter_std, center=self.config.center, + central_rotation_range=(-horizontal_range, horizontal_range), ) ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten() diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/stable_diffusion_utils.py index d7a94805e..ec5b29544 100644 --- a/nerfstudio/generative/stable_diffusion_utils.py +++ b/nerfstudio/generative/stable_diffusion_utils.py @@ -46,7 +46,7 @@ class PositionalTextEmbeddings: back_prompt: str, front_prompt: str, stable_diffusion: StableDiffusion, - positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete", + positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete" ): self.positional_prompting = positional_prompting self.sd_device = stable_diffusion.device diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 86e45e0e1..230de5f43 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -77,8 +77,10 @@ class DreamFusionModelConfig(ModelConfig): """Randomizes light source per output.""" initialize_density: bool = True """Initialize density in center of scene.""" - init_density_strength: float = 10 - """Initial strength of center density""" + taper_range: Tuple[int, int] = (0, 1000) + """Range of step values for the density tapering""" + taper_strength: Tuple[float, float] = (1.0, 0.0) + """Strength schedule of center density""" sphere_collider: bool = True """Use spherical collider instead of box""" random_background: bool = True @@ -126,7 +128,7 @@ class DreamFusionModelConfig(ModelConfig): """start training with lambertian shading after this many iterations""" opacity_penalty: bool = True """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" - opacity_loss_mult: float = 1 + opacity_loss_mult: float = 1e-3 """scale for opacity penalty""" max_res: int = 256 """Maximum resolution of the density field.""" @@ -135,7 +137,7 @@ class DreamFusionModelConfig(ModelConfig): """enables location based prompting""" interpolated_prompting: bool = False """enables interpolated location prompting""" - positional_prompting: Literal["discrete", "interpolated", "base"] = "discrete" + positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete" """ how to incorporate position into prompt""" top_prompt: str = ", overhead view" """appended to prompt for overhead view""" @@ -237,7 +239,7 @@ class DreamFusionModel(Model): single_jitter=self.config.use_single_jitter, update_sched=update_schedule, initial_sampler=UniformSampler(single_jitter=self.config.use_single_jitter), - pdf_histogram_padding=0.001, + pdf_histogram_padding=0.005, ) # renderers @@ -267,7 +269,7 @@ class DreamFusionModel(Model): def taper_density( self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument ): - self.density_strength -= 0.25 + self.density_strength = np.interp(step, self.config.taper_range, self.config.taper_strength) def start_training_normals( self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument @@ -303,6 +305,7 @@ class DreamFusionModel(Model): where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], iters=(150, 300, 450, 600), func=taper_density, + update_every_num_iters=1, args=[self, training_callback_attributes], ), TrainingCallback( @@ -352,11 +355,9 @@ class DreamFusionModel(Model): if self.initialize_density: pos = ray_samples.frustums.get_positions() density_blob = ( - self.config.init_density_strength - * self.density_strength - * torch.exp(-torch.norm(pos, dim=-1) * 10)[..., None] + self.density_strength * self.density_strength * torch.exp(-torch.norm(pos, dim=-1) * 10)[..., None] ) - density += density_blob + density = density + density_blob weights = ray_samples.get_weights(density) weights_list.append(weights) From 7a140a1f5b008f6212dbde86b7453da42a439d10 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Tue, 21 Feb 2023 17:35:32 -0800 Subject: [PATCH 8/9] tune --- nerfstudio/configs/method_configs.py | 4 ++-- nerfstudio/models/dreamfusion.py | 32 +++++++++++----------------- 2 files changed, 15 insertions(+), 21 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index e16f2c302..a03b8bcf1 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -361,7 +361,7 @@ method_configs["dreamfusion"] = TrainerConfig( initialize_density=False, random_background=True, proposal_warmup=500, - proposal_update_every=2, + proposal_update_every=0, proposal_weights_anneal_max_num_iters=100, start_lambertian_training=500, start_normals_training=1000, @@ -372,7 +372,7 @@ method_configs["dreamfusion"] = TrainerConfig( ), optimizers={ "proposal_networks": { - "optimizer": AdamOptimizerConfig(lr=5e-2, eps=1e-15), + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), "scheduler": None, }, "fields": { diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 6d94ed01d..4dd4b6b1f 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -128,7 +128,7 @@ class DreamFusionModelConfig(ModelConfig): """start training with lambertian shading after this many iterations""" opacity_penalty: bool = True """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" - opacity_loss_mult: float = 1e-3 + opacity_loss_mult: float = 1 """scale for opacity penalty""" max_res: int = 256 """Maximum resolution of the density field.""" @@ -303,7 +303,6 @@ class DreamFusionModel(Model): callbacks = [ TrainingCallback( where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - iters=(150, 300, 450, 600), func=taper_density, update_every_num_iters=1, args=[self, training_callback_attributes], @@ -354,9 +353,7 @@ class DreamFusionModel(Model): if self.initialize_density: pos = ray_samples.frustums.get_positions() - density_blob = ( - self.density_strength * self.density_strength * torch.exp(-torch.norm(pos, dim=-1) * 10)[..., None] - ) + density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] density = density + density_blob weights = ray_samples.get_weights(density) @@ -426,22 +423,19 @@ class DreamFusionModel(Model): else: outputs["train_output"] = outputs["rgb"] - if self.training: - outputs["rendered_orientation_loss"] = orientation_loss( - weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions - ) + outputs["rendered_orientation_loss"] = orientation_loss( + weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions + ) - outputs["rendered_pred_normal_loss"] = pred_normal_loss( - weights.detach(), - field_outputs[FieldHeadNames.NORMALS].detach(), - field_outputs[FieldHeadNames.PRED_NORMALS], - ) + outputs["rendered_pred_normal_loss"] = pred_normal_loss( + weights.detach(), + field_outputs[FieldHeadNames.NORMALS].detach(), + field_outputs[FieldHeadNames.PRED_NORMALS], + ) - assert weights.shape[-1] == 1 - if self.config.opacity_penalty: - outputs["opacity_loss"] = ( - torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult - ) + assert weights.shape[-1] == 1 + if self.config.opacity_penalty: + outputs["opacity_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult return outputs From b5ba3d97220f525a17ea0a46f9c00493ba7a7b18 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Tue, 21 Feb 2023 17:54:09 -0800 Subject: [PATCH 9/9] lint --- nerfstudio/configs/method_configs.py | 4 +++- nerfstudio/data/datamanagers/dreamfusion_datamanager.py | 2 +- nerfstudio/generative/stable_diffusion_utils.py | 4 ++-- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index a03b8bcf1..b4f7ec868 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -377,7 +377,9 @@ method_configs["dreamfusion"] = TrainerConfig( }, "fields": { "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15), - "scheduler": ExponentialDecaySchedulerConfig(warmup_steps=2000, lr_final=1e-6, max_steps=20000), + "scheduler": ExponentialDecaySchedulerConfig( + warmup_steps=2000, lr_final=1e-6, max_steps=20000, ramp="linear" + ), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py index 9d7af9d3f..48a686df1 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py @@ -248,7 +248,7 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me # return ray_bundle, {"initialization": False} horizontal_range = min((step / self.config.horizontal_rotation_warmup), 1) * 180 - + cameras, vertical_rotation, central_rotation = random_train_pose( self.config.train_images_per_batch, self.config.train_resolution, diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/stable_diffusion_utils.py index bf32f7ec7..141df766f 100644 --- a/nerfstudio/generative/stable_diffusion_utils.py +++ b/nerfstudio/generative/stable_diffusion_utils.py @@ -46,7 +46,7 @@ class PositionalTextEmbeddings: back_prompt: str, front_prompt: str, stable_diffusion: StableDiffusion, - positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete" + positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete", ): self.positional_prompting = positional_prompting self.sd_device = stable_diffusion.device @@ -105,7 +105,7 @@ class PositionalTextEmbeddings: text_embedding = text_embedding / 90.0 text_embedding = (vert * text_embedding + (90 - vert) * self.top_embed) / 90.0 else: - print('here') + print("here") text_embedding = self.base_embed return text_embedding