diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 8bfdedef7..d64bfdfeb 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -16,9 +16,9 @@ import sys from pathlib import Path -from typing import List, Optional, Union +from typing import List, Optional, Tuple, Union -import matplotlib.pyplot as plt +import mediapy import numpy as np import torch import torch.nn.functional as F @@ -48,14 +48,20 @@ CLIP_SOURCE = "openai/clip-vit-large-patch14" class StableDiffusion(nn.Module): - """Stable Diffusion implementation""" + """Stable Diffusion implementation - def __init__(self, device) -> None: + Args: + device: device to use + num_train_timesteps: number of training timesteps + + """ + + def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000) -> None: super().__init__() self.device = device + self.num_train_timesteps = num_train_timesteps - self.num_train_timesteps = 1000 self.min_step = int(self.num_train_timesteps * 0.02) self.max_step = int(self.num_train_timesteps * 0.98) @@ -65,7 +71,7 @@ class StableDiffusion(nn.Module): beta_schedule="scaled_linear", num_train_timesteps=self.num_train_timesteps, ) - self.alphas = self.scheduler.alphas_cumprod.to(self.device) + self.alphas = self.scheduler.alphas_cumprod.to(self.device) # type: ignore pipe = StableDiffusionPipeline.from_pretrained(SD_SOURCE, torch_dtype=torch.float16) assert pipe is not None @@ -127,65 +133,82 @@ class StableDiffusion(nn.Module): return text_embeddings - def sds_loss(self, text_embeddings, nerf_output, guidance_scale=100): - nerf_output_np = nerf_output.squeeze(0).permute(1, 2, 0).detach().cpu().numpy() - nerf_output_np = np.clip(nerf_output_np, 0.0, 1.0) - plt.imsave("nerf_output.png", nerf_output_np) - nerf_output = F.interpolate(nerf_output, (IMG_DIM, IMG_DIM), mode="bilinear") + def sds_loss( + self, + text_embeddings: TensorType["N", "max_length", "embed_dim"], + image: TensorType["BS", 3, "H", "W"], + guidance_scale: float = 100.0, + ) -> Tuple[torch.Tensor, TensorType["BS", "2N", "H", "W"], TensorType["BS", "2N", "H", "W"]]: + """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 + + Returns: + A tuple of Loss, latent, and gradient + """ + 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) - latents = self.imgs_to_latent(nerf_output) + latents = self.imgs_to_latent(image) + # predict the noise residual with unet, NO grad! with torch.no_grad(): # add noise noise = torch.randn_like(latents) - latents_noisy = self.scheduler.add_noise(latents, noise, t) + latents_noisy = self.scheduler.add_noise(latents, noise, t) # type: ignore # pred noise latent_model_input = torch.cat([latents_noisy] * 2) noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample - # perform guidance (high scale from paper!) + # 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) - # with torch.no_grad(): - # diffused_img = self.latents_to_img(latents_noisy - noise_pred) - # diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy().reshape((512, 512, 3)) - # diffused_img = (diffused_img * 255).round().astype("uint8") - # plt.imsave("sdimg.png", diffused_img) # w(t), sigma_t^2 w = 1 - self.alphas[t] - # w = self.alphas[t] ** 0.5 * (1 - self.alphas[t]) - grad = w * (noise_pred - noise) - # clip grad for stable training? - # grad = grad.clamp(-10, 10) + grad = w * (noise_pred - noise) grad = torch.nan_to_num(grad) - # noise_difference = self.latents_to_img(grad) - # plt.imsave('noise_difference.png', noise_difference.view(3, 512, 512).permute(1, 2, 0).detach().cpu().numpy()) - - # manually backward, since we omitted an item in grad and cannot simply autodiff. - # _t = time.time() - noise_loss = torch.mean(torch.nan_to_num(torch.square(noise_pred - noise))) - # print('SDS LOSS:', noise_loss) - return noise_loss, latents, grad def produce_latents( - self, text_embeddings, height=IMG_DIM, width=IMG_DIM, num_inference_steps=50, guidance_scale=7.5, latents=None - ): + self, + text_embeddings: TensorType["N", "max_length", "embed_dim"], + height: int = IMG_DIM, + width: int = IMG_DIM, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + latents: Optional[TensorType["BS", "2N", "H", "W"]] = None, + ) -> TensorType["BS", "2N", "H", "W"]: + """Produce latents for a given text embedding + + Args: + text_embeddings: Text embeddings + height: Height of the image + width: Width of the image + num_inference_steps: Number of inference steps + guidance_scale: How much to weigh the guidance + latents: Latents to start with + + Returns: + Latents + """ if latents is None: latents = torch.randn( (text_embeddings.shape[0] // 2, self.unet.in_channels, height // 8, width // 8), device=self.device ) - self.scheduler.set_timesteps(num_inference_steps) + self.scheduler.set_timesteps(num_inference_steps) # type: ignore with torch.autocast("cuda"): - for t in self.scheduler.timesteps: + for t in self.scheduler.timesteps: # type: ignore + assert latents is not None # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. latent_model_input = torch.cat([latents] * 2) @@ -198,11 +221,18 @@ class StableDiffusion(nn.Module): noise_pred = noise_pred_uncond + 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"] - + latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"] # type: ignore return latents - def latents_to_img(self, latents): + def latents_to_img(self, latents: TensorType["BS", "2N", "H", "W"]) -> TensorType["BS", 3, "H", "W"]: + """Convert latents to images + + Args: + latents: Latents to convert + + Returns: + Images + """ latents = 1 / CONST_SCALE * latents @@ -213,8 +243,15 @@ class StableDiffusion(nn.Module): return imgs - def imgs_to_latent(self, imgs): - # imgs: [B, 3, H, W] + def imgs_to_latent(self, imgs: TensorType["BS", 3, "H", "W"]) -> TensorType["BS", "2N", "H", "W"]: + """Convert images to latents + + Args: + imgs: Images to convert + + Returns: + Latents + """ imgs = 2 * imgs - 1 posterior = self.auto_encoder.encode(imgs).latent_dist @@ -297,7 +334,7 @@ def generate_image( with torch.no_grad(): sd = StableDiffusion(cuda_device) imgs = sd.prompt_to_img(prompt, negative, steps) - plt.imsave(str(save_path), imgs[0]) + mediapy.write_image(str(save_path), imgs[0]) if __name__ == "__main__":