From 36514797fda7c2d9e875e162351149588fbe3d9d Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Mon, 20 Feb 2023 10:12:41 -0800 Subject: [PATCH] Support multiple versions of stable diffusion (#1461) Co-authored-by: David McAllister --- nerfstudio/generative/stable_diffusion.py | 18 ++-- scripts/generative/trace_stable_diffusion.py | 105 ++++++++++++------- 2 files changed, 76 insertions(+), 47 deletions(-) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index fec3aa588..a3955712b 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -48,9 +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 @@ -87,7 +89,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 +106,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 +137,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/scripts/generative/trace_stable_diffusion.py b/scripts/generative/trace_stable_diffusion.py index a501bd424..eee069eb3 100644 --- a/scripts/generative/trace_stable_diffusion.py +++ b/scripts/generative/trace_stable_diffusion.py @@ -2,6 +2,7 @@ import functools import sys +from dataclasses import dataclass from pathlib import Path import appdirs @@ -9,6 +10,8 @@ import torch import tyro from rich.console import Console +from nerfstudio.generative.stable_diffusion import SD_IDENTIFIERS + CONSOLE = Console(width=120) try: @@ -20,55 +23,77 @@ except ImportError: sys.exit(1) -def jit_unet(save_dir: Path = Path(appdirs.user_data_dir("nerfstudio"))): - """Trace Stable Diffusion UNet for speed improvement +@dataclass +class TraceSD: + """Traces the Stable Diffusion U-Net for better performance.""" - Args: - save_dir: directory to save the traced model - """ + sd_version: str + """Stable diffusion version.""" + save_dir: Path = Path(appdirs.user_data_dir("nerfstudio")) + """Directory to save the traced model.""" - # torch disable grad - torch.set_grad_enabled(False) + def main(self): + """Trace the Stable Diffusion U-Net and save it to disk.""" + # torch disable grad + torch.set_grad_enabled(False) - # 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 + if self.sd_version not in SD_IDENTIFIERS: + CONSOLE.print(f"[bold red]Invalid Stable Diffusion version; choose from {set(SD_IDENTIFIERS.keys())}") + sys.exit(1) - pipe = StableDiffusionPipeline.from_pretrained( - "runwayml/stable-diffusion-v1-5", - torch_dtype=torch.float16, - ).to("cuda") + sd_id = SD_IDENTIFIERS[self.sd_version] + pipe = StableDiffusionPipeline.from_pretrained( + sd_id, + torch_dtype=torch.float16, + ) + assert pipe is not None + pipe = pipe.to("cuda") + pipe.enable_attention_slicing() - pipe.enable_attention_slicing() + # load inputs + def generate_inputs(): + sample = torch.randn(2, 4, 64, 64).half().cuda() + timestep = torch.rand(1).half().cuda() * 999 - 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 + # 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 - # warmup - inputs = None - for _ in range(3): - with torch.inference_mode(): - inputs = generate_inputs() - _ = unet(*inputs) + encoder_hidden_states = torch.randn(2, w, d).half().cuda() + return sample, timestep, encoder_hidden_states - # trace - print("tracing..") - assert inputs is not None - unet_traced = torch.jit.trace(unet, inputs) - unet_traced.eval() - print("done tracing") + 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 - # save the traced model - unet_traced_filename = save_dir / "sd_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) + # warmup + inputs = None + for _ in range(3): + with torch.inference_mode(): + inputs = generate_inputs() + _ = unet(*inputs) + + # trace + print("tracing..") + assert inputs is not None + unet_traced = torch.jit.trace(unet, inputs) + unet_traced.eval() # type: ignore + print("done tracing") + + # save the traced model + filename_sd_id = sd_id.split("/")[-1] + unet_traced_filename = self.save_dir / 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) # type: ignore + + +def entrypoint(): + """Entrypoint for use with pyproject scripts.""" + tyro.extras.set_accent_color("bright_yellow") + tyro.cli(TraceSD).main() if __name__ == "__main__": - tyro.cli(jit_unet) + entrypoint()