diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index b93a8eb5f..ca5d387f9 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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, }, }, diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 5338846d9..22e0652ec 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/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] diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index f64aa0854..3bfe1372d 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/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: diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 35cc49045..3fe6cecc4 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/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)