update to onnx

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.mid filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,48 +1,72 @@
 import argparse
 import glob
-
-import PIL
+import os
+import os.path
+from sys import exit
+import shutil
 import gradio as gr
 import numpy as np
-import torch
-
-import torch.nn.functional as F
+import onnxruntime as rt
+import PIL
+import PIL.ImageColor
 import tqdm
+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
+
+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):
+    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([np.random.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
+
 
-@torch.inference_mode()
 def generate(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):
+             disable_patch_change=False, disable_control_change=False, disable_channels=None):
     if disable_channels is not None:
         disable_channels = [tokenizer.parameter_ids["channel"][c] for c in disable_channels]
     else:
         disable_channels = []
     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)
-    with bar, torch.cuda.amp.autocast(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_base.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:
@@ -56,9 +80,9 @@ def generate(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)
+                logits = model_token.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)
                 if i == 0:
                     next_token_seq = sample
                     eid = sample.item()
@@ -67,29 +91,30 @@ def generate(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
 
 
-def run(tab, instruments, drum_kit, mid, midi_events, gen_events, temp, top_p, top_k, allow_cc, amp):
+def run(tab, instruments, drum_kit, mid, midi_events, gen_events, temp, top_p, top_k, allow_cc):
     mid_seq = []
     max_len = int(gen_events)
     img_len = 1024
     img = np.full((128 * 2, img_len, 3), 255, dtype=np.uint8)
     state = {"t1": 0, "t": 0, "cur_pos": 0}
-    rand = np.random.RandomState(0)
-    colors = {(i, j): rand.randint(0, 200, 3) for i in range(128) for j in range(16)}
+    colors = ['navy', 'blue', 'deepskyblue', 'teal', 'green', 'lightgreen', 'lime', 'orange',
+              'brown', 'grey', 'red', 'pink', 'aqua', 'orchid', 'bisque', 'coral']
+    colors = [PIL.ImageColor.getrgb(color) for color in colors]
 
     def draw_event(tokens):
         if tokens[0] in tokenizer.id_events:
@@ -112,7 +137,7 @@ def run(tab, instruments, drum_kit, mid, midi_events, gen_events, temp, top_p, t
                     img[:, -shift:] = 255
                     state["cur_pos"] += shift
                 t = t - state["cur_pos"]
-                img[p * 2:(p + 1) * 2, t: t + d] = colors[(tr, c)]
+                img[p * 2:(p + 1) * 2, t: t + d] = colors[c]
 
     def get_img():
         t = state["t"] - state["cur_pos"]
@@ -135,7 +160,7 @@ def run(tab, instruments, drum_kit, mid, midi_events, gen_events, temp, top_p, t
             mid.append(tokenizer.event2tokens(["patch_change", 0, 0, i, c, p]))
         mid_seq = mid
         mid = np.asarray(mid, dtype=np.int64)
-        if len(instruments) > 0 or drum_kit != "None":
+        if len(instruments) > 0:
             disable_patch_change = True
             disable_channels = [i for i in range(16) if i not in patches]
     elif mid is not None:
@@ -148,7 +173,7 @@ def run(tab, instruments, drum_kit, mid, midi_events, gen_events, temp, top_p, t
             draw_event(token_seq)
     generator = generate(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)
+                         disable_channels=disable_channels)
     for token_seq in generator:
         mid_seq.append(token_seq)
         draw_event(token_seq)
@@ -179,17 +204,16 @@ if __name__ == "__main__":
     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")
-    parser.add_argument("--device", type=str, default="cpu", help="device to run model")
-    parser.add_argument("--max-gen", type=int, default=512, help="max")
-    soundfont_path = hf_hub_download(repo_id="skytnt/midi-model", filename="soundfont.sf2")
-    model_path = hf_hub_download(repo_id="skytnt/midi-model", filename="model.ckpt")
+    parser.add_argument("--max-gen", type=int, default=256, help="max")
     opt = parser.parse_args()
+    soundfont_path = hf_hub_download(repo_id="skytnt/midi-model", filename="soundfont.sf2")
+    model_base_path = hf_hub_download(repo_id="skytnt/midi-model", filename="onnx/model_base.onnx")
+    model_token_path = hf_hub_download(repo_id="skytnt/midi-model", filename="onnx/model_token.onnx")
+
     tokenizer = MIDITokenizer()
-    model = MIDIModel(tokenizer).to(device=opt.device)
-    ckpt = torch.load(model_path, map_location="cpu")
-    state_dict = ckpt.get("state_dict", ckpt)
-    model.load_state_dict(state_dict, strict=False)
-    model.eval()
+    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+    model_base = rt.InferenceSession(model_base_path, providers=providers)
+    model_token = rt.InferenceSession(model_token_path, providers=providers)
 
     app = gr.Blocks()
     with app:
@@ -199,39 +223,52 @@ if __name__ == "__main__":
                     "Demo for [SkyTNT/midi-model](https://github.com/SkyTNT/midi-model)\n\n"
                     "[Open In Colab]"
                     "(https://colab.research.google.com/github/SkyTNT/midi-model/blob/main/demo.ipynb)"
-                    " for faster running")
+                    " for faster running and longer generation"
+                    )
 
         tab_select = gr.Variable(value=0)
         with gr.Tabs():
             with gr.TabItem("instrument prompt") as tab1:
                 input_instruments = gr.Dropdown(label="instruments (auto if empty)", choices=list(patch2number.keys()),
-                                                multiselect=True, max_choices=10, type="value")
+                                                multiselect=True, max_choices=15, type="value")
                 input_drum_kit = gr.Dropdown(label="drum kit", choices=list(drum_kits2number.keys()), type="value",
                                              value="None")
