Update app.py

app.py

@@ -68,6 +68,57 @@ 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 = 4096
+
+if model_selector == 'with velocity - 3 epochs':
+    PAD_IDX = 512
+
+else:
+    PAD_IDX = 384
+
+model = TransformerWrapper(
+        num_tokens = PAD_IDX+1,
+        max_seq_len = SEQ_LEN,
+        attn_layers = Decoder(dim = 2048,
+                              depth = 8,
+                              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/Godzilla-Piano-Transformer', 
+                                   filename='Godzilla_Piano_Transformer_No_Velocity_Trained_Model_14903_steps_0.4874_loss_0.8571_acc.pth')
+
+model.load_state_dict(torch.load(model_checkpoint, map_location='cuda', weights_only=True))
+
+model = torch.compile(model, mode='max-autotune')
+
+print('=' * 70)
+print('Done!')
+print('=' * 70)
+print('Model will use', dtype, 'precision...')
+print('=' * 70)
+
+model.cuda()
+model.eval()
+
+#==================================================================================
+
 def load_model(model_selector):
 
     return [[], []]
@@ -176,55 +227,6 @@ def generate_music(prime,
     else:
         inputs = prime[-num_mem_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 = 4096
-
-    if model_selector == 'with velocity - 3 epochs':
-        PAD_IDX = 512
-
-    else:
-        PAD_IDX = 384
-    
-    model = TransformerWrapper(
-            num_tokens = PAD_IDX+1,
-            max_seq_len = SEQ_LEN,
-            attn_layers = Decoder(dim = 2048,
-                                  depth = 8,
-                                  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/Godzilla-Piano-Transformer', 
-                                       filename='Godzilla_Piano_Transformer_No_Velocity_Trained_Model_14903_steps_0.4874_loss_0.8571_acc.pth')
-    
-    model.load_state_dict(torch.load(model_checkpoint, map_location='cuda', weights_only=True))
-    
-    model = torch.compile(model, mode='max-autotune')
-    
-    print('=' * 70)
-    print('Done!')
-    print('=' * 70)
-    print('Model will use', dtype, 'precision...')
-    print('=' * 70)
-    
-    model.cuda()
-    model.eval()
-
     print('Generating...')
     
     inp = [inputs] * num_gen_batches