add undo button

app.py

@@ -121,7 +121,7 @@ def send_msgs(msgs):
     return json.dumps(msgs)
 
 
-def run(model_name, tab, mid_seq, instruments, drum_kit, bpm, time_sig, key_sig, mid, midi_events,
+def run(model_name, tab, mid_seq, continuation_state, instruments, drum_kit, bpm, time_sig, key_sig, 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):
     model = models[model_name]
@@ -187,8 +187,10 @@ def run(model_name, tab, mid_seq, instruments, drum_kit, bpm, time_sig, key_sig,
         for token_seq in mid:
             mid_seq.append(token_seq.tolist())
     elif tab == 2 and mid_seq is not None:
+        continuation_state.append(len(mid_seq))
         mid = np.asarray(mid_seq, dtype=np.int64)
     else:
+        continuation_state = [0]
         mid_seq = []
         mid = None
 
@@ -196,12 +198,11 @@ def run(model_name, tab, mid_seq, instruments, drum_kit, bpm, time_sig, key_sig,
         max_len += len(mid)
 
     events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
-    if tab == 2:
-        init_msgs = [create_msg("visualizer_continue", tokenizer.version)]
-    else:
-        init_msgs = [create_msg("visualizer_clear", tokenizer.version),
-                     create_msg("visualizer_append", events)]
-    yield mid_seq, None, None, seed, send_msgs(init_msgs)
+    init_msgs = [create_msg("progress", [0, gen_events])]
+    if tab != 2:
+        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)
@@ -213,29 +214,51 @@ def run(model_name, tab, mid_seq, instruments, drum_kit, bpm, time_sig, key_sig,
         events.append(tokenizer.tokens2event(token_seq))
         ct = time.time()
         if ct - t > 0.5:
-            yield mid_seq, None, None, seed, send_msgs(
-                [create_msg("visualizer_append", events), create_msg("progress", [i + 1, gen_events])])
+            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)
+    audio = synthesizer.synthesis(MIDI.score2opus(mid))
     with open(f"output.mid", 'wb') as f:
         f.write(MIDI.score2midi(mid))
-    audio = synthesizer.synthesis(MIDI.score2opus(mid))
-    events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
-    yield mid_seq, "output.mid", (44100, audio), seed, send_msgs([create_msg("visualizer_end", events)])
+    end_msgs = [create_msg("visualizer_clear", tokenizer.version),
+                create_msg("visualizer_append", events),
+                create_msg("visualizer_end", None),
+                create_msg("progress", [0, 0])]
+    yield mid_seq, continuation_state, "output.mid", (44100, audio), seed, send_msgs(end_msgs)
 
 
 def cancel_run(model_name, mid_seq):
     if mid_seq is None:
         return None, None, []
     tokenizer = models[model_name][2]
+    events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
     mid = tokenizer.detokenize(mid_seq)
+    audio = synthesizer.synthesis(MIDI.score2opus(mid))
     with open(f"output.mid", 'wb') as f:
         f.write(MIDI.score2midi(mid))
-    audio = synthesizer.synthesis(MIDI.score2opus(mid))
+    end_msgs = [create_msg("visualizer_clear", tokenizer.version),
+                create_msg("visualizer_append", events),
+                create_msg("visualizer_end", None),
+                create_msg("progress", [0, 0])]
+    return "output.mid", (44100, audio), send_msgs(end_msgs)
+
+
+def undo_continuation(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]
     events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
-    return "output.mid", (44100, audio), send_msgs([create_msg("visualizer_end", events)])
+    end_msgs = [create_msg("visualizer_clear", tokenizer.version),
+                create_msg("visualizer_append", events),
+                create_msg("visualizer_end", None),
+                create_msg("progress", [0, 0])]
+    return mid_seq, continuation_state, send_msgs(end_msgs)
 
 
 def load_javascript(dir="javascript"):
@@ -341,7 +364,7 @@ if __name__ == "__main__":
                                   type="value", value=list(models.keys())[0])
         tab_select = gr.State(value=0)
         with gr.Tabs():
-            with gr.TabItem("instrument prompt") as tab1:
+            with gr.TabItem("custom prompt") as tab1:
                 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",
@@ -388,6 +411,7 @@ if __name__ == "__main__":
                                        [input_midi, input_midi_events])
             with gr.TabItem("last output prompt") as tab3:
                 gr.Markdown("Continue generating on the last output. Just click the generate button")
+                undo_btn = gr.Button("undo the last continuation")
 
         tab1.select(lambda: 0, None, tab_select, queue=False)
         tab2.select(lambda: 1, None, tab_select, queue=False)
