mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
75 lines
2.1 KiB
Python
75 lines
2.1 KiB
Python
"""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)
|