update sd

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)
 

nerfstudio/models/dreamfusion.py

@@ -192,7 +192,7 @@ class DreamFusionModel(Model):
         """Set the fields and modules"""
         super().populate_modules()
 
-        self.sd = StableDiffusion(self.sd_device)
+        self.sd = StableDiffusion(self.sd_device, version=config.sd_version)
         self.text_embeddings = {
             "top_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", ""),
             "front_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", ""),

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,8 +87,13 @@ class DreamfusionPipeline(VanillaPipeline):
         self.side_prompt = config.side_prompt
         self.back_prompt = config.back_prompt 
         self.front_prompt = config.front_prompt
+<<<<<<< HEAD
         self.sd = StableDiffusion(self.sd_device)
 
+=======
+        self.sd = StableDiffusion(self.sd_device, version=config.sd_version)
+        
+>>>>>>> c81e24bd5ad0240bd94c9f1d98af55348cb2f27c
         if config.location_based_prompting:
             self.top_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", "")
             self.front_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", "")

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"

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