speed up rendering audio

app.py

@@ -1,3 +1,5 @@
+from concurrent.futures import ThreadPoolExecutor
+
 import spaces
 import random
 import argparse
@@ -240,7 +242,8 @@ def run(model_name, tab, mid_seq, continuation_state, continuation_select, instr
 
 def finish_run(model_name, mid_seq):
  if mid_seq is None:
- return None, None, []
+ outputs = [None] * OUTPUT_BATCH_SIZE
+ return *outputs, []
  tokenizer = models[model_name].tokenizer
  outputs = []
  end_msgs = [create_msg("progress", [0, 0])]
@@ -249,16 +252,36 @@ def finish_run(model_name, mid_seq):
  for i in range(OUTPUT_BATCH_SIZE):
  events = [tokenizer.tokens2event(tokens) for tokens in mid_seq[i]]
  mid = tokenizer.detokenize(mid_seq[i])
- audio = synthesizer.synthesis(MIDI.score2opus(mid))
  with open(f"outputs/output{i + 1}.mid", 'wb') as f:
  f.write(MIDI.score2midi(mid))
- outputs += [(44100, audio), f"outputs/output{i + 1}.mid"]
+ outputs.append(f"outputs/output{i + 1}.mid")
  end_msgs += [create_msg("visualizer_clear", [i, tokenizer.version]),
  create_msg("visualizer_append", [i, events]),
  create_msg("visualizer_end", i)]
  return *outputs, send_msgs(end_msgs)
 
 
+def synthesis_task(mid):
+ return synthesizer.synthesis(MIDI.score2opus(mid))
+
+def render_audio(model_name, mid_seq, should_render_audio):
+ if (not should_render_audio) or mid_seq is None:
+ outputs = [None] * OUTPUT_BATCH_SIZE
+ return tuple(outputs)
+ tokenizer = models[model_name].tokenizer
+ outputs = []
+ if not os.path.exists("outputs"):
+ os.mkdir("outputs")
+ audio_futures = []
+ for i in range(OUTPUT_BATCH_SIZE):
+ mid = tokenizer.detokenize(mid_seq[i])
+ audio_future = thread_pool.submit(synthesis_task, mid)
+ audio_futures.append(audio_future)
+ for future in audio_futures:
+ outputs.append((44100, future.result()))
+ return tuple(outputs)
+
+
 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([])
@@ -324,6 +347,7 @@ if __name__ == "__main__":
  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")
+ thread_pool = ThreadPoolExecutor(max_workers=OUTPUT_BATCH_SIZE)
  synthesizer = MidiSynthesizer(soundfont_path)
  models_info = {
  "generic pretrain model (tv2o-medium) by skytnt": ["skytnt/midi-model-tv2o-medium", "", "tv2o-medium"],
@@ -442,20 +466,23 @@ if __name__ == "__main__":
  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=30)
  input_allow_cc = gr.Checkbox(label="allow midi cc event", value=True)
+ input_render_audio = gr.Checkbox(label="render audio after generation", value=True)
  example3 = gr.Examples([[1, 0.95, 128], [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_continuation_state = gr.State([0])
- batch_outputs = []
+ midi_outputs = []
+ audio_outputs = []
  with gr.Tabs(elem_id="output_tabs"):
  for i in range(OUTPUT_BATCH_SIZE):
  with gr.TabItem(f"output {i + 1}") as tab1:
  output_midi_visualizer = gr.HTML(elem_id=f"midi_visualizer_container_{i}")
  output_audio = gr.Audio(label="output audio", format="mp3", elem_id=f"midi_audio_{i}")
  output_midi = gr.File(label="output midi", file_types=[".mid"])
- batch_outputs += [output_audio, output_midi]
+ midi_outputs.append(output_midi)
+ audio_outputs.append(output_audio)
  run_event = run_btn.click(run, [input_model, tab_select, output_midi_seq, output_continuation_state,
  input_continuation_select, input_instruments, input_drum_kit, input_bpm,
  input_time_sig, input_key_sig, input_midi, input_midi_events,
@@ -465,12 +492,17 @@ if __name__ == "__main__":
  input_top_k, input_allow_cc],
  [output_midi_seq, output_continuation_state, input_seed, js_msg],
  concurrency_limit=10, queue=True)
- run_event.then(fn=finish_run,
- inputs=[input_model, output_midi_seq],
- outputs=batch_outputs + [js_msg],
- queue=False)
+ finish_run_event = run_event.then(fn=finish_run,
+ inputs=[input_model, output_midi_seq],
+ outputs=midi_outputs + [js_msg],
+ queue=False)
+ finish_run_event.then(fn=render_audio,
+ inputs=[input_model, output_midi_seq, input_render_audio],
+ outputs=audio_outputs,
+ queue=False)
  # stop_btn.click(None, [], [], cancels=run_event,
  # queue=False)
  undo_btn.click(undo_continuation, [input_model, output_midi_seq, output_continuation_state],
  [output_midi_seq, output_continuation_state, js_msg], queue=False)
  app.queue().launch(server_port=opt.port, share=opt.share, inbrowser=True)
+ thread_pool.shutdown()

midi_synthesizer.py

@@ -1,3 +1,5 @@
+from threading import Lock
+
 import fluidsynth
 import numpy as np
 
@@ -9,14 +11,16 @@ class MidiSynthesizer:
  fl = fluidsynth.Synth(samplerate=float(sample_rate))
  sfid = fl.sfload(soundfont_path)
  self.devices = [[fl, sfid, False]]
+ self.file_lock = Lock()
 
  def get_fluidsynth(self):
  for device in self.devices:
  if not device[2]:
  device[2] = True
  return device
- fl = fluidsynth.Synth(samplerate=float(self.sample_rate))
- sfid = fl.sfload(self.soundfont_path)
+ with self.file_lock:
+ fl = fluidsynth.Synth(samplerate=float(self.sample_rate))
+ sfid = fl.sfload(self.soundfont_path)
  device = [fl, sfid, True]
  self.devices.append(device)
  return device