From 390b91033919666ac4e4b53b7b64dd599b42ce3b Mon Sep 17 00:00:00 2001
From: Matt Tancik <mattt4001@gmail.com>
Date: Thu, 2 Feb 2023 16:59:07 -0800
Subject: [PATCH] Cleanup

---
 nerfstudio/fields/dreamfusion_field.py       | 16 +++-------------
 nerfstudio/pipelines/dreamfusion_pipeline.py |  8 +++++---
 2 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py
index ce3656b03..af2b073c1 100644
--- a/nerfstudio/fields/dreamfusion_field.py
+++ b/nerfstudio/fields/dreamfusion_field.py
@@ -28,7 +28,6 @@ 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.spatial_distortions import SpatialDistortion
 from nerfstudio.fields.base_field import Field
 from nerfstudio.utils import math
 
@@ -58,7 +57,6 @@ class DreamFusionField(Field):
         geo_feat_dim: output geo feat dimensions
         num_layers_color: number of hidden layers for color network
         hidden_dim_color: dimension of hidden layers for color network
-        spatial_distortion: spatial distortion to apply to the scene
     """
 
     def __init__(
@@ -69,16 +67,13 @@ class DreamFusionField(Field):
         geo_feat_dim: int = 15,
         num_layers_color: int = 3,
         hidden_dim_color: int = 64,
-        hidden_dim_transient: int = 64,
-        spatial_distortion: Optional[SpatialDistortion] = None,
+        hidden_dim_normal: int = 32,
     ) -> None:
         super().__init__()
 
         self.aabb = Parameter(aabb, requires_grad=False)
         self.geo_feat_dim = geo_feat_dim
 
-        self.spatial_distortion = spatial_distortion
-
         num_levels = 16
         max_res = 1024
         base_res = 16
@@ -122,7 +117,7 @@ class DreamFusionField(Field):
         # predicted normals
         self.mlp_pred_normals = tcnn.Network(
             n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims,
-            n_output_dims=hidden_dim_transient,
+            n_output_dims=hidden_dim_normal,
             network_config={
                 "otype": "FullyFusedMLP",
                 "activation": "ReLU",
@@ -159,12 +154,7 @@ class DreamFusionField(Field):
 
     def get_density(self, ray_samples: RaySamples):
         """Computes and returns the densities."""
-        if self.spatial_distortion is not None:
-            positions = ray_samples.frustums.get_positions()
-            positions = self.spatial_distortion(positions)
-            positions = (positions + 2.0) / 4.0
-        else:
-            positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
+        positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
         self._sample_locations = positions
         if not self._sample_locations.requires_grad:
             self._sample_locations.requires_grad = True
diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py
index cfb8a967c..e4d1f3def 100644
--- a/nerfstudio/pipelines/dreamfusion_pipeline.py
+++ b/nerfstudio/pipelines/dreamfusion_pipeline.py
@@ -130,11 +130,13 @@ class DreamfusionPipeline(VanillaPipeline):
                 train_output.to(self.sd_device),
                 guidance_scale=int(self.config.guidance_scale),
             )
-            loss_dict["sds_loss"] = sds_loss.to(self.device)
+            sds_loss = sds_loss.to(self.device)
+
         # TODO: opacity penalty using transmittance, not accumultation
         if self.config.opacity_penalty:
-            accum_mean = np.mean(1.0 - accumulation)
-            sds_loss *= np.min((0.5, accum_mean))
+            accum_mean = torch.mean(1.0 - model_outputs["accumulation"])
+            sds_loss *= torch.clamp(accum_mean, min=0.5)
+        loss_dict["sds_loss"] = sds_loss.to(self.device)
 
         model_outputs["latents"] = latents
         model_outputs["grad"] = grad