Update app.py

app.py

@@ -41,16 +41,14 @@ from huggingface_hub import hf_hub_download
 print('=' * 70)
 print('Loading models...')
 print('=' * 70)
-print('Loading small model...')
+print('Loading chords texturing model...')
 print('=' * 70)
 
 SEQ_LEN = 2048
 PAD_IDX = 721
-DEVICE = 'cuda' # 'cpu'
-
-# instantiate the model
+DEVICE = 'cuda'
 
-model = TransformerWrapper(
+tex_model = TransformerWrapper(
     num_tokens = PAD_IDX+1,
     max_seq_len = SEQ_LEN,
     attn_layers = Decoder(dim = 2048,
@@ -61,9 +59,9 @@ model = TransformerWrapper(
                          )
     )
 
-model = AutoregressiveWrapper(model, ignore_index=PAD_IDX)
+tex_model = AutoregressiveWrapper(tex_model, ignore_index=PAD_IDX)
 
-model.to(DEVICE)
+tex_model.to(DEVICE)
 print('=' * 70)
 
 print('Loading model checkpoint...')
@@ -73,9 +71,11 @@ checkpoint = hf_hub_download(
     filename='Chordified_Piano_Transformer_Texturing_Trained_Model_18092_steps_0.7058_loss_0.7977_acc.pth'
 )
 
-model.load_state_dict(torch.load(checkpoint, map_location=DEVICE, weights_only=True))
+tex_model.load_state_dict(torch.load(checkpoint, map_location=DEVICE, weights_only=True))
 
-model.eval()
+tex_model.eval()
+
+tex_model = torch.compile(tex_model)
 
 print('=' * 70)
 print('Done!')
@@ -83,10 +83,49 @@ print('=' * 70)
 
 # =================================================================================================
 
-if DEVICE == 'cpu':
-    dtype = torch.bfloat16
-else:
-    dtype = torch.bfloat16
+print('Loading chords progressions model...')
+print('=' * 70)
+
+SEQ_LEN = 380
+PAD_IDX = 324
+DEVICE = 'cuda' # 'cpu'
+
+# instantiate the model
+
+prg_model = TransformerWrapper(
+    num_tokens = PAD_IDX+1,
+    max_seq_len = SEQ_LEN,
+    attn_layers = Decoder(dim = 2048,
+                          depth = 6,
+                          heads = 16,
+                          rotary_pos_emb = True,
+                          attn_flash = True
+                         )
+    )
+
+prg_model = AutoregressiveWrapper(prg_model, ignore_index=PAD_IDX)
+
+prg_model.to(DEVICE)
+print('=' * 70)
+
+print('Loading model checkpoint...')
+
+checkpoint = hf_hub_download(
+    repo_id='asigalov61/Chordified-Piano-Transformer',
+    filename='Chordified_Piano_Transformer_Texturing_Trained_Model_18092_steps_0.7058_loss_0.7977_acc.pth'
+)
+
+prg_model.load_state_dict(torch.load(checkpoint, map_location=DEVICE, weights_only=True))
+
+prg_model.eval()
+
+prg_model = torch.compile(prg_model)
+
+print('=' * 70)
+
+# =================================================================================================
+
+dtype = torch.bfloat16
 
 ctx = torch.amp.autocast(device_type=DEVICE, dtype=dtype)