From c81e24bd5ad0240bd94c9f1d98af55348cb2f27c Mon Sep 17 00:00:00 2001 From: mcallisterdavid Date: Thu, 16 Feb 2023 02:37:31 +0000 Subject: [PATCH] Choose your own diffusion --- nerfstudio/generative/stable_diffusion.py | 19 ++-- nerfstudio/pipelines/dreamfusion_pipeline.py | 4 +- pyproject.toml | 1 + scripts/trace_stablediff.py | 96 +++++++++++++------- 4 files changed, 77 insertions(+), 43 deletions(-) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index fec3aa588..0af0f23f6 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -48,10 +48,11 @@ except ImportError: 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" - +SD_IDENTIFIERS = { + '1-5':'runwayml/stable-diffusion-v1-5', + '2-0':'stabilityai/stable-diffusion-2-base', + '2-1':'stabilityai/stable-diffusion-2-1-base' +} @dataclass class UNet2DConditionOutput: @@ -87,7 +88,7 @@ class StableDiffusion(nn.Module): num_train_timesteps: number of training timesteps """ - def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000) -> None: + def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000, version='1-5') -> None: super().__init__() self.device = device @@ -104,14 +105,16 @@ class StableDiffusion(nn.Module): ) self.alphas = self.scheduler.alphas_cumprod.to(self.device) # type: ignore - pipe = StableDiffusionPipeline.from_pretrained(SD_SOURCE, torch_dtype=torch.float16) + sd_id = SD_IDENTIFIERS[version] + pipe = StableDiffusionPipeline.from_pretrained(sd_id, torch_dtype=torch.float16) assert pipe is not None pipe = pipe.to(self.device) pipe.enable_attention_slicing() # use jitted unet - unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / "sd_unet_traced.pt" + filename_sd_id = sd_id.split('/')[-1] + unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / f'{filename_sd_id}_unet_traced.pt' if unet_traced_filename.exists(): CONSOLE.print("Loading traced UNet.") unet_traced = torch.jit.load(unet_traced_filename) @@ -133,7 +136,7 @@ class StableDiffusion(nn.Module): del pipe.unet else: CONSOLE.print("[bold yellow] Warning: Loading UNet without JIT acceleration.") - CONSOLE.print(r"Run [yellow]python scripts/generative/trace_stable_diffusion.py[/yellow] for a speedup!") + CONSOLE.print(f"Run [yellow]ns-trace-sd --sd-version {version} [/yellow] for a speedup!") self.unet = pipe.unet self.unet.to(memory_format=torch.channels_last) diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 692a92a71..63f86816b 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -53,6 +53,8 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig): """appended to prompt for back view""" guidance_scale: float = 100 """guidance scale for sds loss""" + sd_version: str = "1-5" + """Stable Diffusion version for dreamfusion""" stablediffusion_device: Optional[str] = None @@ -85,7 +87,7 @@ class DreamfusionPipeline(VanillaPipeline): self.side_prompt = config.side_prompt self.back_prompt = config.back_prompt self.front_prompt = config.front_prompt - self.sd = StableDiffusion(self.sd_device) + self.sd = StableDiffusion(self.sd_device, version=config.sd_version) if config.location_based_prompting: self.top_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", "") diff --git a/pyproject.toml b/pyproject.toml index 8901aa282..b29515d30 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,6 +100,7 @@ ns-install-cli = "scripts.completions.install:entrypoint" ns-process-data = "scripts.process_data:entrypoint" ns-download-data = "scripts.downloads.download_data:entrypoint" ns-train = "scripts.train:entrypoint" +ns-trace-sd = "scripts.trace_stablediff:entrypoint" ns-viewer = "scripts.viewer.run_viewer:entrypoint" ns-eval = "scripts.eval:entrypoint" ns-render = "scripts.render:entrypoint" diff --git a/scripts/trace_stablediff.py b/scripts/trace_stablediff.py index 6395f9bf8..44193307d 100644 --- a/scripts/trace_stablediff.py +++ b/scripts/trace_stablediff.py @@ -1,54 +1,82 @@ """Trace Stable Diffusion UNet for speed improvement""" import functools +import sys +from dataclasses import dataclass from pathlib import Path import appdirs import torch +import tyro from diffusers import StableDiffusionPipeline +from rich.console import Console +from nerfstudio.generative.stable_diffusion import SD_IDENTIFIERS -def jit_unet(): - # torch disable grad - torch.set_grad_enabled(False) +CONSOLE = Console(width=120) - # load inputs - def generate_inputs(): - sample = torch.randn(2, 4, 64, 64).half().cuda() - timestep = torch.rand(1).half().cuda() * 999 - encoder_hidden_states = torch.randn(2, 77, 768).half().cuda() - return sample, timestep, encoder_hidden_states +@dataclass +class TraceSD: + """Traces the Stable Diffusion U-Net for better performance.""" - pipe = StableDiffusionPipeline.from_pretrained( - "runwayml/stable-diffusion-v1-5", - torch_dtype=torch.float16, - ).to("cuda") + sd_version: str + """Stable diffusion version.""" - pipe.enable_attention_slicing() + def main(self): + # torch disable grad + torch.set_grad_enabled(False) - unet = pipe.unet - unet.eval() - unet.to(memory_format=torch.channels_last) # use channels_last memory format - unet.forward = functools.partial(unet.forward, return_dict=False) # set return_dict=False as default + if self.sd_version not in SD_IDENTIFIERS.keys(): + CONSOLE.print(f"[bold red]Invalid Stable Diffusion version; choose from {set(SD_IDENTIFIERS.keys())}") + sys.exit(1) - # warmup - for _ in range(3): - with torch.inference_mode(): - inputs = generate_inputs() - orig_output = unet(*inputs) + sd_id = SD_IDENTIFIERS[self.sd_version] + pipe = StableDiffusionPipeline.from_pretrained( + sd_id, + torch_dtype=torch.float16, + ).to("cuda") + pipe.enable_attention_slicing() - # trace - print("tracing..") - unet_traced = torch.jit.trace(unet, inputs) - unet_traced.eval() - print("done tracing") + # load inputs + def generate_inputs(): + sample = torch.randn(2, 4, 64, 64).half().cuda() + timestep = torch.rand(1).half().cuda() * 999 - # save the traced model - unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / "unet_traced.pt" - if not unet_traced_filename.exists(): - unet_traced_filename.parent.mkdir(parents=True, exist_ok=True) - unet_traced.save(unet_traced_filename) + # get shape of encoder hidden state (varies by SD version) + text_input = pipe.tokenizer("", return_tensors="pt") + _, w, d = pipe.text_encoder(text_input.input_ids.cuda())[0].shape + encoder_hidden_states = torch.randn(2, w, d).half().cuda() + return sample, timestep, encoder_hidden_states + + unet = pipe.unet + unet.eval() + unet.to(memory_format=torch.channels_last) # use channels_last memory format + unet.forward = functools.partial(unet.forward, return_dict=False) # set return_dict=False as default + + # warmup + for _ in range(3): + with torch.inference_mode(): + inputs = generate_inputs() + _ = unet(*inputs) + + # trace + print("tracing..") + unet_traced = torch.jit.trace(unet, inputs) + unet_traced.eval() + print("done tracing") + + # save the traced model + filename_sd_id = sd_id.split('/')[-1] + unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / f'{filename_sd_id}_unet_traced.pt' + if not unet_traced_filename.exists(): + unet_traced_filename.parent.mkdir(parents=True, exist_ok=True) + unet_traced.save(unet_traced_filename) + +def entrypoint(): + """Entrypoint for use with pyproject scripts.""" + tyro.extras.set_accent_color("bright_yellow") + tyro.cli(TraceSD).main() if __name__ == "__main__": - jit_unet() + entrypoint()