up

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import torch
 from diffusers import QwenImagePipeline
 from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
-from optimization import get_compiled_transformer
+from optimization import compile_transformer
 from hub_utils import _push_compiled_graph_to_hub
 from huggingface_hub import whoami
 import spaces
@@ -16,7 +16,7 @@ pipe = QwenImagePipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=dtype).t
 pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 
 # --- Ahead-of-time compilation ---
-compiled_transformer = get_compiled_transformer(pipe, prompt="prompt")
+compiled_transformer = compile_transformer(pipe, prompt="prompt")
 
 @spaces.GPU(duration=120)
 def push_to_hub(repo_id, filename, oauth_token: gr.OAuthToken):

optimization.py

@@ -38,23 +38,18 @@ INDUCTOR_CONFIGS = {
 }
 
 
-def get_compiled_transformer(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
-
-    @spaces.GPU(duration=1500)
-    def compile_transformer():
-        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)
-
-    compiled = compile_transformer()
-    return compiled
+@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)