Add spaces support

app.py

@@ -4,7 +4,7 @@ torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 
 import os
-
+import spaces
 from diffusers.models import AutoencoderKL
 from models import FLAV_models
 
@@ -35,6 +35,7 @@ def setup_models():
     
     device = "cuda"
     vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-ema")
+    vae.eval()
     
     model = FLAV_models["FLAV-B/1"](
         latent_size= 256//8,
@@ -46,7 +47,7 @@ def setup_models():
 
     ckpt_path = hf_hub_download(repo_id="MaverickAlex/R-FLAV", filename="aist-ema.pth")
 
-    state_dict = torch.load(ckpt_path)
+    state_dict = torch.load(ckpt_path, map_location="cpu")
 
     ema = EMAModel(model.parameters())
     ema.load_state_dict(state_dict)
@@ -63,7 +64,7 @@ def setup_models():
     vocoder.to(device)
     
 
-
+@spaces.GPU
 def generate_video(num_frames=10, steps=2, seed=42):
     global vae, model, vocoder
     # Setup device