Update app.py

app.py

@@ -57,6 +57,46 @@ NUM_OUT_BATCHES = 8
 
 #==================================================================================
 
+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 = 8192
+PAD_IDX = 19463
+
+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_path = 'Giant_Music_Transformer_Medium_Trained_Model_10446_steps_0.7202_loss_0.8233_acc.pth'
+
+model.load_state_dict(torch.load(model_path, map_location='cpu'))
+
+print('=' * 70)
+print('Done!')
+print('=' * 70)
+print('Model will use', dtype, 'precision...')
+print('=' * 70)
+
+#==================================================================================
+
 def load_midi(input_midi):
 
     raw_score = TMIDIX.midi2single_track_ms_score(input_midi.name)
@@ -219,52 +259,8 @@ def generate_music(prime,
                     model_sampling_top_p
                   ):
 
-
-    #==============================================================================
-    
-    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 = 8192
-    PAD_IDX = 19463
-    
-    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_path = 'Giant_Music_Transformer_Medium_Trained_Model_10446_steps_0.7202_loss_0.8233_acc.pth'
-    
-    model.load_state_dict(torch.load(model_path))
-    
-    print('=' * 70)
-    
     model.cuda()
     model.eval()
-    
-    print('Done!')
-    print('=' * 70)
-    print('Model will use', dtype, 'precision...')
-    print('=' * 70)
-
-    #==============================================================================
 
     print('Generating...')