torch seems slow, bring back onnx

app.py

@@ -1,54 +1,79 @@
 import argparse
 import glob
-import json
-import os
+import os.path
 import time
+import uuid
 
 import gradio as gr
 import numpy as np
-import torch
-
-import torch.nn.functional as F
+import onnxruntime as rt
 import tqdm
+import json
+from huggingface_hub import hf_hub_download
 
 import MIDI
-from midi_model import MIDIModel
-from midi_tokenizer import MIDITokenizer
 from midi_synthesizer import synthesis
-from huggingface_hub import hf_hub_download
+from midi_tokenizer import MIDITokenizer
 
 MAX_SEED = np.iinfo(np.int32).max
 in_space = os.getenv("SYSTEM") == "spaces"
 
 
-@torch.inference_mode()
+def softmax(x, axis):
+    x_max = np.amax(x, axis=axis, keepdims=True)
+    exp_x_shifted = np.exp(x - x_max)
+    return exp_x_shifted / np.sum(exp_x_shifted, axis=axis, keepdims=True)
+
+
+def sample_top_p_k(probs, p, k, generator=None):
+    if generator is None:
+        generator = np.random
+    probs_idx = np.argsort(-probs, axis=-1)
+    probs_sort = np.take_along_axis(probs, probs_idx, -1)
+    probs_sum = np.cumsum(probs_sort, axis=-1)
+    mask = probs_sum - probs_sort > p
+    probs_sort[mask] = 0.0
+    mask = np.zeros(probs_sort.shape[-1])
+    mask[:k] = 1
+    probs_sort = probs_sort * mask
+    probs_sort /= np.sum(probs_sort, axis=-1, keepdims=True)
+    shape = probs_sort.shape
+    probs_sort_flat = probs_sort.reshape(-1, shape[-1])
+    probs_idx_flat = probs_idx.reshape(-1, shape[-1])
+    next_token = np.stack([generator.choice(idxs, p=pvals) for pvals, idxs in zip(probs_sort_flat, probs_idx_flat)])
+    next_token = next_token.reshape(*shape[:-1])
+    return next_token
+
+
 def generate(model, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
-             disable_patch_change=False, disable_control_change=False, disable_channels=None, amp=True, generator=None):
+             disable_patch_change=False, disable_control_change=False, disable_channels=None, generator=None):
     if disable_channels is not None:
         disable_channels = [tokenizer.parameter_ids["channel"][c] for c in disable_channels]
     else:
         disable_channels = []
+    if generator is None:
+        generator = np.random
     max_token_seq = tokenizer.max_token_seq
     if prompt is None:
-        input_tensor = torch.full((1, max_token_seq), tokenizer.pad_id, dtype=torch.long, device=model.device)
+        input_tensor = np.full((1, max_token_seq), tokenizer.pad_id, dtype=np.int64)
         input_tensor[0, 0] = tokenizer.bos_id  # bos
     else:
         prompt = prompt[:, :max_token_seq]
         if prompt.shape[-1] < max_token_seq:
             prompt = np.pad(prompt, ((0, 0), (0, max_token_seq - prompt.shape[-1])),
                             mode="constant", constant_values=tokenizer.pad_id)
-        input_tensor = torch.from_numpy(prompt).to(dtype=torch.long, device=model.device)
-    input_tensor = input_tensor.unsqueeze(0)
+        input_tensor = prompt
+    input_tensor = input_tensor[None, :, :]
     cur_len = input_tensor.shape[1]
     bar = tqdm.tqdm(desc="generating", total=max_len - cur_len, disable=in_space)
-    with bar, torch.amp.autocast(device_type=model.device, enabled=amp):
+    with bar:
         while cur_len < max_len:
             end = False
-            hidden = model.forward(input_tensor)[0, -1].unsqueeze(0)
-            next_token_seq = None
+            hidden = model[0].run(None, {'x': input_tensor})[0][:, -1]
+            next_token_seq = np.empty((1, 0), dtype=np.int64)
             event_name = ""
             for i in range(max_token_seq):
-                mask = torch.zeros(tokenizer.vocab_size, dtype=torch.int64, device=model.device)
+                mask = np.zeros(tokenizer.vocab_size, dtype=np.int64)
                 if i == 0:
                     mask_ids = list(tokenizer.event_ids.values()) + [tokenizer.eos_id]
                     if disable_patch_change:
