Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import os
 import sys
+import torch
+import spaces
 import shutil
 import uuid
 import subprocess
@@ -15,6 +17,13 @@ snapshot_download(
     local_dir = "./checkpoints"  
 )
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+if torch.cuda.is_available():
+    torch_dtype = torch.float16
+else:
+    torch_dtype = torch.float32
+
 import tempfile
 from moviepy.editor import VideoFileClip
 from pydub import AudioSegment
@@ -78,7 +87,7 @@ from diffusers.utils.import_utils import is_xformers_available
 from accelerate.utils import set_seed
 from latentsync.whisper.audio2feature import Audio2Feature
 
-
+@spaces.GPU
 def main(video_path, audio_path, progress=gr.Progress(track_tqdm=True)):
     inference_ckpt_path = "checkpoints/latentsync_unet.pt"
     unet_config_path = "configs/unet/second_stage.yaml"
@@ -118,7 +127,7 @@ def main(video_path, audio_path, progress=gr.Progress(track_tqdm=True)):
     unet, _ = UNet3DConditionModel.from_pretrained(
         OmegaConf.to_container(config.model),
         inference_ckpt_path,  # load checkpoint
-        device="cpu",
+        device=device,
     )
 
     unet = unet.to(dtype=torch.float16)