everything but draemfusion works kinda

nerfstudio/configs/method_configs.py

@@ -447,7 +447,7 @@ method_configs["dreamfusion"] = TrainerConfig(
     steps_per_eval_batch=50,
     steps_per_eval_image=50,
     steps_per_save=200,
-    max_num_iterations=30000,
+    max_num_iterations=20000,
     mixed_precision=True,
     pipeline=VanillaPipelineConfig(
         generative=True,
@@ -501,7 +501,7 @@ method_configs["dreamngp"] = TrainerConfig(
     pipeline=VanillaPipelineConfig(
         generative=True,
         datamanager=DreamFusionDataManagerConfig(
-            horizontal_rotation_warmup=0,
+            horizontal_rotation_warmup=2000,
         ),
         model=DreamfusionNGPModelConfig(
             eval_num_rays_per_chunk=8192,
@@ -547,7 +547,7 @@ method_configs["dreamtensor"] = TrainerConfig(
     pipeline=VanillaPipelineConfig(
         generative=True,
         datamanager=DreamFusionDataManagerConfig(
-            horizontal_rotation_warmup=0,
+            horizontal_rotation_warmup=2000,
         ),
         model=DreamfusionTensorModelConfig(
             initialize_density=True,
@@ -557,14 +557,23 @@ method_configs["dreamtensor"] = TrainerConfig(
         ),
     ),
     optimizers={
-        "encodings": {
-            "optimizer": AdamOptimizerConfig(lr=0.0002),
-            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=30000),
-        },
+        # "fields": {
+        #     "optimizer": AdamOptimizerConfig(lr=5e-5),
+        #     "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-5, max_steps=20000),
+        # },        
+        # "encodings": {
+        #     "optimizer": AdamOptimizerConfig(lr=0.0002),
+        #     "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=20000),
+        # },
         "fields": {
             "optimizer": AdamOptimizerConfig(lr=0.001),
             "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000),
         },
+        "encodings": {
+            "optimizer": AdamOptimizerConfig(lr=0.02),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=30000),
+        },
+
     },
     viewer=ViewerConfig(num_rays_per_chunk=32000),
     vis="viewer",
@@ -672,7 +681,7 @@ method_configs["dreamembed"] = TrainerConfig(
     pipeline=VanillaPipelineConfig(
         generative=True,
         datamanager=DreamFusionDataManagerConfig(
-            horizontal_rotation_warmup=0,
+            horizontal_rotation_warmup=2000,
         ),
         model=DreamEmbeddingModelConfig(
             eval_num_rays_per_chunk=1 << 15,

nerfstudio/fields/dreamtensor_field.py

@@ -113,8 +113,8 @@ class DreamFusionTensorField(TensoRFField):
         )
 
     def get_density(self, ray_samples: RaySamples) -> TensorType:
-        # positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
-        # positions = positions * 2 - 1
+        positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
+        positions = positions * 2 - 1
         positions = ray_samples.frustums.get_positions()
         density = self.density_encoding(positions)
         density_enc = torch.sum(density, dim=-1)[:, :, None]
@@ -124,8 +124,8 @@ class DreamFusionTensorField(TensoRFField):
 
     def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None) -> TensorType:
         d = ray_samples.frustums.directions
-        # positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
-        # positions = positions * 2 - 1
+        positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
+        positions = positions * 2 - 1
         positions = ray_samples.frustums.get_positions()
 
         rgb_features = self.color_encoding(positions)

nerfstudio/models/dreamembedding.py

@@ -180,7 +180,7 @@ class DreamEmbeddingModel(DreamFusionModel):
             pos = ray_samples.frustums.get_positions()
             # density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
             # density = density + density_blob
-            density_blob = (-torch.exp(torch.norm(pos, dim=-1) / 0.6) + 2)[..., None]
+            density_blob = self.density_strength * (-torch.exp(torch.norm(pos, dim=-1) / 0.6) + 2)[..., None]
             density = torch.max(density + density_blob, torch.tensor([0.], device=self.device))
 
         weights = ray_samples.get_weights(density)

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 realistic pineapple with green leaves"
+    prompt: str = "A high-quality photo of a pineapple"
     # prompt: str = "a high-quality photo of a banana"
     """prompt for stable dreamfusion"""
 
@@ -78,7 +78,7 @@ class DreamFusionModelConfig(ModelConfig):
     """Randomizes light source per output."""
     initialize_density: bool = True
     """Initialize density in center of scene."""
-    taper_range: Tuple[int, int] = (0, 500)
+    taper_range: Tuple[int, int] = (0, 1000)
     """Range of step values for the density tapering"""
     taper_strength: Tuple[float, float] = (1.0, 0.0)
     """Strength schedule of center density"""
@@ -358,8 +358,8 @@ class DreamFusionModel(Model):
             pos = ray_samples.frustums.get_positions()
             # density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
             # density = density + density_blob
-            density_blob = (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 0.2)[..., None]
-            density = torch.max(density + density_blob, torch.tensor([0.], device=self.device))
+            density_blob = self.density_strength * (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 1.0)[..., None]
+            density = torch.max(density + density_blob, torch.tensor([1e-15], device=self.device))
 
         weights = ray_samples.get_weights(density)
         weights_list.append(weights)
@@ -391,7 +391,7 @@ class DreamFusionModel(Model):
             outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])
 
         normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
-        pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
+        # pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
 
         # lambertian shading
         if self.config.random_light_source:  # and self.training:
@@ -410,33 +410,42 @@ class DreamFusionModel(Model):
         )
 
         outputs["normals"] = self.shader_normals(normals, weights=accum_mask)
-        outputs["pred_normals"] = self.shader_normals(pred_normals, weights=accum_mask)
+        # outputs["pred_normals"] = self.shader_normals(pred_normals, weights=accum_mask)
         outputs["shaded"] = accum_mask * shaded
         outputs["other_train_output"] = accum_mask * shaded_albedo + background
         outputs["shaded_albedo"] = accum_mask * shaded_albedo
         outputs["rgb"] = accum_mask * rgb + background
 
-        if shading_weight > 0:
-            samp = np.random.random_sample()
-            if samp > 0.5 and not self.training:
-                outputs["train_output"] = outputs["shaded"]
-            elif samp < 0.2 and self.random_background:
-                rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device)
-                outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv
-            else:
-                outputs["train_output"] = accum_mask * shaded_albedo + background
+        # if shading_weight > 0:
+        #     samp = np.random.random_sample()
+        #     if samp > 0.5 and self.training:
+        #         outputs["train_output"] = outputs["shaded"]
+        #     elif samp < 0.2 and self.random_background and self.training:
+        #         rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device)
+        #         outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv
+        #     else:
+        #         outputs["train_output"] = accum_mask * shaded_albedo + background
+        # else:
+        #     outputs["train_output"] = outputs["rgb"]
+
+        samp = np.random.random_sample()
+        if samp < 0.4:
+            rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device)
+            train_output = accum_mask * rgb + rand_bg * accum_mask_inv
         else:
-            outputs["train_output"] = outputs["rgb"]
+            train_output = accum_mask * rgb + background
+
+        outputs["train_output"] = train_output
 
         outputs["rendered_orientation_loss"] = orientation_loss(
             weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
         )
 
-        outputs["rendered_pred_normal_loss"] = pred_normal_loss(
-            weights.detach(),
-            field_outputs[FieldHeadNames.NORMALS].detach(),
-            field_outputs[FieldHeadNames.PRED_NORMALS],
-        )
+        # outputs["rendered_pred_normal_loss"] = pred_normal_loss(
+        #     weights.detach(),
+        #     field_outputs[FieldHeadNames.NORMALS].detach(),
+        #     field_outputs[FieldHeadNames.PRED_NORMALS],
+        # )
 
         assert weights.shape[-1] == 1
         if self.config.opacity_penalty:
@@ -455,9 +464,9 @@ class DreamFusionModel(Model):
                 outputs["rendered_orientation_loss"]
             )
             # ground truth supervision for normals
