diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 46ca5e279..cc5bada47 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -27,7 +27,6 @@ from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.configs.base_config import ViewerConfig from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig from nerfstudio.data.datamanagers.depth_datamanager import DepthDataManagerConfig -from nerfstudio.data.datamanagers.sdf_datamanager import SDFDataManagerConfig from nerfstudio.data.datamanagers.dreamfusion_datamanager import ( DreamFusionDataManagerConfig, ) @@ -448,24 +447,22 @@ method_configs["dreamembed"] = TrainerConfig( pipeline=VanillaPipelineConfig( generative=True, datamanager=DreamFusionDataManagerConfig( - horizontal_rotation_warmup=4000, + horizontal_rotation_warmup=0, ), model=DreamEmbeddingModelConfig( eval_num_rays_per_chunk=1 << 15, distortion_loss_mult=0.1, interlevel_loss_mult=100.0, - orientation_loss_mult=1., - max_res=256, + orientation_loss_mult=1.0, + max_res=128, sphere_collider=True, initialize_density=False, - random_background=True, proposal_warmup=500, proposal_update_every=0, proposal_weights_anneal_max_num_iters=100, - start_normals_training=1000, opacity_loss_mult=1, positional_prompting="discrete", - guidance_scale=100, + guidance_scale=50, ), ), optimizers={ @@ -476,7 +473,7 @@ method_configs["dreamembed"] = TrainerConfig( "fields": { "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15), "scheduler": ExponentialDecaySchedulerConfig( - warmup_steps=2000, lr_final=1e-6, max_steps=20000, ramp="linear" + warmup_steps=1000, lr_final=1e-6, max_steps=20000, ramp="linear" ), }, }, diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py index 2a2e5247e..ddfff95a9 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py @@ -55,9 +55,7 @@ class TrivialDataset(InputDataset): def __getitem__(self, index: int) -> Dict: return { - "image": torch.cat( - [torch.ones(128, 256, 3), torch.zeros(128, 256, 3)], dim=0 - ), + "image": torch.cat([torch.ones(128, 256, 3), torch.zeros(128, 256, 3)], dim=0), "image_idx": index, } @@ -95,13 +93,11 @@ def random_train_pose( ) # This is the uniform sample on the part of the sphere we care about where 0 = 0 degrees and 1 = 360 degrees sampled_uniform = ( - torch.rand(size) * (vertical_rotation_range[1] - vertical_rotation_range[0]) - + vertical_rotation_range[0] + torch.rand(size) * (vertical_rotation_range[1] - vertical_rotation_range[0]) + vertical_rotation_range[0] ) / 180 vertical_rotation = torch.arccos(1 - 2 * sampled_uniform) central_rotation = torch.deg2rad( - torch.rand(size) * (central_rotation_range[1] - central_rotation_range[0]) - + central_rotation_range[0] + torch.rand(size) * (central_rotation_range[1] - central_rotation_range[0]) + central_rotation_range[0] ) c_cos = torch.cos(central_rotation) @@ -131,9 +127,7 @@ def random_train_pose( # Default directions are facing in the -z direction, so origins should face opposite way origins = torch.stack([torch.tensor([0, 0, 1])] * size, dim=0) - origins = (origins * radius_mean) + ( - origins * (torch.randn((origins.shape)) * radius_std) - ) + origins = (origins * radius_mean) + (origins * (torch.randn((origins.shape)) * radius_std)) R = torch.bmm(rot_z, rot_x) # Want to have Rx @ Ry @ origin t = ( torch.bmm(R, origins.unsqueeze(-1)) @@ -220,9 +214,7 @@ class DreamFusionDataManager(DataManager): # pylint: disable=abstract-method self.test_split = "test" if test_mode in ["test", "inference"] else "val" if self.config.data is not None: - CONSOLE.print( - "[red] --data should not be used with the DreamFusionDataManager[/red]" - ) + CONSOLE.print("[red] --data should not be used with the DreamFusionDataManager[/red]") sys.exit(1) cameras, _, _ = random_train_pose( @@ -266,7 +258,7 @@ class DreamFusionDataManager(DataManager): # pylint: disable=abstract-method # ).flatten() # return ray_bundle, {"initialization": False} - horizontal_range = min((step / self.config.horizontal_rotation_warmup), 1) * 180 + horizontal_range = min((step / max(1, self.config.horizontal_rotation_warmup)), 1) * 180 cameras, vertical_rotation, central_rotation = random_train_pose( self.config.train_images_per_batch, @@ -280,9 +272,7 @@ class DreamFusionDataManager(DataManager): # pylint: disable=abstract-method center=self.config.center, central_rotation_range=(-horizontal_range, horizontal_range), ) - ray_bundle = cameras.generate_rays( - torch.tensor(list(range(self.config.train_images_per_batch))) - ).flatten() + ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten() # camera_idx = torch.randint(0, self.eval_cameras.shape[0], [1], dtype=torch.long, device=self.device) # ray_bundle = self.eval_cameras.generate_rays(camera_idx).flatten() diff --git a/nerfstudio/field_components/field_heads.py b/nerfstudio/field_components/field_heads.py index 630ee874c..45c420894 100644 --- a/nerfstudio/field_components/field_heads.py +++ b/nerfstudio/field_components/field_heads.py @@ -42,9 +42,6 @@ class FieldHeadNames(Enum): SDF = "sdf" ALPHA = "alpha" GRADIENT = "gradient" - SDF = "sdf" - ALPHA = "alpha" - GRADIENT = "gradient" class FieldHead(FieldComponent): diff --git a/nerfstudio/fields/dreamembedding_field.py b/nerfstudio/fields/dreamembedding_field.py index 97480fd3a..be69e52ef 100644 --- a/nerfstudio/fields/dreamembedding_field.py +++ b/nerfstudio/fields/dreamembedding_field.py @@ -1,3 +1,17 @@ +# 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. + from typing import Optional @@ -18,6 +32,7 @@ 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] @@ -26,6 +41,7 @@ def get_normalized_directions(directions: TensorType["bs":..., 3]): """ return (directions + 1.0) / 2.0 + class DreamEmbeddingField(Field): """DreamEmbedding Field that uses TCNN @@ -132,7 +148,7 @@ class DreamEmbeddingField(Field): # from smaller internal (float16) parameters. density = trunc_exp(density_before_activation.to(positions)) return density, base_mlp_out - + def get_background_feature(self, ray_bundle: RayBundle): """Predicts background colors at infinity.""" directions = get_normalized_directions(ray_bundle.directions) @@ -142,7 +158,7 @@ class DreamEmbeddingField(Field): background_feature = self.mlp_background_feature(directions_flat).view(*outputs_shape, -1).to(directions) return background_feature - + def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None): assert density_embedding is not None outputs = {} @@ -156,4 +172,4 @@ class DreamEmbeddingField(Field): feature = self.mlp_head(h).view(*outputs_shape, -1).to(directions) outputs.update({FieldHeadNames.FEATURE: feature}) - return outputs \ No newline at end of file + return outputs diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/stable_diffusion_utils.py index d67b5999d..d7b66b7d3 100644 --- a/nerfstudio/generative/stable_diffusion_utils.py +++ b/nerfstudio/generative/stable_diffusion_utils.py @@ -16,9 +16,9 @@ Stable diffusion utils """ +import torch from torchtyping import TensorType from typing_extensions import Literal -import torch from nerfstudio.generative.stable_diffusion import StableDiffusion diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 357119db8..3278ff9e5 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -342,15 +342,16 @@ class NormalsRenderer(nn.Module): if normalize: n = safe_normalize(n) return n - + + class FeatureRenderer(nn.Module): """Calculate features along the ray.""" @classmethod def forward( - cls, + cls, features: TensorType["bs":..., "num_samples", 4], weights: TensorType["bs":..., "num_samples", 1], ) -> TensorType["bs":..., 4]: feat = torch.sum(weights * features, dim=-2) - return feat \ No newline at end of file + return feat diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index 4040c33be..933609884 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -44,13 +44,9 @@ class ModelConfig(InstantiateConfig): """target class to instantiate""" enable_collider: bool = True """Whether to create a scene collider to filter rays.""" - collider_params: Optional[Dict[str, float]] = to_immutable_dict( - {"near_plane": 2.0, "far_plane": 6.0} - ) + collider_params: Optional[Dict[str, float]] = to_immutable_dict({"near_plane": 2.0, "far_plane": 6.0}) """parameters to instantiate scene collider with""" - loss_coefficients: Dict[str, float] = to_immutable_dict( - {"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0} - ) + loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0}) """parameters to instantiate density field with""" eval_num_rays_per_chunk: int = 4096 """specifies number of rays per chunk during eval""" @@ -80,9 +76,7 @@ class Model(nn.Module): super().__init__() self.config = config self.scene_box = scene_box - self.render_aabb = ( - None # the box that we want to render - should be a subset of scene_box - ) + self.render_aabb = None # the box that we want to render - should be a subset of scene_box self.num_train_data = num_train_data self.kwargs = kwargs self.collider = None @@ -160,9 +154,7 @@ class Model(nn.Module): return {} @abstractmethod - def get_loss_dict( - self, outputs, batch, metrics_dict=None - ) -> Dict[str, torch.Tensor]: + def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: """Computes and returns the losses dict. Args: @@ -172,9 +164,7 @@ class Model(nn.Module): """ @torch.no_grad() - def get_outputs_for_camera_ray_bundle( - self, camera_ray_bundle: RayBundle - ) -> Dict[str, torch.Tensor]: + def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle) -> Dict[str, torch.Tensor]: """Takes in camera parameters and computes the output of the model. Args: @@ -187,9 +177,7 @@ class Model(nn.Module): for i in range(0, num_rays, num_rays_per_chunk): start_idx = i end_idx = i + num_rays_per_chunk - ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle( - start_idx, end_idx - ) + ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx) outputs = self.forward(ray_bundle=ray_bundle) for output_name, output in outputs.items(): # type: ignore outputs_lists[output_name].append(output) @@ -224,10 +212,7 @@ class Model(nn.Module): Args: loaded_state: dictionary of pre-trained model states """ - state = { - key.replace("module.", ""): value - for key, value in loaded_state["model"].items() - } + state = {key.replace("module.", ""): value for key, value in loaded_state["model"].items()} self.load_state_dict(state) # type: ignore def update_to_step(self, step: int) -> None: diff --git a/nerfstudio/models/dreamembedding.py b/nerfstudio/models/dreamembedding.py index bb4de6209..be9c76928 100644 --- a/nerfstudio/models/dreamembedding.py +++ b/nerfstudio/models/dreamembedding.py @@ -1,14 +1,26 @@ +# 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. + from __future__ import annotations from collections import defaultdict from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple, Type - -import torch import imageio +import torch import torch.nn.functional as F - from typing_extensions import Literal from nerfstudio.cameras.rays import RayBundle @@ -28,7 +40,6 @@ from nerfstudio.model_components.losses import ( interlevel_loss, orientation_loss, ) - from nerfstudio.model_components.renderers import FeatureRenderer from nerfstudio.models.dreamfusion import DreamFusionModel, DreamFusionModelConfig from nerfstudio.utils import colormaps, colors, math, misc @@ -57,8 +68,6 @@ class DreamEmbeddingModelConfig(DreamFusionModelConfig): """Strength schedule of center density""" sphere_collider: bool = True """Use spherical collider instead of box""" - random_background: bool = True - """Randomly choose between using background mlp and random color for background""" target_transmittance_start: float = 0.4 """target transmittance for opacity penalty. This is the percent of the scene that is background when rendered at the start of training""" @@ -67,7 +76,7 @@ class DreamEmbeddingModelConfig(DreamFusionModelConfig): background when rendered at the end of training""" transmittance_end_schedule: int = 1500 """number of iterations to reach target_transmittance_end""" - num_proposal_samples_per_ray: Tuple[int, ...] = (256, 96) + num_proposal_samples_per_ray: Tuple[int, ...] = (96, 48) """Number of samples per ray for each proposal network.""" num_nerf_samples_per_ray: int = 48 """Number of samples per ray for the nerf network.""" @@ -85,13 +94,13 @@ class DreamEmbeddingModelConfig(DreamFusionModelConfig): "hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, - "max_res": 128, + "max_res": 64, }, { "hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, - "max_res": 256, + "max_res": 128, }, ] ) @@ -106,8 +115,6 @@ class DreamEmbeddingModelConfig(DreamFusionModelConfig): """Proposal loss multiplier.""" distortion_loss_mult: float = 1.0 """Distortion loss multiplier.""" - start_normals_training: int = 1000 - """Start training normals after this many iterations""" start_lambertian_training: int = 1000 """start training with lambertian shading after this many iterations""" opacity_penalty: bool = True @@ -159,34 +166,19 @@ class DreamEmbeddingModel(DreamFusionModel): super().populate_modules() # setting up fields - self.field = DreamEmbeddingField( - self.scene_box.aabb, max_res=self.config.max_res - ) - image = imageio.imread('/home/terrance/nerfstudio/data/generative/redfrogcrop.png') - image = (torch.from_numpy(image).to(self.sd_device).half() / 255.0).unsqueeze(0).permute(0, 3, 1, 2) - self.input_image = F.interpolate(image, (512, 512), mode="bilinear") - - input_latent = self.sd.imgs_to_latent(self.input_image) - self.input_latent = input_latent.permute(0, 2, 3, 1).flatten(end_dim=-2) - + self.field = DreamEmbeddingField(self.scene_box.aabb, max_res=self.config.max_res) self.renderer_feature = FeatureRenderer() - self.l2_loss = MSELoss() def get_outputs(self, ray_bundle: RayBundle, ray_dims: Tuple): # uniform sampling background_feature = self.field.get_background_feature(ray_bundle) - ray_samples, weights_list, ray_samples_list = self.proposal_sampler( - ray_bundle, density_fns=self.density_fns - ) + ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns) field_outputs = self.field(ray_samples, compute_normals=True) density = field_outputs[FieldHeadNames.DENSITY] 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_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] density = density + density_blob weights = ray_samples.get_weights(density) @@ -195,13 +187,9 @@ class DreamEmbeddingModel(DreamFusionModel): accumulation = self.renderer_accumulation(weights) depth = self.renderer_depth(weights, ray_samples) - features = self.renderer_feature( - features=field_outputs[FieldHeadNames.FEATURE], weights=weights - ) + features = self.renderer_feature(features=field_outputs[FieldHeadNames.FEATURE], weights=weights) - accum_mask = torch.clamp( - (torch.nan_to_num(accumulation, nan=0.0)), min=0.0, max=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 background = accum_mask_inv * background_feature @@ -236,9 +224,7 @@ class DreamEmbeddingModel(DreamFusionModel): 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.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 @@ -258,13 +244,9 @@ class DreamEmbeddingModel(DreamFusionModel): outputs["ray_samples_list"] = ray_samples_list 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] - ) + 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 - ) + normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) outputs["normals"] = self.shader_normals(normals, weights=accum_mask) outputs["rendered_orientation_loss"] = orientation_loss( @@ -275,34 +257,25 @@ class DreamEmbeddingModel(DreamFusionModel): 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 - ) + outputs["opacity_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult return outputs - def get_loss_dict( - self, outputs, batch, metrics_dict=None - ) -> Dict[str, torch.Tensor]: + def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: # Scaling metrics by coefficients to create the losses. loss_dict = {} loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) if self.train_normals: # orientation loss for computed normals - loss_dict[ - "orientation_loss" - ] = self.config.orientation_loss_mult * torch.mean( + loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean( outputs["rendered_orientation_loss"] ) else: loss_dict["orientation_loss"] = 0 if self.config.opacity_penalty: - loss_dict["opacity_loss"] = ( - self.config.opacity_loss_mult * outputs["opacity_loss"].mean() - ) + loss_dict["opacity_loss"] = self.config.opacity_loss_mult * outputs["opacity_loss"].mean() if self.prompt != self.cur_prompt: self.cur_prompt = self.prompt @@ -339,23 +312,17 @@ class DreamEmbeddingModel(DreamFusionModel): loss_dict["sds_loss"] = sds_loss.to(self.device) if self.training: - loss_dict[ - "distortion_loss" - ] = self.config.distortion_loss_mult * distortion_loss( + loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss( outputs["weights_list"], outputs["ray_samples_list"] ) - loss_dict[ - "interlevel_loss" - ] = self.config.interlevel_loss_mult * interlevel_loss( + 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]: + def forward(self, ray_bundle: RayBundle, ray_dims: Tuple = ()) -> Dict[str, torch.Tensor]: """Run forward starting with a ray bundle. This outputs different things depending on the configuration of the model and whether or not the batch is provided (whether or not we are training basically) @@ -369,9 +336,7 @@ class DreamEmbeddingModel(DreamFusionModel): return self.get_outputs(ray_bundle, ray_dims) @torch.no_grad() - def get_outputs_for_camera_ray_bundle( - self, camera_ray_bundle: RayBundle - ) -> Dict[str, torch.Tensor]: + def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle) -> Dict[str, torch.Tensor]: """Takes in camera parameters and computes the output of the model. Args: @@ -384,12 +349,8 @@ class DreamEmbeddingModel(DreamFusionModel): for i in range(0, num_rays, num_rays_per_chunk): start_idx = i end_idx = i + num_rays_per_chunk - ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle( - start_idx, end_idx - ) - outputs = self.forward( - ray_bundle=ray_bundle, ray_dims=(image_height, image_width) - ) + ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx) + outputs = self.forward(ray_bundle=ray_bundle, ray_dims=(image_height, image_width)) for output_name, output in outputs.items(): # type: ignore outputs_lists[output_name].append(output) outputs = {}