update fields to use tcnn

nerfstudio/configs/method_configs.py

@@ -481,15 +481,15 @@ method_configs["dreamfusion"] = TrainerConfig(
     ),
     optimizers={
         "proposal_networks": {
-            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15),
             "scheduler": ExponentialDecaySchedulerConfig(
-                warmup_steps=2000, lr_final=1e-6, max_steps=20000, ramp="linear"
+                warmup_steps=4000, lr_final=5e-7, max_steps=20000, ramp="linear"
             ),
         },
         "fields": {
-            "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
+            "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15),
             "scheduler": ExponentialDecaySchedulerConfig(
-                warmup_steps=2000, lr_final=1e-6, max_steps=20000, ramp="linear"
+                warmup_steps=4000, lr_final=5e-7, max_steps=20000, ramp="linear"
             ),
         },
     },

nerfstudio/fields/dreamfusion_field.py

@@ -17,22 +17,24 @@ Field for dreamfusion model
 """
 
 
-from typing import Optional, Tuple, Dict
+from typing import Optional, Tuple, Dict, Literal
 
 import numpy as np
 import torch
 from torch.nn.parameter import Parameter
 from jaxtyping import Float
-from torch import Tensor
+from torch import Tensor, nn
 
+from nerfstudio.field_components.mlp import MLP
 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
+from nerfstudio.field_components.encodings import HashEncoding, SHEncoding
 from nerfstudio.fields.base_field import Field
 
 try:
-    import tinycudann as tcnn
+    pass
 except ImportError:
     # tinycudann module doesn't exist
     pass
@@ -73,6 +75,7 @@ class DreamFusionField(Field):
         log2_hashmap_size: int = 19,
         num_layers_color: int = 3,
         hidden_dim_color: int = 64,
+        implementation: Literal["tcnn", "torch"] = "tcnn",
     ) -> None:
         super().__init__()
 
@@ -81,64 +84,94 @@ class DreamFusionField(Field):
 
         base_res = 16
         features_per_level = 2
-        growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
+        np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
 
-        self.position_encoding = tcnn.Encoding(
-            n_input_dims=3,
-            encoding_config={"otype": "Frequency", "n_frequencies": 2},
+        self.direction_encoding = SHEncoding(
+            levels=4,
+            implementation=implementation,
         )
 
-        self.direction_encoding = tcnn.Encoding(
-            n_input_dims=3,
-            encoding_config={
-                "otype": "SphericalHarmonics",
-                "degree": 4,
-            },
+        encoder = HashEncoding(
+            num_levels=num_levels,
+            min_res=base_res,
+            max_res=max_res,
+            log2_hashmap_size=log2_hashmap_size,
+            features_per_level=features_per_level,
+            implementation=implementation,
+        )
+        network = MLP(
+            in_dim=encoder.get_out_dim(),
+            num_layers=num_layers,
+            layer_width=hidden_dim,
+            out_dim=1 + self.geo_feat_dim,
+            activation=nn.ReLU(),
+            out_activation=None,
+            implementation=implementation,
+        )
+        self.mlp_base = torch.nn.Sequential(encoder, network)
+
+        # self.mlp_base = tcnn.NetworkWithInputEncoding(
+        #     n_input_dims=3,
+        #     n_output_dims=1 + self.geo_feat_dim,
+        #     encoding_config={
+        #         "otype": "HashGrid",
+        #         "n_levels": num_levels,
+        #         "n_features_per_level": features_per_level,
+        #         "log2_hashmap_size": log2_hashmap_size,
+        #         "base_resolution": base_res,
+        #         "per_level_scale": growth_factor,
+        #     },
+        #     network_config={
+        #         "otype": "FullyFusedMLP",
+        #         "activation": "ReLU",
+        #         "output_activation": "None",
+        #         "n_neurons": hidden_dim,
+        #         "n_hidden_layers": num_layers - 1,
+        #     },
+        # )
+
+        # self.mlp_background_color = tcnn.Network(
+        #     n_input_dims=self.direction_encoding.n_output_dims,
+        #     n_output_dims=3,
+        #     network_config={
+        #         "otype": "FullyFusedMLP",
+        #         "activation": "ReLU",
+        #         "output_activation": "Sigmoid",
+        #         "n_neurons": 32,
+        #         "n_hidden_layers": 1,
+        #     },
+        # )
+
+        self.mlp_background_color = MLP(
+            in_dim=self.direction_encoding.get_out_dim(),
+            num_layers=2,
+            layer_width=32,
+            out_dim=3,
+            activation=nn.ReLU(),
+            out_activation=nn.Sigmoid(),
+            implementation=implementation,
         )
 
-        self.mlp_base = tcnn.NetworkWithInputEncoding(
-            n_input_dims=3,
-            n_output_dims=1 + self.geo_feat_dim,
-            encoding_config={
-                "otype": "HashGrid",
-                "n_levels": num_levels,
-                "n_features_per_level": features_per_level,
-                "log2_hashmap_size": log2_hashmap_size,
-                "base_resolution": base_res,
-                "per_level_scale": growth_factor,
-            },
-            network_config={
-                "otype": "FullyFusedMLP",
-                "activation": "ReLU",
-                "output_activation": "None",
-                "n_neurons": hidden_dim,
-                "n_hidden_layers": num_layers - 1,
-            },
-        )
-
-        self.mlp_background_color = tcnn.Network(
-            n_input_dims=self.direction_encoding.n_output_dims,
-            n_output_dims=3,
-            network_config={
-                "otype": "FullyFusedMLP",
-                "activation": "ReLU",
-                "output_activation": "Sigmoid",
-                "n_neurons": 32,
-                "n_hidden_layers": 1,
-            },
-        )
-
-        self.mlp_head = tcnn.Network(
-            n_input_dims=self.geo_feat_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,
-            },
+        self.mlp_head = MLP(
+            in_dim=self.geo_feat_dim,
+            num_layers=num_layers_color,
+            layer_width=hidden_dim_color,
+            out_dim=3,
+            activation=nn.ReLU(),
+            out_activation=nn.Sigmoid(),
+            implementation=implementation,
         )
+        # self.mlp_head = tcnn.Network(
+        #     n_input_dims=self.geo_feat_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) -> Tuple[Tensor, Tensor]:
         """Computes and returns the densities."""

nerfstudio/models/dreamfusion.py

@@ -66,7 +66,7 @@ class DreamFusionModelConfig(ModelConfig):
 
     _target: Type = field(default_factory=lambda: DreamFusionModel)
     """target class to instantiate"""
-    prompt: str = "a high quality photo of a tree frog"
+    prompt: str = "a high quality zoomed out photo of an orange tabby cat sitting on a box"
     """prompt for stable dreamfusion"""
 
     orientation_loss_mult: Tuple[float, float] = (0.01, 100.0)
@@ -159,6 +159,8 @@ class DreamFusionModelConfig(ModelConfig):
     """diffusion model for SDS loss"""
     sd_version: str = "1-5"
     """model version when using stable diffusion"""
+    implementation: Literal["tcnn", "torch"] = "tcnn"
+    """Which implementation to use for the model."""
 
 
 class DreamFusionModel(Model):
@@ -229,8 +231,7 @@ class DreamFusionModel(Model):
         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 = HashMLPDensityField(
-                self.scene_box.aabb,
-                **prop_net_args,
+                self.scene_box.aabb, **prop_net_args, implementation=self.config.implementation
             )
             self.proposal_networks.append(network)
         self.density_fns.extend([network.density_fn for network in self.proposal_networks])