SDS return loss

nerfstudio/generative/stable_diffusion.py

@@ -19,7 +19,7 @@
 import sys
 from dataclasses import dataclass
 from pathlib import Path
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Union
 
 import appdirs
 import mediapy
@@ -29,6 +29,8 @@ import torch.nn.functional as F
 import tyro
 from rich.console import Console
 from torch import nn
+from torch.cuda.amp import custom_bwd, custom_fwd
+from torch.cuda.amp.grad_scaler import GradScaler
 from torchtyping import TensorType
 
 CONSOLE = Console(width=120)
@@ -58,6 +60,26 @@ class UNet2DConditionOutput:
     sample: torch.FloatTensor
 
 
+class _SDSGradient(torch.autograd.Function):  # pylint: disable=abstract-method
+    """Custom gradient function for SDS loss. Since it is already computed, we can just return it."""
+
+    @staticmethod
+    @custom_fwd
+    def forward(ctx, input_tensor, gt_grad): # pylint: disable=arguments-differ
+        del input_tensor
+        ctx.save_for_backward(gt_grad)
+        # Return magniture of gradient, not the actual loss.
+        return torch.mean(gt_grad**2) ** 0.5
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad): # pylint: disable=arguments-differ
+        del grad
+        (gt_grad,) = ctx.saved_tensors
+        batch_size = len(gt_grad)
+        return gt_grad / batch_size, None
+
+
 class StableDiffusion(nn.Module):
     """Stable Diffusion implementation
     Args:
@@ -162,14 +184,16 @@ class StableDiffusion(nn.Module):
         text_embeddings: TensorType["N", "max_length", "embed_dim"],
         image: TensorType["BS", 3, "H", "W"],
         guidance_scale: float = 100.0,
-    ) -> Tuple[torch.Tensor, TensorType["BS", 4, "H", "W"], TensorType["BS", 4, "H", "W"]]:
+        grad_scaler: Optional[GradScaler] = None,
+    ) -> torch.Tensor:
         """Score Distilation Sampling loss proposed in DreamFusion paper (https://dreamfusion3d.github.io/)
         Args:
             text_embeddings: Text embeddings
             image: Rendered image
             guidance_scale: How much to weigh the guidance
+            grad_scaler: Grad scaler
         Returns:
-            A tuple of Loss, latent, and gradient
+            The loss
         """
         image = F.interpolate(image, (IMG_DIM, IMG_DIM), mode="bilinear")
         t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)
@@ -194,9 +218,11 @@ class StableDiffusion(nn.Module):
         grad = w * (noise_pred - noise)
         grad = torch.nan_to_num(grad)
 
-        sds_mse = torch.mean(torch.nan_to_num(torch.square(noise_pred - noise)))
+        if grad_scaler is not None:
+            latents = grad_scaler.scale(latents)
+        loss = _SDSGradient.apply(latents, grad)
 
-        return sds_mse, latents, grad
+        return loss
 
     def produce_latents(
         self,

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -41,10 +41,11 @@ class SpecifyGradient(torch.autograd.Function):
     @staticmethod
     @custom_bwd
     def backward(ctx, grad):
-        gt_grad, = ctx.saved_tensors
+        (gt_grad,) = ctx.saved_tensors
         batch_size = len(gt_grad)
         return gt_grad / batch_size, None
 
+
 @dataclass
 class DreamfusionPipelineConfig(VanillaPipelineConfig):
     """Configuration for pipeline instantiation"""
@@ -169,22 +170,14 @@ class DreamfusionPipeline(VanillaPipeline):
             text_embedding = self.base_text_embedding
 
         with torch.autocast(device_type="cuda", dtype=torch.float16):
-            sds_loss, latents, grad = self.sd.sds_loss(
+            sds_loss = self.sd.sds_loss(
                 text_embedding.to(self.sd_device),
                 train_output.to(self.sd_device),
                 guidance_scale=int(self.config.guidance_scale),
+                grad_scaler=self.grad_scaler,
             )
-            sds_loss = sds_loss.to(self.device)
 
-        # TODO: opacity penalty using transmittance, not accumultation
-        if self.config.opacity_penalty:
-            accum_mean = torch.mean(1.0 - model_outputs["accumulation"])
-            sds_loss *= torch.clamp(accum_mean, min=0.5)
-        metrics_dict["sds_loss"] = sds_loss.to(self.device)
-        loss_dict["latent_loss"] = SpecifyGradient.apply(self.grad_scaler.scale(latents), grad)
-
-        model_outputs["latents"] = latents
-        model_outputs["grad"] = grad
+        loss_dict["sds_loss"] = sds_loss.to(self.device)
 
         return model_outputs, loss_dict, metrics_dict