torch

app.py

@@ -1,79 +1,53 @@
 import argparse
 import glob
 import json
-import os.path
+import os
 import time
 
 import gradio as gr
 import numpy as np
-import onnxruntime as rt
+import torch
+import torch.nn.functional as F
 import tqdm
 from huggingface_hub import hf_hub_download
 
 import MIDI
+from midi_model import MIDIModel, MIDIModelConfig
 from midi_synthesizer import MidiSynthesizer
-from midi_tokenizer import MIDITokenizer
 
 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
-
-
-def generate(model, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
+@torch.inference_mode()
+def generate(model: MIDIModel, 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):
- tokenizer = model[2]
+ tokenizer = model.tokenizer
  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)
+ bar = tqdm.tqdm(desc="generating", total=max_len - cur_len)
  with bar:
  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 +61,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 +72,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
 
@@ -125,7 +99,7 @@ def run(model_name, tab, mid_seq, continuation_state, instruments, drum_kit, bpm
  reduce_cc_st, remap_track_channel, add_default_instr, remove_empty_channels, seed, seed_rand,
  gen_events, temp, top_p, top_k, allow_cc):
  model = models[model_name]
- tokenizer = model[2]
+ tokenizer = model.tokenizer
  bpm = int(bpm)
  if time_sig == "auto":
  time_sig = None
@@ -147,7 +121,7 @@ def run(model_name, tab, mid_seq, continuation_state, instruments, drum_kit, bpm
  max_len = gen_events
  if seed_rand:
  seed = np.random.randint(0, MAX_SEED)
- generator = np.random.RandomState(seed)
+ generator = torch.Generator(opt.device).manual_seed(seed)
  disable_patch_change = False
  disable_channels = None
  if tab == 0:
@@ -203,22 +177,24 @@ def run(model_name, tab, mid_seq, continuation_state, instruments, drum_kit, bpm
  init_msgs += [create_msg("visualizer_clear", tokenizer.version),
  create_msg("visualizer_append", events)]
  yield mid_seq, continuation_state, None, None, seed, send_msgs(init_msgs)
- 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)
- events = []
- t = time.time() + 1
- for i, token_seq in enumerate(midi_generator):
- token_seq = token_seq.tolist()
- mid_seq.append(token_seq)
- events.append(tokenizer.tokens2event(token_seq))
- ct = time.time()
- if ct - t > 0.5:
- yield (mid_seq, continuation_state, None, None, seed,
- send_msgs([create_msg("visualizer_append", events),
- create_msg("progress", [i + 1, gen_events])]))
- t = ct
- events = []
+ ctx = torch.amp.autocast(device_type=opt.device, dtype=torch.bfloat16, enabled=opt.device != "cpu")
+ with ctx:
+ 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)
+ events = []
+ t = time.time() + 1
+ for i, token_seq in enumerate(midi_generator):
+ token_seq = token_seq.tolist()
+ mid_seq.append(token_seq)
+ events.append(tokenizer.tokens2event(token_seq))
+ ct = time.time()
+ if ct - t > 0.5:
+ yield (mid_seq, continuation_state, None, None, seed,
+ send_msgs([create_msg("visualizer_append", events),
+ create_msg("progress", [i + 1, gen_events])]))
+ t = ct
+ events = []
 
  events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
  mid = tokenizer.detokenize(mid_seq)
@@ -235,7 +211,7 @@ def run(model_name, tab, mid_seq, continuation_state, instruments, drum_kit, bpm
 def cancel_run(model_name, mid_seq):
  if mid_seq is None:
  return None, None, []
- tokenizer = models[model_name][2]
+ tokenizer = models[model_name].tokenizer
  events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
  mid = tokenizer.detokenize(mid_seq)
  audio = synthesizer.synthesis(MIDI.score2opus(mid))
@@ -248,11 +224,12 @@ def cancel_run(model_name, mid_seq):
  return "output.mid", (44100, audio), send_msgs(end_msgs)
 
 
-def undo_continuation(mid_seq, continuation_state):
+def undo_continuation(model_name, mid_seq, continuation_state):
  if mid_seq is None or len(continuation_state) < 2:
  return mid_seq, continuation_state, send_msgs([])
  mid_seq = mid_seq[:continuation_state[-1]]
  continuation_state = continuation_state[:-1]
+ tokenizer = models[model_name].tokenizer
  events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
  end_msgs = [create_msg("visualizer_clear", tokenizer.version),
  create_msg("visualizer_append", events),
@@ -293,21 +270,6 @@ def hf_hub_download_retry(repo_id, filename):
  raise err
 
 
-def get_tokenizer(config_name):
- tv, size = config_name.split("-")
- tv = tv[1:]
- if tv[-1] == "o":
- o = True
- tv = tv[:-1]
- else:
- o = False
- if tv not in ["v1", "v2"]:
- raise ValueError(f"Unknown tokenizer version {tv}")
- tokenizer = MIDITokenizer(tv)
- tokenizer.set_optimise_midi(o)
- return tokenizer
-
-
 number2drum_kits = {-1: "None", 0: "Standard", 8: "Room", 16: "Power", 24: "Electric", 25: "TR-808", 32: "Jazz",
  40: "Blush", 48: "Orchestra"}
 patch2number = {v: k for k, v in MIDI.Number2patch.items()}
@@ -319,6 +281,7 @@ 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="cuda", help="device to run model")
  parser.add_argument("--max-gen", type=int, default=1024, help="max")
  opt = parser.parse_args()
  soundfont_path = hf_hub_download_retry(repo_id="skytnt/midi-model", filename="soundfont.sf2")
@@ -331,14 +294,17 @@ if __name__ == "__main__":
  "touhou finetune model (tv1-medium) by skytnt": ["skytnt/midi-model-ft", "touhou/", "tv1-medium"],
  }
  models = {}
- providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+ if opt.device == "cuda":
+ torch.backends.cuda.enable_mem_efficient_sdp(True)
+ torch.backends.cuda.enable_flash_sdp(True)
  for name, (repo_id, path, config) 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)
- tokenizer = get_tokenizer(config)
- models[name] = [model_base, model_token, tokenizer]
+ model_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}model.ckpt")
+ model = MIDIModel(config=MIDIModelConfig.from_name(config))
+ ckpt = torch.load(model_path, map_location="cpu")
+ state_dict = ckpt.get("state_dict", ckpt)
+ model.load_state_dict(state_dict, strict=False)
+ model.to(device=opt.device, dtype=torch.bfloat16 if opt.device == "cuda" else torch.float32).eval()
+ models[name] = model
 
  load_javascript()
  app = gr.Blocks()
@@ -447,6 +413,6 @@ if __name__ == "__main__":
  stop_btn.click(cancel_run, [input_model, output_midi_seq],
  [output_midi, output_audio, js_msg],
  cancels=run_event, queue=False)
- undo_btn.click(undo_continuation, [output_midi_seq, output_continuation_state],
+ undo_btn.click(undo_continuation, [input_model, output_midi_seq, output_continuation_state],
  [output_midi_seq, output_continuation_state, js_msg], queue=False)
  app.launch(server_port=opt.port, share=opt.share, inbrowser=True)

midi_model.py

@@ -1,3 +1,5 @@
+from typing import Union
+
 import numpy as np
 import torch
 import torch.nn as nn
@@ -5,23 +7,61 @@ import torch.nn.functional as F
 import tqdm
 from transformers import LlamaModel, LlamaConfig
 
-from midi_tokenizer import MIDITokenizer
+from midi_tokenizer import MIDITokenizerV1, MIDITokenizerV2, MIDITokenizer
+
+config_name_list = ["tv1-medium", "tv2-medium", "tv2o-medium", "tv2-large", "tv2o-large"]
+
+
+class MIDIModelConfig:
+ def __init__(self, tokenizer: Union[MIDITokenizerV1, MIDITokenizerV2],
+ net_config: LlamaConfig, net_token_config: LlamaConfig):
+ self.tokenizer = tokenizer
+ self.net_config = net_config
+ self.net_token_config = net_token_config
+ self.n_embd = net_token_config.hidden_size
+
+ @staticmethod
+ def get_config(tokenizer_ver="v2", optimise_midi=True, n_layer=12, n_head=16, n_embd=1024, n_inner=4096):
+ tokenizer = MIDITokenizer(tokenizer_ver)
+ tokenizer.set_optimise_midi(optimise_midi)
+ net_config = 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)
+ net_token_config = 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)
+ return MIDIModelConfig(tokenizer, net_config, net_token_config)
+
+ @staticmethod
+ def from_name(name="tv2o-medium"):
+ tv, size = name.split("-")
+ tv = tv[1:]
+ if tv[-1] == "o":
+ o = True
+ tv = tv[:-1]
+ else:
+ o = False
+ if tv not in ["v1", "v2"]:
+ raise ValueError(f"Unknown tokenizer version {tv}")
+ if size == "medium":
+ return MIDIModelConfig.get_config(tokenizer_ver=tv, optimise_midi=o,
+ n_layer=12, n_head=16, n_embd=1024, n_inner=4096)
+ elif size == "large":
+ return MIDIModelConfig.get_config(tokenizer_ver=tv, optimise_midi=o,
+ n_layer=24, n_head=16, n_embd=1024, n_inner=4096)
+ else:
+ raise ValueError(f"Unknown model size {size}")
 
 
 class MIDIModel(nn.Module):
- def __init__(self, tokenizer: MIDITokenizer, n_layer=12, n_head=16, n_embd=1024, n_inner=4096,
- *args, **kwargs):
+ def __init__(self, config: MIDIModelConfig, *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))
- self.lm_head = nn.Linear(n_embd, tokenizer.vocab_size, bias=False)
+ self.tokenizer = config.tokenizer
+ self.net = LlamaModel(config.net_config)
+ self.net_token = LlamaModel(config.net_token_config)
+ self.lm_head = nn.Linear(config.n_embd, self.tokenizer.vocab_size, bias=False)
  self.device = "cpu"
 
  def to(self, *args, **kwargs):
@@ -71,7 +111,7 @@ class MIDIModel(nn.Module):
  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):
+ def generate(self, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20, generator=None):
  tokenizer = self.tokenizer
  max_token_seq = tokenizer.max_token_seq
  if prompt is None:
@@ -86,7 +126,7 @@ class MIDIModel(nn.Module):
  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):
+ with bar:
  while cur_len < max_len:
  end = False
  hidden = self.forward(input_tensor)[0, -1].unsqueeze(0)
@@ -123,4 +163,4 @@ class MIDIModel(nn.Module):
  bar.update(1)
  if end:
  break
- return input_tensor[0].cpu().numpy()
+ return input_tensor[0].cpu().numpy()

requirements.txt

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