@@ -406,18 +430,22 @@ if __name__ == "__main__":
         run_btn = gr.Button("generate", variant="primary")
         stop_btn = gr.Button("stop and output")
         output_midi_seq = gr.State()
+        output_continuation_state = gr.State([0])
         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, output_midi_seq, input_instruments,
-                                        input_drum_kit, input_bpm, input_time_sig, input_key_sig, input_midi,
-                                        input_midi_events, input_reduce_cc_st, input_remap_track_channel,
+        run_event = run_btn.click(run, [input_model, tab_select, output_midi_seq, output_continuation_state,
+                                        input_instruments, input_drum_kit, input_bpm, input_time_sig, input_key_sig,
+                                        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],
+                                  [output_midi_seq, output_continuation_state,
+                                   output_midi, output_audio, input_seed, js_msg],
                                   concurrency_limit=3)
         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],
+                       [output_midi_seq, output_continuation_state, js_msg], queue=False)
     app.launch(server_port=opt.port, share=opt.share, inbrowser=True)

javascript/app.js

@@ -400,6 +400,8 @@ customElements.define('midi-visualizer', MidiVisualizer);
         }
     })
 
+    let hasProgressBar = false;
+
     function createProgressBar(progressbarContainer){
         let parentProgressbar = progressbarContainer.parentNode;
         let divProgress = document.createElement('div');
@@ -421,15 +423,23 @@ customElements.define('midi-visualizer', MidiVisualizer);
         divInner.style.width = "0%";
         divProgress.appendChild(divInner);
         parentProgressbar.insertBefore(divProgress, progressbarContainer);
+        hasProgressBar = true;
     }
 
     function removeProgressBar(progressbarContainer){
         let parentProgressbar = progressbarContainer.parentNode;
         let divProgress = parentProgressbar.querySelector(".progressDiv");
         parentProgressbar.removeChild(divProgress);
+        hasProgressBar = false;
     }
 
     function setProgressBar(progressbarContainer, progress, total){
+        if (!hasProgressBar)
+            createProgressBar(midi_visualizer_container_inited)
+        if (hasProgressBar && total === 0){
+            removeProgressBar(midi_visualizer_container_inited)
+            return
+        }
         let parentProgressbar = progressbarContainer.parentNode;
         let divProgress = parentProgressbar.querySelector(".progressDiv");
         let divInner = parentProgressbar.querySelector(".progress");
@@ -453,31 +463,21 @@ customElements.define('midi-visualizer', MidiVisualizer);
             case "visualizer_clear":
                 midi_visualizer.clearMidiEvents(false);
                 midi_visualizer.version = msg.data
-                createProgressBar(midi_visualizer_container_inited)
-                break;
-            case "visualizer_continue":
-                midi_visualizer.version = msg.data
-                createProgressBar(midi_visualizer_container_inited)
                 break;
             case "visualizer_append":
                 msg.data.forEach( value => {
                     midi_visualizer.appendMidiEvent(value);
                 })
                 break;
+            case "visualizer_end":
+                midi_visualizer.finishAppendMidiEvent()
+                midi_visualizer.setPlayTime(0);
+                break;
             case "progress":
                 let progress = msg.data[0]
                 let total = msg.data[1]
                 setProgressBar(midi_visualizer_container_inited, progress, total)
                 break;
-            case "visualizer_end":
-                midi_visualizer.clearMidiEvents(true);
-                msg.data.forEach( value => {
-                    midi_visualizer.appendMidiEvent(value);
-                })
-                midi_visualizer.finishAppendMidiEvent()
-                midi_visualizer.setPlayTime(0);
-                removeProgressBar(midi_visualizer_container_inited);
-                break;
             default:
         }
     }

midi_synthesizer.py

@@ -1,6 +1,7 @@
 import fluidsynth
 import numpy as np
 
+
 class MidiSynthesizer:
     def __init__(self, soundfont_path, sample_rate=44100):
         self.soundfont_path = soundfont_path
@@ -21,8 +22,8 @@ class MidiSynthesizer:
         return device
 
     def release_fluidsynth(self, device):
-        device[0].system_reset()
         device[0].get_samples(self.sample_rate*5) # wait for silence
+        device[0].system_reset()
         device[2] = False
 
     def synthesis(self, midi_opus):
@@ -73,4 +74,4 @@ class MidiSynthesizer:
             if max_val != 0:
                 ss = (ss / max_val) * np.iinfo(np.int16).max
         ss = ss.astype(np.int16)
-        return ss
+        return ss