Update app.py

app.py

@@ -181,28 +181,30 @@ def benchmark(Batch, Heads, N, D_head, provider):
     return tflops(ms), tflops(max_ms), tflops(min_ms)
 
 
-
-@spaces.GPU(duration=180) # Triton benchmarks can take a minute
-def run_benchmark():
-    # Ensure we are in a clean directory for images
-    output_dir = "./plots"
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
-        
-    # Run the triton benchmark
-    # This will generate several .png files in the save_path
-    bench_flash_attention.run(save_path=output_dir, print_data=True)
+# 2. --- WRAP THE RUN COMMAND IN A DECORATED FUNCTION ---
+@spaces.GPU(duration=150) # High duration for Triton compilation + Benchmarking
+def start_benchmarking():
+    # Triton saves plots to the current directory by default
+    save_path = "./plots"
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
+    
+    # Run your original benchmark function
+    # Note: Ensure bench_flash_attention is defined above this
+    bench_flash_attention.run(save_path=save_path, print_data=True)
     
-    # Collect the generated images
-    images = [os.path.join(output_dir, f) for f in os.listdir(output_dir) if f.endswith('.png')]
+    # Find the .png files generated by Triton
+    images = [os.path.join(save_path, f) for f in os.listdir(save_path) if f.endswith('.png')]
     return images
 
-# Gradio Interface
+# 3. --- CREATE THE GRADIO GUI TO KEEP THE SPACE ALIVE ---
 with gr.Blocks() as demo:
-    gr.Markdown("# Triton Fused Attention Benchmark on ZeroGPU")
-    run_btn = gr.Button("Run Benchmark")
-    out_gallery = gr.Gallery(label="Performance Plots")
-
-    run_btn.click(fn=run_benchmark, outputs=out_gallery)
+    gr.Markdown("# Triton Attention Benchmark")
+    gr.Markdown("Click the button below to trigger the ZeroGPU and run the Triton benchmark.")
+    
+    run_btn = gr.Button("Run Benchmark (H100/H200)", variant="primary")
+    plot_gallery = gr.Gallery(label="Generated Performance Plots", columns=2)
+    
+    run_btn.click(fn=start_benchmarking, outputs=plot_gallery)
 
 demo.launch()