From dd2a50574a8fc573e466a83f7dce33376e3366cc Mon Sep 17 00:00:00 2001 From: mcallisterdavid Date: Tue, 13 Dec 2022 04:55:56 +0000 Subject: [PATCH] Normals prediction training fixed --- nerfstudio/generative/stable_diffusion.py | 4 +- .../model_components/scene_colliders.py | 3 + nerfstudio/models/dreamfusion.py | 60 ++++++++++--------- nerfstudio/pipelines/dreamfusion_pipeline.py | 15 +++-- 4 files changed, 50 insertions(+), 32 deletions(-) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 23471908f..d0f12759d 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -13,7 +13,7 @@ from transformers import CLIPTextModel, CLIPTokenizer, logging CONSOLE = Console(width=120) logging.set_verbosity_error() IMG_DIM = 512 -const_scale = 0.15#8215 +const_scale = 0.18215 class StableDiffusion(nn.Module): """Stable Diffusion implementation @@ -82,6 +82,8 @@ class StableDiffusion(nn.Module): plt.imsave('nerf_output.png', nerf_output_np) nerf_output = F.interpolate(nerf_output, (IMG_DIM, IMG_DIM), mode='bilinear') + print(torch.max(nerf_output), torch.min(nerf_output)) + t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device) latents = self.imgs_to_latent(nerf_output) # predict the noise residual with unet, NO grad! diff --git a/nerfstudio/model_components/scene_colliders.py b/nerfstudio/model_components/scene_colliders.py index c50b41042..fe27cdb8e 100644 --- a/nerfstudio/model_components/scene_colliders.py +++ b/nerfstudio/model_components/scene_colliders.py @@ -147,6 +147,9 @@ class SphereCollider(SceneCollider): nears = torch.clamp(nears, min=near_plane) fars = torch.maximum(fars, nears + 1e-6) + nears = torch.nan_to_num(nears, nan=0.0) + fars = torch.nan_to_num(fars, nan=0.0) + return nears, fars def set_nears_and_fars(self, ray_bundle: RayBundle) -> RayBundle: diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index c00b13c06..c29f03e74 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -68,7 +68,7 @@ class DreamFusionModelConfig(ModelConfig): """Orientation loss multipier on computed normals.""" pred_normal_loss_mult: float = 0.0003 """Predicted normal loss multiplier.""" - random_light_source: bool = False + random_light_source: bool = True """Randomizes light source per output.""" initialize_density: bool = True """Initialize density in center of scene.""" @@ -153,7 +153,7 @@ class DreamFusionModel(Model): # colliders if self.config.enable_collider: self.collider = AABBBoxCollider(scene_box=self.scene_box) - # self.collider = SphereCollider(torch.Tensor([0, 0, 0]), 1.0) + self.collider = SphereCollider(torch.Tensor([0, 0, 0]), 0.8) def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} @@ -178,19 +178,25 @@ class DreamFusionModel(Model): accumulation = self.renderer_accumulation(weights) depth = self.renderer_depth(weights, ray_samples_uniform) rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) + rgb = torch.clamp(rgb, 0.0, 1.0) - background = torch.clamp((1 - torch.nan_to_num(accumulation, nan=0.0)), min=0.0, max=1.0) * background_rgb + accum_mask = torch.clamp((torch.nan_to_num(accumulation, nan=0.0)), min=0.0, max=1.0) + accum_mask_inv = 1.0 - accum_mask + + background = accum_mask_inv * background_rgb outputs = { "rgb": rgb, "background_rgb": background_rgb, "background": background, - "accumulation": accumulation, + "accumulation": accum_mask, "depth": depth, } 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)#.detach() + + # lambertian shading if self.config.random_light_source: @@ -200,28 +206,28 @@ class DreamFusionModel(Model): light_d = math.safe_normalize(light_d) ratio = 0.3 - lambertian = ratio + (1 - ratio) * (normals @ light_d).clamp(min=0) - + # with torch.no_grad(): + lambertian = ratio + (1 - ratio) * (pred_normals.detach() @ light_d).clamp(min=0) shaded = lambertian.unsqueeze(-1).repeat(1, 3) shaded_albedo = rgb * lambertian.unsqueeze(-1) outputs["normals"] = normals outputs["pred_normals"] = pred_normals - outputs["shaded"] = torch.nan_to_num(accumulation, nan=0.0) * shaded + torch.ones_like(shaded) * torch.clamp((1 - torch.nan_to_num(accumulation, nan=0.0)), min=0.0, max=1.0) + outputs["shaded"] = accum_mask * shaded + torch.ones_like(shaded) * accum_mask_inv outputs["shaded_albedo"] = shaded_albedo - outputs["render"] = torch.nan_to_num(accumulation, nan=0.0) * shaded_albedo + background + outputs["render"] = accum_mask * shaded_albedo + background - if self.train_normals: + # if self.train_normals: - 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 + ) - 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], + ) return outputs @@ -230,16 +236,16 @@ class DreamFusionModel(Model): loss_dict = {} loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) - if self.train_normals: - # orientation loss for computed normals - loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean( - outputs["rendered_orientation_loss"] - ) + # if self.train_normals: + # orientation loss for computed normals + loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean( + 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"] - ) + # ground truth supervision for normals + loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( + outputs["rendered_pred_normal_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 b08660d79..8757ec336 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -42,7 +42,7 @@ class DreamfusionPipeline(VanillaPipeline): local_rank ) self.sd = StableDiffusion(device) - self.text_embedding = self.sd.get_text_embeds("A high quality photo of a pineapple.", "") + self.text_embedding = self.sd.get_text_embeds("A high quality photo of a squirrel.", "") @profiler.time_function def custom_step(self, step: int, grad_scaler: GradScaler, optimizers: Optimizers): @@ -77,12 +77,19 @@ class DreamfusionPipeline(VanillaPipeline): sds_loss, latents, grad = self.sd.sds_loss(self.text_embedding, albedo_output) grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True) - optimizers.optimizer_scaler_step_all(grad_scaler) - grad_scaler.update() - optimizers.scheduler_step_all(step) + # optimizers.scheduler_step_all(step) + # optimizers.optimizer_scaler_step_all(grad_scaler) + # grad_scaler.update() metrics_dict = self.model.get_metrics_dict(model_outputs, batch) loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict) loss_dict["sds_loss"] = sds_loss + normals_loss = loss_dict["orientation_loss"] + loss_dict["pred_normal_loss"] + grad_scaler.scale(normals_loss).backward() # type: ignore + optimizers.optimizer_scaler_step_all(grad_scaler) + + grad_scaler.update() + optimizers.scheduler_step_all(step) + return model_outputs, loss_dict, metrics_dict \ No newline at end of file