batch

.gitignore

@@ -151,3 +151,4 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
 output.mid
+/outputs/

app.py

@@ -18,11 +18,12 @@ from midi_model import MIDIModel, MIDIModelConfig
 from midi_synthesizer import MidiSynthesizer
 
 MAX_SEED = np.iinfo(np.int32).max
+OUTPUT_BATCH_SIZE = 4
 in_space = os.getenv("SYSTEM") == "spaces"
 
 
 @torch.inference_mode()
-def generate(model: MIDIModel, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
+def generate(model: MIDIModel, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98, top_k=20,
              disable_patch_change=False, disable_control_change=False, disable_channels=None, generator=None):
     tokenizer = model.tokenizer
     if disable_channels is not None:
@@ -33,49 +34,69 @@ def generate(model: MIDIModel, prompt=None, max_len=512, temp=1.0, top_p=0.98, t
     if prompt is None:
         input_tensor = torch.full((1, max_token_seq), tokenizer.pad_id, dtype=torch.long, device=model.device)
         input_tensor[0, 0] = tokenizer.bos_id  # bos
+        input_tensor = input_tensor.unsqueeze(0)
+        input_tensor = torch.cat([input_tensor] * batch_size, dim=0)
     else:
-        prompt = prompt[:, :max_token_seq]
+        if len(prompt.shape) == 2:
+            prompt = prompt[None, :]
+            prompt = np.repeat(prompt, repeats=batch_size, axis=0)
+        elif prompt.shape[0] == 1:
+            prompt = np.repeat(prompt, repeats=batch_size, axis=0)
+        elif len(prompt.shape) != 3 or prompt.shape[0] != batch_size:
+            raise ValueError(f"invalid shape for prompt, {prompt.shape}")
+        prompt = prompt[..., :max_token_seq]
         if prompt.shape[-1] < max_token_seq:
-            prompt = np.pad(prompt, ((0, 0), (0, max_token_seq - prompt.shape[-1])),
+            prompt = np.pad(prompt, ((0, 0), (0, 0), (0, max_token_seq - prompt.shape[-1])),
                             mode="constant", constant_values=tokenizer.pad_id)
         input_tensor = torch.from_numpy(prompt).to(dtype=torch.long, device=model.device)
-    input_tensor = input_tensor.unsqueeze(0)
     cur_len = input_tensor.shape[1]
     bar = tqdm.tqdm(desc="generating", total=max_len - cur_len)
     with bar:
         while cur_len < max_len:
-            end = False
-            hidden = model.forward(input_tensor)[0, -1].unsqueeze(0)
+            end = [False] * batch_size
+            hidden = model.forward(input_tensor)[:, -1]
             next_token_seq = None
-            event_name = ""
+            event_names = [""] * batch_size
             for i in range(max_token_seq):
-                mask = torch.zeros(tokenizer.vocab_size, dtype=torch.int64, device=model.device)
-                if i == 0:
-                    mask_ids = list(tokenizer.event_ids.values()) + [tokenizer.eos_id]
-                    if disable_patch_change:
-                        mask_ids.remove(tokenizer.event_ids["patch_change"])
-                    if disable_control_change:
-                        mask_ids.remove(tokenizer.event_ids["control_change"])
-                    mask[mask_ids] = 1
-                else:
-                    param_name = tokenizer.events[event_name][i - 1]
-                    mask_ids = tokenizer.parameter_ids[param_name]
-                    if param_name == "channel":
-                        mask_ids = [i for i in mask_ids if i not in disable_channels]
-                    mask[mask_ids] = 1
+                mask = torch.zeros((batch_size, tokenizer.vocab_size), dtype=torch.int64, device=model.device)
+                for b in range(batch_size):
+                    if end[b]:
+                        mask[b, tokenizer.pad_id] = 1
+                        continue
+                    if i == 0:
+                        mask_ids = list(tokenizer.event_ids.values()) + [tokenizer.eos_id]
+                        if disable_patch_change:
+                            mask_ids.remove(tokenizer.event_ids["patch_change"])
+                        if disable_control_change:
+                            mask_ids.remove(tokenizer.event_ids["control_change"])
+                        mask[b, mask_ids] = 1
+                    else:
+                        param_names = tokenizer.events[event_names[b]]
+                        if i > len(param_names):
+                            mask[b, tokenizer.pad_id] = 1
+                            continue
+                        param_name = param_names[i - 1]
+                        mask_ids = tokenizer.parameter_ids[param_name]
+                        if param_name == "channel":
+                            mask_ids = [i for i in mask_ids if i not in disable_channels]
+                        mask[b, mask_ids] = 1
+                mask = mask.unsqueeze(1)
                 logits = model.forward_token(hidden, next_token_seq)[:, -1:]
                 scores = torch.softmax(logits / temp, dim=-1) * mask
-                sample = model.sample_top_p_k(scores, top_p, top_k, generator=generator)
+                samples = model.sample_top_p_k(scores, top_p, top_k, generator=generator)
                 if i == 0:
-                    next_token_seq = sample
-                    eid = sample.item()
-                    if eid == tokenizer.eos_id:
-                        end = True
-                        break
-                    event_name = tokenizer.id_events[eid]
+                    next_token_seq = samples
+                    for b in range(batch_size):
+                        if end[b]:
+                            continue
+                        eid = samples[b].item()
+                        if eid == tokenizer.eos_id:
+                            end[b] = True
+                        else:
+                            event_names[b] = tokenizer.id_events[eid]
                 else:
-                    next_token_seq = torch.cat([next_token_seq, sample], dim=1)
-                    if len(tokenizer.events[event_name]) == i:
+                    next_token_seq = torch.cat([next_token_seq, samples], dim=1)
+                    if all([len(tokenizer.events[event_names[b]]) == i for b in range(batch_size) if not end[b]]):
                         break
             if next_token_seq.shape[1] < max_token_seq:
                 next_token_seq = F.pad(next_token_seq, (0, max_token_seq - next_token_seq.shape[1]),
@@ -84,8 +105,8 @@ def generate(model: MIDIModel, prompt=None, max_len=512, temp=1.0, top_p=0.98, t
             input_tensor = torch.cat([input_tensor, next_token_seq], dim=1)
             cur_len += 1
             bar.update(1)
-            yield next_token_seq.reshape(-1).cpu().numpy()
-            if end:
+            yield next_token_seq[:, 0].cpu().numpy()
+            if all(end):
                 break
 
 
@@ -96,8 +117,9 @@ def create_msg(name, data):
 def send_msgs(msgs):
     return json.dumps(msgs)
 
-def get_duration(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,
+
+def get_duration(model_name, tab, mid_seq, continuation_state, continuation_select, 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):
     if "large" in model_name:
         return gen_events // 10 + 15
@@ -106,9 +128,9 @@ def get_duration(model_name, tab, mid_seq, continuation_state, instruments, drum
 
 
 @spaces.GPU(duration=get_duration)
-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):
+def run(model_name, tab, mid_seq, continuation_state, continuation_select, 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]
     model.to(device=opt.device)
     tokenizer = model.tokenizer
@@ -156,8 +178,8 @@ def run(model_name, tab, mid_seq, continuation_state, instruments, drum_kit, bpm
             patches[9] = drum_kits2number[drum_kit]
         for i, (c, p) in enumerate(patches.items()):
             mid.append(tokenizer.event2tokens(["patch_change", 0, 0, i + 1, c, p]))
-        mid_seq = mid
-        mid = np.asarray(mid, dtype=np.int64)
+        mid = np.asarray([mid] * OUTPUT_BATCH_SIZE, dtype=np.int64)
+        mid_seq = mid.tolist()
         if len(instruments) > 0:
             disable_patch_change = True
             disable_channels = [i for i in range(16) if i not in patches]
@@ -167,84 +189,91 @@ def run(model_name, tab, mid_seq, continuation_state, instruments, drum_kit, bpm
                                  remap_track_channel=remap_track_channel,
                                  add_default_instr=add_default_instr,
                                  remove_empty_channels=remove_empty_channels)
-        mid = np.asarray(mid, dtype=np.int64)
         mid = mid[:int(midi_events)]
-        mid_seq = []
-        for token_seq in mid:
-            mid_seq.append(token_seq.tolist())
+        mid = np.asarray([mid] * OUTPUT_BATCH_SIZE, dtype=np.int64)
+        mid_seq = mid.tolist()
     elif tab == 2 and mid_seq is not None:
-        continuation_state.append(len(mid_seq))
         mid = np.asarray(mid_seq, dtype=np.int64)
+        if continuation_select > 0:
+            continuation_state.append(mid_seq)
+            mid = np.repeat(mid[continuation_select - 1:continuation_select], repeats=OUTPUT_BATCH_SIZE, axis=0)
+            mid_seq = mid.tolist()
+        else:
+            continuation_state.append(mid.shape[1])
     else:
         continuation_state = [0]
-        mid_seq = []
-        mid = None
+        mid = [[tokenizer.bos_id] + [tokenizer.pad_id] * (tokenizer.max_token_seq - 1)]
+        mid = np.asarray([mid] * OUTPUT_BATCH_SIZE, dtype=np.int64)
+        mid_seq = mid.tolist()
 
     if mid is not None:
-        max_len += len(mid)
+        max_len += mid.shape[1]
 
-    events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
     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)
-    events = []
-    t = time.time() + 1
-    for i, token_seq in enumerate(midi_generator):
-        token_seq = token_seq.tolist()
-        mid_seq.append(token_seq)
-        events.append(tokenizer.tokens2event(token_seq))
-        ct = time.time()
-        if ct - t > 0.5:
-            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 = []
+    if not (tab == 2 and continuation_select == 0):
+        for i in range(OUTPUT_BATCH_SIZE):
+            events = [tokenizer.tokens2event(tokens) for tokens in mid_seq[i]]
+            init_msgs += [create_msg("visualizer_clear", [i, tokenizer.version]),
+                          create_msg("visualizer_append", [i, events])]
+    yield mid_seq, continuation_state, seed, send_msgs(init_msgs)
+    midi_generator = generate(model, mid, batch_size=OUTPUT_BATCH_SIZE, 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)
+    events = [list() for i in range(OUTPUT_BATCH_SIZE)]
+    t = time.time()
+    for i, token_seqs in enumerate(midi_generator):
+        token_seqs = token_seqs.tolist()
+        for j in range(OUTPUT_BATCH_SIZE):
+            token_seq = token_seqs[j]
+            mid_seq[j].append(token_seq)
+            events[j].append(tokenizer.tokens2event(token_seq))
+        if time.time() - t > 0.2:
+            msgs = [create_msg("progress", [i + 1, gen_events])]
+            for j in range(OUTPUT_BATCH_SIZE):
+                msgs += [create_msg("visualizer_append", [j, events[j]])]
+                events[j] = list()
+            yield mid_seq, continuation_state, seed, send_msgs(msgs)
+            t = time.time()
+    yield mid_seq, continuation_state, seed, send_msgs([])
 
-    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))
-    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):
+def finish_run(model_name, mid_seq):
     if mid_seq is None:
         return None, None, []
     tokenizer = models[model_name].tokenizer
-    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))
-    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)
+    outputs = []
+    end_msgs = [create_msg("progress", [0, 0])]
+    if not os.path.exists("outputs"):
+        os.mkdir("outputs")
+    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"]
+        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 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([])
-    mid_seq = mid_seq[:continuation_state[-1]]
-    continuation_state = continuation_state[:-1]
     tokenizer = models[model_name].tokenizer
