Add model download functionality

app.py

@@ -12,6 +12,11 @@ from transformers import T5Tokenizer
 from transformer_model import Transformer
 from miditok import REMI, TokenizerConfig
 from pathlib import Path
+from huggingface_hub import hf_hub_download
+
+repo_id = "amaai-lab/text2midi"
+# Download the model.bin file
+model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
 
 
 def save_wav(filepath):
@@ -81,7 +86,7 @@ def generate_midi(caption, temperature=0.9, max_len=500):
     vocab_size = len(r_tokenizer)
     print("Vocab size: ", vocab_size)
     model = Transformer(vocab_size, 768, 8, 2048, 18, 1024, False, 8, device=device)
-    model_path = os.path.join("amaai-lab/text2midi", "pytorch_model.bin")
+    # model_path = os.path.join("amaai-lab/text2midi", "pytorch_model.bin")
     model.load_state_dict(torch.load(model_path, map_location=device))
     model.eval()
     tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")

requirements.txt

@@ -5,4 +5,5 @@ transformers
 st-moe-pytorch
 jsonlines
 miditok==3.0.3
-sentencepiece
+sentencepiece
+huggingface_hub