Update app.py

app.py

@@ -2,6 +2,7 @@ import torch
 import spaces
 import gradio as gr
 from diffusers import DiffusionPipeline
+from optimization import optimize_pipeline_ 
 
 # Load the pipeline once at startup
 print("Loading Z-Image-Turbo pipeline...")
@@ -11,6 +12,16 @@ pipe = DiffusionPipeline.from_pretrained(
     low_cpu_mem_usage=False,
 )
 pipe.to("cuda")
+pipe.transformer.set_attention_backend("_flash_3_hub")
+
+#warmup
+optimize_pipeline_(
+    pipe,
+    prompt="prompt",
+    num_inference_steps=1,
+    guidance_scale=0.0
+)
+
 print("Pipeline loaded!")
 
 @spaces.GPU