@@ -62,9 +87,9 @@ def generate(model, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
                     if param_name == "channel":
                         mask_ids = [i for i in mask_ids if i not in disable_channels]
                     mask[mask_ids] = 1
-                logits = model.forward_token(hidden, next_token_seq)[:, -1:]
-                scores = torch.softmax(logits / temp, dim=-1) * mask
-                sample = model.sample_top_p_k(scores, top_p, top_k, generator=generator)
+                logits = model[1].run(None, {'x': next_token_seq, "hidden": hidden})[0][:, -1:]
+                scores = softmax(logits / temp, -1) * mask
+                sample = sample_top_p_k(scores, top_p, top_k, generator)
                 if i == 0:
                     next_token_seq = sample
                     eid = sample.item()
@@ -73,17 +98,17 @@ def generate(model, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
                         break
                     event_name = tokenizer.id_events[eid]
                 else:
-                    next_token_seq = torch.cat([next_token_seq, sample], dim=1)
+                    next_token_seq = np.concatenate([next_token_seq, sample], axis=1)
                     if len(tokenizer.events[event_name]) == i:
                         break
             if next_token_seq.shape[1] < max_token_seq:
-                next_token_seq = F.pad(next_token_seq, (0, max_token_seq - next_token_seq.shape[1]),
-                                       "constant", value=tokenizer.pad_id)
-            next_token_seq = next_token_seq.unsqueeze(1)
-            input_tensor = torch.cat([input_tensor, next_token_seq], dim=1)
+                next_token_seq = np.pad(next_token_seq, ((0, 0), (0, max_token_seq - next_token_seq.shape[-1])),
+                                        mode="constant", constant_values=tokenizer.pad_id)
+            next_token_seq = next_token_seq[None, :, :]
+            input_tensor = np.concatenate([input_tensor, next_token_seq], axis=1)
             cur_len += 1
             bar.update(1)
-            yield next_token_seq.reshape(-1).cpu().numpy()
+            yield next_token_seq.reshape(-1)
             if end:
                 break
 
@@ -104,7 +129,7 @@ def run(model_name, tab, instruments, drum_kit, bpm, mid, midi_events, midi_opt,
     max_len = gen_events
     if seed_rand:
         seed = np.random.randint(0, MAX_SEED)
-    generator = torch.Generator(device).manual_seed(seed)
+    generator = np.random.RandomState(seed)
     disable_patch_change = False
     disable_channels = None
     if tab == 0:
@@ -135,16 +160,14 @@ def run(model_name, tab, instruments, drum_kit, bpm, mid, midi_events, midi_opt,
         for token_seq in mid:
             mid_seq.append(token_seq.tolist())
     max_len += len(mid)
-
     events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
     init_msgs = [create_msg("visualizer_clear", None), create_msg("visualizer_append", events)]
     t = time.time() + 1
     yield mid_seq, None, None, seed, send_msgs(init_msgs)
     model = models[model_name]
-    amp = device == "cuda"
     midi_generator = generate(model, mid, max_len=max_len, temp=temp, top_p=top_p, top_k=top_k,
                          disable_patch_change=disable_patch_change, disable_control_change=not allow_cc,
-                         disable_channels=disable_channels, amp=amp, generator=generator)
+                         disable_channels=disable_channels, generator=generator)
     events = []
     for i, token_seq in enumerate(midi_generator):
         token_seq = token_seq.tolist()
@@ -222,21 +245,15 @@ if __name__ == "__main__":
                    "j-pop finetune model": ["skytnt/midi-model-ft", "jpop/"],
                    "touhou finetune model": ["skytnt/midi-model-ft", "touhou/"],
                    }
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    if device=="cuda": # flash attn
-        torch.backends.cuda.enable_mem_efficient_sdp(True)
-        torch.backends.cuda.enable_flash_sdp(True)
     models = {}
     tokenizer = MIDITokenizer()
+    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
     for name, (repo_id, path) in models_info.items():
-
-        model_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}model.ckpt")
-        model = MIDIModel(tokenizer).to(device=device)
-        ckpt = torch.load(model_path, weights_only=True)
-        state_dict = ckpt.get("state_dict", ckpt)
-        model.load_state_dict(state_dict, strict=False)
-        model.eval()
-        models[name] = model
+        model_base_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}onnx/model_base.onnx")
+        model_token_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}onnx/model_token.onnx")
+        model_base = rt.InferenceSession(model_base_path, providers=providers)
+        model_token = rt.InferenceSession(model_token_path, providers=providers)
+        models[name] = [model_base, model_token]
 
     load_javascript()
     app = gr.Blocks()
@@ -248,8 +265,7 @@ if __name__ == "__main__":
                     "[Open In Colab]"
                     "(https://colab.research.google.com/github/SkyTNT/midi-model/blob/main/demo.ipynb)"
                     " for faster running and longer generation\n\n"
-                    "**Update v1.2**: Optimise the tokenizer and dataset\n\n"
-                    f"Device: {device}"
+                    "**Update v1.2**: Optimise the tokenizer and dataset"
                     )
         js_msg = gr.Textbox(elem_id="msg_receiver", visible=False)
         js_msg.change(None, [js_msg], [], js="""
@@ -319,4 +335,4 @@ if __name__ == "__main__":
                                   [output_midi_seq, output_midi, output_audio, input_seed, js_msg],
                                   concurrency_limit=3)
         stop_btn.click(cancel_run, [output_midi_seq], [output_midi, output_audio, js_msg], cancels=run_event, queue=False)
-    app.launch(server_port=opt.port, share=opt.share, inbrowser=True)
+    app.launch(server_port=opt.port, share=opt.share, inbrowser=True)

requirements.txt

@@ -1,8 +1,6 @@
---extra-index-url https://download.pytorch.org/whl/cu124
 Pillow
 numpy
-torch
-transformers>=4.36
+onnxruntime-gpu
 gradio==4.43.0
 pyfluidsynth
 tqdm