Update app.py

app.py

@@ -24,6 +24,50 @@ from midi_to_colab_audio import midi_to_colab_audio
 
 from huggingface_hub import hf_hub_download
 
+# =================================================================================================
+
+print('Loading model...')
+
+SEQ_LEN = 1802
+PAD_IDX = 771
+DEVICE = 'cuda' # 'cpu'
+
+# instantiate the model
+
+model = TransformerWrapper(
+    num_tokens = PAD_IDX+1,
+    max_seq_len = SEQ_LEN,
+    attn_layers = Decoder(dim = 1024, 
+                          depth = 8, 
+                          heads = 8, 
+                          rotary_pos_emb=True, 
+                          attn_flash = True
+                         )
+    )
+
+model = AutoregressiveWrapper(model, ignore_index = PAD_IDX)
+
+
+print('=' * 70)
+
+print('Loading model checkpoint...')
+
+model_checkpoint = hf_hub_download(repo_id='asigalov61/Score-2-Performance-Transformer', 
+                                   filename='Score_2_Performance_Transformer_Final_Small_Trained_Model_4496_steps_1.5185_loss_0.5589_acc.pth'
+                                  )
+
+model.load_state_dict(torch.load(model_checkpoint, map_location='cpu', weights_only=True))
+
+model = torch.compile(model, mode='max-autotune')
+
+dtype = torch.bfloat16
+
+ctx = torch.amp.autocast(device_type=DEVICE, dtype=dtype)
+
+print('=' * 70)
+print('Done!')
+print('=' * 70)
+
 # =================================================================================================
                        
 @spaces.GPU
@@ -66,45 +110,7 @@ def Convert_Score_to_Performance(input_midi,
     print('Model sampling top p value:', input_model_top_p)
     print('=' * 70)
 
-    #===============================================================================
-
-    print('Loading model...')
 
-    SEQ_LEN = 4096
-    PAD_IDX = 673
-    DEVICE = 'cuda' # 'cpu'
-
-    # instantiate the model
-
-    model = TransformerWrapper(
-        num_tokens = PAD_IDX+1,
-        max_seq_len = SEQ_LEN,
-        attn_layers = Decoder(dim = 1024, depth = 16, heads = 16, rotary_pos_emb=True, 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('Guided_Rpck_Music_Transformer_Trained_Model_12081_steps_0.4113_loss_0.8747_acc.pth',
-                   map_location=DEVICE))
-    print('=' * 70)
-
-    model.eval()
-
-    if DEVICE == 'cpu':
-        dtype = torch.bfloat16
-    else:
-        dtype = torch.bfloat16
-
-    ctx = torch.amp.autocast(device_type=DEVICE, dtype=dtype)
-
-    print('Done!')
-    print('=' * 70)
 
     #===============================================================================
 
@@ -204,6 +210,9 @@ def Convert_Score_to_Performance(input_midi,
     print('=' * 70)
     print('Generating...')
 
+    model.to(DEVICE)
+    model.eval()
+
     #==================================================================
 
     def generate_continuation(num_prime_tokens, num_gen_tokens):