Update app.py

app.py

@@ -5,6 +5,7 @@ from pyramid_dit import PyramidDiTForVideoGeneration
 from diffusers.utils import load_image, export_to_video
 from huggingface_hub import snapshot_download
 import os
+import spaces
 
 # Download and load the model
 model_path = 'pyramid_flow_model'
@@ -25,6 +26,7 @@ model.dit.to("cuda")
 model.text_encoder.to("cuda")
 model.vae.enable_tiling()
 
+@spaces.GPU(duration=120)
 def generate_video(prompt, height, width, duration, guidance_scale, video_guidance_scale):
     temp = 16 if duration == "5s" else 31