typing + linting

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 = {}
 

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

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."""

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)