Update app.py

app.py

@@ -56,7 +56,10 @@ print('=' * 70)
 
 #==================================================================================
 
-MODEL_CHECKPOINT_VEL = 'Monster_Piano_Transformer_Velocity_Trained_Model_59896_steps_0.9055_loss_0.735_acc.pth'
+MODEL_CHECKPOINTS = {
+                    'with velocity': 'Monster_Piano_Transformer_Velocity_Trained_Model_59896_steps_0.9055_loss_0.735_acc.pth',
+                    'without velocity': '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'
 
@@ -66,45 +69,47 @@ PREVIEW_LENGTH = 120 # in tokens
 
 #==================================================================================
 
-print('=' * 70)
-print('Instantiating model...')
-
-device_type = 'cuda'
-dtype = 'bfloat16'
-
-ptdtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
-ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)
-
-SEQ_LEN = 2048
-PAD_IDX = 512
-
-model = TransformerWrapper(
-        num_tokens = PAD_IDX+1,
-        max_seq_len = SEQ_LEN,
-        attn_layers = Decoder(dim = 2048,
-                              depth = 4,
-                              heads = 32,
-                              rotary_pos_emb = True,
-                              attn_flash = True
-                              )
-)
-
-model = AutoregressiveWrapper(model, ignore_index=PAD_IDX, pad_value=PAD_IDX)
-
-print('=' * 70)
-print('Loading model checkpoint...')
-
-model_checkpoint = hf_hub_download(repo_id='asigalov61/Monster-Piano-Transformer', filename=MODEL_CHECKPOINT_VEL)
+def load_model(model_selector):
 
-model.load_state_dict(torch.load(model_checkpoint, map_location='cpu'))
-
-model = torch.compile(model, mode='max-autotune')
-
-print('=' * 70)
-print('Done!')
-print('=' * 70)
-print('Model will use', dtype, 'precision...')
-print('=' * 70)
+    print('=' * 70)
+    print('Instantiating model...')
+    
+    device_type = 'cuda'
+    dtype = 'bfloat16'
+    
+    ptdtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
+    ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+    
+    SEQ_LEN = 2048
+    PAD_IDX = 512
+    
+    model = TransformerWrapper(
+            num_tokens = PAD_IDX+1,
+            max_seq_len = SEQ_LEN,
+            attn_layers = Decoder(dim = 2048,
+                                  depth = 4,
+                                  heads = 32,
+                                  rotary_pos_emb = True,
+                                  attn_flash = True
+                                  )
+    )
+    
+    model = AutoregressiveWrapper(model, ignore_index=PAD_IDX, pad_value=PAD_IDX)
+    
+    print('=' * 70)
+    print('Loading model checkpoint...')      
+    
+    model_checkpoint = hf_hub_download(repo_id='asigalov61/Monster-Piano-Transformer', filename=MODEL_CHECKPOINTS[model_selector])
+    
+    model.load_state_dict(torch.load(model_checkpoint, map_location='cpu'))
+    
+    model = torch.compile(model, mode='max-autotune')
+    
+    print('=' * 70)
+    print('Done!')
+    print('=' * 70)
+    print('Model will use', dtype, 'precision...')
+    print('=' * 70)
 
 #==================================================================================
 
@@ -304,7 +309,8 @@ def generate_callback_wrapper(input_midi,
                               # model_sampling_top_p,
                               final_composition, 
                               generated_batches, 
-                              block_lines
+                              block_lines,
+                              model_selector
                              ):
 
     print('=' * 70)
@@ -317,6 +323,10 @@ def generate_callback_wrapper(input_midi,
             fn1 = fn.split('.')[0]
             print('Input file name:', fn)
 
+    print('Selected model type:', model_selector)
+
+    load_model(model_selector)
+    
     print('Num prime tokens:', num_prime_tokens)
     print('Num gen tokens:', num_gen_tokens)
     print('Num mem tokens:', num_mem_tokens)
@@ -481,6 +491,13 @@ with gr.Blocks() as demo:
                             [final_composition, generated_batches, block_lines])
     
     gr.Markdown("## Generate")
+
+    model_selector = gr.gr.Dropdown(["with velocity", 
+                                     "Without velocity"
+                                    ], 
+                                    label="Select model", 
+                                    info="Select desired Monster Piano Transformer model"
+                                   )
     
     num_prime_tokens = gr.Slider(15, 1024, value=1024, step=1, label="Number of prime tokens")
     num_gen_tokens = gr.Slider(15, 1024, value=1024, step=1, label="Number of tokens to generate")
@@ -511,7 +528,8 @@ with gr.Blocks() as demo:
                         # model_sampling_top_p,
                         final_composition,
                         generated_batches,
-                        block_lines                        
+                        block_lines,
+                        model_selector
                        ], 
                        outputs
                       )