From e229a430d5cced6a399f96aa9490fbb5ea5cc59b Mon Sep 17 00:00:00 2001 From: terrance Date: Wed, 10 May 2023 21:25:25 +0000 Subject: [PATCH] add deepfloyd option first try --- nerfstudio/generative/deepfloyd.py | 162 +++++++++++++++++++++++++++++ nerfstudio/models/dreamfusion.py | 18 +++- 2 files changed, 176 insertions(+), 4 deletions(-) create mode 100644 nerfstudio/generative/deepfloyd.py diff --git a/nerfstudio/generative/deepfloyd.py b/nerfstudio/generative/deepfloyd.py new file mode 100644 index 000000000..6fbf88667 --- /dev/null +++ b/nerfstudio/generative/deepfloyd.py @@ -0,0 +1,162 @@ +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import List, Optional, Union + +import appdirs +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 torch.cuda.amp import custom_bwd, custom_fwd +from torch.cuda.amp.grad_scaler import GradScaler +from torchtyping import TensorType + + + +from diffusers import IFPipeline +from transformers import T5EncoderModel, T5Tokenizer + +class DeepFloyd(nn.Module): + + # load in model + def __init__(self, device: Union[torch.device, str]): + super().__init__() + self.device = device + + self.text_encoder = T5EncoderModel.from_pretrained( + "DeepFloyd/IF-I-M-v1.0", + subfolder="text_encoder", + load_in_8bit=True, + variant="8bit", + device_map="auto", + ) + + self.pipe = IFPipeline.from_pretrained( + "DeepFloyd/IF-I-M-v1.0", + text_encoder=self.text_encoder, + safety_checker=None, + watermarker=None, + feature_extractor=None, + requires_safety_checker=False, + variant="fp16", + torch_dtype=torch.float16, + ).to(self.device) + + self.pipe.enable_attention_slicing(1) + self.pipe.unet.to(memory_format=torch.channels_last) + + self.unet = self.pipe.unet + for p in self.unet.parameters(): + p.requires_grad_(False) + + self.scheduler = self.pipe.scheduler + + self.num_train_timesteps = self.scheduler.config.num_train_timesteps + self.min_step = int(self.num_train_timesteps * 0.02) + self.max_step = int(self.num_train_timesteps * 0.98) + + + self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to( + self.device + ) + + def get_text_embeds( + self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]] + ) -> TensorType[2, "max_length", "embed_dim"]: + prompt = [prompt] if isinstance(prompt, str) else prompt + negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + prompt_embeds, negative_embeds = self.pipe.encode_prompt(prompt, negative_prompt=negative_prompt) + + return torch.cat([prompt_embeds, negative_embeds]) + + def sds_loss( + self, + text_embeddings, + image, + guidance_scale, + grad_scaler, + ) -> torch.Tensor: + with torch.autocast(device_type="cuda", enabled=False): + image = F.interpolate(image.half(), (IMG_DIM, IMG_DIM), mode="bilinear", align_corners=False) + t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device) + + # predict the noise residual with unet, NO grad! + with torch.no_grad(): + # add noise + noise = torch.randn_like(image) + image_noisy = self.scheduler.add_noise(image, noise, t) # type: ignore + # pred noise + image_model_input = torch.cat([image_noisy] * 2) + noise_pred = self.unet(image_model_input, t, encoder_hidden_states=text_embeddings).sample + + # perform guidance + noise_pred_text, noise_pred_uncond = noise_pred.chunk(2) + noise_pred_text, predicted_variance = noise_pred_text.split(3, dim=1) + noise_pred_uncond, _ = noise_pred_uncond.split(3, dim=1) + + noise_pred = noise_pred_text + 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) + + if grad_scaler is not None: + latents = grad_scaler.scale(latents) + loss = _SDSGradient.apply(latents, grad) + + return loss + + def prompt_to_img( + self, + prompts: Union[str, List[str]], + negative_prompts: Union[str, List[str]] = "", + generator=None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + ) -> 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) + prompt_embeds, negative_embeds = self.pipe.encode_prompt(prompts, negative_prompt=negative_prompts) + image = self.pipe(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator).images + + return image + +def generate_image( + prompt: str, negative: str = "", seed: int = 0, steps: int = 50, save_path: Path = Path("test_deepfloyd.png") +): + """Generate an image from a prompt using DeepFloyd IF. + 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. + """ + generator = torch.manual_seed(seed) + cuda_device = torch.device("cuda") + with torch.no_grad(): + df = DeepFloyd(cuda_device) + imgs = df.prompt_to_img(prompt, negative, generator, steps) + imgs[0].save(save_path) + # mediapy.write_image(str(save_path), imgs[0]) + +if __name__ == "__main__": + tyro.cli(generate_image) \ No newline at end of file diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 4dd4b6b1f..7d68f8a22 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -36,6 +36,7 @@ from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.density_fields import HashMLPDensityField from nerfstudio.fields.dreamfusion_field import DreamFusionField from nerfstudio.generative.stable_diffusion import StableDiffusion +from nerfstudio.generative.deepfloyd import DeepFloyd from nerfstudio.generative.stable_diffusion_utils import PositionalTextEmbeddings from nerfstudio.model_components.losses import ( MSELoss, @@ -198,14 +199,16 @@ class DreamFusionModel(Model): """Set the fields and modules""" super().populate_modules() - self.sd = StableDiffusion(self.sd_device, version=self.sd_version) + # self.sd = StableDiffusion(self.sd_device, version=self.sd_version) + self.df = DeepFloyd(self.sd_device) self.text_embeddings = PositionalTextEmbeddings( base_prompt=self.cur_prompt, top_prompt=self.cur_prompt + self.top_prompt, side_prompt=self.cur_prompt + self.side_prompt, back_prompt=self.cur_prompt + self.back_prompt, front_prompt=self.cur_prompt + self.front_prompt, - stable_diffusion=self.sd, + # stable_diffusion=self.sd, + stable_diffusion=self.df, positional_prompting=self.positional_prompting, ) @@ -480,11 +483,18 @@ class DreamFusionModel(Model): .permute(0, 3, 1, 2) ) - sds_loss = 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.guidance_scale), + # grad_scaler=self.grad_scaler, + # ) + + sds_loss = self.df.sds_loss( text_embedding.to(self.sd_device), train_output.to(self.sd_device), guidance_scale=int(self.guidance_scale), - grad_scaler=self.grad_scaler, + grad_scaler=self.grad_scaler, ) loss_dict["sds_loss"] = sds_loss.to(self.device)