From 3c71429d300550a2f8709b6b5377bb9945b2310e Mon Sep 17 00:00:00 2001 From: terrance Date: Fri, 3 Feb 2023 22:28:03 +0000 Subject: [PATCH] cleaning up jit code --- nerfstudio/generative/stable_diffusion.py | 2 +- .../{trace_unet.py => trace_stablediff.py} | 29 ------------------- 2 files changed, 1 insertion(+), 30 deletions(-) rename scripts/{trace_unet.py => trace_stablediff.py} (60%) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index eb16710f0..1816da32d 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -92,7 +92,7 @@ class StableDiffusion(nn.Module): unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / "unet_traced.pt" if unet_traced_filename.exists(): CONSOLE.print("Loading traced UNet.") - unet_traced = torch.jit.load("unet_traced.pt") + unet_traced = torch.jit.load(unet_traced_filename) class TracedUNet(torch.nn.Module): """Jitted version of UNet""" diff --git a/scripts/trace_unet.py b/scripts/trace_stablediff.py similarity index 60% rename from scripts/trace_unet.py rename to scripts/trace_stablediff.py index 17577037b..6395f9bf8 100644 --- a/scripts/trace_unet.py +++ b/scripts/trace_stablediff.py @@ -1,7 +1,6 @@ """Trace Stable Diffusion UNet for speed improvement""" import functools -import time from pathlib import Path import appdirs @@ -10,14 +9,9 @@ from diffusers import StableDiffusionPipeline def jit_unet(): - # torch disable grad torch.set_grad_enabled(False) - # set variables - n_experiments = 2 - unet_runs_per_experiment = 50 - # load inputs def generate_inputs(): sample = torch.randn(2, 4, 64, 64).half().cuda() @@ -49,29 +43,6 @@ def jit_unet(): unet_traced.eval() print("done tracing") - # warmup and optimize graph - for _ in range(5): - with torch.inference_mode(): - inputs = generate_inputs() - orig_output = unet_traced(*inputs) - - # benchmarking - with torch.inference_mode(): - for _ in range(n_experiments): - torch.cuda.synchronize() - start_time = time.time() - for _ in range(unet_runs_per_experiment): - orig_output = unet_traced(*inputs) - torch.cuda.synchronize() - print(f"unet traced inference took {time.time() - start_time:.2f} seconds") - for _ in range(n_experiments): - torch.cuda.synchronize() - start_time = time.time() - for _ in range(unet_runs_per_experiment): - orig_output = unet(*inputs) - torch.cuda.synchronize() - print(f"unet inference took {time.time() - start_time:.2f} seconds") - # save the traced model unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / "unet_traced.pt" if not unet_traced_filename.exists():