Update app.py

app.py

@@ -93,21 +93,21 @@ def GenerateMIDI(num_tok, idrums, iinstr):
     model = TransformerWrapper(
         num_tokens=3088,
         max_seq_len=SEQ_LEN,
-        attn_layers=Decoder(dim=1024, depth=16, heads=8)
+        attn_layers=Decoder(dim=1024, depth=16, heads=8, attn_flash=True)
     )
 
     model = AutoregressiveWrapper(model)
 
     model = torch.nn.DataParallel(model)
 
-    model.cpu()
+    model.cuda()
     print('=' * 70)
 
     print('Loading model checkpoint...')
 
     model.load_state_dict(
         torch.load('Allegro_Music_Transformer_Tiny_Trained_Model_80000_steps_0.9457_loss_0.7443_acc.pth',
-                   map_location='cpu'))
+                   map_location='cuda'))
     print('=' * 70)
 
     model.eval()
@@ -125,13 +125,14 @@ def GenerateMIDI(num_tok, idrums, iinstr):
 
     for i in range(max(1, min(512, num_tok))):
         
-        inp = torch.LongTensor([outy]).cpu()
+        inp = torch.LongTensor([outy]).cuda()
 
-        out = model.module.generate(inp,
-                                    1,
-                                    temperature=0.9,
-                                    return_prime=False,
-                                    verbose=False)
+        with torch.amp.autocast(device_type='cuda', dtype=torch.float16)
+            out = model.module.generate(inp,
+                                        1,
+                                        temperature=0.9,
+                                        return_prime=False,
+                                        verbose=False)
 
         out0 = out[0].tolist()
         outy.extend(out0)
@@ -253,7 +254,7 @@ if __name__ == "__main__":
         run_btn = gr.Button("generate", variant="primary")
         interrupt_btn = gr.Button("interrupt")
 
-        output_midi_seq = gr.JSON()
+        output_midi_seq = gr.HTML()
         output_midi_visualizer = gr.HTML(elem_id="midi_visualizer_container")
         output_audio = gr.Audio(label="output audio", format="mp3", elem_id="midi_audio")
         output_midi = gr.File(label="output midi", file_types=[".mid"])