up

app.py

@@ -29,8 +29,6 @@ def push_to_hub(repo_id, filename, oauth_token: gr.OAuthToken, progress=gr.Progr
         # --- Ahead-of-time compilation ---
         start = time.perf_counter()
         compiled_transformer = compile_transformer(pipe, prompt="prompt")
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
         end = time.perf_counter()
         print(f"Compilation took: {start - end} seconds.")
 

optimization.py

@@ -31,12 +31,14 @@ def compile_transformer(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.k
         with spaces.aoti_capture(pipeline.transformer) as call:
             pipeline(*args, **kwargs)
 
+        print("Inputs captured.")
         dynamic_shapes = tree_map(lambda v: None, call.kwargs)
         dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
 
         exported = torch.export.export(
             mod=pipeline.transformer, args=call.args, kwargs=call.kwargs, dynamic_shapes=dynamic_shapes
         )
+        print("Export done.")
         return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
 
     compiled_transformer = f()