Update app.py

app.py

@@ -82,6 +82,7 @@
 import gradio as gr
 import torch
 from diffusers import FluxPipeline
+import spaces
 from torchao.quantization import swap_conv2d_1x1_to_linear, apply_dynamic_quant
 
 # Step 1: Enable PyTorch 2-specific optimizations
@@ -123,6 +124,7 @@ pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True)
 pipe.vae.decode = torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True)
 
 # Image generation function
+@spaces.GPU(duration=120)
 def generate_image(prompt, guidance_scale, num_inference_steps):
     # Generate the image using the optimized pipeline
     image = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images[0]