Apply optimize_pipeline_

app.py

@@ -19,6 +19,7 @@ from diffusers import ZImagePipeline
 from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel
 
 from pe import prompt_template
+from optimization import optimize_pipeline_ 
 
 # ==================== Environment Variables ==================================
 MODEL_PATH = os.environ.get("MODEL_PATH", "Tongyi-MAI/Z-Image-Turbo")
@@ -409,6 +410,18 @@ def generate(
 
 init_app()
 
+#Apply the scheduler for the compilation
+scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0)
+pipe.scheduler = scheduler
+
+#Perform the one-time AoTI compilation
+optimize_pipeline_(
+    pipe,
+    prompt="prompt",
+    num_inference_steps=1,
+    guidance_scale=0.0
+)
+
 with gr.Blocks(title="Z-Image Demo") as demo:
     gr.Markdown(
         """<div align="center">