Trace Stable Diffusion (#1362)

* Trace Stable Diffusion

* pylint

Co-authored by Terrance Wang

nerfstudio/generative/stable_diffusion.py

@@ -17,9 +17,11 @@
 # Modified from https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py
 
 import sys
+from dataclasses import dataclass
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
 
+import appdirs
 import mediapy
 import numpy as np
 import torch
@@ -49,13 +51,18 @@ SD_SOURCE = "runwayml/stable-diffusion-v1-5"
 CLIP_SOURCE = "openai/clip-vit-large-patch14"
 
 
+@dataclass
+class UNet2DConditionOutput:
+    """Class to hold traced model"""
+
+    sample: torch.FloatTensor
+
+
 class StableDiffusion(nn.Module):
     """Stable Diffusion implementation
-
     Args:
         device: device to use
         num_train_timesteps: number of training timesteps
-
     """
 
     def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000) -> None:
@@ -81,24 +88,46 @@ class StableDiffusion(nn.Module):
 
         pipe.enable_attention_slicing()
 
+        # use jitted unet
+        unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / "sd_unet_traced.pt"
+        if unet_traced_filename.exists():
+            CONSOLE.print("Loading traced UNet.")
+            unet_traced = torch.jit.load(unet_traced_filename)
+
+            class TracedUNet(torch.nn.Module):
+                """Jitted version of UNet"""
+
+                def __init__(self):
+                    super().__init__()
+                    self.in_channels = pipe.unet.in_channels
+                    self.device = pipe.unet.device
+
+                def forward(self, latent_model_input, t, encoder_hidden_states):  # pylint: disable=no-self-use
+                    """Forward pass"""
+                    sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0]
+                    return UNet2DConditionOutput(sample=sample)
+
+            self.unet = TracedUNet()
+            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!")
+            self.unet = pipe.unet
+            self.unet.to(memory_format=torch.channels_last)
+
         self.tokenizer = pipe.tokenizer
-        self.unet = pipe.unet
         self.text_encoder = pipe.text_encoder
         self.auto_encoder = pipe.vae
 
-        self.unet.to(memory_format=torch.channels_last)
-
         CONSOLE.print("Stable Diffusion loaded!")
 
     def get_text_embeds(
         self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]]
     ) -> TensorType[2, "max_length", "embed_dim"]:
         """Get text embeddings for prompt and negative prompt
-
         Args:
             prompt: Prompt text
             negative_prompt: Negative prompt text
-
         Returns:
             Text embeddings
         """
@@ -135,12 +164,10 @@ class StableDiffusion(nn.Module):
         guidance_scale: float = 100.0,
     ) -> Tuple[torch.Tensor, TensorType["BS", 4, "H", "W"], TensorType["BS", 4, "H", "W"]]:
         """Score Distilation Sampling loss proposed in DreamFusion paper (https://dreamfusion3d.github.io/)
-
         Args:
             text_embeddings: Text embeddings
             image: Rendered image
             guidance_scale: How much to weigh the guidance
-
         Returns:
             A tuple of Loss, latent, and gradient
         """
@@ -181,7 +208,6 @@ class StableDiffusion(nn.Module):
         latents: Optional[TensorType["BS", 4, "H", "W"]] = None,
     ) -> TensorType["BS", 4, "H", "W"]:
         """Produce latents for a given text embedding
-
         Args:
             text_embeddings: Text embeddings
             height: Height of the image
@@ -189,7 +215,6 @@ class StableDiffusion(nn.Module):
             num_inference_steps: Number of inference steps
             guidance_scale: How much to weigh the guidance
             latents: Latents to start with
-
         Returns:
             Latents
         """
@@ -209,7 +234,9 @@ class StableDiffusion(nn.Module):
 
                 # predict the noise residual
                 with torch.no_grad():
-                    noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]
+                    noise_pred = self.unet(
+                        latent_model_input, t.to(self.device), encoder_hidden_states=text_embeddings
+                    ).sample
 
                 # perform guidance
                 noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
@@ -221,10 +248,8 @@ class StableDiffusion(nn.Module):
 
     def latents_to_img(self, latents: TensorType["BS", 4, "H", "W"]) -> TensorType["BS", 3, "H", "W"]:
         """Convert latents to images
-
         Args:
             latents: Latents to convert
-
         Returns:
             Images
         """
@@ -240,10 +265,8 @@ class StableDiffusion(nn.Module):
 
     def imgs_to_latent(self, imgs: TensorType["BS", 3, "H", "W"]) -> TensorType["BS", 4, "H", "W"]:
         """Convert images to latents
-
         Args:
             imgs: Images to convert
-
         Returns:
             Latents
         """
@@ -263,14 +286,12 @@ class StableDiffusion(nn.Module):
         latents=None,
     ) -> np.ndarray:
         """Generate an images from a prompts.
-
         Args:
             prompts: The prompt to generate an image from.
             negative_prompts: The negative prompt to generate an image from.
             num_inference_steps: The number of inference steps to perform.
             guidance_scale: The scale of the guidance.
             latents: The latents to start from, defaults to random.
-
         Returns:
             The generated image.
         """
@@ -297,14 +318,12 @@ class StableDiffusion(nn.Module):
         self, prompts, negative_prompts="", num_inference_steps=50, guidance_scale=7.5, latents=None
     ) -> np.ndarray:
         """Generate an image from a prompt.
-
         Args:
             prompts: The prompt to generate an image from.
             negative_prompts: The negative prompt to generate an image from.
             num_inference_steps: The number of inference steps to perform.
             guidance_scale: The scale of the guidance.
             latents: The latents to start from, defaults to random.
-
         Returns:
             The generated image.
         """
@@ -315,7 +334,6 @@ def generate_image(
     prompt: str, negative: str = "", seed: int = 0, steps: int = 50, save_path: Path = Path("test_sd.png")
 ):
     """Generate an image from a prompt using Stable Diffusion.
-
     Args:
         prompt: The prompt to use.
         negative: The negative prompt to use.

scripts/generative/trace_stable_diffusion.py

@@ -0,0 +1,74 @@
+"""Trace Stable Diffusion UNet for speed improvement"""
+
+import functools
+import sys
+from pathlib import Path
+
+import appdirs
+import torch
+import tyro
+from rich.console import Console
+
+CONSOLE = Console(width=120)
+
+try:
+    from diffusers import StableDiffusionPipeline
+except ImportError:
+    CONSOLE.print("[bold red]Missing Stable Diffusion packages.")
+    CONSOLE.print(r"Install using [yellow]pip install nerfstudio\[gen][/yellow]")
+    CONSOLE.print(r"or [yellow]pip install -e .\[gen][/yellow] if installing from source.")
+    sys.exit(1)
+
+
+def jit_unet(save_dir: Path = Path(appdirs.user_data_dir("nerfstudio"))):
+    """Trace Stable Diffusion UNet for speed improvement
+
+    Args:
+        save_dir: directory to save the traced model
+    """
+
+    # 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
+
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        torch_dtype=torch.float16,
+    ).to("cuda")
+
+    pipe.enable_attention_slicing()
+
+    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
+    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()
+    print("done tracing")
+
+    # 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)
+
+
+if __name__ == "__main__":
+    tyro.cli(jit_unet)