-            loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean(
-                outputs["rendered_pred_normal_loss"]
-            )
+            # loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean(
+            #     outputs["rendered_pred_normal_loss"]
+            # )
         else:
             loss_dict["orientation_loss"] = 0
             loss_dict["pred_normal_loss"] = 0
@@ -528,6 +537,6 @@ class DreamFusionModel(Model):
             "prop_depth_0": prop_depth_0,
             "prop_depth_1": prop_depth_1,
             "normals": outputs["normals"],
-            "pred_normals": outputs["pred_normals"],
+            # "pred_normals": outputs["pred_normals"],
         }
         return metrics_dict, images_dict

nerfstudio/models/dreamfusion_ngp.py

@@ -224,13 +224,7 @@ class DreamfusionNGPModel(DreamFusionModel):
                 where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
                 update_every_num_iters=1,
                 func=update_occupancy_grid,
-            ),            
-            TrainingCallback(
-                where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
-                iters=(self.config.start_normals_training,),
-                func=start_training_normals,
-                args=[self, training_callback_attributes],
-            ),
+            ),   
             TrainingCallback(
                 where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
                 func=taper_density,
@@ -263,7 +257,7 @@ class DreamfusionNGPModel(DreamFusionModel):
 
         if self.initialize_density:
             pos = ray_samples.frustums.get_positions()
-            density_blob = (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 0.5)[..., None]
+            density_blob = self.density_strength * (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 0.5)[..., None]
             density = torch.max(density + density_blob, torch.tensor([0.], device=self.device))
 
         packed_info = nerfacc.pack_info(ray_indices, num_rays)
@@ -303,48 +297,54 @@ class DreamfusionNGPModel(DreamFusionModel):
             "rgb": accum_mask * rgb + background
         }
 
-        normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
-        outputs["normals"] = self.shader_normals(normals, weights=accum_mask)
-
-        # lambertian shading
-        if self.config.random_light_source:  # and self.training:
-            light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals)
+        samp = np.random.random_sample()
+        if samp < 0.4:
+            rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device)
+            train_output = accum_mask * rgb + rand_bg * accum_mask_inv
         else:
-            light_d = ray_bundle.origins[0]
-        light_d = math.safe_normalize(light_d)
+            train_output = accum_mask * rgb + background
+        
+        outputs["train_output"] = train_output
 
-        if (self.train_shaded and np.random.random_sample() > 0.75) or not self.training:
-            shading_weight = 0.9
-        else:
-            shading_weight = 0.0
+        # normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
+        # outputs["normals"] = self.shader_normals(normals, weights=accum_mask)
 
-        shaded, shaded_albedo = self.shader_lambertian(
-            rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False
-        )
+        # # lambertian shading
+        # if self.config.random_light_source:  # and self.training:
+        #     light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals)
+        # else:
+        #     light_d = ray_bundle.origins[0]
+        # light_d = math.safe_normalize(light_d)
 
-        outputs["shaded"] = accum_mask * shaded
+        # if (self.train_shaded and np.random.random_sample() > 0.75) or not self.training:
+        #     shading_weight = 0.9
+        # else:
+        #     shading_weight = 0.0
 
-        if shading_weight > 0:
-            samp = np.random.random_sample()
-            if samp > 0.5 and not self.training:
-                outputs["train_output"] = outputs["shaded"]
-            elif samp < 0.2 and self.random_background:
-                rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device)
-                outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv
-            else:
-                outputs["train_output"] = accum_mask * shaded_albedo + background
-        else:
-            outputs["train_output"] = outputs["rgb"]
+        # shaded, shaded_albedo = self.shader_lambertian(
+        #     rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False
+        # )
 
-        # print(outputs["rgb"].shape)
-        # outputs["train_output"] = outputs["rgb"]
+        # outputs["shaded"] = accum_mask * shaded
+
+        # if shading_weight > 0:
+        #     samp = np.random.random_sample()
+        #     if samp > 0.5 and not self.training:
+        #         outputs["train_output"] = outputs["shaded"]
+        #     elif samp < 0.2 and self.random_background:
+        #         rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device)
+        #         outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv
+        #     else:
+        #         outputs["train_output"] = accum_mask * shaded_albedo + background
+        # else:
+        #     outputs["train_output"] = outputs["rgb"]
 
 
-        outputs["rendered_orientation_loss"] = orientation_loss(
-            weights.detach(),
-            field_outputs[FieldHeadNames.NORMALS],
-            ray_bundle.directions,
-        )
+        # outputs["rendered_orientation_loss"] = orientation_loss(
+        #     weights.detach(),
+        #     field_outputs[FieldHeadNames.NORMALS],
+        #     ray_bundle.directions,
+        # )
 
         assert weights.shape[-1] == 1
         if self.config.opacity_penalty:

nerfstudio/models/dreamfusion_tensor.py

@@ -104,12 +104,12 @@ class DreamfusionTensorModelConfig(DreamFusionModelConfig):
     transmittance_end_schedule: int = 1500
     """number of iterations to reach target_transmittance_end"""
 
-    init_resolution: int = 300
+    init_resolution: int = 128
     """initial render resolution"""
     final_resolution: int = 300
     """final render resolution"""
     # upsampling_iters: Tuple[int, ...] = (2000, 3000, 4000, 5500, 7000)
