Normals prediction training fixed

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!

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:

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(

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