-    events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
-    end_msgs = [create_msg("visualizer_clear", tokenizer.version),
-                create_msg("visualizer_append", events),
-                create_msg("visualizer_end", None),
-                create_msg("progress", [0, 0])]
+    if isinstance(continuation_state[-1], list):
+        mid_seq = continuation_state[-1]
+    else:
+        mid_seq = [ms[:continuation_state[-1]] for ms in mid_seq]
+    continuation_state = continuation_state[:-1]
+    end_msgs = [create_msg("progress", [0, 0])]
+    for i in range(OUTPUT_BATCH_SIZE):
+        events = [tokenizer.tokens2event(tokens) for tokens in mid_seq[i]]
+        end_msgs += [create_msg("visualizer_clear", [i, tokenizer.version]),
+                     create_msg("visualizer_append", [i, events]),
+                     create_msg("visualizer_end", i)]
     return mid_seq, continuation_state, send_msgs(end_msgs)
 
 
@@ -296,13 +325,14 @@ if __name__ == "__main__":
     opt = parser.parse_args()
     soundfont_path = hf_hub_download_retry(repo_id="skytnt/midi-model", filename="soundfont.sf2")
     synthesizer = MidiSynthesizer(soundfont_path)
-    models_info = {"generic pretrain model (tv2o-medium) by skytnt": ["skytnt/midi-model-tv2o-medium", "", "tv2o-medium"],
-                   "generic pretrain model (tv2o-large) by asigalov61": ["asigalov61/Music-Llama", "", "tv2o-large"],
-                   "generic pretrain model (tv2o-medium) by asigalov61": ["asigalov61/Music-Llama-Medium", "", "tv2o-medium"],
-                   "generic pretrain model (tv1-medium) by skytnt": ["skytnt/midi-model", "", "tv1-medium"],
-                   "j-pop finetune model (tv2o-medium) by skytnt": ["skytnt/midi-model-ft", "jpop-tv2o-medium/", "tv2o-medium"],
-                   "touhou finetune model (tv2o-medium) by skytnt": ["skytnt/midi-model-ft", "touhou-tv2o-medium/", "tv2o-medium"],
-                   }
+    models_info = {
+        "generic pretrain model (tv2o-medium) by skytnt": ["skytnt/midi-model-tv2o-medium", "", "tv2o-medium"],
+        "generic pretrain model (tv2o-large) by asigalov61": ["asigalov61/Music-Llama", "", "tv2o-large"],
+        "generic pretrain model (tv2o-medium) by asigalov61": ["asigalov61/Music-Llama-Medium", "", "tv2o-medium"],
+        "generic pretrain model (tv1-medium) by skytnt": ["skytnt/midi-model", "", "tv1-medium"],
+        "j-pop finetune model (tv2o-medium) by skytnt": ["skytnt/midi-model-ft", "jpop-tv2o-medium/", "tv2o-medium"],
+        "touhou finetune model (tv2o-medium) by skytnt": ["skytnt/midi-model-ft", "touhou-tv2o-medium/", "tv2o-medium"],
+    }
     models = {}
     if opt.device == "cuda":
         torch.backends.cudnn.deterministic = True
