Update app.py

app.py

@@ -1,4 +1,3 @@
-import argparse
 import glob
 import json
 import os.path
@@ -11,6 +10,7 @@ import torch
 import torch.nn.functional as F
 
 import gradio as gr
+import spaces
 
 from x_transformer import *
 import tqdm
@@ -24,7 +24,7 @@ in_space = os.getenv("SYSTEM") == "spaces"
 
 # =================================================================================================
 
-@torch.no_grad()
+@spaces.GPU
 def GenerateMIDI(num_tok, idrums, iinstr):
     print('=' * 70)
     print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
@@ -83,6 +83,38 @@ def GenerateMIDI(num_tok, idrums, iinstr):
 
     yield output, None, None, [create_msg("visualizer_clear", None)]
 
+
+    print('Loading model...')
+
+    SEQ_LEN = 2048
+
+    # instantiate the model
+
+    model = TransformerWrapper(
+        num_tokens=3088,
+        max_seq_len=SEQ_LEN,
+        attn_layers=Decoder(dim=1024, depth=16, heads=8)
+    )
+
+    model = AutoregressiveWrapper(model)
+
+    model = torch.nn.DataParallel(model)
+
+    model.cpu()
+    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'))
+    print('=' * 70)
+
+    model.eval()
+
+    print('Done!')
+    print('=' * 70)
+
     outy = start_tokens
     
     ctime = 0
@@ -201,42 +233,6 @@ if __name__ == "__main__":
     print('App start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
     print('=' * 70)
     
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--share", action="store_true", default=False, help="share gradio app")
-    parser.add_argument("--port", type=int, default=7860, help="gradio server port")
-    opt = parser.parse_args()
-    
-    print('Loading model...')
-
-    SEQ_LEN = 2048
-
-    # instantiate the model
-
-    model = TransformerWrapper(
-        num_tokens=3088,
-        max_seq_len=SEQ_LEN,
-        attn_layers=Decoder(dim=1024, depth=16, heads=8)
-    )
-
-    model = AutoregressiveWrapper(model)
-
-    model = torch.nn.DataParallel(model)
-
-    model.cpu()
-    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'))
-    print('=' * 70)
-
-    model.eval()
-
-    print('Done!')
-    print('=' * 70)
-
     load_javascript()
     app = gr.Blocks()
     with app:
@@ -267,4 +263,4 @@ if __name__ == "__main__":
                                   [output_midi_seq, output_midi, output_audio, js_msg])
         interrupt_btn.click(cancel_run, output_midi_seq, [output_midi, output_audio, js_msg],
                             cancels=run_event, queue=False)
-        app.queue(concurrency_count=1).launch(server_port=opt.port, share=opt.share, inbrowser=True)
+        app.queue().launch()