diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 7bf71354f..d0c76262d 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -341,7 +341,7 @@ method_configs["dreamfusion"] = TrainerConfig( method_name="dreamfusion", steps_per_eval_batch=50, steps_per_eval_image=50, - steps_per_save=2000, + steps_per_save=200, max_num_iterations=30000, mixed_precision=True, pipeline=DreamfusionPipelineConfig( @@ -352,7 +352,9 @@ method_configs["dreamfusion"] = TrainerConfig( ), model=DreamFusionModelConfig( eval_num_rays_per_chunk=1 << 15, - distortion_loss_mult=10, + distortion_loss_mult=10.0, + orientation_loss_mult=0.1, + max_res=256, sphere_collider=True, initialize_density=False, random_background=True, @@ -362,9 +364,9 @@ method_configs["dreamfusion"] = TrainerConfig( proposal_weights_anneal_max_num_iters=100, start_lambertian_training=1000, start_normals_training=500, + opacity_loss_mult=0.001, ), - alphas_loss_mult=1.0, - opacity_loss_mult=1.0, + interpolated_prompting=False, guidance_scale=100, ), optimizers={ diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py index f3b61e58f..4960fbe48 100644 --- a/nerfstudio/fields/dreamfusion_field.py +++ b/nerfstudio/fields/dreamfusion_field.py @@ -65,6 +65,9 @@ class DreamFusionField(Field): num_layers: int = 2, hidden_dim: int = 64, geo_feat_dim: int = 15, + num_levels: int = 16, + max_res: int = 256, + log2_hashmap_size: int = 19, num_layers_color: int = 3, hidden_dim_color: int = 64, hidden_dim_normal: int = 32, @@ -74,10 +77,7 @@ class DreamFusionField(Field): self.aabb = Parameter(aabb, requires_grad=False) self.geo_feat_dim = geo_feat_dim - num_levels = 16 - max_res = 1024 base_res = 16 - log2_hashmap_size = 19 features_per_level = 2 growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index 5df52d2fa..08ee1fd95 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -509,7 +509,16 @@ class VolumetricSampler(Sampler): class ProposalNetworkSampler(Sampler): - """Sampler that uses a proposal network to generate samples.""" + """Sampler that uses a proposal network to generate samples. + + Args: + num_proposal_samples_per_ray: Number of samples to generate per ray for each proposal step. + num_nerf_samples_per_ray: Number of samples to generate per ray for the NERF model. + num_proposal_network_iterations: Number of proposal network iterations to run. + single_jitter: Use a same random jitter for all samples along a ray. + update_sched: A function that takes the iteration number of steps between updates. + initial_sampler: Sampler to use for the first iteration. Uses UniformLinDispPiecewise if not set. + """ def __init__( self, @@ -518,6 +527,7 @@ class ProposalNetworkSampler(Sampler): num_proposal_network_iterations: int = 2, single_jitter: bool = False, update_sched: Callable = lambda x: 1, + initial_sampler: Optional[Sampler] = None, ) -> None: super().__init__() self.num_proposal_samples_per_ray = num_proposal_samples_per_ray @@ -528,7 +538,10 @@ class ProposalNetworkSampler(Sampler): raise ValueError("num_proposal_network_iterations must be >= 1") # samplers - self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter) + if initial_sampler is None: + self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter) + else: + self.initial_sampler = initial_sampler self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter) self._anneal = 1.0 diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 4cc56fb19..c2a716946 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -41,7 +41,10 @@ from nerfstudio.model_components.losses import ( orientation_loss, pred_normal_loss, ) -from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler +from nerfstudio.model_components.ray_samplers import ( + ProposalNetworkSampler, + UniformSampler, +) from nerfstudio.model_components.renderers import ( AccumulationRenderer, DepthRenderer, @@ -118,6 +121,12 @@ class DreamFusionModelConfig(ModelConfig): """Start training normals after this many iterations""" start_lambertian_training: int = 1000 """start training with lambertian shading after this many iterations""" + opacity_penalty: bool = True + """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" + opacity_loss_mult: float = 1 + """scale for opacity penalty""" + max_res: int = 256 + """Maximum resolution of the density field.""" class DreamFusionModel(Model): @@ -148,7 +157,7 @@ class DreamFusionModel(Model): super().populate_modules() # setting up fields - self.field = DreamFusionField(self.scene_box.aabb) + self.field = DreamFusionField(self.scene_box.aabb, max_res=self.config.max_res) # samplers self.density_fns = [] @@ -176,10 +185,9 @@ class DreamFusionModel(Model): num_proposal_network_iterations=self.config.num_proposal_iterations, single_jitter=self.config.use_single_jitter, update_sched=update_schedule, + initial_sampler=UniformSampler(single_jitter=self.config.use_single_jitter), ) - # self.sampler_uniform = UniformSampler(num_samples=self.num_samples, single_jitter=True) - # renderers self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) self.renderer_accumulation = AccumulationRenderer() @@ -357,11 +365,7 @@ class DreamFusionModel(Model): if samp > 0.5 and not self.training: outputs["train_output"] = outputs["shaded"] elif samp < 0.2 and self.random_background: - rand_bg = torch.zeros(background.shape) - rand_bg.index_fill_(1, torch.tensor([0]), np.random.random_sample()) - rand_bg.index_fill_(1, torch.tensor([1]), np.random.random_sample()) - rand_bg.index_fill_(1, torch.tensor([2]), np.random.random_sample()) - rand_bg = rand_bg.to(self.device) + rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device) outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv else: outputs["train_output"] = accum_mask * shaded_albedo + background @@ -370,7 +374,7 @@ class DreamFusionModel(Model): if self.training or True: outputs["rendered_orientation_loss"] = orientation_loss( - weights, field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions + weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions ) outputs["rendered_pred_normal_loss"] = pred_normal_loss( @@ -380,7 +384,10 @@ class DreamFusionModel(Model): ) assert weights.shape[-1] == 1 - outputs["alphas_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) + if self.config.opacity_penalty: + outputs["opacity_loss"] = ( + torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult + ) return outputs @@ -402,11 +409,9 @@ class DreamFusionModel(Model): loss_dict["orientation_loss"] = 0 loss_dict["pred_normal_loss"] = 0 - loss_dict["alphas_loss"] = outputs["alphas_loss"].mean() + if self.config.opacity_penalty: + loss_dict["opacity_loss"] = self.config.opacity_loss_mult * outputs["opacity_loss"].mean() - loss_dict["opacity_loss"] = -torch.minimum( - (1 - outputs["accumulation"]).mean(), torch.tensor(self.target_transmittance) - ) if self.training: loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss( outputs["weights_list"], outputs["ray_samples_list"] diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index b228c7027..692a92a71 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -51,14 +51,6 @@ 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)""" - alphas_loss_mult: float = 1 - """scale for alphas penalty""" - opacity_penalty: bool = True - """enables penalty to encourage transparent scenes, as in "dreamfields" paper""" - opacity_loss_mult: float = 1 - """scale for opacity penalty""" guidance_scale: float = 100 """guidance scale for sds loss"""