-    upsampling_iters: Tuple[int, ...] = (500, 750, 1000)
+    upsampling_iters: Tuple[int, ...] = (1000, 3000, 7000)
     """specifies a list of iteration step numbers to perform upsampling"""
     num_samples: int = 100
     """Number of samples in field evaluation"""
@@ -247,7 +247,32 @@ class DreamfusionTensorModel(DreamFusionModel):
             self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
         ):
             self.density_strength = np.interp(step, self.config.taper_range, self.config.taper_strength)
+        def reinitialize_optimizer(
+            self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
+        ):
+            index = self.upsampling_iters.index(step)
+            resolution = self.upsampling_steps[index]
 
+            # upsample the position and direction grids
+            self.field.density_encoding.upsample_grid(resolution)
+            self.field.color_encoding.upsample_grid(resolution)
+
+            # reinitialize the encodings optimizer
+            optimizers_config = training_callback_attributes.optimizers.config
+            enc = training_callback_attributes.pipeline.get_param_groups()["encodings"]
+            lr_init = optimizers_config["encodings"]["optimizer"].lr
+
+            training_callback_attributes.optimizers.optimizers["encodings"] = optimizers_config["encodings"][
+                "optimizer"
+            ].setup(params=enc)
+            if optimizers_config["encodings"]["scheduler"]:
+                training_callback_attributes.optimizers.schedulers["encodings"] = (
+                    optimizers_config["encodings"]["scheduler"]
+                    .setup()
+                    .get_scheduler(
+                        optimizer=training_callback_attributes.optimizers.optimizers["encodings"], lr_init=lr_init
+                    )
+                )
         return [
             TrainingCallback(
                 where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
@@ -255,6 +280,12 @@ class DreamfusionTensorModel(DreamFusionModel):
                 update_every_num_iters=1,
                 args=[self, training_callback_attributes],
             ),
+            TrainingCallback(
+                where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
+                iters=self.upsampling_iters,
+                func=reinitialize_optimizer,
+                args=[self, training_callback_attributes],
+            )
         ]
         
     def get_param_groups(self) -> Dict[str, List[Parameter]]:
@@ -292,7 +323,7 @@ class DreamfusionTensorModel(DreamFusionModel):
 
         if self.initialize_density:
             pos = ray_samples_pdf.frustums.get_positions()
-            density_blob = (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 1.0)[..., None]
+            density_blob = self.density_strength * (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 1.0)[..., None]
             density = torch.max(density + density_blob, torch.tensor([0.], device=self.device))
 
         weights_fine = ray_samples_pdf.get_weights(density)
@@ -312,14 +343,24 @@ class DreamfusionTensorModel(DreamFusionModel):
         accum_mask_inv = 1.0 - accum_mask
         background = accum_mask_inv * background_rgb
 
+        samp = np.random.random_sample()
+        if samp < 0.4:
+            rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device)
+            train_output = accum_mask * rgb + rand_bg * accum_mask_inv
+        else:
+            train_output = accum_mask * rgb + background
+
         outputs = {
             "rgb_only": rgb,
             "background_rgb": background_rgb,
             "background": background,
             "accumulation": accum_mask,
             "depth": depth,
-            "train_output": accum_mask * rgb + background
+            "train_output": train_output
         } 
+
+        outputs["opacity_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * 0.5
+
         return outputs
 
 
@@ -359,6 +400,8 @@ class DreamfusionTensorModel(DreamFusionModel):
             grad_scaler=self.grad_scaler,
         )
 
+        loss_dict["opacity_loss"] = outputs["opacity_loss"].mean()
+
         loss_dict["sds_loss"] = sds_loss.to(self.device)
         return loss_dict
     

nerfstudio/pipelines/base_pipeline.py

@@ -285,16 +285,6 @@ class VanillaPipeline(Pipeline):
         """
         ray_bundle, batch = self.datamanager.next_train(step)
         model_outputs = self.model(ray_bundle)
-
-        # if step % 100 == 0:
-        #     image = model_outputs["train_output"]
-        #     image = image.view(int(image.shape[0] ** 0.5), int(image.shape[0] ** 0.5), 3)
-        #     image = image.detach().cpu().numpy()
-        #     image = (image * 255).round().astype("uint8")
-        #     mediapy.write_image(str('dreamngp.png'), image)
-        #     # mediapy.write_image(str('dreamfusion.png'), image)
-
-
         metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
 
         if self.config.datamanager.camera_optimizer is not None: