diff --git a/nerfstudio/engine/dreamfusion_trainer.py b/nerfstudio/engine/dreamfusion_trainer.py index 1f062318d..de456152d 100644 --- a/nerfstudio/engine/dreamfusion_trainer.py +++ b/nerfstudio/engine/dreamfusion_trainer.py @@ -60,6 +60,10 @@ class DreamfusionTrainer(Trainer): grad = model_outputs["grad"] self.grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True) normals_loss = loss_dict["orientation_loss"] + loss_dict["pred_normal_loss"] + if self.pipeline.config.alphas_penalty: + normals_loss += loss_dict["alphas_loss"] + if self.pipeline.config.opacity_penalty: + normals_loss += loss_dict["opacity_loss"] self.grad_scaler.scale(normals_loss).backward() # type: ignore self.optimizers.optimizer_scaler_step_all(self.grad_scaler) diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 6ef860aac..38211bacd 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -259,6 +259,9 @@ class DreamFusionModel(Model): field_outputs[FieldHeadNames.PRED_NORMALS], ) + assert weights.shape[-1] == 1 + outputs["alphas_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01).mean() + return outputs def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: @@ -276,6 +279,9 @@ class DreamFusionModel(Model): loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( outputs["rendered_pred_normal_loss"] ) + + loss_dict["alphas_loss"] = outputs["alphas_loss"] + return loss_dict def get_image_metrics_and_images( diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 667ac7072..b8477ab3c 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -48,10 +48,11 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig): """appended to prompt for front view""" back_prompt: str = ", back view" """appended to prompt for back view""" - + alphas_penalty: bool = True + """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" opacity_penalty: bool = True - """enables penalty to encourage transparent scenes""" - target_transmittance: float = 0.5 + """enables penalty to encourage transparent scenes, as in "dreamfields" paper""" + target_transmittance: float = 0.8 """target transmittance for opacity penalty""" guidance_scale: float = 100 @@ -144,6 +145,9 @@ class DreamfusionPipeline(VanillaPipeline): if self.config.opacity_penalty: accum_mean = np.mean(1.0 - accumulation) sds_loss *= np.min((0.5, accum_mean)) + loss_dict["opacity_loss"] = -torch.minimum( + (1 - model_outputs["accumulation"]).mean(), torch.tensor(self.config.target_transmittance) + ) model_outputs["latents"] = latents model_outputs["grad"] = grad