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