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