interpolated location prompting option, lowering init density, transmittance loss interpolation flipped

nerfstudio/configs/method_configs.py

@@ -351,7 +351,7 @@ method_configs["dreamfusion"] = DreamfusionTrainerConfig(
             eval_num_rays_per_batch=4096,
         ),
         model=DreamFusionModelConfig(eval_num_rays_per_chunk=1 << 15, distortion_loss_mult=10),
-        alphas_loss_mult=1.0,
+        alphas_loss_mult=1.5,
         opacity_loss_mult=1.0,
         guidance_scale=100,
     ),

nerfstudio/models/dreamfusion.py

@@ -64,8 +64,8 @@ class DreamFusionModelConfig(ModelConfig):
     _target: Type = field(default_factory=lambda: DreamFusionModel)
     """target class to instantiate"""
 
-    num_samples: int = 128
-    """Number of samples in field evaluation"""
+    # num_samples: int = 128
+    # """Number of samples in field evaluation"""
     orientation_loss_mult: float = 0.01
     """Orientation loss multipier on computed normals."""
     pred_normal_loss_mult: float = 0.001
@@ -74,17 +74,17 @@ class DreamFusionModelConfig(ModelConfig):
     """Randomizes light source per output."""
     initialize_density: bool = True
     """Initialize density in center of scene."""
-    init_density_strength: float = 50.0
+    init_density_strength: float = 0.1
     """Initial strength of center density"""
     sphere_collider: bool = True
     """Use spherical collider instead of box"""
-    target_transmittance_start: float = 0.8
+    target_transmittance_start: float = 0.4
     """target transmittance for opacity penalty. This is the percent of the scene that is
     background when rendered at the start of training"""
-    target_transmittance_end: float = 0.4
+    target_transmittance_end: float = 0.7
     """target transmittance for opacity penalty. This is the percent of the scene that is
     background when rendered at the end of training"""
-    transmittance_end_schedule: int = 1000
+    transmittance_end_schedule: int = 1500
     """number of iterations to reach target_transmittance_end"""
     num_proposal_samples_per_ray: Tuple[int, ...] = (256, 96)
     """Number of samples per ray for each proposal network."""
@@ -131,7 +131,6 @@ class DreamFusionModel(Model):
         config: DreamFusionModelConfig,
         **kwargs,
     ) -> None:
-        self.num_samples = config.num_samples
         self.initialize_density = config.initialize_density
         self.train_normals = False
         self.train_shaded = False
@@ -203,7 +202,7 @@ class DreamFusionModel(Model):
         def taper_density(
             self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
         ):
-            self.density_strength -= 0.2
+            self.density_strength -= 0.25
 
         def start_training_normals(
             self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
@@ -218,7 +217,10 @@ class DreamFusionModel(Model):
         def update_target_transmittance(
             self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
         ):
-            if step < self.config.transmittance_end_schedule:
+            if (
+                step < self.config.transmittance_end_schedule
+                and self.target_transmittance > self.config.target_transmittance_end
+            ):
                 self.target_transmittance -= (1 / self.config.transmittance_end_schedule) * (
                     self.config.target_transmittance_start - self.config.target_transmittance_end
                 )
@@ -236,19 +238,19 @@ class DreamFusionModel(Model):
         callbacks = [
             TrainingCallback(
                 where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
-                iters=(200, 400, 600, 800, 1000),
+                iters=(150, 300, 450, 600),
                 func=taper_density,
                 args=[self, training_callback_attributes],
             ),
             TrainingCallback(
                 where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
-                iters=(1000,),
+                iters=(600,),
                 func=start_training_normals,
                 args=[self, training_callback_attributes],
             ),
             TrainingCallback(
                 where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
-                iters=(1000,),
+                iters=(600,),
                 func=start_shaded_training,
                 args=[self, training_callback_attributes],
             ),
@@ -387,10 +389,6 @@ class DreamFusionModel(Model):
 
         loss_dict["alphas_loss"] = outputs["alphas_loss"].mean()
 
-        assert (
-            self.target_transmittance >= self.config.target_transmittance_end
-            and self.target_transmittance <= self.config.target_transmittance_start
-        )
         loss_dict["opacity_loss"] = -torch.minimum(
             (1 - outputs["accumulation"]).mean(), torch.tensor(self.target_transmittance)
         )

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -37,7 +37,7 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig):
     """target class to instantiate"""
     datamanager: DreamFusionDataManagerConfig = DreamFusionDataManagerConfig()
     """specifies the datamanager config"""
-    prompt: str = "A zoomed out high-quality photo of Temple of Heaven"
+    prompt: str = "A high-quality photo of a tree frog on a stump"
     """prompt for stable dreamfusion"""
 
     location_based_prompting: bool = True
@@ -51,10 +51,10 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig):
     front_prompt: str = ", front view"
     """appended to prompt for front view"""
     back_prompt: str = ", back view"
-    """appended to prompt for back view"""
+    """appended to prompt for back view"""  
     alphas_penalty: bool = True
     """enables penalty to encourage sparse weights (penalizing for uniform density along ray)"""
-    alphas_loss_mult: float = 1
+    alphas_scale: float = 1
     """scale for alphas penalty"""
     opacity_penalty: bool = True
     """enables penalty to encourage transparent scenes, as in "dreamfields" paper"""
@@ -114,18 +114,27 @@ class DreamfusionPipeline(VanillaPipeline):
 
         if self.config.location_based_prompting:
             if self.config.interpolated_prompting:
-                horiz = batch["central"]
-                vert = batch["vertical"]
+                horiz = batch["central"].to(self.sd_device)
+                vert = max(batch["vertical"].to(self.sd_device), 0)
 
-                if batch["central"] > 315 or batch["central"] <= 45:
-                    text_embedding = self.front_text_embedding
-                elif batch["central"] > 45 and batch["central"] <= 135:
-                    text_embedding = self.side_text_embedding
-                elif batch["central"] > 135 and batch["central"] <= 225:
-                    text_embedding = self.back_text_embedding
-                else:  # batch["central"] > 225 and batch["central"] <= 315:
-                    text_embedding = self.side_text_embedding
+                if batch["central"] > 0 or batch["central"] <= 90:
+                    text_embedding = (
+                        (horiz) * self.side_text_embedding + (90 - horiz) * self.front_text_embedding
+                    ) / 90.0
+                elif batch["central"] > 90 and batch["central"] <= 180:
+                    text_embedding = (
+                        (horiz - 90) * self.back_text_embedding + (180 - horiz) * self.side_text_embedding
+                    ) / 90.0
+                elif batch["central"] > 180 and batch["central"] <= 270:
+                    text_embedding = (
+                        (horiz - 180) * self.side_text_embedding + (270 - horiz) * self.back_text_embedding
+                    ) / 90.0
+                else:  # batch["central"] > 270 and batch["central"] <= 360:
+                    text_embedding = (
+                        (horiz - 270) * self.front_text_embedding + (360 - horiz) * self.side_text_embedding
+                    ) / 90.0
 
+                text_embedding = (vert * text_embedding + (90 - vert) * self.top_text_embedding) / 90.0
 
             else:
                 if batch["vertical"] < 30: