Update app.py

app.py

@@ -2,6 +2,9 @@ import argparse
 import glob
 import os.path
 
+import torch
+import torch.nn.functional as F
+
 import gradio as gr
 import numpy as np
 import onnxruntime as rt
@@ -13,7 +16,54 @@ import TMIDIX
 
 in_space = os.getenv("SYSTEM") == "spaces"
 
-providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+providers = ['CPUExecutionProvider']
+
+#=================================================================================================
+
+def generate(
+            start_tokens,
+            seq_len,
+            max_seq_len = 2048,
+            temperature = 0.9,
+            verbose=True,
+            return_prime=False,
+            ):
+
+    out = torch.LongTensor([start_tokens])
+
+    st = len(start_tokens)
+
+    if verbose:
+      print("Generating sequence of max length:", seq_len)
+
+    for s in range(seq_len):
+        x = out[:, -max_seq_len:]
+
+        torch_in = x.tolist()[0]
+
+        logits = torch.FloatTensor(session.run(None, {'input': [torch_in]})[0])[:, -1]
+        
+        filtered_logits = logits
+
+        probs = F.softmax(filtered_logits / temperature, dim=-1)
+
+        sample = torch.multinomial(probs, 1)
+
+        out = torch.cat((out, sample), dim=-1)
+
+        if verbose:
+          if s % 32 == 0:
+            print(s, '/', seq_len)
+
+    if return_prime:
+      return out[:, :]
+    
+    else:
+      return out[:, st:]
+
+
+#=================================================================================================
+
 
 def load_javascript(dir="javascript"):
     scripts_list = glob.glob(f"{dir}/*.js")
@@ -55,7 +105,7 @@ if __name__ == "__main__":
     opt = parser.parse_args()
 
 
-    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+    providers = ['CPUExecutionProvider']
 
     session = rt.InferenceSession('Allegro_Music_Transformer_Small_Trained_Model_56000_steps_0.9399_loss_0.7374_acc.onnx', providers=providers)