Update app.py

app.py

@@ -31,19 +31,24 @@ def GenerateAccompaniment(input_midi, input_num_tokens, input_conditioning_type,
 
     print('Loading model...')
 
-    SEQ_LEN = 8192 # Models seq len
-    PAD_IDX = 707 # Models pad index
-    DEVICE = 'cuda' # 'cuda'
+    SEQ_LEN = 8192
+    PAD_IDX = 2239
+    DEVICE = 'cuda' # 'cpu'
 
     # instantiate the model
 
     model = TransformerWrapper(
         num_tokens = PAD_IDX+1,
         max_seq_len = SEQ_LEN,
-        attn_layers = Decoder(dim = 2048, depth = 4, heads = 16, attn_flash = True)
+        attn_layers = Decoder(dim = 2048, 
+                              depth = 8, 
+                              heads = 32, 
+                              rotary_pos_emb = True,  
+                              attn_flash = True
+                             )
         )
     
-    model = AutoregressiveWrapper(model, ignore_index = PAD_IDX)
+    model = AutoregressiveWrapper(model, ignore_index = PAD_IDX, pad_value=PAD_IDX)
 
     model.to(DEVICE)
     print('=' * 70)
@@ -51,7 +56,7 @@ def GenerateAccompaniment(input_midi, input_num_tokens, input_conditioning_type,
     print('Loading model checkpoint...')
 
     model.load_state_dict(
-        torch.load('Chords_Progressions_Transformer_Small_2048_Trained_Model_12947_steps_0.9316_loss_0.7386_acc.pth',
+        torch.load('Ultimate_Chords_Progressions_Transformer_Trained_Model_LAX_5858_steps_0.4506_loss_0.8724_acc.pth',
                    map_location=DEVICE))
     print('=' * 70)
 
@@ -62,7 +67,7 @@ def GenerateAccompaniment(input_midi, input_num_tokens, input_conditioning_type,
     else:
         dtype = torch.float16
 
-    ctx = torch.amp.autocast(device_type=DEVICE, dtype=dtype)
+    ctx = torch.amp.autocast(device_type='cuda', dtype=dtype)
 
     print('Done!')
     print('=' * 70)