From 5b1c4a33dc057d182e1d0cc8e05f726369ce5ab1 Mon Sep 17 00:00:00 2001 From: terrance Date: Sun, 30 Apr 2023 20:54:28 +0000 Subject: [PATCH] everything but draemfusion works kinda --- nerfstudio/configs/method_configs.py | 25 +++++--- nerfstudio/fields/dreamtensor_field.py | 8 +-- nerfstudio/models/dreamembedding.py | 2 +- nerfstudio/models/dreamfusion.py | 59 +++++++++-------- nerfstudio/models/dreamfusion_ngp.py | 84 ++++++++++++------------- nerfstudio/models/dreamfusion_tensor.py | 51 +++++++++++++-- nerfstudio/pipelines/base_pipeline.py | 10 --- 7 files changed, 145 insertions(+), 94 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 43a41a4f6..b4e697585 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -447,7 +447,7 @@ method_configs["dreamfusion"] = TrainerConfig( steps_per_eval_batch=50, steps_per_eval_image=50, steps_per_save=200, - max_num_iterations=30000, + max_num_iterations=20000, mixed_precision=True, pipeline=VanillaPipelineConfig( generative=True, @@ -501,7 +501,7 @@ method_configs["dreamngp"] = TrainerConfig( pipeline=VanillaPipelineConfig( generative=True, datamanager=DreamFusionDataManagerConfig( - horizontal_rotation_warmup=0, + horizontal_rotation_warmup=2000, ), model=DreamfusionNGPModelConfig( eval_num_rays_per_chunk=8192, @@ -547,7 +547,7 @@ method_configs["dreamtensor"] = TrainerConfig( pipeline=VanillaPipelineConfig( generative=True, datamanager=DreamFusionDataManagerConfig( - horizontal_rotation_warmup=0, + horizontal_rotation_warmup=2000, ), model=DreamfusionTensorModelConfig( initialize_density=True, @@ -557,14 +557,23 @@ method_configs["dreamtensor"] = TrainerConfig( ), ), optimizers={ - "encodings": { - "optimizer": AdamOptimizerConfig(lr=0.0002), - "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=30000), - }, + # "fields": { + # "optimizer": AdamOptimizerConfig(lr=5e-5), + # "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-5, max_steps=20000), + # }, + # "encodings": { + # "optimizer": AdamOptimizerConfig(lr=0.0002), + # "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=20000), + # }, "fields": { "optimizer": AdamOptimizerConfig(lr=0.001), "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000), }, + "encodings": { + "optimizer": AdamOptimizerConfig(lr=0.02), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=30000), + }, + }, viewer=ViewerConfig(num_rays_per_chunk=32000), vis="viewer", @@ -672,7 +681,7 @@ method_configs["dreamembed"] = TrainerConfig( pipeline=VanillaPipelineConfig( generative=True, datamanager=DreamFusionDataManagerConfig( - horizontal_rotation_warmup=0, + horizontal_rotation_warmup=2000, ), model=DreamEmbeddingModelConfig( eval_num_rays_per_chunk=1 << 15, diff --git a/nerfstudio/fields/dreamtensor_field.py b/nerfstudio/fields/dreamtensor_field.py index 271311855..2019a3ae1 100644 --- a/nerfstudio/fields/dreamtensor_field.py +++ b/nerfstudio/fields/dreamtensor_field.py @@ -113,8 +113,8 @@ class DreamFusionTensorField(TensoRFField): ) def get_density(self, ray_samples: RaySamples) -> TensorType: - # positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) - # positions = positions * 2 - 1 + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + positions = positions * 2 - 1 positions = ray_samples.frustums.get_positions() density = self.density_encoding(positions) density_enc = torch.sum(density, dim=-1)[:, :, None] @@ -124,8 +124,8 @@ class DreamFusionTensorField(TensoRFField): def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None) -> TensorType: d = ray_samples.frustums.directions - # positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) - # positions = positions * 2 - 1 + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + positions = positions * 2 - 1 positions = ray_samples.frustums.get_positions() rgb_features = self.color_encoding(positions) diff --git a/nerfstudio/models/dreamembedding.py b/nerfstudio/models/dreamembedding.py index e8b41599d..e1955acd6 100644 --- a/nerfstudio/models/dreamembedding.py +++ b/nerfstudio/models/dreamembedding.py @@ -180,7 +180,7 @@ class DreamEmbeddingModel(DreamFusionModel): pos = ray_samples.frustums.get_positions() # density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] # density = density + density_blob - density_blob = (-torch.exp(torch.norm(pos, dim=-1) / 0.6) + 2)[..., None] + density_blob = self.density_strength * (-torch.exp(torch.norm(pos, dim=-1) / 0.6) + 2)[..., None] density = torch.max(density + density_blob, torch.tensor([0.], device=self.device)) weights = ray_samples.get_weights(density) diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 62c9d9d47..2abf72974 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -66,7 +66,7 @@ class DreamFusionModelConfig(ModelConfig): _target: Type = field(default_factory=lambda: DreamFusionModel) """target class to instantiate""" - prompt: str = "A high-quality photo of a realistic pineapple with green leaves" + prompt: str = "A high-quality photo of a pineapple" # prompt: str = "a high-quality photo of a banana" """prompt for stable dreamfusion""" @@ -78,7 +78,7 @@ class DreamFusionModelConfig(ModelConfig): """Randomizes light source per output.""" initialize_density: bool = True """Initialize density in center of scene.""" - taper_range: Tuple[int, int] = (0, 500) + taper_range: Tuple[int, int] = (0, 1000) """Range of step values for the density tapering""" taper_strength: Tuple[float, float] = (1.0, 0.0) """Strength schedule of center density""" @@ -358,8 +358,8 @@ class DreamFusionModel(Model): pos = ray_samples.frustums.get_positions() # density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] # density = density + density_blob - density_blob = (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 0.2)[..., None] - density = torch.max(density + density_blob, torch.tensor([0.], device=self.device)) + density_blob = self.density_strength * (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 1.0)[..., None] + density = torch.max(density + density_blob, torch.tensor([1e-15], device=self.device)) weights = ray_samples.get_weights(density) weights_list.append(weights) @@ -391,7 +391,7 @@ class DreamFusionModel(Model): outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i]) normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) - pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) + # pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) # lambertian shading if self.config.random_light_source: # and self.training: @@ -410,33 +410,42 @@ class DreamFusionModel(Model): ) outputs["normals"] = self.shader_normals(normals, weights=accum_mask) - outputs["pred_normals"] = self.shader_normals(pred_normals, weights=accum_mask) + # outputs["pred_normals"] = self.shader_normals(pred_normals, weights=accum_mask) outputs["shaded"] = accum_mask * shaded outputs["other_train_output"] = accum_mask * shaded_albedo + background outputs["shaded_albedo"] = accum_mask * shaded_albedo outputs["rgb"] = accum_mask * rgb + background - if shading_weight > 0: - samp = np.random.random_sample() - if samp > 0.5 and not self.training: - outputs["train_output"] = outputs["shaded"] - elif samp < 0.2 and self.random_background: - 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 + # if shading_weight > 0: + # samp = np.random.random_sample() + # if samp > 0.5 and self.training: + # outputs["train_output"] = outputs["shaded"] + # elif samp < 0.2 and self.random_background and self.training: + # 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 + # else: + # outputs["train_output"] = outputs["rgb"] + + samp = np.random.random_sample() + if samp < 0.4: + rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device) + train_output = accum_mask * rgb + rand_bg * accum_mask_inv else: - outputs["train_output"] = outputs["rgb"] + train_output = accum_mask * rgb + background + + outputs["train_output"] = train_output outputs["rendered_orientation_loss"] = orientation_loss( weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions ) - outputs["rendered_pred_normal_loss"] = pred_normal_loss( - weights.detach(), - field_outputs[FieldHeadNames.NORMALS].detach(), - field_outputs[FieldHeadNames.PRED_NORMALS], - ) + # outputs["rendered_pred_normal_loss"] = pred_normal_loss( + # weights.detach(), + # field_outputs[FieldHeadNames.NORMALS].detach(), + # field_outputs[FieldHeadNames.PRED_NORMALS], + # ) assert weights.shape[-1] == 1 if self.config.opacity_penalty: @@ -455,9 +464,9 @@ class DreamFusionModel(Model): outputs["rendered_orientation_loss"] ) # ground truth supervision for normals - loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( - outputs["rendered_pred_normal_loss"] - ) + # loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( + # outputs["rendered_pred_normal_loss"] + # ) else: loss_dict["orientation_loss"] = 0 loss_dict["pred_normal_loss"] = 0 @@ -528,6 +537,6 @@ class DreamFusionModel(Model): "prop_depth_0": prop_depth_0, "prop_depth_1": prop_depth_1, "normals": outputs["normals"], - "pred_normals": outputs["pred_normals"], + # "pred_normals": outputs["pred_normals"], } return metrics_dict, images_dict diff --git a/nerfstudio/models/dreamfusion_ngp.py b/nerfstudio/models/dreamfusion_ngp.py index 3811e8ea0..bcbdb79e9 100644 --- a/nerfstudio/models/dreamfusion_ngp.py +++ b/nerfstudio/models/dreamfusion_ngp.py @@ -224,13 +224,7 @@ class DreamfusionNGPModel(DreamFusionModel): where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], update_every_num_iters=1, func=update_occupancy_grid, - ), - TrainingCallback( - where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - iters=(self.config.start_normals_training,), - func=start_training_normals, - args=[self, training_callback_attributes], - ), + ), TrainingCallback( where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], func=taper_density, @@ -263,7 +257,7 @@ class DreamfusionNGPModel(DreamFusionModel): if self.initialize_density: pos = ray_samples.frustums.get_positions() - density_blob = (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 0.5)[..., None] + density_blob = self.density_strength * (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 0.5)[..., None] density = torch.max(density + density_blob, torch.tensor([0.], device=self.device)) packed_info = nerfacc.pack_info(ray_indices, num_rays) @@ -303,48 +297,54 @@ class DreamfusionNGPModel(DreamFusionModel): "rgb": accum_mask * rgb + background } - normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) - outputs["normals"] = self.shader_normals(normals, weights=accum_mask) - - # lambertian shading - if self.config.random_light_source: # and self.training: - light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals) + samp = np.random.random_sample() + if samp < 0.4: + rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device) + train_output = accum_mask * rgb + rand_bg * accum_mask_inv else: - light_d = ray_bundle.origins[0] - light_d = math.safe_normalize(light_d) + train_output = accum_mask * rgb + background + + outputs["train_output"] = train_output - if (self.train_shaded and np.random.random_sample() > 0.75) or not self.training: - shading_weight = 0.9 - else: - shading_weight = 0.0 + # normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) + # outputs["normals"] = self.shader_normals(normals, weights=accum_mask) - shaded, shaded_albedo = self.shader_lambertian( - rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False - ) + # # lambertian shading + # if self.config.random_light_source: # and self.training: + # light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals) + # else: + # light_d = ray_bundle.origins[0] + # light_d = math.safe_normalize(light_d) - outputs["shaded"] = accum_mask * shaded + # if (self.train_shaded and np.random.random_sample() > 0.75) or not self.training: + # shading_weight = 0.9 + # else: + # shading_weight = 0.0 - if shading_weight > 0: - samp = np.random.random_sample() - if samp > 0.5 and not self.training: - outputs["train_output"] = outputs["shaded"] - elif samp < 0.2 and self.random_background: - 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 - else: - outputs["train_output"] = outputs["rgb"] + # shaded, shaded_albedo = self.shader_lambertian( + # rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False + # ) - # print(outputs["rgb"].shape) - # outputs["train_output"] = outputs["rgb"] + # outputs["shaded"] = accum_mask * shaded + + # if shading_weight > 0: + # samp = np.random.random_sample() + # if samp > 0.5 and not self.training: + # outputs["train_output"] = outputs["shaded"] + # elif samp < 0.2 and self.random_background: + # 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 + # else: + # outputs["train_output"] = outputs["rgb"] - outputs["rendered_orientation_loss"] = orientation_loss( - weights.detach(), - field_outputs[FieldHeadNames.NORMALS], - ray_bundle.directions, - ) + # outputs["rendered_orientation_loss"] = orientation_loss( + # weights.detach(), + # field_outputs[FieldHeadNames.NORMALS], + # ray_bundle.directions, + # ) assert weights.shape[-1] == 1 if self.config.opacity_penalty: diff --git a/nerfstudio/models/dreamfusion_tensor.py b/nerfstudio/models/dreamfusion_tensor.py index 2ce5410ae..6dafa71eb 100644 --- a/nerfstudio/models/dreamfusion_tensor.py +++ b/nerfstudio/models/dreamfusion_tensor.py @@ -104,12 +104,12 @@ class DreamfusionTensorModelConfig(DreamFusionModelConfig): transmittance_end_schedule: int = 1500 """number of iterations to reach target_transmittance_end""" - init_resolution: int = 300 + init_resolution: int = 128 """initial render resolution""" final_resolution: int = 300 """final render resolution""" # upsampling_iters: Tuple[int, ...] = (2000, 3000, 4000, 5500, 7000) - upsampling_iters: Tuple[int, ...] = (500, 750, 1000) + upsampling_iters: Tuple[int, ...] = (1000, 3000, 7000) """specifies a list of iteration step numbers to perform upsampling""" num_samples: int = 100 """Number of samples in field evaluation""" @@ -247,7 +247,32 @@ class DreamfusionTensorModel(DreamFusionModel): self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument ): self.density_strength = np.interp(step, self.config.taper_range, self.config.taper_strength) + def reinitialize_optimizer( + self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument + ): + index = self.upsampling_iters.index(step) + resolution = self.upsampling_steps[index] + # upsample the position and direction grids + self.field.density_encoding.upsample_grid(resolution) + self.field.color_encoding.upsample_grid(resolution) + + # reinitialize the encodings optimizer + optimizers_config = training_callback_attributes.optimizers.config + enc = training_callback_attributes.pipeline.get_param_groups()["encodings"] + lr_init = optimizers_config["encodings"]["optimizer"].lr + + training_callback_attributes.optimizers.optimizers["encodings"] = optimizers_config["encodings"][ + "optimizer" + ].setup(params=enc) + if optimizers_config["encodings"]["scheduler"]: + training_callback_attributes.optimizers.schedulers["encodings"] = ( + optimizers_config["encodings"]["scheduler"] + .setup() + .get_scheduler( + optimizer=training_callback_attributes.optimizers.optimizers["encodings"], lr_init=lr_init + ) + ) return [ TrainingCallback( where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], @@ -255,6 +280,12 @@ class DreamfusionTensorModel(DreamFusionModel): update_every_num_iters=1, args=[self, training_callback_attributes], ), + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + iters=self.upsampling_iters, + func=reinitialize_optimizer, + args=[self, training_callback_attributes], + ) ] def get_param_groups(self) -> Dict[str, List[Parameter]]: @@ -292,7 +323,7 @@ class DreamfusionTensorModel(DreamFusionModel): if self.initialize_density: pos = ray_samples_pdf.frustums.get_positions() - density_blob = (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 1.0)[..., None] + density_blob = self.density_strength * (-0.05 * torch.exp(5 * torch.norm(pos, dim=-1)) + 1.0)[..., None] density = torch.max(density + density_blob, torch.tensor([0.], device=self.device)) weights_fine = ray_samples_pdf.get_weights(density) @@ -312,14 +343,24 @@ class DreamfusionTensorModel(DreamFusionModel): accum_mask_inv = 1.0 - accum_mask background = accum_mask_inv * background_rgb + samp = np.random.random_sample() + if samp < 0.4: + rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device) + train_output = accum_mask * rgb + rand_bg * accum_mask_inv + else: + train_output = accum_mask * rgb + background + outputs = { "rgb_only": rgb, "background_rgb": background_rgb, "background": background, "accumulation": accum_mask, "depth": depth, - "train_output": accum_mask * rgb + background + "train_output": train_output } + + outputs["opacity_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * 0.5 + return outputs @@ -359,6 +400,8 @@ class DreamfusionTensorModel(DreamFusionModel): grad_scaler=self.grad_scaler, ) + loss_dict["opacity_loss"] = outputs["opacity_loss"].mean() + loss_dict["sds_loss"] = sds_loss.to(self.device) return loss_dict \ No newline at end of file diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index b3c7e6f14..00b9fb7f8 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -285,16 +285,6 @@ class VanillaPipeline(Pipeline): """ ray_bundle, batch = self.datamanager.next_train(step) model_outputs = self.model(ray_bundle) - - # if step % 100 == 0: - # image = model_outputs["train_output"] - # image = image.view(int(image.shape[0] ** 0.5), int(image.shape[0] ** 0.5), 3) - # image = image.detach().cpu().numpy() - # image = (image * 255).round().astype("uint8") - # mediapy.write_image(str('dreamngp.png'), image) - # # mediapy.write_image(str('dreamfusion.png'), image) - - metrics_dict = self.model.get_metrics_dict(model_outputs, batch) if self.config.datamanager.camera_optimizer is not None: