Update app.py

app.py

@@ -58,7 +58,8 @@ else:
         )
         pipe.save_pretrained("./checkpoints", variant="fp16")
 
-pipe.to("cuda")
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+# pipe.to(device)
 # pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
 #pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
 
@@ -77,6 +78,8 @@ def generate_video(
     output_folder: str = "outputs",
 ):
     global pipe
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pipe.to(device)
     # note julian: normally we should resize input images, but normally they are already in 1024x576, so..
 
     # also, I would like to experiment with vertical videos, and 1024x512 videos