@@ -391,7 +421,12 @@ if __name__ == "__main__":
                 example2 = gr.Examples([[file, 128] for file in glob.glob("example/*.mid")],
                                        [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")
+                gr.Markdown("Continue generating on the last output.")
+                input_continuation_select = gr.Radio(label="select output to continue generating", value="all",
+                                                     choices=["all"] + [f"output{i + 1}" for i in
+                                                                        range(OUTPUT_BATCH_SIZE)],
+                                                     type="index"
+                                                     )
                 undo_btn = gr.Button("undo the last continuation")
 
         tab1.select(lambda: 0, None, tab_select, queue=False)
@@ -413,21 +448,29 @@ if __name__ == "__main__":
         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"])
+        batch_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]
         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_continuation_select, 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_continuation_state,
-                                   output_midi, output_audio, input_seed, js_msg],
-                                  concurrency_limit=10)
-        stop_btn.click(cancel_run, [input_model, output_midi_seq],
-                       [output_midi, output_audio, js_msg],
-                       cancels=run_event, queue=False)
+                                  [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)
+        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)

javascript/app.js

@@ -1,3 +1,5 @@
+const MIDI_OUTPUT_BATCH_SIZE=4;
+
 /**
  * 自动绕过 shadowRoot 的 querySelector
  * @param {string} selector - 要查询的 CSS 选择器
@@ -594,33 +596,49 @@ class MidiVisualizer extends HTMLElement{
 customElements.define('midi-visualizer', MidiVisualizer);
 
 (()=>{
-    let midi_visualizer_container_inited = null
-    let midi_audio_audio_inited = null;
-    let midi_audio_cursor_inited = null;
-    let midi_visualizer = document.createElement('midi-visualizer')
-    onUiUpdate((m)=>{
-        let app = gradioApp()
-        let midi_visualizer_container = app.querySelector("#midi_visualizer_container");
-        if(!!midi_visualizer_container && midi_visualizer_container_inited!== midi_visualizer_container){
-            midi_visualizer_container.appendChild(midi_visualizer)
-            midi_visualizer_container_inited = midi_visualizer_container;
-        }
-        let midi_audio = app.querySelector("#midi_audio");
-        if (!!midi_audio){
-            let midi_audio_cursor = midi_audio.deepQuerySelector(".cursor");
-            if(!!midi_audio_cursor && midi_audio_cursor_inited!==midi_audio_cursor){
-                midi_visualizer.bindWaveformCursor(midi_audio_cursor)
-                midi_audio_cursor_inited = midi_audio_cursor
+    function midi_visualizer_setup(idx, midi_visualizer){
+        let midi_visualizer_container_inited = null
+        let midi_audio_audio_inited = null;
+        let midi_audio_cursor_inited = null;
+        onUiUpdate((m)=>{
+            let app = gradioApp()
+            let midi_visualizer_container = app.querySelector(`#midi_visualizer_container_${idx}`);
+            if(!!midi_visualizer_container && midi_visualizer_container_inited!== midi_visualizer_container){
+                midi_visualizer_container.appendChild(midi_visualizer)
+                midi_visualizer_container_inited = midi_visualizer_container;
             }
-            let midi_audio_audio = midi_audio.deepQuerySelector("audio");
-            if(!!midi_audio_audio && midi_audio_audio_inited!==midi_audio_audio){
-                midi_visualizer.bindAudioPlayer(midi_audio_audio)
-                midi_audio_audio_inited = midi_audio_audio
+            let midi_audio = app.querySelector(`#midi_audio_${idx}`);
+            if (!!midi_audio){
+                let midi_audio_cursor = midi_audio.deepQuerySelector(".cursor");
+                if(!!midi_audio_cursor && midi_audio_cursor_inited!==midi_audio_cursor){
+                    midi_visualizer.bindWaveformCursor(midi_audio_cursor)
+                    midi_audio_cursor_inited = midi_audio_cursor
+                }
+                let midi_audio_audio = midi_audio.deepQuerySelector("audio");
+                if(!!midi_audio_audio && midi_audio_audio_inited!==midi_audio_audio){
+                    midi_visualizer.bindAudioPlayer(midi_audio_audio)
+                    midi_audio_audio_inited = midi_audio_audio
+                }
             }
-        }
-    })
+        });
+    }
+
+    let midi_visualizers = []
+    for (let i = 0; i < MIDI_OUTPUT_BATCH_SIZE ; i++){
+        let midi_visualizer = document.createElement('midi-visualizer');
+        midi_visualizers.push(midi_visualizer);
+        midi_visualizer_setup(i, midi_visualizer)
+    }
 
     let hasProgressBar = false;
+    let output_tabs_inited = null;
+    onUiUpdate((m)=>{
+        let app = gradioApp()
+        let output_tabs = app.querySelector("#output_tabs");
+        if(!!output_tabs && output_tabs_inited!== output_tabs){
+            output_tabs_inited = output_tabs;
+        }
+    });
 
     function createProgressBar(progressbarContainer){
         let parentProgressbar = progressbarContainer.parentNode;
@@ -653,15 +671,15 @@ customElements.define('midi-visualizer', MidiVisualizer);
         hasProgressBar = false;
     }
 
-    function setProgressBar(progressbarContainer, progress, total){
+    function setProgressBar(progress, total){
         if (!hasProgressBar)
-            createProgressBar(midi_visualizer_container_inited)
+            createProgressBar(output_tabs_inited)
         if (hasProgressBar && total === 0){
-            removeProgressBar(midi_visualizer_container_inited)
+            removeProgressBar(output_tabs_inited)
             return
         }
-        let parentProgressbar = progressbarContainer.parentNode;
-        let divProgress = parentProgressbar.querySelector(".progressDiv");
+        let parentProgressbar = output_tabs_inited.parentNode;
+        // let divProgress = parentProgressbar.querySelector(".progressDiv");
         let divInner = parentProgressbar.querySelector(".progress");
         if(total===0)
             total = 1;
@@ -679,24 +697,30 @@ customElements.define('midi-visualizer', MidiVisualizer);
         }
     })
     function handleMsg(msg){
+        let idx;
         switch (msg.name) {
             case "visualizer_clear":
-                midi_visualizer.clearMidiEvents(false);
-                midi_visualizer.version = msg.data
+                idx = msg.data[0];
+                let ver = msg.data[1];
+                midi_visualizers[idx].clearMidiEvents(false);
+                midi_visualizers[idx].version = ver;
                 break;
             case "visualizer_append":
-                msg.data.forEach( value => {
-                    midi_visualizer.appendMidiEvent(value);
+                idx = msg.data[0];
+                let events = msg.data[1];
+                events.forEach( value => {
+                    midi_visualizers[idx].appendMidiEvent(value);
                 })
                 break;
             case "visualizer_end":
-                midi_visualizer.finishAppendMidiEvent()
-                midi_visualizer.setPlayTime(0);
+                idx = msg.data;
+                midi_visualizers[idx].finishAppendMidiEvent()
+                midi_visualizers[idx].setPlayTime(0);
                 break;
             case "progress":
                 let progress = msg.data[0]
                 let total = msg.data[1]
-                setProgressBar(midi_visualizer_container_inited, progress, total)
+                setProgressBar(progress, total)
                 break;
             default:
         }

midi_model.py

@@ -111,49 +111,69 @@ class MIDIModel(nn.Module):
         return next_token
 
     @torch.inference_mode()
-    def generate(self, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20, generator=None):
+    def generate(self, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98, top_k=20, generator=None):
         tokenizer = self.tokenizer
         max_token_seq = tokenizer.max_token_seq
         if prompt is None:
             input_tensor = torch.full((1, max_token_seq), tokenizer.pad_id, dtype=torch.long, device=self.device)
             input_tensor[0, 0] = tokenizer.bos_id  # bos
+            input_tensor = input_tensor.unsqueeze(0)
+            input_tensor = torch.cat([input_tensor] * batch_size, dim=0)
         else:
-            prompt = prompt[:, :max_token_seq]
+            if len(prompt.shape) == 2:
+                prompt = prompt[None, :]
+                prompt = np.repeat(prompt, repeats=batch_size, axis=0)
+            elif prompt.shape[0] == 1:
+                prompt = np.repeat(prompt, repeats=batch_size, axis=0)
+            else:
+                raise ValueError(f"invalid shape for prompt, {prompt.shape}")
+            prompt = prompt[..., :max_token_seq]
             if prompt.shape[-1] < max_token_seq:
-                prompt = np.pad(prompt, ((0, 0), (0, max_token_seq - prompt.shape[-1])),
+                prompt = np.pad(prompt, ((0, 0), (0, 0), (0, max_token_seq - prompt.shape[-1])),
                                 mode="constant", constant_values=tokenizer.pad_id)
             input_tensor = torch.from_numpy(prompt).to(dtype=torch.long, device=self.device)
-        input_tensor = input_tensor.unsqueeze(0)
+
         cur_len = input_tensor.shape[1]
         bar = tqdm.tqdm(desc="generating", total=max_len - cur_len)
         with bar:
             while cur_len < max_len:
-                end = False
-                hidden = self.forward(input_tensor)[0, -1].unsqueeze(0)
+                end = [False] * batch_size
+                hidden = self.forward(input_tensor)[:, -1]
                 next_token_seq = None
-                event_name = ""
+                event_names = [""] * batch_size
                 for i in range(max_token_seq):
-                    mask = torch.zeros(tokenizer.vocab_size, dtype=torch.int64, device=self.device)
-                    if i == 0:
-                        mask[list(tokenizer.event_ids.values()) + [tokenizer.eos_id]] = 1
-                    else:
-                        param_name = tokenizer.events[event_name][i - 1]
-                        mask[tokenizer.parameter_ids[param_name]] = 1
-
+                    mask = torch.zeros((batch_size, tokenizer.vocab_size), dtype=torch.int64, device=self.device)
+                    for b in range(batch_size):
+                        if end[b]:
+                            mask[b, tokenizer.pad_id] = 1
+                            continue
+                        if i == 0:
+                            mask[b, list(tokenizer.event_ids.values()) + [tokenizer.eos_id]] = 1
+                        else:
+                            param_names = tokenizer.events[event_names[b]]
+                            if i > len(param_names):
+                                mask[b, tokenizer.pad_id] = 1
+                                continue
+                            mask[b, tokenizer.parameter_ids[param_names[i - 1]]] = 1
+                    mask = mask.unsqueeze(1)
                     logits = self.forward_token(hidden, next_token_seq)[:, -1:]
                     scores = torch.softmax(logits / temp, dim=-1) * mask
-                    sample = self.sample_top_p_k(scores, top_p, top_k, generator=generator)
+                    samples = self.sample_top_p_k(scores, top_p, top_k, generator=generator)
                     if i == 0:
-                        next_token_seq = sample
-                        eid = sample.item()
-                        if eid == tokenizer.eos_id:
-                            end = True
-                            break
-                        event_name = tokenizer.id_events[eid]
+                        next_token_seq = samples
+                        for b in range(batch_size):
+                            if end[b]:
+                                continue
+                            eid = samples[b].item()
+                            if eid == tokenizer.eos_id:
+                                end[b] = True
+                            else:
+                                event_names[b] = tokenizer.id_events[eid]
                     else:
-                        next_token_seq = torch.cat([next_token_seq, sample], dim=1)
-                        if len(tokenizer.events[event_name]) == i:
+                        next_token_seq = torch.cat([next_token_seq, samples], dim=1)
+                        if all([len(tokenizer.events[event_names[b]]) == i for b in range(batch_size) if not end[b]]):
                             break
+
                 if next_token_seq.shape[1] < max_token_seq:
                     next_token_seq = F.pad(next_token_seq, (0, max_token_seq - next_token_seq.shape[1]),
                                            "constant", value=tokenizer.pad_id)
@@ -161,6 +181,7 @@ class MIDIModel(nn.Module):
                 input_tensor = torch.cat([input_tensor, next_token_seq], dim=1)
                 cur_len += 1
                 bar.update(1)
-                if end:
+
+                if all(end):
                     break
-        return input_tensor[0].cpu().numpy()
+        return input_tensor.cpu().numpy()