Update app.py

app.py

@@ -7,6 +7,8 @@ from huggingface_hub import snapshot_download
 from pyramid_dit import PyramidDiTForVideoGeneration
 from diffusers.utils import export_to_video
 
+import spaces 
+
 # Constants
 MODEL_PATH = "pyramid-flow-model"
 MODEL_REPO = "rain1011/pyramid-flow-sd3"
@@ -35,6 +37,7 @@ def load_model():
 model = load_model()
 
 # Text-to-video generation function
+@spaces.GPU(duration=240)
 def generate_video(prompt, duration, guidance_scale, video_guidance_scale):
  temp = int(duration * 2.4) # Convert seconds to temp value (assuming 24 FPS)
  torch_dtype = torch.bfloat16 if MODEL_DTYPE == "bf16" else torch.float32
@@ -58,6 +61,7 @@ def generate_video(prompt, duration, guidance_scale, video_guidance_scale):
  return output_path
 
 # Image-to-video generation function
+@spaces.GPU(duration=240)
 def generate_video_from_image(image, prompt, duration, video_guidance_scale):
  temp = int(duration * 2.4) # Convert seconds to temp value (assuming 24 FPS)
  torch_dtype = torch.bfloat16 if MODEL_DTYPE == "bf16" else torch.float32