Update app.py

app.py

@@ -6,8 +6,10 @@ torch.backends.cudnn.allow_tf32 = True
 import os
 import spaces
 from diffusers.models import AutoencoderKL
-from models import FLAV_models
-
+from models import FLAV
+from huggingface_hub import hf_hub_download
+import torch
+    
 from diffusion.rectified_flow import RectifiedFlow
 
 from diffusers.training_utils import EMAModel
@@ -27,7 +29,7 @@ AUDIO_T_PER_FRAME = 1600 // 160
 vae = None
 model = None
 vocoder = None
-audio_scale = 3.50
+audio_scale = 3.5009668382765917
 
 
 def setup_models():
@@ -37,25 +39,13 @@ def setup_models():
     vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-ema")
     vae.eval()
     
-    model = FLAV_models["FLAV-B/1"](
-        latent_size= 256//8,
-        in_channels = 4,
-        num_classes = 0,
-        predict_frames = 10,
-        causal_attn = True,
-    )
-
-    ckpt_path = hf_hub_download(repo_id="MaverickAlex/R-FLAV", filename="aist-ema.pth")
-
-    state_dict = torch.load(ckpt_path, map_location="cpu")
 
-    ema = EMAModel(model.parameters())
-    ema.load_state_dict(state_dict)
-    ema.copy_to(model.parameters())
+    model_ckpt = "MaverickAlex/R-FLAV-B-1-AIST" # MaverickAlex/R-FLAV-B-1-LS
 
-    hf_hub_download(repo_id="MaverickAlex/R-FLAV", filename="vocoder-aist/config.json")
-    vocoder_path = hf_hub_download(repo_id="MaverickAlex/R-FLAV", filename="vocoder-aist/vocoder.pt")
+    model = FLAV.from_pretrained(model_ckpt)
 
+    hf_hub_download(repo_id=model_ckpt, filename="vocoder/config.json")
+    vocoder_path = hf_hub_download(repo_id=model_ckpt, filename="vocoder/vocoder.pt")
     vocoder_path = vocoder_path.replace("vocoder.pt", "")
     vocoder = Generator.from_pretrained(vocoder_path)