+                example1 = gr.Examples([
+                    [[], "None"],
+                    [["Acoustic Grand"], "None"],
+                    [["Acoustic Grand", "Violin", "Viola", "Cello", "Contrabass", "Timpani"], "Orchestra"],
+                    [["Acoustic Guitar(nylon)", "Acoustic Guitar(steel)", "Electric Guitar(jazz)",
+                      "Electric Guitar(clean)", "Electric Guitar(muted)", "Overdriven Guitar", "Distortion Guitar",
+                      "Electric Bass(finger)"], "Standard"],
+                    [["Acoustic Grand", "String Ensemble 1", "Trombone", "Tuba", "Muted Trumpet", "French Horn", "Oboe",
+                      "English Horn", "Bassoon", "Clarinet"], "Orchestra"]
+
+                ], [input_instruments, input_drum_kit])
             with gr.TabItem("midi prompt") as tab2:
                 input_midi = gr.File(label="input midi", file_types=[".midi", ".mid"], type="binary")
                 input_midi_events = gr.Slider(label="use first n midi events as prompt", minimum=1, maximum=512,
                                               step=1,
                                               value=128)
+                example2 = gr.Examples([[file, 128] for file in glob.glob("example/*.mid")],
+                                       [input_midi, input_midi_events])
 
         tab1.select(lambda: 0, None, tab_select, queue=False)
         tab2.select(lambda: 1, None, tab_select, queue=False)
         input_gen_events = gr.Slider(label="generate n midi events", minimum=1, maximum=opt.max_gen,
                                      step=1, value=opt.max_gen)
         input_temp = gr.Slider(label="temperature", minimum=0.1, maximum=1.2, step=0.01, value=1)
-        input_top_p = gr.Slider(label="top p", minimum=0.1, maximum=1, step=0.01, value=0.97)
+        input_top_p = gr.Slider(label="top p", minimum=0.1, maximum=1, step=0.01, value=0.98)
         input_top_k = gr.Slider(label="top k", minimum=1, maximum=50, step=1, value=20)
-        input_allow_cc = gr.Checkbox(label="allow control change event", value=True)
-        input_amp = gr.Checkbox(label="enable amp", value=True)
+        input_allow_cc = gr.Checkbox(label="allow midi cc event", value=True)
         run_btn = gr.Button("generate", variant="primary")
         stop_btn = gr.Button("stop")
         output_midi_seq = gr.Variable()
         output_midi_img = gr.Image(label="output image")
         output_midi = gr.File(label="output midi", file_types=[".mid"])
-        output_audio = gr.Audio(label="output audio", format="mp3")
+        output_audio = gr.Audio(label="output audio", format="wav")
         run_event = run_btn.click(run, [tab_select, input_instruments, input_drum_kit, input_midi, input_midi_events,
                                         input_gen_events, input_temp, input_top_p, input_top_k,
-                                        input_allow_cc, input_amp],
+                                        input_allow_cc],
                                   [output_midi_seq, output_midi_img, output_midi, output_audio])
         stop_btn.click(cancel_run, output_midi_seq, [output_midi, output_audio], cancels=run_event, queue=False)
-    app.queue(1).launch(server_port=opt.port, share=opt.share)
+    app.queue(2).launch(server_port=opt.port, share=opt.share, inbrowser=True)

example/Bach--Fugue-in-D-Minor.mid

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1398121eb86a33e73f90ec84be71dac6abc0ddf11372ea7cdd9e01586938a56b
+size 7720

example/Beethoven--Symphony-No5-in-C-Minor-Fate-Opus-67.mid

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28ff6fdcd644e781d36411bf40ab7a1f4849adddbcd1040eaec22751c5ca99d2
+size 87090

example/Chopin--Nocturne No. 9 in B Major, Opus 32 No.1, Andante Sostenuto.mid

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3a236e647ad9f5d0af680d3ca19d3b60f334c4bde6b4f86310f63405245c476e
+size 13484

example/Mozart--Requiem, No.1..mid

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa49bf4633401e16777fe47f6f53a494c2166f5101af6dafc60114932a59b9bd
+size 14695

example/castle_in_the_sky.mid

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa14aec6f1be15c4fddd0decc6d9152204f160d4e07e05d8d1dc9f209c309ff7
+size 7957

example/eva-残酷な天使のテーゼ.mid

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e513487543d7e27ec5dc30f027302d2a3b5a3aaf9af554def1e5cd6a7a8d355a
+size 17671

midi_model.py

