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