integrate model download

app.py

@@ -18,7 +18,7 @@ from PIL import Image
 
 
 st.set_page_config(
-    page_title="TastyPiano",
+    page_title="Test",
     page_icon="🎹",
 )
 

src/music2cocktailrep/training/latent_translation/setup_trained_model.py

@@ -5,11 +5,22 @@ from src.music2cocktailrep.training.latent_translation.vae_model import get_gml_
 from src.music.config import TRANSLATION_VAE_CHKP_PATH
 from src.cocktails.utilities.cocktail_utilities import get_bunch_of_rep_keys
 import os
+from huggingface_hub import hf_hub_download
+from shutil import copy
 
-
+TOKEN = os.environ['secret']
 rep_keys = get_bunch_of_rep_keys()['custom']
 
 def setup_trained_model(checkpoint_path=TRANSLATION_VAE_CHKP_PATH):
+    # download translation model
+    repo_id = "ccolas/translation_vae"
+    filename = "checkpoints_best_eval_old.save"
+    downloaded_path = hf_hub_download(repo_id=repo_id,
+                                      filename=filename,
+                                      repo_type='model',
+                                      use_auth_token=TOKEN)
+    model_path = checkpoint_path + 'checkpoints_best_eval.save'
+    copy(downloaded_path, model_path)
     with open(checkpoint_path + 'params.json', 'r') as f:
         params = json.load(f)
 
@@ -28,7 +39,6 @@ def setup_trained_model(checkpoint_path=TRANSLATION_VAE_CHKP_PATH):
     def normalize_music_input(input):
         return (input - stats_music[0]) / stats_music[1]
 
-    model_path = checkpoint_path + 'checkpoints_best_eval.save'
     model.load_state_dict(torch.load(model_path))
     model.eval()