to torch

README.md

@@ -3,8 +3,8 @@ title: Midi Music Generator
 emoji: 🎼🎶
 colorFrom: red
 colorTo: indigo
-sdk: docker
-sdk_version: 4.36.1
+sdk: gradio
+sdk_version: 4.43.0
 app_file: app.py
 pinned: true
 license: apache-2.0

app.py

@@ -1,79 +1,54 @@
 import argparse
 import glob
-import os.path
+import json
+import os
 import time
-import uuid
 
 import gradio as gr
 import numpy as np
-import onnxruntime as rt
+import torch
+
+import torch.nn.functional as F
 import tqdm
-import json
-from huggingface_hub import hf_hub_download
 
 import MIDI
-from midi_synthesizer import synthesis
+from midi_model import MIDIModel
 from midi_tokenizer import MIDITokenizer
+from midi_synthesizer import synthesis
+from huggingface_hub import hf_hub_download
 
 MAX_SEED = np.iinfo(np.int32).max
 in_space = os.getenv("SYSTEM") == "spaces"
 
 
-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
-
-
+@torch.inference_mode()
 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, generator=None):
+             disable_patch_change=False, disable_control_change=False, disable_channels=None, amp=True, 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 = np.full((1, max_token_seq), tokenizer.pad_id, dtype=np.int64)
+        input_tensor = torch.full((1, max_token_seq), tokenizer.pad_id, dtype=torch.long, device=model.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 = prompt
-    input_tensor = input_tensor[None, :, :]
+        input_tensor = torch.from_numpy(prompt).to(dtype=torch.long, device=model.device)
+    input_tensor = input_tensor.unsqueeze(0)
     cur_len = input_tensor.shape[1]
     bar = tqdm.tqdm(desc="generating", total=max_len - cur_len, disable=in_space)
-    with bar:
+    with bar, torch.amp.autocast(device_type=model.device, enabled=amp):
         while cur_len < max_len:
             end = False
-            hidden = model[0].run(None, {'x': input_tensor})[0][:, -1]
-            next_token_seq = np.empty((1, 0), dtype=np.int64)
+            hidden = model.forward(input_tensor)[0, -1].unsqueeze(0)
+            next_token_seq = None
             event_name = ""
             for i in range(max_token_seq):
-                mask = np.zeros(tokenizer.vocab_size, dtype=np.int64)
+                mask = torch.zeros(tokenizer.vocab_size, dtype=torch.int64, device=model.device)
                 if i == 0:
                     mask_ids = list(tokenizer.event_ids.values()) + [tokenizer.eos_id]
                     if disable_patch_change:
@@ -87,9 +62,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[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)
+                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)
                 if i == 0:
                     next_token_seq = sample
                     eid = sample.item()
@@ -98,17 +73,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 = np.concatenate([next_token_seq, sample], axis=1)
+                    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 = 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)
+                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)
-            yield next_token_seq.reshape(-1)
+            yield next_token_seq.reshape(-1).cpu().numpy()
             if end:
                 break
 
@@ -129,7 +104,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 = np.random.RandomState(seed)
+    generator = torch.Generator(device).manual_seed(seed)
     disable_patch_change = False
     disable_channels = None
     if tab == 0:
@@ -160,14 +135,16 @@ 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, generator=generator)
+                         disable_channels=disable_channels, amp=amp, generator=generator)
     events = []
     for i, token_seq in enumerate(midi_generator):
         token_seq = token_seq.tolist()
@@ -245,15 +222,18 @@ 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"
     models = {}
     tokenizer = MIDITokenizer()
-    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
     for name, (repo_id, path) in models_info.items():
-        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]
+
+        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)
+        state_dict = ckpt.get("state_dict", ckpt)
+        model.load_state_dict(state_dict, strict=False)
+        model.eval()
+        models[name] = model
 
     load_javascript()
     app = gr.Blocks()
@@ -265,7 +245,8 @@ 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"
+                    "**Update v1.2**: Optimise the tokenizer and dataset\n\n"
+                    f"Device: {device}"
                     )
         js_msg = gr.Textbox(elem_id="msg_receiver", visible=False)
         js_msg.change(None, [js_msg], [], js="""

midi_model.py

@@ -0,0 +1,129 @@
+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 midi_tokenizer import MIDITokenizer
+
+
+class MIDIModel(nn.Module):
+    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:
+            torch.backends.cuda.enable_mem_efficient_sdp(True)
+            torch.backends.cuda.enable_flash_sdp(True)
+        self.lm_head = nn.Linear(n_embd, tokenizer.vocab_size, bias=False)
+        self.device = "cpu"
+
+    def to(self, *args, **kwargs):
+        if "device" in kwargs:
+            self.device = kwargs["device"]
+        return super(MIDIModel, self).to(*args, **kwargs)
+
+    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, midi_sequence_length, token_sequence_length)
+        :return: hidden (batch_size, midi_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, generator=None):
+        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, generator=generator).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, generator=None):
+        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, generator=generator)
+                    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,5 +1,8 @@
 Pillow
 numpy
-onnxruntime-gpu
+torch
+transformers>=4.36
 gradio==4.43.0
 pyfluidsynth
+tqdm
+huggingface_hub