Update src/music2cocktailrep/training/latent_translation/setup_trained_model.py

src/music2cocktailrep/training/latent_translation/setup_trained_model.py

@@ -10,19 +10,19 @@ from shutil import copy
 import hashlib
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-# TOKEN = os.environ['token']
+TOKEN = os.environ['token']
 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)
+    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)
+    copy(downloaded_path, model_path)
     with open(checkpoint_path + 'params.json', 'r') as f:
         params = json.load(f)