Update app.py

app.py

@@ -334,21 +334,26 @@ def CompareMIDIs(input_src_midi, input_trg_midi, input_sampling_resolution, inpu
     model = TransformerWrapper(
         num_tokens = PAD_IDX+1,
         max_seq_len = SEQ_LEN,
-        attn_layers = Decoder(dim = 1024, depth = 32, heads = 32, attn_flash = True)
+        attn_layers = Decoder(dim = 2048, depth = 8, heads = 32, attn_flash = True)
         )
     
     model = AutoregressiveWrapper(model, ignore_index = PAD_IDX)
 
-    model.to(DEVICE)
     print('=' * 70)
 
     print('Loading model checkpoint...')
 
-    model.load_state_dict(
-        torch.load('Giant_Music_Transformer_Large_Trained_Model_36074_steps_0.3067_loss_0.927_acc.pth',
-                   map_location=DEVICE))
+    model_checkpoint = hf_hub_download(repo_id='asigalov61/Giant-Music-Transformer', 
+                                       filename='Giant_Music_Transformer_Medium_Trained_Model_42174_steps_0.5211_loss_0.8542_acc.pth'
+                                      )
+    
+    model.load_state_dict(torch.load(kar_model_checkpoint, map_location='cpu', weights_only=True))
+    
+    model = torch.compile(model, mode='max-autotune')
+    
     print('=' * 70)
-
+    
+    model.to(DEVICE)
     model.eval()
 
     if DEVICE == 'cpu':