From 71fd9f724596dc7c437a43002cc611632b914649 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Tue, 14 Feb 2023 10:14:48 -0800 Subject: [PATCH] Update SDS loss (#1426) --- nerfstudio/generative/stable_diffusion.py | 40 +++++++++++++++++++---- 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index a91c3b49f..fec3aa588 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/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) @@ -186,7 +210,7 @@ class StableDiffusion(nn.Module): # perform guidance noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + noise_pred = noise_pred_text + guidance_scale * (noise_pred_text - noise_pred_uncond) # w(t), sigma_t^2 w = 1 - self.alphas[t] @@ -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, @@ -240,7 +266,7 @@ class StableDiffusion(nn.Module): # perform guidance noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + noise_pred = noise_pred_text + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"] # type: ignore