v1.2

app.py

@@ -111,16 +111,19 @@ def create_msg(name, data):
     return {"name": name, "data": data, "uuid": uuid.uuid4().hex}
 
 
-def send_msgs(msgs, msgs_history):
+def send_msgs(msgs, msgs_history=None):
+    if msgs_history is None:
+        msgs_history = []
     msgs_history.append(msgs)
-    if len(msgs_history) > 50:
-        msgs_history.pop(0)
+    if len(msgs_history) > 25:
+        msgs_history= msgs_history[1:]
     return json.dumps(msgs_history)
 
 
-def run(model_name, tab, instruments, drum_kit, mid, midi_events, gen_events, temp, top_p, top_k, allow_cc):
+def run(model_name, tab, instruments, drum_kit, bpm, mid, midi_events, gen_events, temp, top_p, top_k, allow_cc):
     msgs_history = []
     mid_seq = []
+    bpm = int(bpm)
     gen_events = int(gen_events)
     max_len = gen_events
 
@@ -129,6 +132,8 @@ def run(model_name, tab, instruments, drum_kit, mid, midi_events, gen_events, te
     if tab == 0:
         i = 0
         mid = [[tokenizer.bos_id] + [tokenizer.pad_id] * (tokenizer.max_token_seq - 1)]
+        if bpm != 0:
+            mid.append(tokenizer.event2tokens(["set_tempo",0,0,0, bpm]))
         patches = {}
         if instruments is None:
             instruments = []
@@ -151,10 +156,10 @@ def run(model_name, tab, instruments, drum_kit, mid, midi_events, gen_events, te
         max_len += len(mid)
         for token_seq in mid:
             mid_seq.append(token_seq.tolist())
-    init_msgs = [create_msg("visualizer_clear", None)]
+    init_msgs = [create_msg("visualizer_clear", False)]
     for tokens in mid_seq:
         init_msgs.append(create_msg("visualizer_append", tokenizer.tokens2event(tokens)))
-    yield mid_seq, None, None, send_msgs(init_msgs, msgs_history), msgs_history
+    yield mid_seq, None, None, send_msgs(init_msgs, msgs_history)
     model = models[model_name]
     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,
@@ -163,22 +168,31 @@ def run(model_name, tab, instruments, drum_kit, mid, midi_events, gen_events, te
         token_seq = token_seq.tolist()
         mid_seq.append(token_seq)
         event = tokenizer.tokens2event(token_seq)
-        yield mid_seq, None, None, send_msgs([create_msg("visualizer_append", event), create_msg("progress", [i + 1, gen_events])], msgs_history), msgs_history
+        yield mid_seq, None, None, send_msgs([create_msg("visualizer_append", event), create_msg("progress", [i + 1, gen_events])], msgs_history)
     mid = tokenizer.detokenize(mid_seq)
     with open(f"output.mid", 'wb') as f:
         f.write(MIDI.score2midi(mid))
     audio = synthesis(MIDI.score2opus(mid), soundfont_path)
-    yield mid_seq, "output.mid", (44100, audio), send_msgs([create_msg("visualizer_end", None)], msgs_history), msgs_history
+    # resend all msgs
+    msgs = [create_msg("visualizer_end", None), create_msg("visualizer_clear", True)]
+    for tokens in mid_seq:
+        msgs.append(create_msg("visualizer_append", tokenizer.tokens2event(tokens)))
+    msgs.append(create_msg("visualizer_end", None))
+    yield mid_seq, "output.mid", (44100, audio), send_msgs(msgs)
 
 
-def cancel_run(mid_seq, msgs_history):
+def cancel_run(mid_seq):
     if mid_seq is None:
         return None, None, []
     mid = tokenizer.detokenize(mid_seq)
     with open(f"output.mid", 'wb') as f:
         f.write(MIDI.score2midi(mid))
     audio = synthesis(MIDI.score2opus(mid), soundfont_path)
-    return "output.mid", (44100, audio), send_msgs([create_msg("visualizer_end", None)], msgs_history)
+    msgs = [create_msg("visualizer_end", None), create_msg("visualizer_clear", True)]
+    for tokens in mid_seq:
+        msgs.append(create_msg("visualizer_append", tokenizer.tokens2event(tokens)))
+    msgs.append(create_msg("visualizer_end", None))
+    return "output.mid", (44100, audio), send_msgs(msgs)
 
 
 def load_javascript(dir="javascript"):
@@ -200,6 +214,7 @@ def load_javascript(dir="javascript"):
 
 
 def hf_hub_download_retry(repo_id, filename):
+    print(f"downloading {repo_id} {filename}")
     retry = 0
     err = None
     while retry < 30:
@@ -246,9 +261,9 @@ 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 and longer generation"
+                    " for faster running and longer generation\n\n"
+                    "**Update v1.2**: Optimise the tokenizer and dataset"
                     )
-        js_msg_history_state = gr.State(value=[])
         js_msg = gr.Textbox(elem_id="msg_receiver", visible=False)
         js_msg.change(None, [js_msg], [], js="""
         (msg_json) =>{
@@ -262,19 +277,25 @@ if __name__ == "__main__":
         tab_select = gr.State(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()),
+                input_instruments = gr.Dropdown(label="🪗instruments (auto if empty)", choices=list(patch2number.keys()),
                                                 multiselect=True, max_choices=15, type="value")
-                input_drum_kit = gr.Dropdown(label="drum kit", choices=list(drum_kits2number.keys()), type="value",
+                input_drum_kit = gr.Dropdown(label="🥁drum kit", choices=list(drum_kits2number.keys()), type="value",
                                              value="None")
+                input_bpm = gr.Slider(label="BPM (beats per minute, auto if 0)", minimum=0, maximum=255,
+                                              step=1,
+                                              value=0)
                 example1 = gr.Examples([
                     [[], "None"],
                     [["Acoustic Grand"], "None"],
-                    [["Acoustic Grand", "Violin", "Viola", "Cello", "Contrabass"], "Orchestra"],
-                    [["Flute", "Cello", "Bassoon", "Tuba"], "None"],
-                    [["Violin", "Viola", "Cello", "Contrabass", "Trumpet", "French Horn", "Brass Section",
-                      "Flute", "Piccolo", "Tuba", "Trombone", "Timpani"], "Orchestra"],
-                    [["Acoustic Guitar(nylon)", "Acoustic Guitar(steel)", "Electric Guitar(jazz)",
-                      "Electric Guitar(clean)", "Electric Guitar(muted)", "Overdriven Guitar", "Distortion Guitar",
+                    [['Acoustic Grand', 'SynthStrings 2', 'SynthStrings 1', 'Pizzicato Strings',
+                      'Pad 2 (warm)', 'Tremolo Strings', 'String Ensemble 1'], "Orchestra"],
+                    [['Trumpet', 'Oboe', 'Trombone', 'String Ensemble 1', 'Clarinet',
+                      'French Horn', 'Pad 4 (choir)', 'Bassoon', 'Flute'], "None"],
+                    [['Flute', 'French Horn', 'Clarinet', 'String Ensemble 2', 'English Horn', 'Bassoon',
+                      'Oboe', 'Pizzicato Strings'], "Orchestra"],
+                    [['Electric Piano 2', 'Lead 5 (charang)', 'Electric Bass(pick)', 'Lead 2 (sawtooth)',
+                      'Pad 1 (new age)', 'Orchestra Hit', 'Cello', 'Electric Guitar(clean)'], "Standard"],
+                    [["Electric Guitar(clean)", "Electric Guitar(muted)", "Overdriven Guitar", "Distortion Guitar",
                       "Electric Bass(finger)"], "Standard"]
                 ], [input_instruments, input_drum_kit])
             with gr.TabItem("midi prompt") as tab2:
@@ -292,19 +313,19 @@ if __name__ == "__main__":
         with gr.Accordion("options", open=False):
             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.98)
-            input_top_k = gr.Slider(label="top k", minimum=1, maximum=128, step=1, value=48)
+            input_top_k = gr.Slider(label="top k", minimum=1, maximum=128, step=1, value=20)
             input_allow_cc = gr.Checkbox(label="allow midi cc event", value=True)
-            example3 = gr.Examples([[1, 0.98, 12], [1.2, 0.95, 8]], [input_temp, input_top_p, input_top_k])
+            example3 = gr.Examples([[1, 0.98, 20], [1, 0.98, 12]], [input_temp, input_top_p, input_top_k])
         run_btn = gr.Button("generate", variant="primary")
         stop_btn = gr.Button("stop and output")
         output_midi_seq = gr.State()
         output_midi_visualizer = gr.HTML(elem_id="midi_visualizer_container")
         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_midi,
-                                        input_midi_events, input_gen_events, input_temp, input_top_p, input_top_k,
-                                        input_allow_cc],
-                                  [output_midi_seq, output_midi, output_audio, js_msg, js_msg_history_state],
+        run_event = run_btn.click(run, [input_model, tab_select, input_instruments, input_drum_kit, input_bpm,
+                                        input_midi, input_midi_events, input_gen_events, input_temp,
+                                        input_top_p, input_top_k, input_allow_cc],
+                                  [output_midi_seq, output_midi, output_audio, js_msg],
                                   concurrency_limit=3)
-        stop_btn.click(cancel_run, [output_midi_seq, js_msg_history_state], [output_midi, output_audio, js_msg], cancels=run_event, queue=False)
+        stop_btn.click(cancel_run, [output_midi_seq], [output_midi, output_audio, js_msg], cancels=run_event, queue=False)
     app.launch(server_port=opt.port, share=opt.share, inbrowser=True)

javascript/app.js

@@ -146,13 +146,14 @@ class MidiVisualizer extends HTMLElement{
         this.setPlayTime(0);
     }
 
-    clearMidiEvents(){
+    clearMidiEvents(keepColor=false){
         this.pause()
         this.midiEvents = [];
         this.activeNotes = [];
         this.midiTimes = [];
         this.t1 = 0
-        this.colorMap.clear()
+        if (!keepColor)
+            this.colorMap.clear()
         this.setPlayTime(0);
         this.totalTimeMs = 0;
         this.playTimeMs = 0
@@ -426,7 +427,7 @@ customElements.define('midi-visualizer', MidiVisualizer);
         handled_msgs.push(msg.uuid);
         switch (msg.name) {
             case "visualizer_clear":
-                midi_visualizer.clearMidiEvents();
+                midi_visualizer.clearMidiEvents(msg.data);
                 createProgressBar(midi_visualizer_container_inited)
                 break;
             case "visualizer_append":

midi_tokenizer.py

@@ -42,22 +42,48 @@ class MIDITokenizer:
         tempo = int((60 / bpm) * 10 ** 6)
         return tempo
 
-    def tokenize(self, midi_score, add_bos_eos=True):
+    def tokenize(self, midi_score, add_bos_eos=True, cc_eps=4, tempo_eps=4):
         ticks_per_beat = midi_score[0]
         event_list = {}
         for track_idx, track in enumerate(midi_score[1:129]):
             last_notes = {}
+            patch_dict = {}
+            control_dict = {}
+            last_tempo = 0
             for event in track:
+                if event[0] not in self.events:
+                    continue
                 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":
                     new_event[4] = max(1, round(16 * new_event[4] / ticks_per_beat))
                 elif event[0] == "set_tempo":
-                    new_event[4] = int(self.tempo2bpm(new_event[4]))
+                    if new_event[4] == 0: # invalid tempo
+                        continue
+                    bpm = int(self.tempo2bpm(new_event[4]))
+                    new_event[4] = min(bpm, 255)
                 if event[0] == "note":
                     key = tuple(new_event[:4] + new_event[5:-1])
                 else:
                     key = tuple(new_event[:-1])
+                if event[0] == "patch_change":
+                    c, p = event[2:]
+                    last_p = patch_dict.setdefault(c, None)
+                    if last_p == p:
+                        continue
+                    patch_dict[c] = p
+                elif event[0] == "control_change":
+                    c, cc, v = event[2:]
+                    last_v = control_dict.setdefault((c, cc), 0)
+                    if abs(last_v - v) < cc_eps:
+                        continue
+                    control_dict[(c, cc)] = v
+                elif event[0] == "set_tempo":
+                    tempo = new_event[-1]
+                    if abs(last_tempo - tempo) < tempo_eps:
+                        continue
+                    last_tempo = tempo
+
                 if event[0] == "note":  # to eliminate note overlap due to quantization
                     cp = tuple(new_event[5:7])
                     if cp in last_notes:
@@ -71,21 +97,39 @@ class MIDITokenizer:
         event_list = list(event_list.values())
         event_list = sorted(event_list, key=lambda e: e[1:4])
         midi_seq = []
+        setup_events = {}
+        notes_in_setup = False
+        for i, event in enumerate(event_list):  # optimise setup
+            new_event = [*event]
+            if event[0] != "note":
+                new_event[1] = 0
+                new_event[2] = 0
+            has_next = False
+            has_pre = False
+            if i < len(event_list) - 1:
+                next_event = event_list[i + 1]
+                has_next = event[1] + event[2] == next_event[1] + next_event[2]
+            if notes_in_setup and i > 0:
+                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:]
+                break
+            else:
+                if event[0] == "note":
+                    notes_in_setup = True
+                key = tuple(event[3:-1])
+            setup_events[key] = new_event
 
         last_t1 = 0
         for event in event_list:
-            name = event[0]
-            if name in self.event_ids:
-                params = event[1:]
-                cur_t1 = params[0]
-                params[0] = params[0] - last_t1
-                if not all([0 <= params[i] < self.event_parameters[p] for i, p in enumerate(self.events[name])]):
-                    continue
-                tokens = [self.event_ids[name]] + [self.parameter_ids[p][params[i]]
-                                                   for i, p in enumerate(self.events[name])]
-                tokens += [self.pad_id] * (self.max_token_seq - len(tokens))
-                midi_seq.append(tokens)
-                last_t1 = cur_t1
+            cur_t1 = event[1]
+            event[1] = event[1] - last_t1
+            tokens = self.event2tokens(event)
+            if not tokens:
+                continue
+            midi_seq.append(tokens)
+            last_t1 = cur_t1
 
         if add_bos_eos:
             bos = [self.bos_id] + [self.pad_id] * (self.max_token_seq - 1)
@@ -96,6 +140,8 @@ class MIDITokenizer:
     def event2tokens(self, event):
         name = event[0]
         params = event[1:]
+        if not all([0 <= params[i] < self.event_parameters[p] for i, p in enumerate(self.events[name])]):
+            return []
         tokens = [self.event_ids[name]] + [self.parameter_ids[p][params[i]]
                                            for i, p in enumerate(self.events[name])]
         tokens += [self.pad_id] * (self.max_token_seq - len(tokens))
@@ -120,14 +166,10 @@ class MIDITokenizer:
         t1 = 0
         for tokens in midi_seq:
             if tokens[0] in self.id_events:
-                name = self.id_events[tokens[0]]
-                if len(tokens) <= len(self.events[name]):
+                event = self.tokens2event(tokens)
+                if not event:
                     continue
-                params = tokens[1:]
-                params = [params[i] - self.parameter_ids[p][0] for i, p in enumerate(self.events[name])]
-                if not all([0 <= params[i] < self.event_parameters[p] for i, p in enumerate(self.events[name])]):
-                    continue
-                event = [name] + params
+                name = event[0]
                 if name == "set_tempo":
                     event[4] = self.bpm2tempo(event[4])
                 if event[0] == "note":
@@ -183,7 +225,7 @@ class MIDITokenizer:
         return img
 
     def augment(self, midi_seq, max_pitch_shift=4, max_vel_shift=10, max_cc_val_shift=10, max_bpm_shift=10,
-                max_track_shift=128, max_channel_shift=16):
+                max_track_shift=0, max_channel_shift=16):
         pitch_shift = random.randint(-max_pitch_shift, max_pitch_shift)
         vel_shift = random.randint(-max_vel_shift, max_vel_shift)
         cc_val_shift = random.randint(-max_cc_val_shift, max_cc_val_shift)
@@ -239,16 +281,85 @@ class MIDITokenizer:
             midi_seq_new.append(tokens_new)
         return midi_seq_new
 
-    def check_alignment(self, midi_seq, threshold=0.3):
-        total = 0
-        hist = [0] * 16
-        for tokens in midi_seq:
-            if tokens[0] in self.id_events and self.id_events[tokens[0]] == "note":
-                t2 = tokens[2] - self.parameter_ids["time2"][0]
-                total += 1
-                hist[t2] += 1
-        if total == 0:
-            return False
-        hist = sorted(hist, reverse=True)
-        p = sum(hist[:2]) / total
-        return p > threshold
+    def check_quality(self, midi_seq, alignment_min=0.4, tonality_min=0.8, piano_max=0.7, notes_bandwidth_min=3, notes_density_max=30, notes_density_min=2.5, total_notes_max=10000, total_notes_min=500, note_window_size=16):
+        total_notes = 0
+        channels = []
+        time_hist = [0] * 16
+        note_windows = {}
+        notes_sametime = []
+        notes_density_list = []
+        tonality_list = []
+        notes_bandwidth_list = []
+        instruments = {}
+        piano_channels = []
+        undef_instrument = False
+        abs_t1 = 0
+        last_t = 0
+        for tsi, tokens in enumerate(midi_seq):
+            event = self.tokens2event(tokens)
+            if not event:
+                continue
+            t1, t2, tr = event[1:4]
+            abs_t1 += t1
+            t = abs_t1 * 16 + t2
+            c = None
+            if event[0] == "note":
+                d, c, p, v = event[4:]
+                total_notes += 1
+                time_hist[t2] += 1
+                if c != 9:  # ignore drum channel
+                    if c not in instruments:
+                        undef_instrument = True
+                    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]
+                    notes_sametime_p = [p_ for _, p_ in notes_sametime]
+                    if len(notes_sametime) > 0:
+                        notes_bandwidth_list.append(max(notes_sametime_p) - min(notes_sametime_p))
+                notes_sametime.append((t + d - 1, p))
+            elif event[0] == "patch_change":
+                c, p = event[4:]
+                instruments[c] = p
+                if p == 0 and c not in piano_channels:
+                    piano_channels.append(c)
+            if c is not None and c not in channels:
+                channels.append(c)
+            last_t = t
+        reasons = []
+        if total_notes < total_notes_min:
+            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:
+            return False, reasons
+        time_hist = sorted(time_hist, reverse=True)
+        alignment = sum(time_hist[:2]) / total_notes
+        for notes in note_windows.values():
+            key_hist = [0] * 12
+            for p in notes:
+                key_hist[p % 12] += 1
+            key_hist = sorted(key_hist, reverse=True)
+            tonality_list.append(sum(key_hist[:7]) / len(notes))
+            notes_density_list.append(len(notes) / note_window_size)
+        tonality_list = sorted(tonality_list)
+        tonality = sum(tonality_list)/len(tonality_list)
+        notes_bandwidth = sum(notes_bandwidth_list)/len(notes_bandwidth_list) if notes_bandwidth_list else 0
+        notes_density = max(notes_density_list) if notes_density_list else 0
+        piano_ratio = len(piano_channels) / len(channels)
+        if len(channels) <=3:  # ignore piano threshold if it is a piano solo midi
+            piano_max = 1
+        if alignment < alignment_min: # check weather the notes align to the bars (because some midi files are recorded)
+            reasons.append("alignment")
+        if tonality < tonality_min:  # check whether the music is tonal
+            reasons.append("tonality")
+        if notes_bandwidth < notes_bandwidth_min:  # check whether music is melodic line only
+            reasons.append("bandwidth")
+        if not notes_density_min < notes_density < notes_density_max:
+            reasons.append("density")
+        if piano_ratio > piano_max: # check whether most instruments is piano (because some midi files don't have instruments assigned correctly)
+            reasons.append("piano")
+        return not reasons, reasons