Update app.py

app.py

@@ -10,8 +10,11 @@ from skyreelsinfer import TaskType
 from skyreelsinfer.skyreels_video_infer import SkyReelsVideoInfer
 from diffusers.utils import export_to_video
 from diffusers.utils import load_image
-import torch
-device="cuda:0" if torch.cuda.is_available() else "cpu" # Pass device to the constructor
+
+@spaces.GPU(duration=5)
+def init_gpu():
+    import torch
+    device="cuda:0" if torch.cuda.is_available() else "cpu" # Pass device to the constructor
 
 predictor = None
 task_type = None
@@ -79,6 +82,7 @@ def create_gradio_interface():
 if __name__ == "__main__":
     #import multiprocessing
     #multiprocessing.freeze_support()
+    init_gpu()
     init_predictor()
     demo = create_gradio_interface()
     demo.launch()