cleaning up jit code

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"""

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