Location based prompting, density initialization fades, LR change

nerfstudio/configs/method_configs.py

@@ -271,12 +271,8 @@ method_configs["dreamfusion"] = Config(
         model=DreamFusionModelConfig(),
     ),
     optimizers={
-        "proposal_networks": {
-            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
-            "scheduler": None,
-        },
         "fields": {
-            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15),
             "scheduler": None,
         },
     },

nerfstudio/generative/stable_diffusion.py

@@ -113,6 +113,11 @@ class StableDiffusion(nn.Module):
         # perform guidance (high scale from paper!)
         noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
         noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        # with torch.no_grad():
+        #     diffused_img = self.latents_to_img(latents_noisy - noise_pred)
+        #     diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy().reshape((512, 512, 3))
+        #     diffused_img = (diffused_img * 255).round().astype("uint8")
+        #     plt.imsave("sd_img.png", diffused_img)
 
         # w(t), sigma_t^2
         w = 1 - self.alphas[t]

nerfstudio/models/dreamfusion.py

@@ -94,6 +94,7 @@ class DreamFusionModel(Model):
         self.initialize_density = config.initialize_density
         self.train_normals = False
         self.train_shaded = False
+        self.density_strength = 1.0
         super().__init__(config=config, **kwargs)
 
     def get_training_callbacks(
@@ -101,10 +102,10 @@ class DreamFusionModel(Model):
     ) -> List[TrainingCallback]:
 
         # the callback that we want to run every X iterations after the training iteration
-        def stop_initialize_density(
+        def taper_density(
             self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
         ):
-            self.initialize_density = False
+            self.density_strength -= 0.2
 
         def start_training_normals(
             self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
@@ -119,8 +120,8 @@ class DreamFusionModel(Model):
         callbacks = [
             TrainingCallback(
                 where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
-                iters=(1000,),
-                func=stop_initialize_density,
+                iters=(200, 400, 600, 800, 1000),
+                func=taper_density,
                 args=[self, training_callback_attributes],
             ),
             TrainingCallback(
@@ -185,7 +186,7 @@ class DreamFusionModel(Model):
         density = field_outputs[FieldHeadNames.DENSITY]
         if self.initialize_density:
             pos = ray_samples_uniform.frustums.get_positions()
-            density_blob = 100 * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
+            density_blob = 50 * self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
             density += density_blob
 
         weights = ray_samples_uniform.get_weights(density)
@@ -231,15 +232,15 @@ class DreamFusionModel(Model):
         outputs["pred_normals"] = pred_normals
         outputs["shaded"] = accum_mask * shaded + torch.ones_like(shaded) * accum_mask_inv
         outputs["shaded_albedo"] = shaded_albedo
-        outputs["render"] = accum_mask * shaded_albedo + background
+        outputs["render"] = accum_mask * rgb + background
 
         if ratio > 0:
             if np.random.random_sample() > 0.5:
-                outputs["train_output"] = accum_mask * shaded + background
+                outputs["train_output"] = shaded
             else:
                 outputs["train_output"] = accum_mask * shaded_albedo + background
         else:
-            outputs["train_output"] = accum_mask * rgb + background
+            outputs["train_output"] = outputs["render"]
 
         # if self.train_normals:
 

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -25,6 +25,10 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig):
     """specifies the datamanager config"""
     prompt: str = "A high quality photo of a pineapple"
     """prompt for stable dreamfusion"""
+    location_based_prompting: bool = True
+    """enables location based prompting"""
+    opacity_penalty: bool = True
+    """enables penalty to encourage transparent scenes"""
 
 
 class DreamfusionPipeline(VanillaPipeline):
@@ -38,7 +42,11 @@ class DreamfusionPipeline(VanillaPipeline):
     ):
         super().__init__(config, device, test_mode, world_size, local_rank)
         self.sd = StableDiffusion(device)
-        self.text_embedding = self.sd.get_text_embeds(config.prompt, "")
+        self.base_text_embedding = self.sd.get_text_embeds(config.prompt, "")
+        if config.location_based_prompting:
+            self.front_text_embedding = self.sd.get_text_embeds(f"{config.prompt}, front view", "")
+            self.side_text_embedding = self.sd.get_text_embeds(f"{config.prompt}, side view", "")
+            self.back_text_embedding = self.sd.get_text_embeds(f"{config.prompt}, back view", "")
 
     @profiler.time_function
     def custom_step(self, step: int, grad_scaler: GradScaler, optimizers: Optimizers):
@@ -55,8 +63,7 @@ class DreamfusionPipeline(VanillaPipeline):
         ray_bundle, batch = self.datamanager.next_train(step)
         model_outputs = self.model(ray_bundle)
 
-        # Just uses albedo for now
-        albedo_output = model_outputs["train_output"].view(1, 64, 64, 3).permute(0, 3, 1, 2)
+        train_output = model_outputs["train_output"].view(1, 64, 64, 3).permute(0, 3, 1, 2)
 
         accumulation = model_outputs["accumulation"].view(64, 64).detach().cpu().numpy()
         accumulation = np.clip(accumulation, 0.0, 1.0)
@@ -70,10 +77,23 @@ class DreamfusionPipeline(VanillaPipeline):
         shaded = np.clip(shaded, 0.0, 1.0)
         plt.imsave("nerf_textureless.jpg", shaded)
 
-        sds_loss, latents, grad = self.sd.sds_loss(self.text_embedding, albedo_output)
+        if self.config.location_based_prompting:
+            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
+            elif batch["central"] > 225 and batch["central"] <= 315:
+                text_embedding = self.side_text_embedding
+        else:
+            text_embedding = self.base_text_embedding
+
+        sds_loss, latents, grad = self.sd.sds_loss(text_embedding, train_output)
         # TODO: opacity penalty using transmittance, not accumultation
-        accum_mean = np.mean(1.0 - accumulation)
-        sds_loss *= np.min((0.8, accum_mean))
+        if self.config.opacity_penalty:
+            accum_mean = np.mean(1.0 - accumulation)
+            sds_loss *= np.min((0.5, accum_mean))
 
         grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True)
         # optimizers.scheduler_step_all(step)