Support multiple versions of stable diffusion (#1461)

Co-authored-by: David McAllister <davidrmc2@gmail.com>

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)
 

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()