Merge branch 'generative' of https://github.com/plenoptix/nerfactory into generative

nerfstudio/generative/stable_diffusion.py

@@ -117,7 +117,7 @@ class StableDiffusion(nn.Module):
         #     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)
+        #     plt.imsave("sdimg.png", diffused_img)
 
         # w(t), sigma_t^2
         w = 1 - self.alphas[t]

nerfstudio/models/dreamfusion.py

@@ -66,12 +66,14 @@ class DreamFusionModelConfig(ModelConfig):
     """Number of samples in field evaluation"""
     orientation_loss_mult: float = 0.01
     """Orientation loss multipier on computed normals."""
-    pred_normal_loss_mult: float = 0.0003
+    pred_normal_loss_mult: float = 0.001
     """Predicted normal loss multiplier."""
     random_light_source: bool = True
     """Randomizes light source per output."""
     initialize_density: bool = True
     """Initialize density in center of scene."""
+    init_density_strength: float = 50.0
+    """Initial strength of center density"""
     sphere_collider: bool = True
     """Use spherical collider instead of box"""
 
@@ -183,9 +185,13 @@ class DreamFusionModel(Model):
         field_outputs = self.field(ray_samples_uniform, compute_normals=True)
 
         density = field_outputs[FieldHeadNames.DENSITY]
-        if self.initialize_density:
+        if self.initialize_density and self.training:
             pos = ray_samples_uniform.frustums.get_positions()
-            density_blob = 50 * self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
+            density_blob = (
+                self.config.init_density_strength
+                * self.density_strength
+                * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
+            )
             density += density_blob
 
         weights = ray_samples_uniform.get_weights(density)
@@ -212,7 +218,7 @@ class DreamFusionModel(Model):
         pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)  # .detach()
 
         # lambertian shading
-        if self.config.random_light_source:
+        if self.config.random_light_source and self.training:  # and if training
             light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals)
         else:
             light_d = ray_bundle.origins[0]

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -25,12 +25,9 @@ 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"""
-    target_transmittance: float = 0.25
-    """target transmittance for opacity penalty"""
     side_prompt: str = ", side view"
     """appended to prompt for side view"""
     front_prompt: str = ", front view"
@@ -38,6 +35,14 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig):
     back_prompt: str = ", back view"
     """appended to prompt for back view"""
 
+    opacity_penalty: bool = True
+    """enables penalty to encourage transparent scenes"""
+    target_transmittance: float = 0.5
+    """target transmittance for opacity penalty"""
+
+    guidance_scale: float = 100
+    """guidance scale for sds loss"""
+
 
 class DreamfusionPipeline(VanillaPipeline):
     def __init__(
@@ -51,11 +56,12 @@ class DreamfusionPipeline(VanillaPipeline):
         super().__init__(config, device, test_mode, world_size, local_rank)
         self.generative = True
         self.sd = StableDiffusion(device)
-        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}{config.front_prompt}", "")
             self.side_text_embedding = self.sd.get_text_embeds(f"{config.prompt}{config.side_prompt}", "")
             self.back_text_embedding = self.sd.get_text_embeds(f"{config.prompt}{config.back_prompt}", "")
+        else:
+            self.base_text_embedding = self.sd.get_text_embeds(config.prompt, "")
 
     @profiler.time_function
     def custom_step(self, step: int, grad_scaler: GradScaler, optimizers: Optimizers):
@@ -78,13 +84,13 @@ class DreamfusionPipeline(VanillaPipeline):
         accumulation = np.clip(accumulation, 0.0, 1.0)
         plt.imsave("nerf_accumulation.jpg", accumulation)
 
-        background = model_outputs["background"].view(64, 64, 3).detach().cpu().numpy()
-        background = np.clip(background, 0.0, 1.0)
-        plt.imsave("nerf_background.jpg", background)
+        # background = model_outputs["background"].view(64, 64, 3).detach().cpu().numpy()
+        # background = np.clip(background, 0.0, 1.0)
+        # plt.imsave("nerf_background.jpg", background)
 
-        shaded = model_outputs["shaded"].view(64, 64, 3).detach().cpu().numpy()
-        shaded = np.clip(shaded, 0.0, 1.0)
-        plt.imsave("nerf_textureless.jpg", shaded)
+        # shaded = model_outputs["shaded"].view(64, 64, 3).detach().cpu().numpy()
+        # shaded = np.clip(shaded, 0.0, 1.0)
+        # plt.imsave("nerf_textureless.jpg", shaded)
 
         if self.config.location_based_prompting:
             if batch["central"] > 315 or batch["central"] <= 45:
@@ -98,7 +104,9 @@ class DreamfusionPipeline(VanillaPipeline):
         else:
             text_embedding = self.base_text_embedding
 
-        sds_loss, latents, grad = self.sd.sds_loss(text_embedding, train_output)
+        sds_loss, latents, grad = self.sd.sds_loss(
+            text_embedding, train_output, guidance_scale=self.config.guidance_scale
+        )
         # TODO: opacity penalty using transmittance, not accumultation
         if self.config.opacity_penalty:
             accum_mean = np.mean(1.0 - accumulation)