Update app.py

app.py

@@ -29,6 +29,8 @@ torch.backends.cuda.enable_math_sdp(True)
 torch.backends.cuda.enable_flash_sdp(True)
 torch.backends.cuda.enable_cudnn_sdp(True)
 
+from huggingface_hub import hf_hub_download
+
 import TMIDIX
 
 from midi_to_colab_audio import midi_to_colab_audio
@@ -54,11 +56,11 @@ print('=' * 70)
 
 #==================================================================================
 
-MODEL_CHECKPOINT = 'Giant_Music_Transformer_Medium_Trained_Model_42174_steps_0.5211_loss_0.8542_acc.pth'
+MODEL_CHECKPOINT = 'Monster_Piano_Transformer_Velocity_Trained_Model_59896_steps_0.9055_loss_0.735_acc.pth'
 
 SOUDFONT_PATH = 'SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2'
 
-NUM_OUT_BATCHES = 8
+NUM_OUT_BATCHES = 16
 
 PREVIEW_LENGTH = 120 # in tokens
 
@@ -73,14 +75,14 @@ dtype = 'bfloat16'
 ptdtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
 ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)
 
-SEQ_LEN = 8192
-PAD_IDX = 19463
+SEQ_LEN = 2048
+PAD_IDX = 512
 
 model = TransformerWrapper(
         num_tokens = PAD_IDX+1,
         max_seq_len = SEQ_LEN,
         attn_layers = Decoder(dim = 2048,
-                              depth = 8,
+                              depth = 4,
                               heads = 32,
                               rotary_pos_emb = True,
                               attn_flash = True
@@ -92,7 +94,9 @@ model = AutoregressiveWrapper(model, ignore_index=PAD_IDX, pad_value=PAD_IDX)
 print('=' * 70)
 print('Loading model checkpoint...')
 
-model.load_state_dict(torch.load(MODEL_CHECKPOINT, map_location='cpu'))
+model_checkpoint = hf_hub_download(repo_id='asigalov61/Monster-Piano-Transformer', filename=MODEL_CHECKPOINT_VEL)
+
+model.load_state_dict(torch.load(model_checkpoint, map_location='cpu'))
 
 model = torch.compile(model, mode='max-autotune')