From 51e7846a6184a9aa12dca486642b1ca1fd9da92b Mon Sep 17 00:00:00 2001 From: terrance Date: Wed, 3 May 2023 23:47:25 +0000 Subject: [PATCH] ngp and tensor work well now, changed shading loss for dreamfusion --- nerfstudio/configs/method_configs.py | 42 ++--- nerfstudio/fields/dreamembedding_field.py | 2 +- nerfstudio/fields/dreamngp_field.py | 76 +++----- nerfstudio/fields/dreamtensor_field.py | 33 ---- nerfstudio/generative/stable_diffusion.py | 7 +- nerfstudio/models/dreamembedding.py | 68 +++----- nerfstudio/models/dreamfusion.py | 76 ++++---- nerfstudio/models/dreamfusion_ngp.py | 164 +++++++++++------- nerfstudio/models/dreamfusion_tensor.py | 26 ++- .../viewer/server/render_state_machine.py | 4 +- scripts/viewer/run_viewer.py | 60 +------ 11 files changed, 245 insertions(+), 313 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index ba645c8c3..2efac5bc8 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -447,16 +447,16 @@ method_configs["phototourism"] = TrainerConfig( method_configs["dreamfusion"] = TrainerConfig( method_name="dreamfusion", - experiment_name="", + experiment_name="dreamfusion", steps_per_eval_batch=50, steps_per_eval_image=50, steps_per_save=200, - max_num_iterations=20000, + max_num_iterations=15000, mixed_precision=True, pipeline=VanillaPipelineConfig( generative=True, datamanager=DreamFusionDataManagerConfig( - horizontal_rotation_warmup=2000, + horizontal_rotation_warmup=4000, ), model=DreamFusionModelConfig( eval_num_rays_per_chunk=1 << 15, @@ -480,11 +480,11 @@ method_configs["dreamfusion"] = TrainerConfig( optimizers={ "proposal_networks": { "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + # "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), "scheduler": None, }, "fields": { "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15), - # "scheduler": None, "scheduler": ExponentialDecaySchedulerConfig( warmup_steps=2000, lr_final=1e-6, max_steps=20000, ramp="linear" ), @@ -500,37 +500,32 @@ method_configs["dreamngp"] = TrainerConfig( steps_per_eval_batch=50, steps_per_eval_image=50, steps_per_save=200, - max_num_iterations=20000, + max_num_iterations=15000, mixed_precision=True, pipeline=VanillaPipelineConfig( generative=True, datamanager=DreamFusionDataManagerConfig( - horizontal_rotation_warmup=2000, + horizontal_rotation_warmup=4000, ), model=DreamfusionNGPModelConfig( eval_num_rays_per_chunk=8192, - contraction_type=ContractionType.AABB, - render_step_size=0.001, - max_num_samples_per_ray=48, orientation_loss_mult=1.0, - sphere_collider=True, initialize_density=True, - proposal_warmup=500, - proposal_update_every=0, - proposal_weights_anneal_max_num_iters=100, - opacity_loss_mult=1, positional_prompting="discrete", guidance_scale=50, + grid_levels=1, + alpha_thre=0.0, + cone_angle=0.0, ), ), optimizers={ "proposal_networks": { "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + # "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), "scheduler": None, }, "fields": { "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15), - # "scheduler": None, "scheduler": ExponentialDecaySchedulerConfig( warmup_steps=2000, lr_final=1e-6, max_steps=20000, ramp="linear" ), @@ -546,12 +541,12 @@ method_configs["dreamtensor"] = TrainerConfig( steps_per_eval_batch=50, steps_per_eval_image=50, steps_per_save=200, - max_num_iterations=20000, + max_num_iterations=15000, mixed_precision=True, pipeline=VanillaPipelineConfig( generative=True, datamanager=DreamFusionDataManagerConfig( - horizontal_rotation_warmup=2000, + horizontal_rotation_warmup=4000, ), model=DreamfusionTensorModelConfig( initialize_density=True, @@ -561,14 +556,6 @@ method_configs["dreamtensor"] = TrainerConfig( ), ), optimizers={ - # "fields": { - # "optimizer": AdamOptimizerConfig(lr=5e-5), - # "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-5, max_steps=20000), - # }, - # "encodings": { - # "optimizer": AdamOptimizerConfig(lr=0.0002), - # "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=20000), - # }, "fields": { "optimizer": AdamOptimizerConfig(lr=0.001), "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000), @@ -680,12 +667,12 @@ method_configs["dreamembed"] = TrainerConfig( steps_per_eval_batch=50, steps_per_eval_image=50, steps_per_save=200, - max_num_iterations=20000, + max_num_iterations=15000, mixed_precision=True, pipeline=VanillaPipelineConfig( generative=True, datamanager=DreamFusionDataManagerConfig( - horizontal_rotation_warmup=2000, + horizontal_rotation_warmup=4000, ), model=DreamEmbeddingModelConfig( eval_num_rays_per_chunk=1 << 15, @@ -706,6 +693,7 @@ method_configs["dreamembed"] = TrainerConfig( optimizers={ "proposal_networks": { "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + # "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), "scheduler": None, }, "fields": { diff --git a/nerfstudio/fields/dreamembedding_field.py b/nerfstudio/fields/dreamembedding_field.py index be69e52ef..fd183cf1a 100644 --- a/nerfstudio/fields/dreamembedding_field.py +++ b/nerfstudio/fields/dreamembedding_field.py @@ -77,7 +77,7 @@ class DreamEmbeddingField(Field): self.position_encoding = tcnn.Encoding( n_input_dims=3, - encoding_config={"otype": "Frequency", "n_frequencies": 2}, + encoding_config={"otype": "Frequency", "n_frequencies": 3}, ) self.direction_encoding = tcnn.Encoding( diff --git a/nerfstudio/fields/dreamngp_field.py b/nerfstudio/fields/dreamngp_field.py index e91d00339..1d1d78f61 100644 --- a/nerfstudio/fields/dreamngp_field.py +++ b/nerfstudio/fields/dreamngp_field.py @@ -21,7 +21,6 @@ from typing import Dict, Optional, Tuple import numpy as np import torch -from nerfacc import ContractionType, contract from torch.nn.parameter import Parameter from torchtyping import TensorType @@ -30,8 +29,12 @@ 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.spatial_distortions import ( + SceneContraction, + SpatialDistortion, +) from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn -from nerfstudio.fields.instant_ngp_field import TCNNInstantNGPField +from nerfstudio.fields.nerfacto_field import TCNNNerfactoField try: import tinycudann as tcnn except ImportError: @@ -47,7 +50,7 @@ def get_normalized_directions(directions: TensorType["bs":..., 3]): return (directions + 1.0) / 2.0 -class DreamNGPField(TCNNInstantNGPField): +class DreamNGPField(TCNNNerfactoField): """TCNN implementation of the Instant-NGP field. Args: @@ -69,34 +72,35 @@ class DreamNGPField(TCNNInstantNGPField): def __init__( self, aabb: TensorType, + num_images: int, 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: Optional[bool] = False, - num_images: Optional[int] = None, - appearance_embedding_dim: int = 32, - contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE, num_levels: int = 16, - log2_hashmap_size: int = 19, max_res: int = 2048, + 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, + pass_semantic_gradients: bool = False, + use_pred_normals: bool = False, + use_average_appearance_embedding: bool = False, + spatial_distortion: Optional[SpatialDistortion] = None, ) -> None: super().__init__( - aabb, - num_layers, - hidden_dim, - geo_feat_dim, - num_layers_color, - hidden_dim_color, - use_appearance_embedding, - num_images, - appearance_embedding_dim, - contraction_type, - num_levels, - log2_hashmap_size, - max_res + aabb, + num_images, + log2_hashmap_size=log2_hashmap_size, + max_res=max_res, + spatial_distortion=None ) + self.mlp_background_color = tcnn.Network( n_input_dims=self.direction_encoding.n_output_dims, n_output_dims=3, @@ -109,32 +113,6 @@ class DreamNGPField(TCNNInstantNGPField): }, ) - def get_density(self, ray_samples: RaySamples): - """Computes and returns the densities""" - # positions = ray_samples.frustums.get_positions() - positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) - positions_flat = positions.view(-1, 3) - positions_flat = contract(x=positions_flat, roi=self.aabb, type=self.contraction_type) - - self._sample_locations = positions_flat - if not self._sample_locations.requires_grad: - self._sample_locations.requires_grad = True - - h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1) - density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) - self._density_before_activation = density_before_activation - - # import pdb - # pdb.set_trace() - - # 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_background_rgb(self, ray_bundle: RayBundle): """Predicts background colors at infinity.""" directions = get_normalized_directions(ray_bundle.directions) diff --git a/nerfstudio/fields/dreamtensor_field.py b/nerfstudio/fields/dreamtensor_field.py index 2019a3ae1..2c9340288 100644 --- a/nerfstudio/fields/dreamtensor_field.py +++ b/nerfstudio/fields/dreamtensor_field.py @@ -21,7 +21,6 @@ from typing import Dict, Optional, Tuple import numpy as np import torch -from nerfacc import ContractionType, contract from torch.nn.parameter import Parameter from torchtyping import TensorType @@ -112,38 +111,6 @@ class DreamFusionTensorField(TensoRFField): }, ) - def get_density(self, ray_samples: RaySamples) -> TensorType: - positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) - positions = positions * 2 - 1 - positions = ray_samples.frustums.get_positions() - density = self.density_encoding(positions) - density_enc = torch.sum(density, dim=-1)[:, :, None] - relu = torch.nn.ReLU() - density_enc = relu(density_enc) - return density_enc - - def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None) -> TensorType: - d = ray_samples.frustums.directions - positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) - positions = positions * 2 - 1 - positions = ray_samples.frustums.get_positions() - - rgb_features = self.color_encoding(positions) - rgb_features = self.B(rgb_features) - - d_encoded = self.direction_encoding(d) - rgb_features_encoded = self.feature_encoding(rgb_features) - - if self.use_sh: - sh_mult = self.sh(d)[:, :, None] - rgb_sh = rgb_features.view(sh_mult.shape[0], sh_mult.shape[1], 3, sh_mult.shape[-1]) - rgb = torch.relu(torch.sum(sh_mult * rgb_sh, dim=-1) + 0.5) - else: - out = self.mlp_head(torch.cat([rgb_features, d, rgb_features_encoded, d_encoded], dim=-1)) # type: ignore - rgb = self.field_output_rgb(out) - - return rgb - def get_background_rgb(self, ray_bundle: RayBundle): """Predicts background colors at infinity.""" directions = get_normalized_directions(ray_bundle.directions) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index ce030e677..a30993066 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -20,6 +20,8 @@ import sys from dataclasses import dataclass from pathlib import Path from typing import List, Optional, Union +from PIL import Image + import appdirs import mediapy @@ -46,6 +48,7 @@ CONSOLE = Console(width=120) IMG_DIM = 512 CONST_SCALE = 0.18215 SD_IDENTIFIERS = { + "1-4": "CompVis/stable-diffusion-v1-4", "1-5": "runwayml/stable-diffusion-v1-5", "2-0": "stabilityai/stable-diffusion-2-base", "2-1": "stabilityai/stable-diffusion-2-1-base", @@ -263,6 +266,7 @@ class StableDiffusion(nn.Module): image: TensorType["BS", 3, "H", "W"], guidance_scale: float = 100.0, grad_scaler: Optional[GradScaler] = None, + max_step_scale: float = 1.0 ) -> torch.Tensor: """Score Distilation Sampling loss proposed in DreamFusion paper (https://dreamfusion3d.github.io/) Args: @@ -276,7 +280,8 @@ class StableDiffusion(nn.Module): # Disable autocast when using diffusers pipeline. with torch.autocast(device_type="cuda", enabled=False): image = F.interpolate(image.half(), (IMG_DIM, IMG_DIM), mode="bilinear") - t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device) + max_step = int((self.max_step + 1) * max_step_scale) + t = torch.randint(self.min_step, max_step, [1], dtype=torch.long, device=self.device) latents = self.imgs_to_latent(image) # predict the noise residual with unet, NO grad! diff --git a/nerfstudio/models/dreamembedding.py b/nerfstudio/models/dreamembedding.py index e1955acd6..5c6bedfab 100644 --- a/nerfstudio/models/dreamembedding.py +++ b/nerfstudio/models/dreamembedding.py @@ -18,6 +18,7 @@ from collections import defaultdict from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple, Type +import numpy as np import imageio import torch import torch.nn.functional as F @@ -51,7 +52,8 @@ class DreamEmbeddingModelConfig(DreamFusionModelConfig): _target: Type = field(default_factory=lambda: DreamEmbeddingModel) """target class to instantiate""" - prompt: str = "A high quality zoomed out photo of a teddy bear" + # prompt: str = "A high quality zoomed out photo of a teddy bear" + prompt: str = "a high-quality photo of a pineapple" """prompt for stable dreamfusion""" orientation_loss_mult: float = 0.0001 @@ -124,26 +126,6 @@ class DreamEmbeddingModelConfig(DreamFusionModelConfig): max_res: int = 256 """Maximum resolution of the density field.""" - location_based_prompting: bool = True - """enables location based prompting""" - interpolated_prompting: bool = False - """enables interpolated location prompting""" - positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete" - """ how to incorporate position into prompt""" - top_prompt: str = ", overhead view" - """appended to prompt for overhead view""" - side_prompt: str = ", side view" - """appended to prompt for side view""" - front_prompt: str = ", front view" - """appended to prompt for front view""" - back_prompt: str = ", back view" - """appended to prompt for back view""" - guidance_scale: float = 100 - """guidance scale for sds loss""" - stablediffusion_device: Optional[str] = None - """device for stable diffusion""" - sd_version: str = "1-5" - class DreamEmbeddingModel(DreamFusionModel): """DreamEmbeddingModel Model @@ -165,6 +147,10 @@ class DreamEmbeddingModel(DreamFusionModel): """Set the fields and modules""" super().populate_modules() + self.feature_to_rgb_matrix = torch.tensor([[0.298, 0.187, -0.158, -0.184], + [0.207, 0.286, 0.189, -0.271], + [0.208, 0.173, 0.264, -0.473]]).to(self.sd_device) + # setting up fields self.field = DreamEmbeddingField(self.scene_box.aabb, max_res=self.config.max_res) self.renderer_feature = FeatureRenderer() @@ -178,9 +164,7 @@ class DreamEmbeddingModel(DreamFusionModel): if self.initialize_density: pos = ray_samples.frustums.get_positions() - # density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] - # density = density + density_blob - density_blob = self.density_strength * (-torch.exp(torch.norm(pos, dim=-1) / 0.6) + 2)[..., None] + density_blob = self.density_strength * (-torch.exp(torch.norm(pos, dim=-1) / 0.6) + 1.0)[..., None] density = torch.max(density + density_blob, torch.tensor([0.], device=self.device)) weights = ray_samples.get_weights(density) @@ -196,17 +180,24 @@ class DreamEmbeddingModel(DreamFusionModel): background = accum_mask_inv * background_feature + feature_to_rgb = features @ self.feature_to_rgb_matrix.T + + samp = np.random.random_sample() + if samp < 0.2: + rand_bg = torch.ones_like(background) * torch.rand(4, device=self.device) + latents = accum_mask * features + rand_bg * accum_mask_inv + else: + latents = accum_mask * features + background + outputs = { - "feature_only": features, + "feature_to_rgb": feature_to_rgb, "background_feature": background_feature, "background": background, "accumulation": accum_mask, "depth": depth, + "latents": latents } - latents = accum_mask * features + background - outputs["latents"] = latents - if ray_dims: height = ray_dims[0] width = ray_dims[1] @@ -219,9 +210,7 @@ class DreamEmbeddingModel(DreamFusionModel): white_bg = torch.tensor([1.0, 1.0, 1.0], device=self.device) a_mask = accum_mask.reshape(1, 1, height, width) - a_mask = torch.nn.functional.interpolate(a_mask, size=(height * 8, width * 8), mode="bilinear") - a_mask = torch.flatten(a_mask.permute(0, 2, 3, 1), end_dim=-2) a_mask_inv = 1 - a_mask @@ -229,8 +218,6 @@ class DreamEmbeddingModel(DreamFusionModel): rgb = torch.flatten(rgb, end_dim=-2) rgb = a_mask * rgb + (a_mask_inv * white_bg) rgb = rgb.to(torch.float32) - # print("rgb", rgb) - # print(rgb.shape) outputs["rgb"] = rgb @@ -314,8 +301,6 @@ class DreamEmbeddingModel(DreamFusionModel): loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( outputs["weights_list"], outputs["ray_samples_list"] ) - # if 45 < batch["central"] <= 135: - # loss_dict["input_latent_loss"] = self.l2_loss(self.input_latent.detach(), outputs["latents"]) return loss_dict def forward(self, ray_bundle: RayBundle, ray_dims: Tuple = ()) -> Dict[str, torch.Tensor]: @@ -378,16 +363,19 @@ class DreamEmbeddingModel(DreamFusionModel): accumulation=outputs["accumulation"], ) - size = outputs["latents"].shape[0] + # print(outputs.keys()) - latents_input = outputs["latents"].view(1, size, size, 4).half() - latents_input = latents_input.permute(0, 3, 1, 2) - rgb = self._sd.latents_to_img(latents_input) - rgb = rgb.permute(0, 2, 3, 1)[0].reshape(512, 512, 3).cpu().to(torch.float32) + # size = outputs["latents"].shape[0] + + # latents_input = outputs["latents"].view(1, size, size, 4).half() + # latents_input = latents_input.permute(0, 3, 1, 2) + # rgb = self._sd.latents_to_img(latents_input) + # rgb = rgb.permute(0, 2, 3, 1)[0].reshape(512, 512, 3).cpu().to(torch.float32) metrics_dict = {} images_dict = { - "img": rgb, + "img": outputs["rgb"], + "features": outputs["feature_to_rgb"], "accumulation": acc, "depth": depth, "prop_depth_0": prop_depth_0, diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 2abf72974..1d10674ad 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -134,10 +134,6 @@ class DreamFusionModelConfig(ModelConfig): max_res: int = 256 """Maximum resolution of the density field.""" - location_based_prompting: bool = True - """enables location based prompting""" - interpolated_prompting: bool = False - """enables interpolated location prompting""" positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete" """ how to incorporate position into prompt""" top_prompt: str = ", overhead view" @@ -154,6 +150,11 @@ class DreamFusionModelConfig(ModelConfig): """device for stable diffusion""" sd_version: str = "1-5" + t_taper_range: Tuple[int, int] = (0, 10000) + """Range of step values for the density tapering""" + t_taper_to: float = 1.0 + """Strength schedule of center density""" + class DreamFusionModel(Model): """DreamFusionModel Model @@ -274,6 +275,11 @@ class DreamFusionModel(Model): ): self.density_strength = np.interp(step, self.config.taper_range, self.config.taper_strength) + def taper_sds( + self, training_callback_attributes: TrainingCallbackAttributes, step: int + ): + self.sds_t_max_step = np.interp(step, self.config.t_taper_range, (1.0, self.config.t_taper_to)) + def start_training_normals( self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument ): @@ -310,6 +316,12 @@ class DreamFusionModel(Model): update_every_num_iters=1, args=[self, training_callback_attributes], ), + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + func=taper_sds, + update_every_num_iters=1, + args=[self, training_callback_attributes], + ), TrainingCallback( where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], iters=(self.config.start_normals_training,), @@ -356,10 +368,10 @@ class DreamFusionModel(Model): if self.initialize_density: pos = ray_samples.frustums.get_positions() - # density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] - # density = density + density_blob - density_blob = self.density_strength * (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 1.0)[..., None] - density = torch.max(density + density_blob, torch.tensor([1e-15], device=self.device)) + density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] + density = density + density_blob + # density_blob = self.density_strength * (-torch.exp(torch.norm(pos, dim=-1) / 0.6) + 2)[..., None] + # density = torch.max(density + density_blob, torch.tensor([1e-15], device=self.device)) weights = ray_samples.get_weights(density) weights_list.append(weights) @@ -391,7 +403,6 @@ class DreamFusionModel(Model): outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i]) normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) - # pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) # lambertian shading if self.config.random_light_source: # and self.training: @@ -408,45 +419,31 @@ class DreamFusionModel(Model): shaded, shaded_albedo = self.shader_lambertian( rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False ) + shaded, shaded_albedo = accum_mask * shaded, accum_mask * shaded_albedo outputs["normals"] = self.shader_normals(normals, weights=accum_mask) - # outputs["pred_normals"] = self.shader_normals(pred_normals, weights=accum_mask) - outputs["shaded"] = accum_mask * shaded - outputs["other_train_output"] = accum_mask * shaded_albedo + background - outputs["shaded_albedo"] = accum_mask * shaded_albedo + outputs["shaded"] = shaded + outputs["other_train_output"] = shaded_albedo + background + outputs["shaded_albedo"] = shaded_albedo outputs["rgb"] = accum_mask * rgb + background - # if shading_weight > 0: - # samp = np.random.random_sample() - # if samp > 0.5 and self.training: - # outputs["train_output"] = outputs["shaded"] - # elif samp < 0.2 and self.random_background and self.training: - # rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device) - # outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv - # else: - # outputs["train_output"] = accum_mask * shaded_albedo + background - # else: - # outputs["train_output"] = outputs["rgb"] + # while training 20% of the time use a random background + if np.random.random_sample() < 0.2 and self.random_background and self.training: + background = torch.ones_like(background) * torch.rand(3, device=self.device) * accum_mask_inv - samp = np.random.random_sample() - if samp < 0.4: - rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device) - train_output = accum_mask * rgb + rand_bg * accum_mask_inv + if shading_weight > 0: + samp = np.random.random_sample() + if samp > 0.5: + outputs["train_output"] = outputs["shaded"] + else: + outputs["train_output"] = shaded_albedo + background else: - train_output = accum_mask * rgb + background - - outputs["train_output"] = train_output + outputs["train_output"] = accum_mask * rgb + background 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], - # ) - 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 @@ -463,10 +460,6 @@ class DreamFusionModel(Model): 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"] - # ) else: loss_dict["orientation_loss"] = 0 loss_dict["pred_normal_loss"] = 0 @@ -537,6 +530,5 @@ class DreamFusionModel(Model): "prop_depth_0": prop_depth_0, "prop_depth_1": prop_depth_1, "normals": outputs["normals"], - # "pred_normals": outputs["pred_normals"], } return metrics_dict, images_dict diff --git a/nerfstudio/models/dreamfusion_ngp.py b/nerfstudio/models/dreamfusion_ngp.py index bcbdb79e9..fba0f81f7 100644 --- a/nerfstudio/models/dreamfusion_ngp.py +++ b/nerfstudio/models/dreamfusion_ngp.py @@ -22,7 +22,6 @@ import nerfacc import torch import numpy as np import torch.nn.functional as F -from nerfacc import ContractionType from torch.nn import Parameter from typing_extensions import Literal @@ -99,28 +98,26 @@ class DreamfusionNGPModelConfig(DreamFusionModelConfig): background when rendered at the end of training""" transmittance_end_schedule: int = 1500 """number of iterations to reach target_transmittance_end""" - - contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE - """Contraction type used for spatial deformation of the field.""" - grid_resolution: int = 64 + grid_resolution: int = 128 """Resolution of the grid used for the field.""" - # max_res: int = 2048 - max_res = 256 + grid_levels: int = 4 + """Levels of the grid used for the field.""" + max_res: int = 512 """Maximum resolution of the hashmap for the base mlp.""" log2_hashmap_size: int = 19 """Size of the hashmap for the base mlp""" + alpha_thre: float = 0.01 + """Threshold for opacity skipping.""" cone_angle: float = 0.0 """Should be set to 0.0 for blender scenes but 1./256 for real scenes.""" - render_step_size: float = 0.01 + render_step_size: float = None """Minimum step size for rendering.""" near_plane: float = 0.05 """How far along ray to start sampling.""" far_plane: float = 1e3 """How far along ray to stop sampling.""" - max_num_samples_per_ray: int = 24 - start_normals_training: int = 1000 """Start training normals after this many iterations""" start_lambertian_training: int = 1000 @@ -155,25 +152,26 @@ class DreamfusionNGPModel(DreamFusionModel): # setting up fields self.field = DreamNGPField( aabb=self.scene_box.aabb, - contraction_type=self.config.contraction_type, - use_appearance_embedding=False, num_images=self.num_train_data, log2_hashmap_size=self.config.log2_hashmap_size, max_res=self.config.max_res, + spatial_distortion=None ) self.initialize_density=self.config.initialize_density self.scene_aabb = Parameter(self.scene_box.aabb.flatten(), requires_grad=False) - # Occupancy Grid - self.occupancy_grid = nerfacc.OccupancyGrid( + if self.config.render_step_size is None: + # auto step size: ~1000 samples in the base level grid + self.config.render_step_size = ((self.scene_aabb[3:] - self.scene_aabb[:3]) ** 2).sum().sqrt().item() / 1000 + # Occupancy Grid. + self.occupancy_grid = nerfacc.OccGridEstimator( roi_aabb=self.scene_aabb, resolution=self.config.grid_resolution, - contraction_type=self.config.contraction_type, + levels=self.config.grid_levels, ) # samplers self.sampler = VolumetricSampler( - scene_aabb=self.scene_box.aabb, occupancy_grid=self.occupancy_grid, density_fn=self.field.density_fn, ) @@ -202,22 +200,25 @@ class DreamfusionNGPModel(DreamFusionModel): self, training_callback_attributes: TrainingCallbackAttributes ) -> List[TrainingCallback]: def update_occupancy_grid(step: int): - # TODO: needs to get access to the sampler, on how the step size is determinated at each x. See - # https://github.com/KAIR-BAIR/nerfacc/blob/127223b11401125a9fce5ce269bb0546ee4de6e8/examples/train_ngp_nerf.py#L190-L213 - self.occupancy_grid.every_n_step( + self.occupancy_grid.update_every_n_steps( step=step, - occ_eval_fn=lambda x: self.field.get_opacity(x, self.config.render_step_size), + occ_eval_fn=lambda x: self.field.density_fn(x) * self.config.render_step_size, ) - def start_training_normals( - self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument - ): - self.train_normals = True - def taper_density( self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument ): 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 + ): + self.train_normals = True + + def start_shaded_training( + self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument + ): + self.train_shaded = True return [ TrainingCallback( @@ -225,6 +226,18 @@ class DreamfusionNGPModel(DreamFusionModel): update_every_num_iters=1, func=update_occupancy_grid, ), + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + iters=(self.config.start_normals_training,), + func=start_training_normals, + args=[self, training_callback_attributes], + ), + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + iters=(self.config.start_lambertian_training,), + func=start_shaded_training, + args=[self, training_callback_attributes], + ), TrainingCallback( where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], func=taper_density, @@ -248,6 +261,7 @@ class DreamfusionNGPModel(DreamFusionModel): near_plane=self.config.near_plane, far_plane=self.config.far_plane, render_step_size=self.config.render_step_size, + alpha_thre=self.config.alpha_thre, cone_angle=self.config.cone_angle, ) # field_outputs = self.field(ray_samples, compute_normals=True) @@ -257,16 +271,17 @@ class DreamfusionNGPModel(DreamFusionModel): if self.initialize_density: pos = ray_samples.frustums.get_positions() - density_blob = self.density_strength * (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 0.5)[..., None] + density_blob = self.density_strength * (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 1.0)[..., None] density = torch.max(density + density_blob, torch.tensor([0.], device=self.device)) packed_info = nerfacc.pack_info(ray_indices, num_rays) weights = nerfacc.render_weight_from_density( + t_starts=ray_samples.frustums.starts[..., 0], + t_ends=ray_samples.frustums.ends[..., 0], + sigmas=density[..., 0], packed_info=packed_info, - sigmas=density, - t_starts=ray_samples.frustums.starts, - t_ends=ray_samples.frustums.ends, - ) + )[0] + weights = weights[..., None] rgb = self.renderer_rgb( rgb=field_outputs[FieldHeadNames.RGB], @@ -281,7 +296,6 @@ class DreamfusionNGPModel(DreamFusionModel): 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 accum_mask = torch.clamp((torch.nan_to_num(accumulation, nan=0.0)), min=0.0, max=1.0) accum_mask_inv = 1.0 - accum_mask @@ -306,45 +320,48 @@ class DreamfusionNGPModel(DreamFusionModel): outputs["train_output"] = train_output - # normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) - # outputs["normals"] = self.shader_normals(normals, weights=accum_mask) + normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) + outputs["normals"] = self.shader_normals(normals, weights=accum_mask) - # # lambertian shading - # if self.config.random_light_source: # and self.training: - # light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals) - # else: - # light_d = ray_bundle.origins[0] - # light_d = math.safe_normalize(light_d) + # lambertian shading + if self.config.random_light_source: # and self.training: + light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals) + else: + light_d = ray_bundle.origins[0] + light_d = math.safe_normalize(light_d) - # if (self.train_shaded and np.random.random_sample() > 0.75) or not self.training: - # shading_weight = 0.9 - # else: - # shading_weight = 0.0 + if (self.train_shaded and np.random.random_sample() > 0.75) or not self.training: + shading_weight = 0.9 + else: + shading_weight = 0.0 - # shaded, shaded_albedo = self.shader_lambertian( - # rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False - # ) + shaded, shaded_albedo = self.shader_lambertian( + rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False + ) + shaded, shaded_albedo = accum_mask * shaded, accum_mask * shaded_albedo - # outputs["shaded"] = accum_mask * shaded + outputs["shaded"] = shaded + outputs["other_train_output"] = shaded_albedo + background + outputs["shaded_albedo"] = shaded_albedo - # if shading_weight > 0: - # samp = np.random.random_sample() - # if samp > 0.5 and not self.training: - # outputs["train_output"] = outputs["shaded"] - # elif samp < 0.2 and self.random_background: - # rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device) - # outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv - # else: - # outputs["train_output"] = accum_mask * shaded_albedo + background - # else: - # outputs["train_output"] = outputs["rgb"] + # while training 20% of the time use a random background + if np.random.random_sample() < 0.2 and self.random_background and self.training: + background = torch.ones_like(background) * torch.rand(3, device=self.device) * accum_mask_inv + if shading_weight > 0: + samp = np.random.random_sample() + if samp > 0.5: + outputs["train_output"] = outputs["shaded"] + else: + outputs["train_output"] = shaded_albedo + background + else: + outputs["train_output"] = accum_mask * rgb + background - # 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, + ) assert weights.shape[-1] == 1 if self.config.opacity_penalty: @@ -397,4 +414,25 @@ class DreamfusionNGPModel(DreamFusionModel): loss_dict["sds_loss"] = sds_loss.to(self.device) return loss_dict + + 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]]: + + rgb = outputs["rgb"] + acc = colormaps.apply_colormap(outputs["accumulation"]) + depth = colormaps.apply_depth_colormap( + outputs["depth"], + accumulation=outputs["accumulation"], + ) + + metrics_dict = {} + + images_dict = { + "img": rgb, + "accumulation": acc, + "depth": depth, + } + + return metrics_dict, images_dict \ No newline at end of file diff --git a/nerfstudio/models/dreamfusion_tensor.py b/nerfstudio/models/dreamfusion_tensor.py index 6dafa71eb..7d065e127 100644 --- a/nerfstudio/models/dreamfusion_tensor.py +++ b/nerfstudio/models/dreamfusion_tensor.py @@ -22,7 +22,6 @@ import nerfacc import torch import numpy as np import torch.nn.functional as F -from nerfacc import ContractionType from torch.nn import Parameter from typing_extensions import Literal @@ -77,7 +76,6 @@ class DreamfusionTensorModelConfig(DreamFusionModelConfig): _target: Type = field(default_factory=lambda: DreamfusionTensorModel) """target class to instantiate""" - # prompt: str = "A high quality zoomed out photo of a teddy bear" prompt: str = "A high-quality photo of a pineapple" """prompt for stable dreamfusion""" @@ -356,7 +354,8 @@ class DreamfusionTensorModel(DreamFusionModel): "background": background, "accumulation": accum_mask, "depth": depth, - "train_output": train_output + "train_output": train_output, + "rgb": accum_mask * rgb + background } outputs["opacity_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * 0.5 @@ -404,4 +403,25 @@ class DreamfusionTensorModel(DreamFusionModel): loss_dict["sds_loss"] = sds_loss.to(self.device) return loss_dict + + 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]]: + + rgb = outputs["rgb"] + acc = colormaps.apply_colormap(outputs["accumulation"]) + depth = colormaps.apply_depth_colormap( + outputs["depth"], + accumulation=outputs["accumulation"], + ) + + metrics_dict = {} + + images_dict = { + "img": rgb, + "accumulation": acc, + "depth": depth, + } + + return metrics_dict, images_dict \ No newline at end of file diff --git a/nerfstudio/viewer/server/render_state_machine.py b/nerfstudio/viewer/server/render_state_machine.py index bd601bfd4..f0c5ec21f 100644 --- a/nerfstudio/viewer/server/render_state_machine.py +++ b/nerfstudio/viewer/server/render_state_machine.py @@ -123,7 +123,7 @@ class RenderStateMachine(threading.Thread): ) image_height, image_width = self._calculate_image_res(cam_msg.aspect) - # image_height, image_width = self._calculate_image_res(1.0) + # image_height, image_width = self._calculate_image_res(1.0) #!!!!!!!!!!!!111 intrinsics_matrix, camera_to_world_h = get_intrinsics_matrix_and_camera_to_world_h( @@ -253,7 +253,7 @@ class RenderStateMachine(threading.Thread): image_width: the maximum image width that can be rendered in the time budget """ max_res = self.viewer.control_panel.max_res - # max_res = 128 + # max_res = 128 #!!!!!!!!!!!!!!!!!!!111111 if self.state == "high": # high res is always static image_height = max_res diff --git a/scripts/viewer/run_viewer.py b/scripts/viewer/run_viewer.py index 67d96ae2a..07ceacbe5 100644 --- a/scripts/viewer/run_viewer.py +++ b/scripts/viewer/run_viewer.py @@ -54,52 +54,7 @@ class RunViewer: config.viewer = self.viewer.as_viewer_config() config.viewer.num_rays_per_chunk = num_rays_per_chunk - self._start_viewer(config, pipeline) - - def _start_viewer(self, config: TrainerConfig, pipeline: Pipeline): - base_dir = config.get_base_dir() - viewer_log_path = base_dir / config.viewer.relative_log_filename - - datapath = pipeline.datamanager.get_datapath() - if datapath is None: - datapath = base_dir - - viewer_state, banner_messages = viewer_utils.setup_viewer( - config.viewer, - log_filename=viewer_log_path, - datapath=datapath, - ) - - # We don't need logging, but writer.GLOBAL_BUFFER needs to be populated - config.logging.local_writer.enable = False - writer.setup_local_writer(config.logging, max_iter=config.max_num_iterations, banner_messages=banner_messages) - - assert viewer_state and pipeline.datamanager.train_dataset - viewer_state.init_scene( - dataset=pipeline.datamanager.train_dataset, - start_train=False, - ) - while True: - viewer_state.vis["renderingState/isTraining"].write(False) - self._update_viewer_state(viewer_state, pipeline) - - def _update_viewer_state(self, viewer_state: viewer_utils.ViewerState, pipeline: Pipeline): - """Updates the viewer state by rendering out scene with current pipeline - Returns the time taken to render scene. - - """ - # NOTE: step must be > 0 otherwise the rendering would not happen - step = 1 - num_rays_per_batch = pipeline.datamanager.get_train_rays_per_batch() - with TimeWriter(writer, EventName.ITER_VIS_TIME) as _: - try: - viewer_state.update_scene(self, step, pipeline.model, num_rays_per_batch) - except RuntimeError: - time.sleep(0.03) # sleep to allow buffer to reset - assert viewer_state.vis is not None - viewer_state.vis["renderingState/log_errors"].write( - "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing." - ) + _start_viewer(config, pipeline, step) def save_checkpoint(self, *args, **kwargs): """ @@ -109,7 +64,6 @@ class RunViewer: def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int): """Starts the viewer - Args: config: Configuration of pipeline to load pipeline: Pipeline instance of which to load weights @@ -117,10 +71,15 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int): """ base_dir = config.get_base_dir() viewer_log_path = base_dir / config.viewer.relative_log_filename + + datapath = pipeline.datamanager.get_datapath() + if datapath is None: + datapath = base_dir + viewer_state = ViewerState( config.viewer, log_filename=viewer_log_path, - datapath=pipeline.datamanager.get_datapath(), + datapath=datapath, pipeline=pipeline, ) banner_messages = [f"Viewer at: {viewer_state.viewer_url}"] @@ -147,7 +106,4 @@ def entrypoint(): if __name__ == "__main__": - entrypoint() - -# For sphinx docs -get_parser_fn = lambda: tyro.extras.get_parser(RunViewer) # noqa + entrypoint() \ No newline at end of file