diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 346a3b5c4..d2eb439f2 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -1,8 +1,15 @@ + +# From https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py + import torch from diffusers import AutoencoderKL, PNDMScheduler, UNet2DConditionModel +from rich.console import Console from torch import nn from transformers import CLIPTextModel, CLIPTokenizer, logging +CONSOLE = Console(width=120) +logging.set_verbosity_error() +IMG_DIM = 512 class StableDiffusion(nn.Module): """Stable Diffusion implementation @@ -13,3 +20,177 @@ class StableDiffusion(nn.Module): device ) -> None: super().__init__() + + self.device = device + + try: + with open('./HF_TOKEN', 'r') as f: + self.token = f.read().replace('\n', '') + CONSOLE.print(f'Hugging face access token loaded.') + except FileNotFoundError as e: + raise AssertionError(f'Cannot read hugging face token, make sure to save your Hugging Face token at `./HF_TOKEN`') + + SD_SOURCE = "runwayml/stable-diffusion-v1-5" + CLIP_SOURCE = "openai/clip-vit-large-patch14" + + 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) + + CONSOLE.print(f'Loading Stable Diffusion model from {SD_SOURCE}.') + + self.auto_encoder = AutoencoderKL.from_pretrained(SD_SOURCE, subfolder="vae", use_auth_token=self.token).to(self.device) + self.unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet", use_auth_token=self.token).to(self.device) + + CONSOLE.print(f'Loading CLIP model from {CLIP_SOURCE}.') + + self.tokenizer = CLIPTokenizer.from_pretrained(CLIP_SOURCE) + self.text_encoder = CLIPTextModel.from_pretrained(CLIP_SOURCE).to(self.device) + + 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) + + CONSOLE.print(f'Stable Diffusion loaded!') + + def get_text_embeds(self, prompt, negative_prompt): + # prompt, negative_prompt: [str] + + # 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 train_step(self, text_embeddings, pred_nerf, guidance_scale=100): + + # Upscale NeRF render for Stable Diffusion + pred_nerf = F.interpolate(pred_nerf, (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(pred_nerf) + + # 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) + # 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!) + 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]) + # 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 = torch.nan_to_num(grad) + + # manually backward, since we omitted an item in grad and cannot simply autodiff. + # _t = time.time() + latents.backward(gradient=grad, retain_graph=True) + # torch.cuda.synchronize(); print(f'[TIME] guiding: backward {time.time() - _t:.4f}s') + + return 0 # dummy loss value + + def produce_latents(self, text_embeddings, height=IMG_DIM, width=IMG_DIM, num_inference_steps=50, guidance_scale=7.5, latents=None): + + 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) + + with torch.autocast('cuda'): + for t in self.scheduler.timesteps: + # 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'] + + return latents + + def latents_to_img(self, latents): + + latents = 1 / 0.18215 * 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): + # imgs: [B, 3, H, W] + + imgs = 2 * imgs - 1 + + posterior = self.auto_encoder.encode(imgs).latent_dist + latents = posterior.sample() * 0.18215 + + return latents + + def prompt_to_img(self, prompts, negative_prompts='', num_inference_steps=50, guidance_scale=7.5, latents=None): + + 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, 64, 64] + + diffused_img = self.latents_to_img(latents) + diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy() + diffused_img = (diffused_img * 255).round().astype('uint8') + + return diffused_img + +def seed_everything(seed): + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + +if __name__ == '__main__': + + import argparse + + import matplotlib.pyplot as plt + + parser = argparse.ArgumentParser() + parser.add_argument('prompt', type=str) + parser.add_argument('--negative', default='', type=str) + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--steps', type=int, default=50) + opt = parser.parse_args() + + seed_everything(opt.seed) + + device = torch.device('cuda') + + sd = StableDiffusion(device) + + imgs = sd.prompt_to_img(opt.prompt, opt.negative, opt.steps) + + plt.imsave('test_sd.png', imgs[0]) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 164997d05..ebd472649 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "nerfstudio" -version = "0.1.10" +version = "0.1.11" description = "All-in-one repository for state-of-the-art NeRFs" readme = "README.md" license = { text="Apache 2.0"} @@ -41,6 +41,8 @@ dependencies = [ "rich==12.5.1", "tensorboard==2.9.0", "torch==1.12.1", + "diffusers==0.9.0", + "transformers==4.24.0", "torchmetrics[image]>=0.9.3", "torchtyping>=0.1.4", "torchvision>=0.13.0",