up

app.py

@@ -5,6 +5,7 @@ from diffusers import DiffusionPipeline
 from optimization import compile_transformer
 from hub_utils import _push_compiled_graph_to_hub
 from huggingface_hub import whoami
+import time
 
 # --- Model Loading ---
 dtype = torch.bfloat16
@@ -14,7 +15,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/Flux.1-Dev", torch_dtype=dtype).to(device)
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU
 def push_to_hub(repo_id, filename, oauth_token: gr.OAuthToken, progress=gr.Progress(track_tqdm=True)):
     if not filename.endswith(".pt2"):
         raise NotImplementedError("The filename must end with a `.pt2` extension.")
@@ -24,7 +25,12 @@ def push_to_hub(repo_id, filename, oauth_token: gr.OAuthToken, progress=gr.Progr
         _ = whoami(oauth_token.token)
 
         # --- 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 - time} seconds.")
 
         token = oauth_token.token
         out = _push_compiled_graph_to_hub(