up

app.py

@@ -14,6 +14,8 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Load the model pipeline
 pipe = QwenImagePipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=dtype).to(device)
 pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+# --- Ahead-of-time compilation ---
+compiled_transformer = compile_transformer(pipe, prompt="prompt")
 
 @spaces.GPU(duration=120)
 def push_to_hub(repo_id, filename, oauth_token: gr.OAuthToken):
@@ -23,9 +25,6 @@ def push_to_hub(repo_id, filename, oauth_token: gr.OAuthToken):
     # this will throw if token is invalid
     _ = whoami(oauth_token.token)
 
-    # --- Ahead-of-time compilation ---
-    compiled_transformer = compile_transformer(pipe, prompt="prompt")
-
     token = oauth_token.token 
     out = _push_compiled_graph_to_hub(
         compiled_transformer.archive_file,

optimization.py

@@ -38,18 +38,20 @@ INDUCTOR_CONFIGS = {
 }
 
 
-@spaces.GPU(duration=1500)
 def compile_transformer(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
-    with spaces.aoti_capture(pipeline.transformer) as call:
-        pipeline(*args, **kwargs)
-
-    dynamic_shapes = tree_map(lambda t: 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,
-    )
-    return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
+    @spaces.GPU(duration=1500)
+    def f():
+        with spaces.aoti_capture(pipeline.transformer) as call:
+            pipeline(*args, **kwargs)
+
+        dynamic_shapes = tree_map(lambda t: 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,
+        )
+        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
+    return f