update tokenizer

app.py

@@ -121,7 +121,8 @@ def send_msgs(msgs):
  return json.dumps(msgs)
 
 
-def run(model_name, tab, instruments, drum_kit, bpm, mid, midi_events, midi_opt, seed, seed_rand,
+def run(model_name, tab, instruments, drum_kit, bpm, mid, midi_events,
+ reduce_cc_st, remap_track_channel, add_default_instr, remove_empty_channels, seed, seed_rand,
  gen_events, temp, top_p, top_k, allow_cc):
  mid_seq = []
  bpm = int(bpm)
@@ -153,8 +154,11 @@ def run(model_name, tab, instruments, drum_kit, bpm, mid, midi_events, midi_opt,
  disable_patch_change = True
  disable_channels = [i for i in range(16) if i not in patches]
  elif mid is not None:
- eps = 4 if midi_opt else 0
- mid = tokenizer.tokenize(MIDI.midi2score(mid), cc_eps=eps, tempo_eps=eps)
+ eps = 4 if reduce_cc_st else 0
+ mid = tokenizer.tokenize(MIDI.midi2score(mid), cc_eps=eps, tempo_eps=eps,
+ remap_track_channel=remap_track_channel,
+ add_default_instr=add_default_instr,
+ remove_empty_channels=remove_empty_channels)
  mid = np.asarray(mid, dtype=np.int64)
  mid = mid[:int(midi_events)]
  for token_seq in mid:
@@ -306,7 +310,10 @@ if __name__ == "__main__":
  input_midi_events = gr.Slider(label="use first n midi events as prompt", minimum=1, maximum=512,
  step=1,
  value=128)
- input_midi_opt = gr.Checkbox(label="optimise midi (uncheck if your midi is generate from this model)", value=True)
+ input_reduce_cc_st = gr.Checkbox(label="reduce control_change and set_tempo events", value=True)
+ input_remap_track_channel = gr.Checkbox(label="remap tracks and channels to have only one channel per track", value=True)
+ input_add_default_instr = gr.Checkbox(label="add a default instrument to channels that don't have an instrument", value=True)
+ input_remove_empty_channels = gr.Checkbox(label="remove channels without notes", value=False)
  example2 = gr.Examples([[file, 128] for file in glob.glob("example/*.mid")],
  [input_midi, input_midi_events])
 
@@ -330,8 +337,10 @@ if __name__ == "__main__":
  output_audio = gr.Audio(label="output audio", format="mp3", elem_id="midi_audio")
  output_midi = gr.File(label="output midi", file_types=[".mid"])
  run_event = run_btn.click(run, [input_model, tab_select, input_instruments, input_drum_kit, input_bpm,
- input_midi, input_midi_events, input_midi_opt, input_seed, input_seed_rand,
- input_gen_events, input_temp, input_top_p, input_top_k, input_allow_cc],
+ input_midi, input_midi_events, input_reduce_cc_st, input_remap_track_channel,
+ input_add_default_instr, input_remove_empty_channels, input_seed,
+ input_seed_rand, input_gen_events, input_temp, input_top_p, input_top_k,
+ input_allow_cc],
  [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)

midi_tokenizer.py

@@ -42,9 +42,16 @@ class MIDITokenizer:
  tempo = int((60 / bpm) * 10 ** 6)
  return tempo
 
- def tokenize(self, midi_score, add_bos_eos=True, cc_eps=4, tempo_eps=4):
+ def tokenize(self, midi_score, add_bos_eos=True, cc_eps=4, tempo_eps=4,
+ remap_track_channel=False, add_default_instr=False, remove_empty_channels=False):
  ticks_per_beat = midi_score[0]
  event_list = {}
+ track_idx_map = {i: dict() for i in range(16)}
+ track_idx_dict = {}
+ channels = []
+ patch_channels = []
+ empty_channels = [True]*16
+ channel_note_tracks = {i: list() for i in range(16)}
  for track_idx, track in enumerate(midi_score[1:129]):
  last_notes = {}
  patch_dict = {}
@@ -53,9 +60,18 @@ class MIDITokenizer:
  for event in track:
  if event[0] not in self.events:
  continue
+ c = -1
  t = round(16 * event[1] / ticks_per_beat) # quantization
  new_event = [event[0], t // 16, t % 16, track_idx] + event[2:]
  if event[0] == "note":
+ c = event[3]
+ if c > 15 or c < 0:
+ continue
+ empty_channels[c] = False
+ track_idx_dict.setdefault(c, track_idx)
+ note_tracks = channel_note_tracks[c]
+ if track_idx not in note_tracks:
+ note_tracks.append(track_idx)
  new_event[4] = max(1, round(16 * new_event[4] / ticks_per_beat))
  elif event[0] == "set_tempo":
  if new_event[4] == 0: # invalid tempo
@@ -68,12 +84,18 @@ class MIDITokenizer:
  key = tuple(new_event[:-1])
  if event[0] == "patch_change":
  c, p = event[2:]
+ if c > 15 or c < 0:
+ continue
  last_p = patch_dict.setdefault(c, None)
  if last_p == p:
  continue
  patch_dict[c] = p
+ if c not in patch_channels:
+ patch_channels.append(c)
  elif event[0] == "control_change":
  c, cc, v = event[2:]
+ if c > 15 or c < 0:
+ continue
  last_v = control_dict.setdefault((c, cc), 0)
  if abs(last_v - v) < cc_eps:
  continue
@@ -84,6 +106,13 @@ class MIDITokenizer:
  continue
  last_tempo = tempo
 
+ if c != -1:
+ if c not in channels:
+ channels.append(c)
+ tr_map = track_idx_map[c]
+ if track_idx not in tr_map:
+ tr_map[track_idx] = 0
+
  if event[0] == "note": # to eliminate note overlap due to quantization
  cp = tuple(new_event[5:7])
  if cp in last_notes:
@@ -95,8 +124,79 @@ class MIDITokenizer:
  last_notes[cp] = (key, new_event)
  event_list[key] = new_event
  event_list = list(event_list.values())
- event_list = sorted(event_list, key=lambda e: e[1:4])
- midi_seq = []
+
+ empty_channels = [c for c in channels if empty_channels[c]]
+
+ if remap_track_channel:
+ patch_channels = []
+ channels_count = 0
+ channels_map = {9: 9} if 9 in channels else {}
+ for c in channels:
+ if c == 9:
+ continue
+ channels_map[c] = channels_count
+ channels_count += 1
+ if channels_count == 9:
+ channels_count = 10
+ channels = list(channels_map.values())
+
+ track_count = 0
+ track_idx_map_order = [k for k,v in sorted(list(channels_map.items()), key=lambda x: x[1])]
+ for c in track_idx_map_order: # tracks not to remove
+ if remove_empty_channels and c in empty_channels:
+ continue
+ tr_map = track_idx_map[c]
+ for track_idx in tr_map:
+ note_tracks = channel_note_tracks[c]
+ if len(note_tracks) != 0 and track_idx not in note_tracks:
+ continue
+ track_count += 1
+ tr_map[track_idx] = track_count
+ for c in track_idx_map_order: # tracks to remove
+ if not (remove_empty_channels and c in empty_channels):
+ continue
+ tr_map = track_idx_map[c]
+ for track_idx in tr_map:
+ note_tracks = channel_note_tracks[c]
+ if not (len(note_tracks) != 0 and track_idx not in note_tracks):
+ continue
+ track_count += 1
+ tr_map[track_idx] = track_count
+
+ empty_channels = [channels_map[c] for c in empty_channels]
+
+ for event in event_list:
+ name = event[0]
+ track_idx = event[3]
+ if name == "note":
+ c = event[5]
+ event[5] = channels_map[c]
+ event[3] = track_idx_map[c][track_idx]
+ track_idx_dict[event[5]] = event[3]
+ elif name == "set_tempo":
+ event[3] = 0
+ elif name == "control_change" or name == "patch_change":
+ c = event[4]
+ event[4] = channels_map[c]
+ tr_map = track_idx_map[c]
+ # move the event to first track of the channel if it's original track is empty
+ note_tracks = channel_note_tracks[c]
+ if len(note_tracks) != 0 and track_idx not in note_tracks:
+ track_idx = channel_note_tracks[c][0]
+ new_track_idx = tr_map.setdefault(track_idx, next(iter(tr_map.values())))
+ event[3] = new_track_idx
+ if name == "patch_change" and event[4] not in patch_channels:
+ patch_channels.append(event[4])
+
+ if add_default_instr:
+ for c in channels:
+ if c not in patch_channels:
+ event_list.append(["patch_change", 0,0, track_idx_dict[c], c, 0])
+
+ events_name_order = {"set_tempo":0, "patch_change":1, "control_change":2, "note":3}
+ events_order = lambda e: e[1:4] + [events_name_order[e[0]]]
+ event_list = sorted(event_list, key=events_order)
+
  setup_events = {}
  notes_in_setup = False
  for i, event in enumerate(event_list): # optimise setup
@@ -113,7 +213,7 @@ class MIDITokenizer:
  pre_event = event_list[i - 1]
  has_pre = event[1] + event[2] == pre_event[1] + pre_event[2]
  if (event[0] == "note" and not has_next) or (notes_in_setup and not has_pre) :
- event_list = sorted(setup_events.values(), key=lambda e: 1 if e[0] == "note" else 0) + event_list[i:]
+ event_list = sorted(setup_events.values(), key=events_order) + event_list[i:]
  break
  else:
  if event[0] == "note":
@@ -122,7 +222,10 @@ class MIDITokenizer:
  setup_events[key] = new_event
 
  last_t1 = 0
+ midi_seq = []
  for event in event_list:
+ if remove_empty_channels and event[0] in ["control_change", "patch_change"] and event[4] in empty_channels:
+ continue
  cur_t1 = event[1]
  event[1] = event[1] - last_t1
  tokens = self.event2tokens(event)
@@ -181,7 +284,7 @@ class MIDITokenizer:
  if track_idx not in tracks_dict:
  tracks_dict[track_idx] = []
  tracks_dict[track_idx].append([event[0], t] + event[4:])
- tracks = list(tracks_dict.values())
+ tracks = [tr for idx, tr in sorted(list(tracks_dict.items()), key=lambda it: it[0])]
 
  for i in range(len(tracks)): # to eliminate note overlap
  track = tracks[i]
@@ -292,7 +395,6 @@ class MIDITokenizer:
  notes_bandwidth_list = []
  instruments = {}
  piano_channels = []
- undef_instrument = False
  abs_t1 = 0
  last_t = 0
  for tsi, tokens in enumerate(midi_seq):
@@ -309,7 +411,9 @@ class MIDITokenizer:
  time_hist[t2] += 1
  if c != 9: # ignore drum channel
  if c not in instruments:
- undef_instrument = True
+ instruments[c] = 0
+ if c not in piano_channels:
+ piano_channels.append(c)
  note_windows.setdefault(abs_t1 // note_window_size, []).append(p)
  if last_t != t:
  notes_sametime = [(et, p_) for et, p_ in notes_sametime if et > last_t]
@@ -330,8 +434,6 @@ class MIDITokenizer:
  reasons.append("total_min")
  if total_notes > total_notes_max:
  reasons.append("total_max")
- if undef_instrument:
- reasons.append("undef_instr")
  if len(note_windows) == 0 and total_notes > 0:
  reasons.append("drum_only")
  if reasons: