From 11e7fa87e04f90ad94598d6b41cacfa91fff4751 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Wed, 1 Feb 2023 10:29:08 -0800 Subject: [PATCH] Add Stable Diffusion helper code (#1320) * Add stable diffusion * Add dependencies fix lint * Update types * Update comment * fix comments --------- Co-authored-by: David McAllister Co-authored-by: Abhik Ahuja --- nerfstudio/generative/__init__.py | 13 + nerfstudio/generative/stable_diffusion.py | 336 ++++++++++++++++++++++ pyproject.toml | 8 + 3 files changed, 357 insertions(+) create mode 100644 nerfstudio/generative/__init__.py create mode 100644 nerfstudio/generative/stable_diffusion.py diff --git a/nerfstudio/generative/__init__.py b/nerfstudio/generative/__init__.py new file mode 100644 index 000000000..d926cf95c --- /dev/null +++ b/nerfstudio/generative/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py new file mode 100644 index 000000000..af3aa2196 --- /dev/null +++ b/nerfstudio/generative/stable_diffusion.py @@ -0,0 +1,336 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Stable Diffusion helpers""" + +# Modified from https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py + +import sys +from pathlib import Path +from typing import List, Optional, Tuple, Union + +import mediapy +import numpy as np +import torch +import torch.nn.functional as F +import tyro +from rich.console import Console +from torch import nn +from torchtyping import TensorType + +CONSOLE = Console(width=120) + +try: + from diffusers import PNDMScheduler, StableDiffusionPipeline + from transformers import logging + +except ImportError: + CONSOLE.print("[bold red]Missing Stable Diffusion packages.") + CONSOLE.print(r"Install using [yellow]pip install nerfstudio\[gen][/yellow]") + CONSOLE.print(r"or [yellow]pip install -e .\[gen][/yellow] if installing from source.") + sys.exit(1) + +logging.set_verbosity_error() +IMG_DIM = 512 +CONST_SCALE = 0.18215 + +SD_SOURCE = "runwayml/stable-diffusion-v1-5" +CLIP_SOURCE = "openai/clip-vit-large-patch14" + + +class StableDiffusion(nn.Module): + """Stable Diffusion implementation + + 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.min_step = int(self.num_train_timesteps * 0.02) + self.max_step = int(self.num_train_timesteps * 0.98) + + self.scheduler = PNDMScheduler( + beta_start=0.00085, + beta_end=0.012, + beta_schedule="scaled_linear", + num_train_timesteps=self.num_train_timesteps, + ) + 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 + pipe = pipe.to(self.device) + + pipe.enable_attention_slicing() + + self.tokenizer = pipe.tokenizer + self.unet = pipe.unet + self.text_encoder = pipe.text_encoder + self.auto_encoder = pipe.vae + + self.unet.to(memory_format=torch.channels_last) + + CONSOLE.print("Stable Diffusion loaded!") + + def get_text_embeds( + self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]] + ) -> TensorType[2, "max_length", "embed_dim"]: + """Get text embeddings for prompt and negative prompt + + Args: + prompt: Prompt text + negative_prompt: Negative prompt text + + Returns: + Text embeddings + """ + + # Tokenize text and get embeddings + text_input = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + + with torch.no_grad(): + text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0] + + # Do the same for unconditional embeddings + uncond_input = self.tokenizer( + negative_prompt, padding="max_length", max_length=self.tokenizer.model_max_length, return_tensors="pt" + ) + + with torch.no_grad(): + uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0] + + # Cat for final embeddings + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + + return text_embeddings + + 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", 4, "H", "W"], TensorType["BS", 4, "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(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) # 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 + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # w(t), sigma_t^2 + w = 1 - self.alphas[t] + + grad = w * (noise_pred - noise) + grad = torch.nan_to_num(grad) + + sds_mse = torch.mean(torch.nan_to_num(torch.square(noise_pred - noise))) + + return sds_mse, latents, grad + + def produce_latents( + 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", 4, "H", "W"]] = None, + ) -> TensorType["BS", 4, "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) # type: ignore + + with torch.autocast("cuda"): + 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) + + # predict the noise residual + with torch.no_grad(): + noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"] + + # 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) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"] # type: ignore + return latents + + def latents_to_img(self, latents: TensorType["BS", 4, "H", "W"]) -> TensorType["BS", 3, "H", "W"]: + """Convert latents to images + + Args: + latents: Latents to convert + + Returns: + Images + """ + + latents = 1 / CONST_SCALE * latents + + with torch.no_grad(): + imgs = self.auto_encoder.decode(latents).sample + + imgs = (imgs / 2 + 0.5).clamp(0, 1) + + return imgs + + def imgs_to_latent(self, imgs: TensorType["BS", 3, "H", "W"]) -> TensorType["BS", 4, "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 + latents = posterior.sample() * CONST_SCALE + + return latents + + def prompt_to_img( + self, + prompts: Union[str, List[str]], + negative_prompts: Union[str, List[str]] = "", + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + latents=None, + ) -> np.ndarray: + """Generate an images from a prompts. + + Args: + prompts: The prompt to generate an image from. + negative_prompts: The negative prompt to generate an image from. + num_inference_steps: The number of inference steps to perform. + guidance_scale: The scale of the guidance. + latents: The latents to start from, defaults to random. + + Returns: + The generated image. + """ + + prompts = [prompts] if isinstance(prompts, str) else prompts + negative_prompts = [negative_prompts] if isinstance(negative_prompts, str) else negative_prompts + text_embeddings = self.get_text_embeds(prompts, negative_prompts) + latents = self.produce_latents( + text_embeddings, + height=IMG_DIM, + width=IMG_DIM, + latents=latents, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + ) # [1, 4, resolution, resolution] + + diffused_img = self.latents_to_img(latents.half()) + diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy() + diffused_img = (diffused_img * 255).round().astype("uint8") + + return diffused_img + + def forward( + self, prompts, negative_prompts="", num_inference_steps=50, guidance_scale=7.5, latents=None + ) -> np.ndarray: + """Generate an image from a prompt. + + Args: + prompts: The prompt to generate an image from. + negative_prompts: The negative prompt to generate an image from. + num_inference_steps: The number of inference steps to perform. + guidance_scale: The scale of the guidance. + latents: The latents to start from, defaults to random. + + Returns: + The generated image. + """ + return self.prompt_to_img(prompts, negative_prompts, num_inference_steps, guidance_scale, latents) + + +def generate_image( + prompt: str, negative: str = "", seed: int = 0, steps: int = 50, save_path: Path = Path("test_sd.png") +): + """Generate an image from a prompt using Stable Diffusion. + + Args: + prompt: The prompt to use. + negative: The negative prompt to use. + seed: The random seed to use. + steps: The number of steps to use. + save_path: The path to save the image to. + """ + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + cuda_device = torch.device("cuda") + with torch.no_grad(): + sd = StableDiffusion(cuda_device) + imgs = sd.prompt_to_img(prompt, negative, steps) + mediapy.write_image(str(save_path), imgs[0]) + + +if __name__ == "__main__": + tyro.cli(generate_image) diff --git a/pyproject.toml b/pyproject.toml index 8b1700d4d..ca44d209e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +59,14 @@ dependencies = [ [project.optional-dependencies] + +# Generative related dependencies +gen = [ + "diffusers==0.9.0", + "transformers==4.24.0", +] + + # Development packages dev = [ "black[jupyter]==22.3.0",