Support multiple versions of stable diffusion (#1461)

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

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