From 0e9e26dc7c8cf224e545a13586e990af7bbce9a5 Mon Sep 17 00:00:00 2001 From: Anson Kao Date: Tue, 31 Jan 2023 15:10:10 +0700 Subject: [PATCH 01/11] Explicitly require NVIDIA graphics cards in README (#1313) --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index d4c29ce9c..ca2bd48b4 100644 --- a/README.md +++ b/README.md @@ -81,7 +81,7 @@ For more complex changes (e.g., running with your own data/setting up a new NeRF ### Prerequisites -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) +You must have an NVIDIA video card with CUDA installed on the system. This library has been tested with version 11.3 of CUDA. You can find more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html) ### Create environment From e73692e599ca7047703874ec2d5bf66fe1f6d703 Mon Sep 17 00:00:00 2001 From: LSong <12553253+lsongx@users.noreply.github.com> Date: Tue, 31 Jan 2023 11:56:28 -0800 Subject: [PATCH 02/11] nerfplayer with ngp and nerfacto (#1277) * nerfplayer with ngp and nerfacto * update comments * update cpp extension build * avoid Windows repeated CUDA compile --------- Co-authored-by: machenmusik --- nerfstudio/configs/method_configs.py | 62 +++ nerfstudio/field_components/cuda/_backend.py | 3 +- nerfstudio/field_components/temporal_grid.py | 2 +- .../fields/nerfplayer_nerfacto_field.py | 415 ++++++++++++++++++ nerfstudio/fields/nerfplayer_ngp_field.py | 236 ++++++++++ nerfstudio/model_components/ray_samplers.py | 2 + nerfstudio/models/nerfplayer_nerfacto.py | 251 +++++++++++ nerfstudio/models/nerfplayer_ngp.py | 241 ++++++++++ nerfstudio/viewer/server/viewer_utils.py | 39 +- tests/test_train.py | 2 + 10 files changed, 1238 insertions(+), 15 deletions(-) create mode 100644 nerfstudio/fields/nerfplayer_nerfacto_field.py create mode 100644 nerfstudio/fields/nerfplayer_ngp_field.py create mode 100644 nerfstudio/models/nerfplayer_nerfacto.py create mode 100644 nerfstudio/models/nerfplayer_ngp.py diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 30682a242..947dad880 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -33,6 +33,7 @@ from nerfstudio.data.datamanagers.variable_res_datamanager import ( ) from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig +from nerfstudio.data.dataparsers.dycheck_dataparser import DycheckDataParserConfig from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( InstantNGPDataParserConfig, @@ -49,6 +50,8 @@ from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig from nerfstudio.models.instant_ngp import InstantNGPModelConfig from nerfstudio.models.mipnerf import MipNerfModel from nerfstudio.models.nerfacto import NerfactoModelConfig +from nerfstudio.models.nerfplayer_nerfacto import NerfplayerNerfactoModelConfig +from nerfstudio.models.nerfplayer_ngp import NerfplayerNGPModelConfig from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig from nerfstudio.models.tensorf import TensoRFModelConfig from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig @@ -67,6 +70,8 @@ descriptions = { "tensorf": "tensorf", "dnerf": "Dynamic-NeRF model. (slow)", "phototourism": "Uses the Phototourism data.", + "nerfplayer-nerfacto": "NeRFPlayer with nerfacto backbone.", + "nerfplayer-ngp": "NeRFPlayer with InstantNGP backbone.", } method_configs["nerfacto"] = TrainerConfig( @@ -325,6 +330,63 @@ method_configs["phototourism"] = TrainerConfig( vis="viewer", ) +method_configs["nerfplayer-nerfacto"] = TrainerConfig( + method_name="nerfplayer-nerfacto", + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, + pipeline=VanillaPipelineConfig( + datamanager=DepthDataManagerConfig( + dataparser=DycheckDataParserConfig(), + train_num_rays_per_batch=4096, + eval_num_rays_per_batch=4096, + camera_optimizer=CameraOptimizerConfig( + mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) + ), + ), + model=NerfplayerNerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), + ), + optimizers={ + "proposal_networks": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + "fields": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", +) + +method_configs["nerfplayer-ngp"] = TrainerConfig( + method_name="nerfplayer-ngp", + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, + pipeline=DynamicBatchPipelineConfig( + datamanager=DepthDataManagerConfig(dataparser=DycheckDataParserConfig(), train_num_rays_per_batch=8192), + model=NerfplayerNGPModelConfig( + eval_num_rays_per_chunk=8192, + contraction_type=ContractionType.AABB, + render_step_size=0.001, + max_num_samples_per_ray=48, + near_plane=0.01, + ), + ), + optimizers={ + "fields": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + } + }, + viewer=ViewerConfig(num_rays_per_chunk=64000), + vis="viewer", +) + AnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[ # Don't show unparseable (fixed) arguments in helptext. tyro.conf.FlagConversionOff[ tyro.extras.subcommand_type_from_defaults(defaults=method_configs, descriptions=descriptions) diff --git a/nerfstudio/field_components/cuda/_backend.py b/nerfstudio/field_components/cuda/_backend.py index d4269b9ee..17732a196 100644 --- a/nerfstudio/field_components/cuda/_backend.py +++ b/nerfstudio/field_components/cuda/_backend.py @@ -27,9 +27,10 @@ BUILD_DIR = _get_build_directory(NAME, verbose=False) _C = None if cuda_toolkit_available(): - if os.path.exists(os.path.join(BUILD_DIR, f"{NAME}.so")): + if os.listdir(BUILD_DIR) != []: # If the build exists, we assume the extension has been built # and we can load it. + print("nerfstudio field components: CUDA set up, loading (should be quick)") _C = load( name=NAME, sources=glob.glob(os.path.join(PATH, "csrc/*.cu")), diff --git a/nerfstudio/field_components/temporal_grid.py b/nerfstudio/field_components/temporal_grid.py index c4f5e20c5..6bccc110a 100644 --- a/nerfstudio/field_components/temporal_grid.py +++ b/nerfstudio/field_components/temporal_grid.py @@ -182,7 +182,7 @@ class TemporalGridEncoder(nn.Module): input_dim: int = 3, num_levels: int = 16, level_dim: int = 2, - per_level_scale: int = 2, + per_level_scale: float = 2.0, base_resolution: int = 16, log2_hashmap_size: int = 19, desired_resolution: Optional[int] = None, diff --git a/nerfstudio/fields/nerfplayer_nerfacto_field.py b/nerfstudio/fields/nerfplayer_nerfacto_field.py new file mode 100644 index 000000000..e87df3b03 --- /dev/null +++ b/nerfstudio/fields/nerfplayer_nerfacto_field.py @@ -0,0 +1,415 @@ +# 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. + +""" +Field implementations for NeRFPlayer (https://arxiv.org/abs/2210.15947) implementation with nerfacto backbone +""" + + +from typing import Optional + +import numpy as np +import torch +from torch.nn.parameter import Parameter +from torchtyping import TensorType + +from nerfstudio.cameras.rays import Frustums, RaySamples +from nerfstudio.data.scene_box import SceneBox +from nerfstudio.field_components.activations import trunc_exp +from nerfstudio.field_components.embedding import Embedding +from nerfstudio.field_components.field_heads import ( + FieldHeadNames, + PredNormalsFieldHead, + SemanticFieldHead, + TransientDensityFieldHead, + TransientRGBFieldHead, + UncertaintyFieldHead, +) +from nerfstudio.field_components.spatial_distortions import SpatialDistortion +from nerfstudio.field_components.temporal_grid import TemporalGridEncoder +from nerfstudio.fields.base_field import Field + +try: + import tinycudann as tcnn +except ImportError: + # tinycudann module doesn't exist + pass + + +def get_normalized_directions(directions: TensorType["bs":..., 3]): + """SH encoding must be in the range [0, 1] + + Args: + directions: batch of directions + """ + return (directions + 1.0) / 2.0 + + +class TemporalHashMLPDensityField(Field): + """A lightweight temporal density field module. + + Args: + aabb: Parameters of scene aabb bounds + temporal_dim: Hashing grid parameter. A higher temporal dim means a higher temporal frequency. + num_layers: Number of hidden layers + hidden_dim: Dimension of hidden layers + spatial_distortion: Spatial distortion module + num_levels: Hashing grid parameter. Used for initialize TemporalGridEncoder class. + max_res: Hashing grid parameter. Used for initialize TemporalGridEncoder class. + base_res: Hashing grid parameter. Used for initialize TemporalGridEncoder class. + log2_hashmap_size: Hashing grid parameter. Used for initialize TemporalGridEncoder class. + features_per_level: Hashing grid parameter. Used for initialize TemporalGridEncoder class. + """ + + def __init__( + self, + aabb, + temporal_dim: int = 64, + num_layers: int = 2, + hidden_dim: int = 64, + spatial_distortion: Optional[SpatialDistortion] = None, + num_levels=8, + max_res=1024, + base_res=16, + log2_hashmap_size=18, + features_per_level=2, + ) -> None: + super().__init__() + # from .temporal_grid import test; test() # DEBUG + self.aabb = Parameter(aabb, requires_grad=False) + self.spatial_distortion = spatial_distortion + growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) + + self.encoding = TemporalGridEncoder( + input_dim=3, + temporal_dim=temporal_dim, + num_levels=num_levels, + level_dim=features_per_level, + per_level_scale=growth_factor, + base_resolution=base_res, + log2_hashmap_size=log2_hashmap_size, + ) + self.linear = tcnn.Network( + n_input_dims=num_levels * features_per_level, + n_output_dims=1, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": hidden_dim, + "n_hidden_layers": num_layers - 1, + }, + ) + + # pylint: disable=arguments-differ + def density_fn(self, positions: TensorType["bs":..., 3], times: TensorType["bs", 1]) -> TensorType["bs":..., 1]: + """Returns only the density. Used primarily with the density grid. + + Args: + positions: the origin of the samples/frustums + times: the time of rays + """ + if len(positions.shape) == 3 and len(times.shape) == 2: + # position is [ray, sample, 3]; times is [ray, 1] + times = times[:, None] # RaySamples can handle the shape + # Need to figure out a better way to descibe positions with a ray. + ray_samples = RaySamples( + frustums=Frustums( + origins=positions, + directions=torch.ones_like(positions), + starts=torch.zeros_like(positions[..., :1]), + ends=torch.zeros_like(positions[..., :1]), + pixel_area=torch.ones_like(positions[..., :1]), + ), + times=times, + ) + density, _ = self.get_density(ray_samples) + return density + + def get_density(self, ray_samples: RaySamples): + if self.spatial_distortion is not None: + positions = self.spatial_distortion(ray_samples.frustums.get_positions()) + positions = (positions + 2.0) / 4.0 + else: + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + positions_flat = positions.view(-1, 3) + time_flat = ray_samples.times.reshape(-1, 1) + x = self.encoding(positions_flat, time_flat).to(positions) + density_before_activation = self.linear(x).view(*ray_samples.frustums.shape, -1) + + # Rectifying the density with an exponential is much more stable than a ReLU or + # softplus, because it enables high post-activation (float32) density outputs + # from smaller internal (float16) parameters. + density = trunc_exp(density_before_activation) + return density, None + + def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None): + return {} + + +class NerfplayerNerfactoField(Field): + """NeRFPlayer (https://arxiv.org/abs/2210.15947) field with nerfacto backbone. + + Args: + aabb: parameters of scene aabb bounds + num_images: number of images in the dataset + num_layers: number of hidden layers + hidden_dim: dimension of hidden layers + geo_feat_dim: output geo feat dimensions + num_levels: number of levels of the hashmap for the base mlp + max_res: maximum resolution of the hashmap for the base mlp + log2_hashmap_size: size of the hashmap for the base mlp + num_layers_color: number of hidden layers for color network + num_layers_transient: number of hidden layers for transient network + hidden_dim_color: dimension of hidden layers for color network + hidden_dim_transient: dimension of hidden layers for transient network + appearance_embedding_dim: dimension of appearance embedding + transient_embedding_dim: dimension of transient embedding + use_transient_embedding: whether to use transient embedding + use_semantics: whether to use semantic segmentation + num_semantic_classes: number of semantic classes + use_pred_normals: whether to use predicted normals + use_average_appearance_embedding: whether to use average appearance embedding or zeros for inference + spatial_distortion: spatial distortion to apply to the scene + """ + + def __init__( + self, + aabb, + num_images: int, + num_layers: int = 2, + hidden_dim: int = 64, + geo_feat_dim: int = 15, + temporal_dim: int = 64, + num_levels: int = 16, + features_per_level: int = 2, + log2_hashmap_size: int = 19, + num_layers_color: int = 3, + num_layers_transient: int = 2, + hidden_dim_color: int = 64, + hidden_dim_transient: int = 64, + appearance_embedding_dim: int = 32, + transient_embedding_dim: int = 16, + use_transient_embedding: bool = False, + use_semantics: bool = False, + num_semantic_classes: int = 100, + use_pred_normals: bool = False, + use_average_appearance_embedding: bool = False, + spatial_distortion: Optional[SpatialDistortion] = None, + ) -> None: + super().__init__() + + self.aabb = Parameter(aabb, requires_grad=False) + self.geo_feat_dim = geo_feat_dim + + self.spatial_distortion = spatial_distortion + self.num_images = num_images + self.appearance_embedding_dim = appearance_embedding_dim + self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim) + self.use_average_appearance_embedding = use_average_appearance_embedding + self.use_transient_embedding = use_transient_embedding + self.use_semantics = use_semantics + self.use_pred_normals = use_pred_normals + + self.direction_encoding = tcnn.Encoding( + n_input_dims=3, + encoding_config={ + "otype": "SphericalHarmonics", + "degree": 4, + }, + ) + + self.position_encoding = tcnn.Encoding( + n_input_dims=3, + encoding_config={"otype": "Frequency", "n_frequencies": 2}, + ) + + self.mlp_base = TemporalGridEncoder( + input_dim=3, + temporal_dim=temporal_dim, + num_levels=num_levels, + level_dim=features_per_level, + log2_hashmap_size=log2_hashmap_size, + desired_resolution=1024 * (self.aabb.max() - self.aabb.min()), + ) + self.mlp_base_decode = tcnn.Network( + n_input_dims=num_levels * features_per_level, + n_output_dims=1 + self.geo_feat_dim, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": hidden_dim, + "n_hidden_layers": num_layers - 1, + }, + ) + + # transients + if self.use_transient_embedding: + self.transient_embedding_dim = transient_embedding_dim + self.embedding_transient = Embedding(self.num_images, self.transient_embedding_dim) + self.mlp_transient = tcnn.Network( + n_input_dims=self.geo_feat_dim + self.transient_embedding_dim, + n_output_dims=hidden_dim_transient, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": hidden_dim_transient, + "n_hidden_layers": num_layers_transient - 1, + }, + ) + self.field_head_transient_uncertainty = UncertaintyFieldHead(in_dim=self.mlp_transient.n_output_dims) + self.field_head_transient_rgb = TransientRGBFieldHead(in_dim=self.mlp_transient.n_output_dims) + self.field_head_transient_density = TransientDensityFieldHead(in_dim=self.mlp_transient.n_output_dims) + + # semantics + if self.use_semantics: + self.mlp_semantics = tcnn.Network( + n_input_dims=self.geo_feat_dim, + n_output_dims=hidden_dim_transient, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": 64, + "n_hidden_layers": 1, + }, + ) + self.field_head_semantics = SemanticFieldHead( + in_dim=self.mlp_semantics.n_output_dims, num_classes=num_semantic_classes + ) + + # predicted normals + if self.use_pred_normals: + self.mlp_pred_normals = tcnn.Network( + n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims, + n_output_dims=hidden_dim_transient, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": 64, + "n_hidden_layers": 2, + }, + ) + self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims) + + self.mlp_head = tcnn.Network( + n_input_dims=self.direction_encoding.n_output_dims + self.geo_feat_dim + self.appearance_embedding_dim, + n_output_dims=3, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "Sigmoid", + "n_neurons": hidden_dim_color, + "n_hidden_layers": num_layers_color - 1, + }, + ) + + def get_density(self, ray_samples: RaySamples): + """Computes and returns the densities.""" + if self.spatial_distortion is not None: + positions = ray_samples.frustums.get_positions() + positions = self.spatial_distortion(positions) + positions = (positions + 2.0) / 4.0 + else: + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + positions_flat = positions.view(-1, 3) + assert ray_samples.times is not None, "Time should be included in the input for NeRFPlayer" + time_flat = ray_samples.times.reshape(-1, 1) + h = self.mlp_base(positions_flat, time_flat) + h = self.mlp_base_decode(h).view(*ray_samples.frustums.shape, -1) + density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) + + # Rectifying the density with an exponential is much more stable than a ReLU or + # softplus, because it enables high post-activation (float32) density outputs + # from smaller internal (float16) parameters. + density = trunc_exp(density_before_activation.to(positions)) + return density, base_mlp_out + + def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None): + assert density_embedding is not None + outputs = {} + if ray_samples.camera_indices is None: + raise AttributeError("Camera indices are not provided.") + camera_indices = ray_samples.camera_indices.squeeze() + directions = get_normalized_directions(ray_samples.frustums.directions) + directions_flat = directions.view(-1, 3) + d = self.direction_encoding(directions_flat) + + outputs_shape = ray_samples.frustums.directions.shape[:-1] + + # appearance + if self.training: + embedded_appearance = self.embedding_appearance(camera_indices) + else: + if self.use_average_appearance_embedding: + embedded_appearance = torch.ones( + (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device + ) * self.embedding_appearance.mean(dim=0) + else: + embedded_appearance = torch.zeros( + (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device + ) + + # transients + if self.use_transient_embedding and self.training: + embedded_transient = self.embedding_transient(camera_indices) + transient_input = torch.cat( + [ + density_embedding.view(-1, self.geo_feat_dim), + embedded_transient.view(-1, self.transient_embedding_dim), + ], + dim=-1, + ) + x = self.mlp_transient(transient_input).view(*outputs_shape, -1).to(directions) + outputs[FieldHeadNames.UNCERTAINTY] = self.field_head_transient_uncertainty(x) + outputs[FieldHeadNames.TRANSIENT_RGB] = self.field_head_transient_rgb(x) + outputs[FieldHeadNames.TRANSIENT_DENSITY] = self.field_head_transient_density(x) + + # semantics + if self.use_semantics: + density_embedding_copy = density_embedding.clone().detach() + semantics_input = torch.cat( + [ + density_embedding_copy.view(-1, self.geo_feat_dim), + ], + dim=-1, + ) + x = self.mlp_semantics(semantics_input).view(*outputs_shape, -1).to(directions) + outputs[FieldHeadNames.SEMANTICS] = self.field_head_semantics(x) + + # predicted normals + if self.use_pred_normals: + positions = ray_samples.frustums.get_positions() + + positions_flat = self.position_encoding(positions.view(-1, 3)) + pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) + + x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions) + outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x) + + h = torch.cat( + [ + d, + density_embedding.view(-1, self.geo_feat_dim), + embedded_appearance.view(-1, self.appearance_embedding_dim), + ], + dim=-1, + ) + rgb = self.mlp_head(h).view(*outputs_shape, -1).to(directions) + outputs.update({FieldHeadNames.RGB: rgb}) + + return outputs diff --git a/nerfstudio/fields/nerfplayer_ngp_field.py b/nerfstudio/fields/nerfplayer_ngp_field.py new file mode 100644 index 000000000..ed1cf250e --- /dev/null +++ b/nerfstudio/fields/nerfplayer_ngp_field.py @@ -0,0 +1,236 @@ +# 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. + +""" +NeRFPlayer (https://arxiv.org/abs/2210.15947) field implementations with InstantNGP backbone +""" + + +from typing import Optional + +import torch +from nerfacc import ContractionType, contract +from torch.nn.parameter import Parameter +from torchtyping import TensorType + +from nerfstudio.cameras.rays import Frustums, RaySamples +from nerfstudio.data.scene_box import SceneBox +from nerfstudio.field_components.activations import trunc_exp +from nerfstudio.field_components.embedding import Embedding +from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.field_components.temporal_grid import TemporalGridEncoder +from nerfstudio.fields.base_field import Field +from nerfstudio.fields.instant_ngp_field import get_normalized_directions + +try: + import tinycudann as tcnn +except ImportError: + # tinycudann module doesn't exist + pass + + +class NerfplayerNGPField(Field): + """NeRFPlayer (https://arxiv.org/abs/2210.15947) field with InstantNGP backbone. + + Args: + aabb: parameters of scene aabb bounds + temporal_dim: the dimension of temporal axis, a higher dimension indicates a higher temporal frequency + please refer to the implementation of TemporalGridEncoder for more details + num_levels: the number of multi-resolution levels; same as InstantNGP + features_per_level: the dim of output feature vector for each level; same as InstantNGP + log2_hashmap_size: the size of the table; same as InstantNGP + base_resolution: base resolution for the table; same as InstantNGP + num_layers: number of hidden layers (occupancy decoder network after sampling) + hidden_dim: dimension of hidden layers (occupancy decoder network after sampling) + geo_feat_dim: output geo feat dimensions + num_layers_color: number of hidden layers for color network + hidden_dim_color: dimension of hidden layers for color network + use_appearance_embedding: whether to use appearance embedding + disable_viewing_dependent: if true, disable the viewing dependent effect (no viewing direction as inputs) + Sometimes we need to disable viewing dependent effects in a dynamic scene, because there is + ambiguity between being dynamic and viewing dependent effects. For example, the shadow of the camera + should be a dynamic effect, but may be reconstructed as viewing dependent effects. + num_images: number of images, requried if use_appearance_embedding is True + appearance_embedding_dim: dimension of appearance embedding + contraction_type: type of contraction + """ + + def __init__( + self, + aabb, + temporal_dim: int = 16, + num_levels: int = 16, + features_per_level: int = 2, + log2_hashmap_size: int = 19, + base_resolution: int = 16, + num_layers: int = 2, + hidden_dim: int = 64, + geo_feat_dim: int = 15, + num_layers_color: int = 3, + hidden_dim_color: int = 64, + use_appearance_embedding: bool = False, + disable_viewing_dependent: bool = False, + num_images: Optional[int] = None, + appearance_embedding_dim: int = 32, + contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE, + ) -> None: + super().__init__() + + self.aabb = Parameter(aabb, requires_grad=False) + self.geo_feat_dim = geo_feat_dim + self.contraction_type = contraction_type + + self.use_appearance_embedding = use_appearance_embedding + if use_appearance_embedding: + assert num_images is not None + self.appearance_embedding_dim = appearance_embedding_dim + self.appearance_embedding = Embedding(num_images, appearance_embedding_dim) + + self.direction_encoding = tcnn.Encoding( + n_input_dims=3, + encoding_config={ + "otype": "SphericalHarmonics", + "degree": 4, + }, + ) + + self.mlp_base = TemporalGridEncoder( + input_dim=3, + temporal_dim=temporal_dim, + num_levels=num_levels, + level_dim=features_per_level, + base_resolution=base_resolution, + log2_hashmap_size=log2_hashmap_size, + desired_resolution=1024 * (self.aabb.max() - self.aabb.min()), + ) + self.mlp_base_decode = tcnn.Network( + n_input_dims=num_levels * features_per_level, + n_output_dims=1 + self.geo_feat_dim, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": hidden_dim, + "n_hidden_layers": num_layers - 1, + }, + ) + + in_dim = self.direction_encoding.n_output_dims + self.geo_feat_dim + if disable_viewing_dependent: + in_dim = self.geo_feat_dim + self.direction_encoding = None + if self.use_appearance_embedding: + in_dim += self.appearance_embedding_dim + self.mlp_head = tcnn.Network( + n_input_dims=in_dim, + n_output_dims=3, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "Sigmoid", + "n_neurons": hidden_dim_color, + "n_hidden_layers": num_layers_color - 1, + }, + ) + + def get_density(self, ray_samples: RaySamples): + positions = ray_samples.frustums.get_positions() + positions_flat = positions.view(-1, 3) + positions_flat = contract(x=positions_flat, roi=self.aabb, type=self.contraction_type) + assert ray_samples.times is not None, "Time should be included in the input for NeRFPlayer" + times_flat = ray_samples.times.view(-1, 1) + + h = self.mlp_base(positions_flat, times_flat) + h = self.mlp_base_decode(h).view(*ray_samples.frustums.shape, -1) + density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) + + # Rectifying the density with an exponential is much more stable than a ReLU or + # softplus, because it enables high post-activation (float32) density outputs + # from smaller internal (float16) parameters. + density = trunc_exp(density_before_activation.to(positions)) + return density, base_mlp_out + + def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None): + directions = get_normalized_directions(ray_samples.frustums.directions) + directions_flat = directions.view(-1, 3) + + if self.direction_encoding is not None: + d = self.direction_encoding(directions_flat) + if density_embedding is None: + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + h = torch.cat([d, positions.view(-1, 3)], dim=-1) + else: + h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) + else: + # viewing direction is disabled + if density_embedding is None: + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + h = positions.view(-1, 3) + else: + h = density_embedding.view(-1, self.geo_feat_dim) + + if self.use_appearance_embedding: + if ray_samples.camera_indices is None: + raise AttributeError("Camera indices are not provided.") + camera_indices = ray_samples.camera_indices.squeeze() + if self.training: + embedded_appearance = self.appearance_embedding(camera_indices) + else: + embedded_appearance = torch.zeros( + (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device + ) + h = torch.cat([h, embedded_appearance.view(-1, self.appearance_embedding_dim)], dim=-1) + + rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions) + return {FieldHeadNames.RGB: rgb} + + # pylint: disable=arguments-differ + def density_fn(self, positions: TensorType["bs":..., 3], times: TensorType["bs":..., 1]) -> TensorType["bs":..., 1]: + """Returns only the density. Used primarily with the density grid. + Overwrite this function since density is time dependent now. + + Args: + positions: the origin of the samples/frustums + times: the time of each position + """ + ray_samples = RaySamples( + frustums=Frustums( + origins=positions, + directions=torch.ones_like(positions), + starts=torch.zeros_like(positions[..., :1]), + ends=torch.zeros_like(positions[..., :1]), + pixel_area=torch.ones_like(positions[..., :1]), + ), + times=times, + ) + density, _ = self.get_density(ray_samples) + return density + + def get_opacity(self, positions: TensorType["bs":..., 3], step_size, time_intervals=10) -> TensorType["bs":..., 1]: + """Returns the opacity for a position and time. Used primarily by the occupancy grid. + This will return the maximum opacity of the points in the space in a dynamic sequence. + + Args: + positions: the positions to evaluate the opacity at. + step_size: the step size to use for the opacity evaluation. + time_intervals: sample density on N time stamps + """ + # TODO: Converting opacity by time intervals is slow, and may lead to temporal artifacts. + # (Maybe random sample time and EMA?) + opacity = [] + for t in range(0, time_intervals): + density = self.density_fn(positions, t / (time_intervals - 1) * torch.ones_like(positions)[..., [0]]) + opacity.append(density * step_size) + opacity = torch.stack(opacity, dim=0).max(dim=0).values + return opacity diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index e41416053..5df52d2fa 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -503,6 +503,8 @@ class VolumetricSampler(Sampler): ), camera_indices=camera_indices, ) + if ray_bundle.times is not None: + ray_samples.times = ray_bundle.times[ray_indices] return ray_samples, ray_indices diff --git a/nerfstudio/models/nerfplayer_nerfacto.py b/nerfstudio/models/nerfplayer_nerfacto.py new file mode 100644 index 000000000..f35a748b1 --- /dev/null +++ b/nerfstudio/models/nerfplayer_nerfacto.py @@ -0,0 +1,251 @@ +# 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. + +""" +NeRFPlayer (https://arxiv.org/abs/2210.15947) implementation with nerfacto backbone. +""" + +from __future__ import annotations + +import functools +from dataclasses import dataclass, field +from typing import Dict, List, Type + +import numpy as np +import torch +from torchmetrics import PeakSignalNoiseRatio +from torchmetrics.functional import structural_similarity_index_measure +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity +from typing_extensions import Literal + +from nerfstudio.cameras.rays import RayBundle +from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.field_components.spatial_distortions import SceneContraction +from nerfstudio.fields.nerfplayer_nerfacto_field import ( + NerfplayerNerfactoField, + TemporalHashMLPDensityField, +) +from nerfstudio.model_components.losses import ( + MSELoss, + interlevel_loss, + orientation_loss, + pred_normal_loss, +) +from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler +from nerfstudio.model_components.renderers import ( + AccumulationRenderer, + DepthRenderer, + NormalsRenderer, + RGBRenderer, +) +from nerfstudio.model_components.scene_colliders import NearFarCollider +from nerfstudio.models.base_model import Model +from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig + + +@dataclass +class NerfplayerNerfactoModelConfig(NerfactoModelConfig): + """Nerfplayer Model Config with Nerfacto backbone""" + + _target: Type = field(default_factory=lambda: NerfplayerNerfactoModel) + near_plane: float = 0.05 + """How far along the ray to start sampling.""" + far_plane: float = 1000.0 + """How far along the ray to stop sampling.""" + background_color: Literal["random", "last_sample"] = "random" + """Whether to randomize the background color. (Random is reported to be better on DyCheck.)""" + num_levels: int = 16 + """Hashing grid parameter.""" + features_per_level: int = 2 + """Hashing grid parameter.""" + log2_hashmap_size: int = 18 + """Hashing grid parameter.""" + temporal_dim: int = 32 + """Hashing grid parameter. A higher temporal dim means a higher temporal frequency.""" + proposal_net_args_list: List[Dict] = field( + default_factory=lambda: [ + {"hidden_dim": 16, "temporal_dim": 32, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 64}, + {"hidden_dim": 16, "temporal_dim": 32, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256}, + ] + ) + """Arguments for the proposal density fields.""" + distortion_loss_mult: float = 1e-2 + """Distortion loss multiplier.""" + temporal_tv_weight: float = 1 + """Temporal TV balancing weight for feature channels.""" + depth_weight: float = 1e-1 + """depth loss balancing weight for feature channels.""" + + +class NerfplayerNerfactoModel(NerfactoModel): + """Nerfplayer model with Nerfacto backbone. + + Args: + config: Nerfplayer configuration to instantiate model + """ + + config: NerfplayerNerfactoModelConfig + + def populate_modules(self): + """Set the fields and modules.""" + Model.populate_modules(self) + + scene_contraction = SceneContraction(order=float("inf")) + + # Fields + self.field = NerfplayerNerfactoField( + self.scene_box.aabb, + temporal_dim=self.config.temporal_dim, + num_levels=self.config.num_levels, + features_per_level=self.config.features_per_level, + log2_hashmap_size=self.config.log2_hashmap_size, + spatial_distortion=scene_contraction, + num_images=self.num_train_data, + use_pred_normals=self.config.predict_normals, + use_average_appearance_embedding=self.config.use_average_appearance_embedding, + ) + + self.density_fns = [] + num_prop_nets = self.config.num_proposal_iterations + # Build the proposal network(s) + self.proposal_networks = torch.nn.ModuleList() + if self.config.use_same_proposal_network: + assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed." + prop_net_args = self.config.proposal_net_args_list[0] + network = TemporalHashMLPDensityField( + self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args + ) + self.proposal_networks.append(network) + self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)]) + else: + for i in range(num_prop_nets): + prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)] + network = TemporalHashMLPDensityField( + self.scene_box.aabb, + spatial_distortion=scene_contraction, + **prop_net_args, + ) + self.proposal_networks.append(network) + self.density_fns.extend([network.density_fn for network in self.proposal_networks]) + + # Samplers + update_schedule = lambda step: np.clip( + np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]), + 1, + self.config.proposal_update_every, + ) + 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, + ) + + # Collider + self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane) + + # renderers + self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) + self.renderer_accumulation = AccumulationRenderer() + self.renderer_depth = DepthRenderer(method="expected") # for depth loss + self.renderer_normals = NormalsRenderer() + + # losses + self.rgb_loss = MSELoss() + + # metrics + self.psnr = PeakSignalNoiseRatio(data_range=1.0) + self.ssim = structural_similarity_index_measure + self.lpips = LearnedPerceptualImagePatchSimilarity() + self.temporal_distortion = True # for viewer + + def get_outputs(self, ray_bundle: RayBundle): + assert ray_bundle.times is not None, "Time not provided." + ray_samples, weights_list, ray_samples_list = self.proposal_sampler( + ray_bundle, density_fns=[functools.partial(f, times=ray_bundle.times) for f in self.density_fns] + ) + field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals) + weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY]) + weights_list.append(weights) + ray_samples_list.append(ray_samples) + + rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) + depth = self.renderer_depth(weights=weights, ray_samples=ray_samples) + accumulation = self.renderer_accumulation(weights=weights) + + outputs = { + "rgb": rgb, + "accumulation": accumulation, + "depth": depth, + } + + if self.config.predict_normals: + outputs["normals"] = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) + outputs["pred_normals"] = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) + + # These use a lot of GPU memory, so we avoid storing them for eval. + if self.training: + outputs["weights_list"] = weights_list + outputs["ray_samples_list"] = ray_samples_list + + if self.training and self.config.predict_normals: + 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], + ) + + for i in range(self.config.num_proposal_iterations): + outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i]) + + return outputs + + def get_loss_dict(self, outputs, batch, metrics_dict=None): + loss_dict = {} + image = batch["image"].to(self.device) + loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"]) + if self.training: + loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( + outputs["weights_list"], outputs["ray_samples_list"] + ) + assert metrics_dict is not None and "distortion" in metrics_dict + loss_dict["distortion_loss"] = self.config.distortion_loss_mult * metrics_dict["distortion"] + if self.config.predict_normals: + # orientation loss for computed normals + loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean( + outputs["rendered_orientation_loss"] + ) + + # ground truth supervision for normals + loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( + outputs["rendered_pred_normal_loss"] + ) + if "depth_image" in batch.keys() and self.config.depth_weight > 0: + mask = (batch["depth_image"] != 0).view([-1]) + loss_dict["depth_loss"] = 0 + l = lambda x: self.config.depth_weight * (x - batch["depth_image"][mask]).pow(2).mean() + loss_dict["depth_loss"] = l(outputs["depth"][mask]) + for i in range(self.config.num_proposal_iterations): + loss_dict["depth_loss"] += l(outputs[f"prop_depth_{i}"][mask]) + if self.config.temporal_tv_weight > 0: + loss_dict["temporal_tv_loss"] = self.field.mlp_base.get_temporal_tv_loss() + for net in self.proposal_networks: + loss_dict["temporal_tv_loss"] += net.encoding.get_temporal_tv_loss() + loss_dict["temporal_tv_loss"] *= self.config.temporal_tv_weight + return loss_dict diff --git a/nerfstudio/models/nerfplayer_ngp.py b/nerfstudio/models/nerfplayer_ngp.py new file mode 100644 index 000000000..39507a538 --- /dev/null +++ b/nerfstudio/models/nerfplayer_ngp.py @@ -0,0 +1,241 @@ +# 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. + +""" +Implementation of NeRFPlayer (https://arxiv.org/abs/2210.15947) with InstantNGP backbone. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Type + +import nerfacc +import torch +from nerfacc import ContractionType +from torch.nn import Parameter +from torchmetrics import PeakSignalNoiseRatio +from torchmetrics.functional import structural_similarity_index_measure +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity +from typing_extensions import Literal + +from nerfstudio.cameras.rays import RayBundle +from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.fields.nerfplayer_ngp_field import NerfplayerNGPField +from nerfstudio.model_components.losses import MSELoss +from nerfstudio.model_components.ray_samplers import VolumetricSampler +from nerfstudio.model_components.renderers import ( + AccumulationRenderer, + DepthRenderer, + RGBRenderer, +) +from nerfstudio.models.base_model import Model +from nerfstudio.models.instant_ngp import InstantNGPModelConfig, NGPModel +from nerfstudio.utils import colors + + +@dataclass +class NerfplayerNGPModelConfig(InstantNGPModelConfig): + """NeRFPlayer Model Config with InstantNGP backbone. + Tips for tuning the performance: + 1. If the scene is flickering, this is caused by unwanted high-freq on the temporal dimension. + Try reducing `temporal_dim` first, but don't be too small, otherwise the dynamic object is blurred. + Then try increasing the `temporal_tv_weight`. This is the loss for promoting smoothness among the + temporal channels. + 2. If a faster rendering is preferred, then try reducing `log2_hashmap_size`. If more details are + wanted, try increasing `log2_hashmap_size`. + 3. If the input cameras are of limited numbers, try reducing `num_levels`. `num_levels` is for + multi-resolution volume sampling, and has a similar behavior to the freq in NeRF. With a small + `num_levels`, a blurred rendering will be generated, but it is unlikely to overfit the training views. + """ + + _target: Type = field(default_factory=lambda: NerfplayerNGPModel) + temporal_dim: int = 64 + """Hashing grid parameter. A higher temporal dim means a higher temporal frequency.""" + num_levels: int = 16 + """Hashing grid parameter.""" + features_per_level: int = 2 + """Hashing grid parameter.""" + log2_hashmap_size: int = 17 + """Hashing grid parameter.""" + base_resolution: int = 16 + """Hashing grid parameter.""" + temporal_tv_weight: float = 1 + """Temporal TV loss balancing weight for feature channels.""" + depth_weight: float = 1e-1 + """depth loss balancing weight for feature channels.""" + train_background_color: Literal["random", "black", "white"] = "random" + """The training background color that is given to untrained areas.""" + eval_background_color: Literal["random", "black", "white"] = "white" + """The training background color that is given to untrained areas.""" + disable_viewing_dependent: bool = True + """Disable viewing dependent effects.""" + + +class NerfplayerNGPModel(NGPModel): + """NeRFPlayer Model with InstantNGP backbone. + + Args: + config: NeRFPlayer NGP configuration to instantiate model + """ + + config: NerfplayerNGPModelConfig + field: NerfplayerNGPField + + def populate_modules(self): + """Set the fields and modules.""" + Model.populate_modules(self) + + self.field = NerfplayerNGPField( + aabb=self.scene_box.aabb, + contraction_type=self.config.contraction_type, + use_appearance_embedding=self.config.use_appearance_embedding, + num_images=self.num_train_data, + temporal_dim=self.config.temporal_dim, + num_levels=self.config.num_levels, + features_per_level=self.config.features_per_level, + log2_hashmap_size=self.config.log2_hashmap_size, + base_resolution=self.config.base_resolution, + disable_viewing_dependent=self.config.disable_viewing_dependent, + ) + + self.scene_aabb = Parameter(self.scene_box.aabb.flatten(), requires_grad=False) + + # Occupancy Grid + self.occupancy_grid = nerfacc.OccupancyGrid( + roi_aabb=self.scene_aabb, + resolution=self.config.grid_resolution, + contraction_type=self.config.contraction_type, + ) + + # Sampler + vol_sampler_aabb = self.scene_box.aabb if self.config.contraction_type == ContractionType.AABB else None + self.sampler = VolumetricSampler( + scene_aabb=vol_sampler_aabb, + occupancy_grid=self.occupancy_grid, + density_fn=self.field.density_fn, + ) # need to update the density_fn later during forward (for input time) + + # renderers + self.train_background_color = self.config.train_background_color + self.eval_background_color = self.config.eval_background_color + if self.config.train_background_color in ["white", "black"]: + self.train_background_color = colors.COLORS_DICT[self.config.train_background_color] + if self.config.eval_background_color in ["white", "black"]: + self.eval_background_color = colors.COLORS_DICT[self.config.eval_background_color] + + self.renderer_rgb = RGBRenderer() # will update bgcolor later during forward + self.renderer_accumulation = AccumulationRenderer() + self.renderer_depth = DepthRenderer(method="expected") + + # losses + self.rgb_loss = MSELoss() + + # metrics + self.psnr = PeakSignalNoiseRatio(data_range=1.0) + self.ssim = structural_similarity_index_measure + self.lpips = LearnedPerceptualImagePatchSimilarity() + self.temporal_distortion = True # for viewer + + def get_outputs(self, ray_bundle: RayBundle): + num_rays = len(ray_bundle) + + # update the density_fn of the sampler so that the density is time aware + self.sampler.density_fn = lambda x: self.field.density_fn(x, ray_bundle.times) + with torch.no_grad(): + ray_samples, ray_indices = self.sampler( + ray_bundle=ray_bundle, + near_plane=self.config.near_plane, + far_plane=self.config.far_plane, + render_step_size=self.config.render_step_size, + cone_angle=self.config.cone_angle, + ) + + field_outputs = self.field(ray_samples) + + # accumulation + packed_info = nerfacc.pack_info(ray_indices, num_rays) + weights = nerfacc.render_weight_from_density( + packed_info=packed_info, + sigmas=field_outputs[FieldHeadNames.DENSITY], + t_starts=ray_samples.frustums.starts, + t_ends=ray_samples.frustums.ends, + ) + + # update bgcolor in the renderer; usually random color for training and fixed color for inference + if self.training: + self.renderer_rgb.background_color = self.train_background_color + else: + self.renderer_rgb.background_color = self.eval_background_color + rgb = self.renderer_rgb( + rgb=field_outputs[FieldHeadNames.RGB], + weights=weights, + ray_indices=ray_indices, + num_rays=num_rays, + ) + depth = self.renderer_depth( + weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=num_rays + ) + accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=num_rays) + alive_ray_mask = accumulation.squeeze(-1) > 0 + + outputs = { + "rgb": rgb, + "accumulation": accumulation, + "depth": depth, + "alive_ray_mask": alive_ray_mask, # the rays we kept from sampler + "num_samples_per_ray": packed_info[:, 1], + } + # adding them to outputs for calculating losses + if self.training and self.config.depth_weight > 0: + outputs["ray_indices"] = ray_indices + outputs["ray_samples"] = ray_samples + outputs["weights"] = weights + outputs["sigmas"] = field_outputs[FieldHeadNames.DENSITY] + return outputs + + def get_loss_dict(self, outputs, batch, metrics_dict=None): + image = batch["image"].to(self.device) + mask = outputs["alive_ray_mask"] + rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask]) + loss_dict = {"rgb_loss": rgb_loss} + if "depth_image" in batch.keys() and self.config.depth_weight > 0: + mask = batch["depth_image"] != 0 + # First we calculate the depth value, just like most of the papers. + loss_dict["depth_loss"] = (outputs["depth"][mask] - batch["depth_image"][mask]).abs().mean() + # But this is not enough -- it will lead to fog like reconstructions, even with depth supervision. + # Because this loss only cares about the mean value. + # (Feel free to try it and see the fog-like effects on DyCheck by commenting out the following loss.) + # > The fog can be effectively penalized by multiview inputs (i.e., from another viewing point). + # > However, it is hard to get multiple views under the setting of dynamic scenes. + # > A new view always indicates another camera, which is expensive and brings sync problems. + # In DyCheck (https://arxiv.org/abs/2210.13445), surface sparsity regularizer + # (i.e., `distortion_loss` in nerfstudio) is used to make surface tight. + # The `distortion_loss` can be used here as well, but personally find it hard to implement... + # (Due to volume sampling, seems that cuda kernels are needed for efficiently computing the loss.) + # (Try nerfplayer with nerfacto backbone for distortion loss.) + # Instead, directly penalizing the empty space is found effective here. Perhaps due to a more + # "clear" loss (as it is directly applied to the network outputs, rather than post-processed values). + # But such a loss also has drawbacks: it tends to overfit wrong (or noise) presented in the depth map. + # Some structures are floating in the air with this loss... + gt_depth_packed = batch["depth_image"][outputs["ray_indices"]] + steps = (outputs["ray_samples"].frustums.starts + outputs["ray_samples"].frustums.ends) / 2 + # empty area should not be too close to the given depth, so lets add a margin to the gt depth + margin = (self.scene_box.aabb.max() - self.scene_box.aabb.min()) / 128 + density_min_mask = (gt_depth_packed - steps > margin) & (gt_depth_packed != 0) + loss_dict["depth_loss"] += (outputs["sigmas"][density_min_mask[..., 0]].pow(2)).mean() * 1e-2 + loss_dict["depth_loss"] *= self.config.depth_weight + if self.config.temporal_tv_weight > 0: + loss_dict["temporal_tv_loss"] = self.config.temporal_tv_weight * self.field.mlp_base.get_temporal_tv_loss() + return loss_dict diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index f62cf9d8f..d27548d20 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -196,11 +196,16 @@ class CheckThread(threading.Thread): while not self.state.check_done_render: # check camera data = self.state.vis["renderingState/camera"].read() + render_time = self.state.vis["renderingState/render_time"].read() if data is not None: camera_object = data["object"] - if self.state.prev_camera_matrix is None or ( - not np.allclose(camera_object["matrix"], self.state.prev_camera_matrix) - and not self.state.prev_moving + if ( + self.state.prev_camera_matrix is None + or ( + not np.allclose(camera_object["matrix"], self.state.prev_camera_matrix) + and not self.state.prev_moving + ) + or (render_time is not None and render_time != self.state.prev_render_time) ): self.state.check_interrupt_vis = True self.state.prev_moving = True @@ -280,6 +285,7 @@ class ViewerState: # viewer specific variables self.prev_camera_matrix = None + self.prev_render_time = 0 self.prev_output_type = OutputTypes.INIT self.prev_colormap_type = ColormapTypes.INIT self.prev_moving = False @@ -553,12 +559,23 @@ class ViewerState: return None camera_object = data["object"] + render_time = self.vis["renderingState/render_time"].read() - if self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix): - self.camera_moving = False + if render_time is not None: + if ( + self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix) + ) and (self.prev_render_time == render_time): + self.camera_moving = False + else: + self.prev_camera_matrix = camera_object["matrix"] + self.prev_render_time = render_time + self.camera_moving = True else: - self.prev_camera_matrix = camera_object["matrix"] - self.camera_moving = True + if self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix): + self.camera_moving = False + else: + self.prev_camera_matrix = camera_object["matrix"] + self.camera_moving = True output_type = self.vis["renderingState/output_choice"].read() if output_type is None: @@ -901,10 +918,6 @@ class ViewerState: dim=0, ) - times = self.vis["renderingState/render_time"].read() - if times is not None: - times = torch.tensor([float(times)]) - camera_type_msg = camera_object["camera_type"] if camera_type_msg == "perspective": camera_type = CameraType.PERSPECTIVE @@ -922,11 +935,11 @@ class ViewerState: cy=intrinsics_matrix[1, 2], camera_type=camera_type, camera_to_worlds=camera_to_world[None, ...], - times=times, + times=torch.tensor([float(self.prev_render_time)]), ) camera = camera.to(graph.device) - camera_ray_bundle = camera.generate_rays(camera_indices=0, aabb_box=graph.render_aabb) + camera_ray_bundle = camera.generate_rays(camera_indices=0) graph.eval() diff --git a/tests/test_train.py b/tests/test_train.py index a3f1b6c60..5c52e50e8 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -23,6 +23,8 @@ BLACKLIST = [ "nerfacto", "phototourism", "depth-nerfacto", + "nerfplayer-ngp", + "nerfplayer-nerfacto", ] From f3b3a517e87c9d65072e5c43f61833d236c9991e Mon Sep 17 00:00:00 2001 From: Terrance Wang Date: Tue, 31 Jan 2023 14:02:57 -0800 Subject: [PATCH 03/11] adding check if camera_path_dir exists (#1319) Co-authored-by: Matthew Tancik --- nerfstudio/viewer/server/viewer_utils.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index d27548d20..c1c335e3f 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -401,13 +401,14 @@ class ViewerState: trainer.save_checkpoint(step) # get all camera paths camera_path_dir = os.path.join(self.datapath, "camera_paths") - camera_path_files = os.listdir(camera_path_dir) - all_path_dict = {} - for i in camera_path_files: - if i[-4:] == "json": - all_path_dict[i[:-5]] = load_from_json(Path(os.path.join(camera_path_dir, i))) - self.vis["renderingState/all_camera_paths"].write(all_path_dict) - self.vis["populate_paths_payload"].delete() + if os.path.exists(camera_path_dir): + camera_path_files = os.listdir(camera_path_dir) + all_path_dict = {} + for i in camera_path_files: + if i[-4:] == "json": + all_path_dict[i[:-5]] = load_from_json(Path(os.path.join(camera_path_dir, i))) + self.vis["renderingState/all_camera_paths"].write(all_path_dict) + self.vis["populate_paths_payload"].delete() def _check_webrtc_offer(self): """Check if there is a webrtc offer to respond to.""" From 79395ecd8d02a60661ddc67b32e8260a22db0564 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Tue, 31 Jan 2023 18:25:30 -0800 Subject: [PATCH 04/11] Fix crop robustness (#1321) --- nerfstudio/cameras/cameras.py | 12 +++++++++--- nerfstudio/viewer/server/viewer_utils.py | 2 +- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index fee7495a8..240c771b6 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -147,6 +147,8 @@ class Cameras(TensorDataclass): self.__post_init__() # This will do the dataclass post_init and broadcast all the tensors + self._use_nerfacc = strtobool(os.environ.get("INTERSECT_WITH_NERFACC", "TRUE")) + def _init_get_fc_xy(self, fc_xy, name): """ Parses the input focal length / principle point x or y and returns a tensor of the correct shape @@ -470,9 +472,13 @@ class Cameras(TensorDataclass): rays_o = rays_o.reshape((-1, 3)) rays_d = rays_d.reshape((-1, 3)) - if strtobool(os.environ.get("INTERSECT_WITH_NERFACC", "TRUE")): - nerfacc = importlib.import_module("nerfacc") - t_min, t_max = nerfacc.ray_aabb_intersect(rays_o, rays_d, tensor_aabb) + if self._use_nerfacc: + try: + nerfacc = importlib.import_module("nerfacc") + t_min, t_max = nerfacc.ray_aabb_intersect(rays_o, rays_d, tensor_aabb) + except: # pylint: disable=bare-except + t_min, t_max = nerfstudio.utils.math.intersect_aabb(rays_o, rays_d, tensor_aabb) + self._use_nerfacc = False else: t_min, t_max = nerfstudio.utils.math.intersect_aabb(rays_o, rays_d, tensor_aabb) diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index c1c335e3f..44056537f 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -940,7 +940,7 @@ class ViewerState: ) camera = camera.to(graph.device) - camera_ray_bundle = camera.generate_rays(camera_indices=0) + camera_ray_bundle = camera.generate_rays(camera_indices=0, aabb_box=graph.render_aabb) graph.eval() From edef6d62f675c3357f5b856960c10c2862add6de Mon Sep 17 00:00:00 2001 From: machenmusik Date: Tue, 31 Jan 2023 21:52:42 -0500 Subject: [PATCH 05/11] bump nerfacc to 0.3.4 (#1322) Co-authored-by: Chen, Michael --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 10942cf21..8b1700d4d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ "mediapy>=1.1.0", "msgpack>=1.0.4", "msgpack_numpy>=0.4.8", - "nerfacc==0.3.3", + "nerfacc==0.3.4", "open3d>=0.16.0", "opencv-python==4.6.0.66", "Pillow>=9.3.0", From f5f424095a05ce565249955f5919f4c06e3c7319 Mon Sep 17 00:00:00 2001 From: machenmusik Date: Tue, 31 Jan 2023 22:27:05 -0500 Subject: [PATCH 06/11] Document localhost URL (fixes #1323) (#1324) Document localhost URL Co-authored-by: Chen, Michael --- docs/developer_guides/viewer/viewer_overview.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/developer_guides/viewer/viewer_overview.md b/docs/developer_guides/viewer/viewer_overview.md index b80463e75..4e10319b6 100644 --- a/docs/developer_guides/viewer/viewer_overview.md +++ b/docs/developer_guides/viewer/viewer_overview.md @@ -50,6 +50,10 @@ yarn install yarn start ``` +The local webserver runs on port 4000 by default, +so when `ns-train` is running, you can connect to the viewer locally at +[http://localhost:4000/?websocket_url=ws://localhost:7007](http://localhost:4000/?websocket_url=ws://localhost:7007) + #### Running the Bridge Server Manually The viewer bridge server runs on the same machine that you use for training. The training code will connect to this server with a lightweight TCP connection using ZMQ. The training job will launch the viewer bridge server is you specify `--viewer.launch-bridge-server` in the terminal. Otherwise, you can launch the bridge server manually with the following script. From 11e7fa87e04f90ad94598d6b41cacfa91fff4751 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Wed, 1 Feb 2023 10:29:08 -0800 Subject: [PATCH 07/11] Add Stable Diffusion helper code (#1320) * Add stable diffusion * Add dependencies fix lint * Update types * Update comment * fix comments --------- Co-authored-by: David McAllister Co-authored-by: Abhik Ahuja --- nerfstudio/generative/__init__.py | 13 + nerfstudio/generative/stable_diffusion.py | 336 ++++++++++++++++++++++ pyproject.toml | 8 + 3 files changed, 357 insertions(+) create mode 100644 nerfstudio/generative/__init__.py create mode 100644 nerfstudio/generative/stable_diffusion.py diff --git a/nerfstudio/generative/__init__.py b/nerfstudio/generative/__init__.py new file mode 100644 index 000000000..d926cf95c --- /dev/null +++ b/nerfstudio/generative/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py new file mode 100644 index 000000000..af3aa2196 --- /dev/null +++ b/nerfstudio/generative/stable_diffusion.py @@ -0,0 +1,336 @@ +# 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. + +"""Stable Diffusion helpers""" + +# Modified from https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py + +import sys +from pathlib import Path +from typing import List, Optional, Tuple, Union + +import mediapy +import numpy as np +import torch +import torch.nn.functional as F +import tyro +from rich.console import Console +from torch import nn +from torchtyping import TensorType + +CONSOLE = Console(width=120) + +try: + from diffusers import PNDMScheduler, StableDiffusionPipeline + from transformers import logging + +except ImportError: + CONSOLE.print("[bold red]Missing Stable Diffusion packages.") + CONSOLE.print(r"Install using [yellow]pip install nerfstudio\[gen][/yellow]") + CONSOLE.print(r"or [yellow]pip install -e .\[gen][/yellow] if installing from source.") + sys.exit(1) + +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" + + +class StableDiffusion(nn.Module): + """Stable Diffusion implementation + + Args: + device: device to use + num_train_timesteps: number of training timesteps + + """ + + def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000) -> None: + super().__init__() + + self.device = device + self.num_train_timesteps = num_train_timesteps + + self.min_step = int(self.num_train_timesteps * 0.02) + self.max_step = int(self.num_train_timesteps * 0.98) + + self.scheduler = PNDMScheduler( + beta_start=0.00085, + beta_end=0.012, + beta_schedule="scaled_linear", + num_train_timesteps=self.num_train_timesteps, + ) + self.alphas = self.scheduler.alphas_cumprod.to(self.device) # type: ignore + + pipe = StableDiffusionPipeline.from_pretrained(SD_SOURCE, torch_dtype=torch.float16) + assert pipe is not None + pipe = pipe.to(self.device) + + pipe.enable_attention_slicing() + + self.tokenizer = pipe.tokenizer + self.unet = pipe.unet + self.text_encoder = pipe.text_encoder + self.auto_encoder = pipe.vae + + self.unet.to(memory_format=torch.channels_last) + + CONSOLE.print("Stable Diffusion loaded!") + + def get_text_embeds( + self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]] + ) -> TensorType[2, "max_length", "embed_dim"]: + """Get text embeddings for prompt and negative prompt + + Args: + prompt: Prompt text + negative_prompt: Negative prompt text + + Returns: + Text embeddings + """ + + # Tokenize text and get embeddings + text_input = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + + with torch.no_grad(): + text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0] + + # Do the same for unconditional embeddings + uncond_input = self.tokenizer( + negative_prompt, padding="max_length", max_length=self.tokenizer.model_max_length, return_tensors="pt" + ) + + with torch.no_grad(): + uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0] + + # Cat for final embeddings + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + + return text_embeddings + + def sds_loss( + self, + text_embeddings: TensorType["N", "max_length", "embed_dim"], + image: TensorType["BS", 3, "H", "W"], + guidance_scale: float = 100.0, + ) -> Tuple[torch.Tensor, TensorType["BS", 4, "H", "W"], TensorType["BS", 4, "H", "W"]]: + """Score Distilation Sampling loss proposed in DreamFusion paper (https://dreamfusion3d.github.io/) + + Args: + text_embeddings: Text embeddings + image: Rendered image + guidance_scale: How much to weigh the guidance + + Returns: + A tuple of Loss, latent, and gradient + """ + image = F.interpolate(image, (IMG_DIM, IMG_DIM), mode="bilinear") + t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device) + latents = self.imgs_to_latent(image) + + # predict the noise residual with unet, NO grad! + with torch.no_grad(): + # add noise + noise = torch.randn_like(latents) + latents_noisy = self.scheduler.add_noise(latents, noise, t) # type: ignore + # pred noise + latent_model_input = torch.cat([latents_noisy] * 2) + noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample + + # perform guidance + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # w(t), sigma_t^2 + w = 1 - self.alphas[t] + + grad = w * (noise_pred - noise) + grad = torch.nan_to_num(grad) + + sds_mse = torch.mean(torch.nan_to_num(torch.square(noise_pred - noise))) + + return sds_mse, latents, grad + + def produce_latents( + self, + text_embeddings: TensorType["N", "max_length", "embed_dim"], + height: int = IMG_DIM, + width: int = IMG_DIM, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + latents: Optional[TensorType["BS", 4, "H", "W"]] = None, + ) -> TensorType["BS", 4, "H", "W"]: + """Produce latents for a given text embedding + + Args: + text_embeddings: Text embeddings + height: Height of the image + width: Width of the image + num_inference_steps: Number of inference steps + guidance_scale: How much to weigh the guidance + latents: Latents to start with + + Returns: + Latents + """ + + if latents is None: + latents = torch.randn( + (text_embeddings.shape[0] // 2, self.unet.in_channels, height // 8, width // 8), device=self.device + ) + + self.scheduler.set_timesteps(num_inference_steps) # type: ignore + + with torch.autocast("cuda"): + for t in self.scheduler.timesteps: # type: ignore + assert latents is not None + # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. + latent_model_input = torch.cat([latents] * 2) + + # predict the noise residual + with torch.no_grad(): + noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"] + + # perform guidance + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"] # type: ignore + return latents + + def latents_to_img(self, latents: TensorType["BS", 4, "H", "W"]) -> TensorType["BS", 3, "H", "W"]: + """Convert latents to images + + Args: + latents: Latents to convert + + Returns: + Images + """ + + latents = 1 / CONST_SCALE * latents + + with torch.no_grad(): + imgs = self.auto_encoder.decode(latents).sample + + imgs = (imgs / 2 + 0.5).clamp(0, 1) + + return imgs + + def imgs_to_latent(self, imgs: TensorType["BS", 3, "H", "W"]) -> TensorType["BS", 4, "H", "W"]: + """Convert images to latents + + Args: + imgs: Images to convert + + Returns: + Latents + """ + imgs = 2 * imgs - 1 + + posterior = self.auto_encoder.encode(imgs).latent_dist + latents = posterior.sample() * CONST_SCALE + + return latents + + def prompt_to_img( + self, + prompts: Union[str, List[str]], + negative_prompts: Union[str, List[str]] = "", + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + latents=None, + ) -> np.ndarray: + """Generate an images from a prompts. + + Args: + prompts: The prompt to generate an image from. + negative_prompts: The negative prompt to generate an image from. + num_inference_steps: The number of inference steps to perform. + guidance_scale: The scale of the guidance. + latents: The latents to start from, defaults to random. + + Returns: + The generated image. + """ + + prompts = [prompts] if isinstance(prompts, str) else prompts + negative_prompts = [negative_prompts] if isinstance(negative_prompts, str) else negative_prompts + text_embeddings = self.get_text_embeds(prompts, negative_prompts) + latents = self.produce_latents( + text_embeddings, + height=IMG_DIM, + width=IMG_DIM, + latents=latents, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + ) # [1, 4, resolution, resolution] + + diffused_img = self.latents_to_img(latents.half()) + diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy() + diffused_img = (diffused_img * 255).round().astype("uint8") + + return diffused_img + + def forward( + self, prompts, negative_prompts="", num_inference_steps=50, guidance_scale=7.5, latents=None + ) -> np.ndarray: + """Generate an image from a prompt. + + Args: + prompts: The prompt to generate an image from. + negative_prompts: The negative prompt to generate an image from. + num_inference_steps: The number of inference steps to perform. + guidance_scale: The scale of the guidance. + latents: The latents to start from, defaults to random. + + Returns: + The generated image. + """ + return self.prompt_to_img(prompts, negative_prompts, num_inference_steps, guidance_scale, latents) + + +def generate_image( + prompt: str, negative: str = "", seed: int = 0, steps: int = 50, save_path: Path = Path("test_sd.png") +): + """Generate an image from a prompt using Stable Diffusion. + + Args: + prompt: The prompt to use. + negative: The negative prompt to use. + seed: The random seed to use. + steps: The number of steps to use. + save_path: The path to save the image to. + """ + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + cuda_device = torch.device("cuda") + with torch.no_grad(): + sd = StableDiffusion(cuda_device) + imgs = sd.prompt_to_img(prompt, negative, steps) + mediapy.write_image(str(save_path), imgs[0]) + + +if __name__ == "__main__": + tyro.cli(generate_image) diff --git a/pyproject.toml b/pyproject.toml index 8b1700d4d..ca44d209e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +59,14 @@ dependencies = [ [project.optional-dependencies] + +# Generative related dependencies +gen = [ + "diffusers==0.9.0", + "transformers==4.24.0", +] + + # Development packages dev = [ "black[jupyter]==22.3.0", From 9966c43811776d4de129068b89f105de18d8a57d Mon Sep 17 00:00:00 2001 From: Nicolas Zunhammer Date: Wed, 1 Feb 2023 21:02:10 +0100 Subject: [PATCH 08/11] [Fix] Docker not supporting older CUDA architectures (#1328) Fix docker not supporting older CUDA architectures. Co-authored-by: Nicolas Zunhammer Co-authored-by: Matthew Tancik --- Dockerfile | 3 ++- docs/quickstart/installation.md | 17 ++++++++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/Dockerfile b/Dockerfile index c2e3ab82b..94587f550 100644 --- a/Dockerfile +++ b/Dockerfile @@ -7,7 +7,8 @@ ENV DEBIAN_FRONTEND=noninteractive ## Set timezone as it is required by some packages. ENV TZ=Europe/Berlin ## CUDA architectures, required by tiny-cuda-nn. -ENV TCNN_CUDA_ARCHITECTURES=86 +## NOTE: All commonly used GPU architectures are included and supported here. To speedup the image build process remove all architectures but the one of your explicit GPU. Find details here: https://developer.nvidia.com/cuda-gpus (8.6 translates to 86 in the line below) or in the docs. +ENV TCNN_CUDA_ARCHITECTURES=90;89;86;80;75;70;61;52;37 ## CUDA Home, required to find CUDA in some packages. ENV CUDA_HOME="/usr/local/cuda" diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md index ae6a17cc9..97c63aac8 100644 --- a/docs/quickstart/installation.md +++ b/docs/quickstart/installation.md @@ -95,13 +95,28 @@ docker run --gpus all -v /folder/of/your/data:/workspace/ -v /home//. ``` ### Note - The container works on Linux and Windows, depending on your OS some additional setup steps might be required to provide access to your GPU inside containers. -- Paths on Windows use backslash '\\\\' while unix based systems use a frontslash '/' for paths, where backslashes might require an escape character depending on where they are used (e.g. C:\\\\folder1\\\\folder2...). Ensure to use the correct paths when mounting folders or providing paths as parameters. +- Paths on Windows use backslash '\\' while unix based systems use a frontslash '/' for paths, where backslashes might require an escape character depending on where they are used (e.g. C:\\\\folder1\\\\folder2...). Alternatively, mounts can be quoted (e.g. ```-v 'C:\local_folder:/docker_folder'```). Ensure to use the correct paths when mounting folders or providing paths as parameters. +- Always use full paths, relative paths are known to create issues when being used in mounts into docker. - Everything inside the container, what is not in a mounted folder (workspace in the above example), will be permanently removed after destroying the container. Always do all your tasks and output folder in workdir! - The user inside the container is called 'user' and is mapped to the local user with ID 1000 (usually the first non-root user on Linux systems). - The container currently is based on nvidia/cuda:11.7.1-devel-ubuntu22.04, consequently it comes with CUDA 11.7 which must be supported by the nvidia driver. No local CUDA installation is required or will be affected by using the docker image. - The docker image (respectively Ubuntu 22.04) comes with Python3.10, no older version of Python is installed. - If you call the container with commands directly, you still might want to add the interactive terminal ('-it') flag to get live log outputs of the nerfstudio scripts. In case the container is used in an automated environment the flag should be discarded. +- The current version of docker is built for multi-architecture (CUDA architectures) use. The target architecture must be defined at build time for tinyCUDNN to be able to compile properly. If your GPU architecture is not covered by the following table you need to replace the number in the line ```ENV TCNN_CUDA_ARCHITECTURES=90;89;86;80;75;70;61;52;37``` to you specific architecture. It also is a good idea to remove all architectures but yours (e.g. ```ENV TCNN_CUDA_ARCHITECTURES=86```) to speedup the docker build a lot. +**Currently supported CUDA architectures in the docker image** + +GPU | CUDA arch +-- | -- +H100 | 90 +40X0 | 89 +30X0 | 86 +A100 | 80 +20X0 | 75 +TITAN V / V100 | 70 +10X0 / TITAN Xp | 61 +9X0 | 52 +K80 | 37 ## Installation FAQ From 567268854bd7822cd98a438015cb3220f959dc48 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Wed, 1 Feb 2023 12:18:09 -0800 Subject: [PATCH 09/11] Cleanup aabb intersection logic (#1329) --- nerfstudio/cameras/cameras.py | 11 +---- nerfstudio/utils/math.py | 92 +++++++++++++++++++++++------------ 2 files changed, 63 insertions(+), 40 deletions(-) diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 240c771b6..d7d228bae 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -16,7 +16,6 @@ Camera Models """ import base64 -import importlib import math import os from dataclasses import dataclass @@ -472,15 +471,7 @@ class Cameras(TensorDataclass): rays_o = rays_o.reshape((-1, 3)) rays_d = rays_d.reshape((-1, 3)) - if self._use_nerfacc: - try: - nerfacc = importlib.import_module("nerfacc") - t_min, t_max = nerfacc.ray_aabb_intersect(rays_o, rays_d, tensor_aabb) - except: # pylint: disable=bare-except - t_min, t_max = nerfstudio.utils.math.intersect_aabb(rays_o, rays_d, tensor_aabb) - self._use_nerfacc = False - else: - t_min, t_max = nerfstudio.utils.math.intersect_aabb(rays_o, rays_d, tensor_aabb) + t_min, t_max = nerfstudio.utils.math.intersect_aabb(rays_o, rays_d, tensor_aabb) t_min = t_min.reshape([shape[0], shape[1], 1]) t_max = t_max.reshape([shape[0], shape[1], 1]) diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py index 2658b1c39..9c2077f4b 100644 --- a/nerfstudio/utils/math.py +++ b/nerfstudio/utils/math.py @@ -15,10 +15,13 @@ """ Math Helper Functions """ from dataclasses import dataclass +from typing import Tuple import torch from torchtyping import TensorType +_USE_NERFACC = True + def components_from_spherical_harmonics(levels: int, directions: TensorType[..., 3]) -> TensorType[..., "components"]: """ @@ -193,47 +196,76 @@ def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor: return torch.exp(-0.5 * x_vars) * torch.sin(x_means) -def intersect_aabb(origins, directions, aabb): +@torch.jit.script +def _intersect_aabb( + origins: torch.Tensor, + directions: torch.Tensor, + aabb: torch.Tensor, + max_bound: float = 1e10, + invalid_value: float = 1e10, +) -> Tuple[torch.Tensor, torch.Tensor]: """ - pytorch implementation of ray intersection with AABB box - :param origin: N,3 tensor of 3d positions - :param direction: N,3 tensor of normalized directions - :param aabb: 6,1 array of aabb box in the form of [x_min, y_min, z_min, x_max, y_max, z_max] - :return: t_min, t_max - two tensors of shapes N,1 representing the intersection signed distance from - the origin. t_min is clipped to 0 in case it is negative, in case of no intersection t_min is 1e10. + Implementation of ray intersection with AABB box + + Args: + origins: [N,3] tensor of 3d positions + directions: [N,3] tensor of normalized directions + aabb: [6] array of aabb box in the form of [x_min, y_min, z_min, x_max, y_max, z_max] + max_bound: Maximum value of t_max + invalid_value: Value to return in case of no intersection + + Returns: + t_min, t_max - two tensors of shapes N representing distance of intersection from the origin. """ - max_bound = 1e10 + tx_min = (aabb[:3] - origins) / directions + tx_max = (aabb[3:] - origins) / directions - tx_min = (aabb[0] - origins[:, 0]) / directions[:, 0] - tx_max = (aabb[3] - origins[:, 0]) / directions[:, 0] t_min = torch.min(tx_min, tx_max) t_max = torch.max(tx_min, tx_max) - ty_min_temp = (aabb[1] - origins[:, 1]) / directions[:, 1] - ty_max_temp = (aabb[4] - origins[:, 1]) / directions[:, 1] - ty_min = torch.min(ty_min_temp, ty_max_temp) - ty_max = torch.max(ty_min_temp, ty_max_temp) + t_min = torch.max(t_min, dim=-1).values + t_max = torch.min(t_max, dim=-1).values - cond = torch.logical_or((t_min > ty_max), (ty_min > t_max)) - t_min = torch.where(cond, max_bound, t_min) - t_max = torch.where(cond, max_bound, t_max) + t_min = torch.clamp(t_min, min=0, max=max_bound) + t_max = torch.clamp(t_max, min=0, max=max_bound) - t_min = torch.max(t_min, ty_min) - t_max = torch.min(t_max, ty_max) + cond = t_max <= t_min + t_min = torch.where(cond, invalid_value, t_min) + t_max = torch.where(cond, invalid_value, t_max) - tz_min_temp = (aabb[2] - origins[:, 2]) / directions[:, 2] - tz_max_temp = (aabb[5] - origins[:, 2]) / directions[:, 2] - tz_min = torch.min(tz_min_temp, tz_max_temp) - tz_max = torch.max(tz_min_temp, tz_max_temp) + return t_min, t_max - cond = torch.logical_or((t_min > tz_max), tz_min > t_max) - t_min = torch.where(cond, max_bound, t_min) - t_max = torch.where(cond, max_bound, t_max) - t_min = torch.max(t_min, tz_min) - t_max = torch.min(t_max, tz_max) - t_min = torch.where(t_min > 0, t_min, 0) - t_max = torch.where(t_min == max_bound, max_bound, t_max) +def intersect_aabb( + origins: TensorType["N", 3], + directions: TensorType["N", 3], + aabb: TensorType[6], +) -> Tuple[TensorType["N"], TensorType["N"]]: + """ + Implementation of ray intersection with AABB box + + Args: + origins: 3d positions + directions: Normalized directions + aabb: array of aabb box in the form of [x_min, y_min, z_min, x_max, y_max, z_max] + max_bound: Maximum value of t_max + invalid_value: Value to return in case of no intersection + + Returns: + t_min, t_max - two tensors of shapes N representing distance of intersection from the origin. + """ + + global _USE_NERFACC # pylint: disable=global-statement + if _USE_NERFACC: + try: + import nerfacc # pylint: disable=import-outside-toplevel + + t_min, t_max = nerfacc.ray_aabb_intersect(origins, directions, aabb) + except: # pylint: disable=bare-except + t_min, t_max = _intersect_aabb(origins, directions, aabb, max_bound=1e10, invalid_value=1e10) + _USE_NERFACC = False + else: + t_min, t_max = _intersect_aabb(origins, directions, aabb, max_bound=1e10, invalid_value=1e10) return t_min, t_max From e979f8920093ade6002d7b4848634a2e51f581d4 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Wed, 1 Feb 2023 12:21:57 -0800 Subject: [PATCH 10/11] Avoid protobuf 3.20.0 (#1330) --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index ca44d209e..21bc69456 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,7 @@ dependencies = [ "opencv-python==4.6.0.66", "Pillow>=9.3.0", "plotly>=5.7.0", - "protobuf<=3.20.3", + "protobuf<=3.20.3,!=3.20.0", "pymeshlab>=2022.2.post2", "pyngrok>=5.1.0", "python-socketio>=5.7.1", From 53a830a6917b511c811b09be8641ba0ca7b488a2 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Wed, 1 Feb 2023 13:22:32 -0800 Subject: [PATCH 11/11] Safe normalize (#1332) --- nerfstudio/utils/math.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py index 9c2077f4b..937866dcc 100644 --- a/nerfstudio/utils/math.py +++ b/nerfstudio/utils/math.py @@ -269,3 +269,19 @@ def intersect_aabb( t_min, t_max = _intersect_aabb(origins, directions, aabb, max_bound=1e10, invalid_value=1e10) return t_min, t_max + + +def safe_normalize( + vectors: TensorType["batch_dim":..., "N"], + eps: float = 1e-10, +) -> TensorType["batch_dim":..., "N"]: + """Normalizes vectors. + + Args: + vectors: Vectors to normalize. + eps: Epsilon value to avoid division by zero. + + Returns: + Normalized vectors. + """ + return vectors / (torch.norm(vectors, dim=-1, keepdim=True) + eps)