@@ -1,123 +0,0 @@
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import tqdm
-from transformers import LlamaModel, LlamaConfig
-from transformers.modeling_utils import ModuleUtilsMixin
-
-from midi_tokenizer import MIDITokenizer
-
-
-class MIDIModel(nn.Module, ModuleUtilsMixin):
-    def __init__(self, tokenizer: MIDITokenizer, n_layer=12, n_head=16, n_embd=1024, n_inner=4096, flash=False,
-                 *args, **kwargs):
-        super(MIDIModel, self).__init__()
-        self.tokenizer = tokenizer
-        self.net = LlamaModel(LlamaConfig(vocab_size=tokenizer.vocab_size,
-                                          hidden_size=n_embd, num_attention_heads=n_head,
-                                          num_hidden_layers=n_layer, intermediate_size=n_inner,
-                                          pad_token_id=tokenizer.pad_id, max_position_embeddings=4096))
-        self.net_token = LlamaModel(LlamaConfig(vocab_size=tokenizer.vocab_size,
-                                                hidden_size=n_embd, num_attention_heads=n_head // 4,
-                                                num_hidden_layers=n_layer // 4, intermediate_size=n_inner // 4,
-                                                pad_token_id=tokenizer.pad_id, max_position_embeddings=4096))
-        if flash:
-            self.net = self.net.to_bettertransformer()
-            self.net_token = self.net_token.to_bettertransformer()
-        self.lm_head = nn.Linear(n_embd, tokenizer.vocab_size, bias=False)
-
-    def forward_token(self, hidden_state, x=None):
-        """
-
-        :param hidden_state: (batch_size, n_embd)
-        :param x: (batch_size, token_sequence_length)
-        :return: (batch_size, 1 + token_sequence_length, vocab_size)
-        """
-        hidden_state = hidden_state.unsqueeze(1)  # (batch_size, 1, n_embd)
-        if x is not None:
-            x = self.net_token.embed_tokens(x)
-            hidden_state = torch.cat([hidden_state, x], dim=1)
-        hidden_state = self.net_token.forward(inputs_embeds=hidden_state).last_hidden_state
-        return self.lm_head(hidden_state)
-
-    def forward(self, x):
-        """
-        :param x: (batch_size, time_sequence_length, token_sequence_length)
-        :return: hidden (batch_size, time_sequence_length, n_embd)
-        """
-
-        # merge token sequence
-        x = self.net.embed_tokens(x)
-        x = x.sum(dim=-2)
-        x = self.net.forward(inputs_embeds=x)
-        return x.last_hidden_state
-
-    def sample_top_p_k(self, probs, p, k):
-        probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
-        probs_sum = torch.cumsum(probs_sort, dim=-1)
-        mask = probs_sum - probs_sort > p
-        probs_sort[mask] = 0.0
-        mask = torch.zeros(probs_sort.shape[-1], device=probs_sort.device)
-        mask[:k] = 1
-        probs_sort = probs_sort * mask
-        probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
-        shape = probs_sort.shape
-        next_token = torch.multinomial(probs_sort.reshape(-1, shape[-1]), num_samples=1).reshape(*shape[:-1], 1)
-        next_token = torch.gather(probs_idx, -1, next_token).reshape(*shape[:-1])
-        return next_token
-
-    @torch.inference_mode()
-    def generate(self, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20, amp=True):
-        tokenizer = self.tokenizer
-        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=self.device)
-            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=self.device)
-        input_tensor = input_tensor.unsqueeze(0)
-        cur_len = input_tensor.shape[1]
-        bar = tqdm.tqdm(desc="generating", total=max_len - cur_len)
-        with bar, torch.cuda.amp.autocast(enabled=amp):
-            while cur_len < max_len:
-                end = False
-                hidden = self.forward(input_tensor)[0, -1].unsqueeze(0)
-                next_token_seq = None
-                event_name = ""
-                for i in range(max_token_seq):
-                    mask = torch.zeros(tokenizer.vocab_size, dtype=torch.int64, device=self.device)
-                    if i == 0:
-                        mask[list(tokenizer.event_ids.values()) + [tokenizer.eos_id]] = 1
-                    else:
-                        param_name = tokenizer.events[event_name][i - 1]
-                        mask[tokenizer.parameter_ids[param_name]] = 1
-
-                    logits = self.forward_token(hidden, next_token_seq)[:, -1:]
-                    scores = torch.softmax(logits / temp, dim=-1) * mask
-                    sample = self.sample_top_p_k(scores, top_p, top_k)
-                    if i == 0:
-                        next_token_seq = sample
-                        eid = sample.item()
-                        if eid == tokenizer.eos_id:
-                            end = True
-                            break
-                        event_name = tokenizer.id_events[eid]
-                    else:
-                        next_token_seq = torch.cat([next_token_seq, sample], dim=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)
-                cur_len += 1
-                bar.update(1)
-                if end:
-                    break
-        return input_tensor[0].cpu().numpy()

requirements.txt

@@ -1,6 +1,5 @@
 Pillow
 numpy
-torch
-transformers
+onnxruntime-gpu
 gradio
 pyfluidsynth