Update app.py

app.py

@@ -2,6 +2,8 @@ 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...")
@@ -10,8 +12,23 @@ pipe = DiffusionPipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=False,
 )
+
+
+pipe.transformer.set_attention_backend(attention_backend)
+
+if enable_compile:
+    print("Compiling transformer...")
+    pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs", fullgraph=False)
+
 pipe.to("cuda")
-print("Pipeline loaded!")
+
+optimize_pipeline_(
+    pipe,
+    prompt="prompt",
+    num_inference_steps=1,
+    guidance_scale=0.0
+)
+
 
 @spaces.GPU
 def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed):