Update app.py

app.py

@@ -62,44 +62,6 @@ def GenerateDrums(input_midi, input_num_tokens):
     print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
     start_time = reqtime.time()
 
-    print('Loading model...')
-
-    SEQ_LEN = 8192 # Models seq len
-    PAD_IDX = 393 # Models pad index
-    DEVICE = 'cuda' # 'cuda'
-
-    # instantiate the model
-
-    model = TransformerWrapper(
-        num_tokens = PAD_IDX+1,
-        max_seq_len = SEQ_LEN,
-        attn_layers = Decoder(dim = 1024, depth = 4, heads = 8, attn_flash = True)
-        )
-    
-    model = AutoregressiveWrapper(model, ignore_index = PAD_IDX)
-
-    model.to(DEVICE)
-    print('=' * 70)
-
-    print('Loading model checkpoint...')
-
-    model.load_state_dict(
-        torch.load('Ultimate_Drums_Transformer_Small_Trained_Model_VER3_VEL_11222_steps_0.5749_loss_0.8085_acc.pth',
-                   map_location=DEVICE))
-    print('=' * 70)
-
-    model.eval()
-
-    if DEVICE == 'cpu':
-        dtype = torch.bfloat16
-    else:
-        dtype = torch.float16
-
-    ctx = torch.amp.autocast(device_type=DEVICE, dtype=dtype)
-
-    print('Done!')
-    print('=' * 70)
-
     fn = os.path.basename(input_midi.name)
     fn1 = fn.split('.')[0]
 
@@ -277,6 +239,44 @@ if __name__ == "__main__":
     opt = parser.parse_args()
 
     soundfont = "SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2"
+
+    print('Loading model...')
+
+    SEQ_LEN = 8192 # Models seq len
+    PAD_IDX = 393 # Models pad index
+    DEVICE = 'cuda' # 'cuda'
+
+    # instantiate the model
+
+    model = TransformerWrapper(
+        num_tokens = PAD_IDX+1,
+        max_seq_len = SEQ_LEN,
+        attn_layers = Decoder(dim = 1024, depth = 4, heads = 8, attn_flash = True)
+        )
+    
+    model = AutoregressiveWrapper(model, ignore_index = PAD_IDX)
+
+    model.to(DEVICE)
+    print('=' * 70)
+
+    print('Loading model checkpoint...')
+
+    model.load_state_dict(
+        torch.load('Ultimate_Drums_Transformer_Small_Trained_Model_VER3_VEL_11222_steps_0.5749_loss_0.8085_acc.pth',
+                   map_location=DEVICE))
+    print('=' * 70)
+
+    model.eval()
+
+    if DEVICE == 'cpu':
+        dtype = torch.bfloat16
+    else:
+        dtype = torch.float16
+
+    ctx = torch.amp.autocast(device_type=DEVICE, dtype=dtype)
+
+    print('Done!')
+    print('=' * 70)
     
     app = gr.Blocks()
     with app:
@@ -310,20 +310,20 @@ if __name__ == "__main__":
                                   [output_midi_title, output_midi_summary, output_midi, output_audio, output_plot])
 
         gr.Examples(
-            [["Ultimate-Drums-Transformer-Melody-Seed-1.mid", 16], 
-             ["Ultimate-Drums-Transformer-Melody-Seed-2.mid", 16], 
-             ["Ultimate-Drums-Transformer-Melody-Seed-3.mid", 16],
-             ["Ultimate-Drums-Transformer-Melody-Seed-4.mid", 16],
-             ["Ultimate-Drums-Transformer-Melody-Seed-5.mid", 16],
-             ["Ultimate-Drums-Transformer-Melody-Seed-6.mid", 16],
-             ["Ultimate-Drums-Transformer-MI-Seed-1.mid", 16],
-             ["Ultimate-Drums-Transformer-MI-Seed-2.mid", 16],
-             ["Ultimate-Drums-Transformer-MI-Seed-3.mid", 16],
-             ["Ultimate-Drums-Transformer-MI-Seed-4.mid", 16]],
+            [["Ultimate-Drums-Transformer-Melody-Seed-1.mid", 128], 
+             ["Ultimate-Drums-Transformer-Melody-Seed-2.mid", 128], 
+             ["Ultimate-Drums-Transformer-Melody-Seed-3.mid", 128],
+             ["Ultimate-Drums-Transformer-Melody-Seed-4.mid", 128],
+             ["Ultimate-Drums-Transformer-Melody-Seed-5.mid", 128],
+             ["Ultimate-Drums-Transformer-Melody-Seed-6.mid", 128],
+             ["Ultimate-Drums-Transformer-MI-Seed-1.mid", 128],
+             ["Ultimate-Drums-Transformer-MI-Seed-2.mid", 128],
+             ["Ultimate-Drums-Transformer-MI-Seed-3.mid", 128],
+             ["Ultimate-Drums-Transformer-MI-Seed-4.mid", 128]],
             [input_midi, input_num_tokens],
             [output_midi_title, output_midi_summary, output_midi, output_audio, output_plot],
             GenerateDrums,
-            cache_examples=False,
+            cache_examples=True,
         )
         
         app.queue().launch()