From ddfe2079188d4531ee84a7b99376edfc71e874de Mon Sep 17 00:00:00 2001 From: terrance Date: Thu, 25 May 2023 20:17:14 +0000 Subject: [PATCH] typing + linting --- nerfstudio/fields/dreamfusion_field.py | 23 ++++++---- ...iffusion_utils.py => dreamfusion_utils.py} | 21 +++++---- nerfstudio/generative/utils.py | 3 +- nerfstudio/models/dreamfusion.py | 46 ++++++++++--------- 4 files changed, 51 insertions(+), 42 deletions(-) rename nerfstudio/generative/{stable_diffusion_utils.py => dreamfusion_utils.py} (90%) diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py index db2d75ddc..a42bd4697 100644 --- a/nerfstudio/fields/dreamfusion_field.py +++ b/nerfstudio/fields/dreamfusion_field.py @@ -17,17 +17,18 @@ Field for dreamfusion model """ -from typing import Optional +from typing import Optional, Tuple, Dict import numpy as np import torch from torch.nn.parameter import Parameter -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor from nerfstudio.cameras.rays import RayBundle, RaySamples from nerfstudio.data.scene_box import SceneBox from nerfstudio.field_components.activations import trunc_exp -from nerfstudio.field_components.field_heads import FieldHeadNames, PredNormalsFieldHead +from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.base_field import Field try: @@ -37,7 +38,7 @@ except ImportError: pass -def get_normalized_directions(directions: TensorType["bs":..., 3]): +def get_normalized_directions(directions: Float[Tensor, "bs 3"]): """SH encoding must be in the range [0, 1] Args: @@ -54,13 +55,16 @@ class DreamFusionField(Field): 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 hidden_dim_color: dimension of hidden layers for color network """ def __init__( self, - aabb, + aabb: Tensor, num_layers: int = 2, hidden_dim: int = 64, geo_feat_dim: int = 15, @@ -69,7 +73,6 @@ class DreamFusionField(Field): log2_hashmap_size: int = 19, num_layers_color: int = 3, hidden_dim_color: int = 64, - hidden_dim_normal: int = 32, ) -> None: super().__init__() @@ -137,7 +140,7 @@ class DreamFusionField(Field): }, ) - def get_density(self, ray_samples: RaySamples): + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: """Computes and returns the densities.""" positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) self._sample_locations = positions @@ -154,7 +157,7 @@ class DreamFusionField(Field): density = trunc_exp(density_before_activation.to(positions)) return density, base_mlp_out - def get_background_rgb(self, ray_bundle: RayBundle): + def get_background_rgb(self, ray_bundle: RayBundle) -> Tensor: """Predicts background colors at infinity.""" directions = get_normalized_directions(ray_bundle.directions) @@ -164,7 +167,9 @@ class DreamFusionField(Field): return background_rgb - def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None): + def get_outputs( + self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None + ) -> Dict[FieldHeadNames, Tensor]: assert density_embedding is not None outputs = {} diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/dreamfusion_utils.py similarity index 90% rename from nerfstudio/generative/stable_diffusion_utils.py rename to nerfstudio/generative/dreamfusion_utils.py index 54a4758bf..4ba62b024 100644 --- a/nerfstudio/generative/stable_diffusion_utils.py +++ b/nerfstudio/generative/dreamfusion_utils.py @@ -1,4 +1,4 @@ -# Copyright 2022 The Nerfstudio Team. All rights reserved. +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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. @@ -12,14 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Stable diffusion utils -""" +"""Utility helper functions for generative 3D models""" -from torchtyping import TensorType +import torch +from torch import Tensor +from jaxtyping import Float from typing_extensions import Literal from typing import Union -import torch from nerfstudio.generative.stable_diffusion import StableDiffusion from nerfstudio.generative.deepfloyd import DeepFloyd @@ -65,7 +64,7 @@ class PositionalTextEmbeddings: side_prompt: Prompt for side view back_prompt: Prompt for back view front_prompt: Prompt for front view - """ + """ self.base_embed = self.diffusion.get_text_embeds(base_prompt, "") self.top_embed = self.diffusion.get_text_embeds(top_prompt, "") self.side_embed = self.diffusion.get_text_embeds(side_prompt, "") @@ -75,7 +74,9 @@ class PositionalTextEmbeddings: if isinstance(self.diffusion, DeepFloyd): self.diffusion.delete_text_encoder() - def get_text_embedding(self, vertical_angle: TensorType[1], horizontal_angle: TensorType[1]): + def get_text_embedding( + self, vertical_angle: Float[Tensor, "1"], horizontal_angle: Float[Tensor, "1"] + ) -> Float[Tensor, "2 max_length embed_dim"]: """Get text embedding based on the position of the camera relative to the scene. This trick is used in Dreamfusion (https://dreamfusion3d.github.io/). @@ -99,8 +100,8 @@ class PositionalTextEmbeddings: else: # horizontal_angle > 225 and horizontal_angle <= 315: text_embedding = self.side_embed elif self.positional_prompting == "interpolated": - horiz = horizontal_angle.to(self.diffusion_device) - vert = max(vertical_angle.to(self.diffusion_device), 0) + horiz = horizontal_angle.detach().numpy()[0] + vert = max(vertical_angle.detach().numpy()[0], 0) if 0 < horizontal_angle <= 90: text_embedding = (horiz) * self.side_embed + (90 - horiz) * self.front_embed diff --git a/nerfstudio/generative/utils.py b/nerfstudio/generative/utils.py index e5996da6d..4dedeed8a 100644 --- a/nerfstudio/generative/utils.py +++ b/nerfstudio/generative/utils.py @@ -12,12 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Utility helper functions for diffusion models""" import torch from torch.cuda.amp import custom_bwd, custom_fwd -"""Gradient helper functions for SDS loss""" - class _SDSGradient(torch.autograd.Function): """Custom gradient function for SDS loss. Since it is already computed, we can just return it.""" diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index de2a3aa9e..52a254b84 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -37,13 +37,12 @@ from nerfstudio.fields.density_fields import HashMLPDensityField from nerfstudio.fields.dreamfusion_field import DreamFusionField from nerfstudio.generative.stable_diffusion import StableDiffusion from nerfstudio.generative.deepfloyd import DeepFloyd -from nerfstudio.generative.stable_diffusion_utils import PositionalTextEmbeddings +from nerfstudio.generative.dreamfusion_utils import PositionalTextEmbeddings from nerfstudio.model_components.losses import ( MSELoss, distortion_loss, interlevel_loss, orientation_loss, - pred_normal_loss, ) from nerfstudio.model_components.ray_samplers import ( ProposalNetworkSampler, @@ -159,6 +158,7 @@ class DreamFusionModelConfig(ModelConfig): diffusion_model: Literal["stablediffusion", "deepfloyd"] = "deepfloyd" """diffusion model for SDS loss""" sd_version: str = "1-5" + """model version when using stable diffusion""" class DreamFusionModel(Model): @@ -186,7 +186,6 @@ class DreamFusionModel(Model): self.target_transmittance = config.target_transmittance_start self.grad_scaler = kwargs["grad_scaler"] - self.positional_prompting = config.positional_prompting self.guidance_scale = config.guidance_scale self.top_prompt = config.top_prompt self.side_prompt = config.side_prompt @@ -194,9 +193,7 @@ class DreamFusionModel(Model): self.front_prompt = config.front_prompt self.diffusion_device = ( - torch.device(kwargs["device"]) - if config.diffusion_device is None - else torch.device(config.diffusion_device) + torch.device(kwargs["device"]) if config.diffusion_device is None else torch.device(config.diffusion_device) ) super().__init__(config=config, **kwargs) @@ -217,7 +214,7 @@ class DreamFusionModel(Model): back_prompt=self.cur_prompt + self.back_prompt, front_prompt=self.cur_prompt + self.front_prompt, diffusion_model=self._diffusion_model, - positional_prompting=self.positional_prompting, + positional_prompting=self.config.positional_prompting, ) # setting up fields @@ -238,11 +235,13 @@ class DreamFusionModel(Model): self.proposal_networks.append(network) self.density_fns.extend([network.density_fn for network in self.proposal_networks]) - 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, - ) + def update_schedule(step): + return 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, @@ -274,7 +273,6 @@ class DreamFusionModel(Model): def get_training_callbacks( self, training_callback_attributes: TrainingCallbackAttributes ) -> List[TrainingCallback]: - # the callback that we want to run every X iterations after the training iteration def taper_density( self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument @@ -296,15 +294,21 @@ class DreamFusionModel(Model): ): if step <= self.config.start_normals_training: self.orientation_loss_mult = 0 - else: - self.orientation_loss_mult = np.interp(step, self.config.orientation_loss_mult, - (self.config.start_normals_training, self.config.orientation_loss_mult_end)) - + else: + self.orientation_loss_mult = np.interp( + step, + self.config.orientation_loss_mult, + (self.config.start_normals_training, self.config.orientation_loss_mult_end), + ) + # anneal the weights of the proposal network before doing PDF sampling def set_anneal(step): # https://arxiv.org/pdf/2111.12077.pdf eq. 18 train_frac = np.clip(step / self.config.proposal_weights_anneal_max_num_iters, 0, 1) - bias = lambda x, b: (b * x) / ((b - 1) * x + 1) + + def bias(x, b): + return b * x / ((b - 1) * x + 1) + anneal = bias(train_frac, self.config.proposal_weights_anneal_slope) self.proposal_sampler.set_anneal(anneal) @@ -420,8 +424,8 @@ class DreamFusionModel(Model): outputs["rgb"] = accum_mask * rgb + background # 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 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() @@ -494,7 +498,7 @@ class DreamFusionModel(Model): text_embedding.to(self.diffusion_device), train_output.to(self.diffusion_device), guidance_scale=int(self.guidance_scale), - grad_scaler=self.grad_scaler, + grad_scaler=self.grad_scaler, ) loss_dict["sds_loss"] = sds_loss.to(self.device)