From b57781ab0ce650b328f03a1c38813a5a64969ebb Mon Sep 17 00:00:00 2001
From: Abhik Ahuja <ahujaabhik@berkeley.edu>
Date: Mon, 28 Nov 2022 14:21:00 -0800
Subject: [PATCH] Added dreamfusion to CLI, removed direction embeddings

---
 nerfstudio/configs/method_configs.py   | 33 ++++++++++++++++++++++++++
 nerfstudio/fields/dreamfusion_field.py | 11 +--------
 nerfstudio/models/dreamfusion.py       |  2 +-
 3 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py
index 92fe1966e..957650053 100644
--- a/nerfstudio/configs/method_configs.py
+++ b/nerfstudio/configs/method_configs.py
@@ -43,6 +43,7 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import (
 )
 from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
 from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
+from nerfstudio.models.dreamfusion import DreamFusionModelConfig
 from nerfstudio.models.instant_ngp import InstantNGPModelConfig
 from nerfstudio.models.mipnerf import MipNerfModel
 from nerfstudio.models.nerfacto import NerfactoModelConfig
@@ -62,6 +63,7 @@ descriptions = {
     "tensorf": "tensorf",
     "dnerf": "Dynamic-NeRF model. (slow)",
     "phototourism": "Uses the Phototourism data.",
+    "dreamfusion": "NeRF Model used in implementation of stable dreamfusion.",
 }
 
 method_configs["nerfacto"] = Config(
@@ -249,6 +251,37 @@ method_configs["phototourism"] = Config(
     vis="viewer",
 )
 
+method_configs["dreamfusion"] = Config(
+    method_name="dreamfusion",
+    trainer=TrainerConfig(
+        steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
+    ),
+    pipeline=VanillaPipelineConfig(
+        datamanager=VanillaDataManagerConfig(
+            dataparser=NerfstudioDataParserConfig(),
+            train_num_rays_per_batch=4096,
+            eval_num_rays_per_batch=4096,
+            camera_optimizer=CameraOptimizerConfig(
+                mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
+            ),
+        ),
+        model=DreamFusionModelConfig(),
+    ),
+    optimizers={
+        "proposal_networks": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+        "fields": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
+
 AnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[  # Don't show unparseable (fixed) arguments in helptext.
     tyro.conf.FlagConversionOff[
         tyro.extras.subcommand_type_from_defaults(defaults=method_configs, descriptions=descriptions)
diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py
index 486b6011c..35b6ac5cb 100644
--- a/nerfstudio/fields/dreamfusion_field.py
+++ b/nerfstudio/fields/dreamfusion_field.py
@@ -103,14 +103,6 @@ class DreamFusionField(Field):
         features_per_level = 2
         growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
 
-        self.direction_encoding = tcnn.Encoding(
-            n_input_dims=3,
-            encoding_config={
-                "otype": "SphericalHarmonics",
-                "degree": 4,
-            },
-        )
-
         self.position_encoding = tcnn.Encoding(
             n_input_dims=3,
             encoding_config={"otype": "Frequency", "n_frequencies": 2},
@@ -191,7 +183,6 @@ class DreamFusionField(Field):
 
         directions = get_normalized_directions(ray_samples.frustums.directions)
         directions_flat = directions.view(-1, 3)
-        d = self.direction_encoding(directions_flat)
 
         outputs_shape = ray_samples.frustums.directions.shape[:-1]
 
@@ -205,7 +196,7 @@ class DreamFusionField(Field):
             x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions)
             outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x)
 
-        h = [density_embedding.view(-1, self.geo_feat_dim)]
+        h = density_embedding.view(-1, self.geo_feat_dim)
 
         rgb = self.mlp_head(h).view(*outputs_shape, -1).to(directions)
         outputs.update({FieldHeadNames.RGB: rgb})
diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py
index 4963e1dfc..e2939ed19 100644
--- a/nerfstudio/models/dreamfusion.py
+++ b/nerfstudio/models/dreamfusion.py
@@ -30,7 +30,7 @@ from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
 
 from nerfstudio.cameras.rays import RayBundle
 
-from nerfstudio.field_components.encodings import NeRFEncoding, 
+# from nerfstudio.field_components.encodings import NeRFEncoding,
 from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.fields.dreamfusion_field import DreamFusionField
 from nerfstudio.model_components.losses import MSELoss