From 85108c390effee29aff679dd6a87a616ab94ba8a Mon Sep 17 00:00:00 2001 From: terrance Date: Wed, 8 Feb 2023 20:21:07 +0000 Subject: [PATCH] interpolated location prompting option, lowering init density, transmittance loss interpolation flipped --- nerfstudio/configs/method_configs.py | 2 +- nerfstudio/models/dreamfusion.py | 30 ++++++++--------- nerfstudio/pipelines/dreamfusion_pipeline.py | 35 ++++++++++++-------- 3 files changed, 37 insertions(+), 30 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 54b555ad5..46cde2cfc 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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, ), diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index d386e63ba..7d85e0e11 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/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) ) diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 5c2aead54..afcd3dd0b 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/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: