Update app.py

app.py

@@ -7,7 +7,7 @@ import torch.nn.functional as F
 
 import gradio as gr
 
-import onnxruntime as rt
+from x_transformer import *
 import tqdm
 
 from midi_synthesizer import synthesis
@@ -32,58 +32,24 @@ def GenerateMIDI(idrums, iinstr, progress=gr.Progress()):
 
     start_tokens = [3087, drums, 3075+first_note_instrument_number]
 
-    seq_len = 512
-    max_seq_len = 2048
-    temperature = 0.9
-    verbose=False
-    return_prime=False
-    
-
-    out = torch.FloatTensor([start_tokens])
-
-    st = len(start_tokens)
-
-    if verbose:
-      print("Generating sequence of max length:", seq_len)
-
-    progress(0, desc="Starting...")
-    step = 0
-    
-    for i in progress.tqdm(range(seq_len)):
-
-        try:
-            x = out[:, -max_seq_len:]
+    print('Selected Improv sequence:')
+    print(start_tokens)
+    print('=' * 70)
     
-            torch_in = x.tolist()[0]
+    inp = [start_tokens] * number_of_batches_to_generate
     
-            logits = torch.FloatTensor(session.run(None, {'input': [torch_in]})[0])[:, -1]
-
-            probs = F.softmax(logits / temperature, dim=-1)
+    inp = torch.LongTensor(inp).cpu()
     
-            sample = torch.multinomial(probs, 1)
+    out = model.module.generate(inp,
+                          number_of_tokens_tp_generate,
+                          temperature=temperature,
+                          return_prime=False,
+                          verbose=True)
     
-            out = torch.cat((out, sample), dim=-1)
-
-            if step % 16 == 0:
-                print(step, '/', seq_len)
-
-            step += 1
-        
-            if step >= seq_len:
-                break
-            
-        except Exception as e:
-            print('Error', e)
-            break
-
-    if return_prime:
-      melody_chords_f = out[:, :]
-    
-    else:
-      melody_chords_f = out[:, st:]
-
-    melody_chords_f = melody_chords_f.tolist()[0]
+    melody_chords_f = out[0].tolist()
     
+    print('=' * 70)
+    print('Done!')
     print('=' * 70)
     print('Sample INTs', melody_chords_f[:12])
     print('=' * 70)
@@ -196,7 +162,31 @@ if __name__ == "__main__":
     opt = parser.parse_args()
 
     print('Loading model...')
-    session = rt.InferenceSession('Allegro_Music_Transformer_Small_Trained_Model_56000_steps_0.9399_loss_0.7374_acc.onnx', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+
+    SEQ_LEN = 2048
+    
+    # instantiate the model
+    
+    model = TransformerWrapper(
+        num_tokens = 3088,
+        max_seq_len = SEQ_LEN,
+        attn_layers = Decoder(dim = 1024, depth = 32, heads = 8)
+    )
+    
+    model = AutoregressiveWrapper(model)
+    
+    model = torch.nn.DataParallel(model)
+    
+    model.cpu()
+    print('=' * 70)
+    
+    print('Loading model checkpoint...')
+    
+    model.load_state_dict(torch.load(full_path_to_model_checkpoint, map_location='cpu'))
+    print('=' * 70)
+    
+    model.eval()    
+    
     print('Done!')
     
     app = gr.Blocks()