add seed

app.py

@@ -14,6 +14,7 @@ import MIDI
 from midi_synthesizer import synthesis
 from midi_tokenizer import MIDITokenizer
 
+MAX_SEED = np.iinfo(np.int32).max
 in_space = os.getenv("SYSTEM") == "spaces"
 
 
@@ -23,7 +24,9 @@ def softmax(x, axis):
     return exp_x_shifted / np.sum(exp_x_shifted, axis=axis, keepdims=True)
 
 
-def sample_top_p_k(probs, p, k):
+def sample_top_p_k(probs, p, k, generator=None):
+    if generator is None:
+        generator = np.random
     probs_idx = np.argsort(-probs, axis=-1)
     probs_sort = np.take_along_axis(probs, probs_idx, -1)
     probs_sum = np.cumsum(probs_sort, axis=-1)
@@ -36,17 +39,19 @@ def sample_top_p_k(probs, p, k):
     shape = probs_sort.shape
     probs_sort_flat = probs_sort.reshape(-1, shape[-1])
     probs_idx_flat = probs_idx.reshape(-1, shape[-1])
-    next_token = np.stack([np.random.choice(idxs, p=pvals) for pvals, idxs in zip(probs_sort_flat, probs_idx_flat)])
+    next_token = np.stack([generator.choice(idxs, p=pvals) for pvals, idxs in zip(probs_sort_flat, probs_idx_flat)])
     next_token = next_token.reshape(*shape[:-1])
     return next_token
 
 
 def generate(model, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
-             disable_patch_change=False, disable_control_change=False, disable_channels=None):
+             disable_patch_change=False, disable_control_change=False, disable_channels=None, generator=None):
     if disable_channels is not None:
         disable_channels = [tokenizer.parameter_ids["channel"][c] for c in disable_channels]
     else:
         disable_channels = []
+    if generator is None:
+        generator = np.random
     max_token_seq = tokenizer.max_token_seq
     if prompt is None:
         input_tensor = np.full((1, max_token_seq), tokenizer.pad_id, dtype=np.int64)
@@ -83,7 +88,7 @@ def generate(model, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
                     mask[mask_ids] = 1
                 logits = model[1].run(None, {'x': next_token_seq, "hidden": hidden})[0][:, -1:]
                 scores = softmax(logits / temp, -1) * mask
-                sample = sample_top_p_k(scores, top_p, top_k)
+                sample = sample_top_p_k(scores, top_p, top_k, generator)
                 if i == 0:
                     next_token_seq = sample
                     eid = sample.item()
@@ -120,13 +125,16 @@ def send_msgs(msgs, msgs_history=None):
     return json.dumps(msgs_history)
 
 
-def run(model_name, tab, instruments, drum_kit, bpm, mid, midi_events, gen_events, temp, top_p, top_k, allow_cc):
+def run(model_name, tab, instruments, drum_kit, bpm, mid, midi_events, seed, seed_rand,
+        gen_events, temp, top_p, top_k, allow_cc):
     msgs_history = []
     mid_seq = []
     bpm = int(bpm)
     gen_events = int(gen_events)
     max_len = gen_events
-
+    if seed_rand:
+        seed = np.random.randint(0, MAX_SEED)
+    generator = np.random.RandomState(seed)
     disable_patch_change = False
     disable_channels = None
     if tab == 0:
@@ -159,22 +167,22 @@ def run(model_name, tab, instruments, drum_kit, bpm, mid, midi_events, gen_event
     init_msgs = [create_msg("visualizer_clear", False)]
     for tokens in mid_seq:
         init_msgs.append(create_msg("visualizer_append", tokenizer.tokens2event(tokens)))
-    yield mid_seq, None, None, send_msgs(init_msgs, msgs_history)
+    yield mid_seq, None, None, seed, send_msgs(init_msgs, msgs_history)
     model = models[model_name]
-    generator = generate(model, mid, max_len=max_len, temp=temp, top_p=top_p, top_k=top_k,
+    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)
-    for i, token_seq in enumerate(generator):
+                         disable_channels=disable_channels, generator=generator)
+    for i, token_seq in enumerate(midi_generator):
         token_seq = token_seq.tolist()
         mid_seq.append(token_seq)
         event = tokenizer.tokens2event(token_seq)
-        yield mid_seq, None, None, send_msgs([create_msg("visualizer_append", event), create_msg("progress", [i + 1, gen_events])], msgs_history)
+        yield mid_seq, None, None, seed, send_msgs([create_msg("visualizer_append", event), create_msg("progress", [i + 1, gen_events])], msgs_history)
     mid = tokenizer.detokenize(mid_seq)
     with open(f"output.mid", 'wb') as f:
         f.write(MIDI.score2midi(mid))
     audio = synthesis(MIDI.score2opus(mid), soundfont_path)
     events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
-    yield mid_seq, "output.mid", (44100, audio), send_msgs([create_msg("visualizer_end", events)])
+    yield mid_seq, "output.mid", (44100, audio), seed, send_msgs([create_msg("visualizer_end", events)])
 
 
 def cancel_run(mid_seq):
@@ -232,8 +240,8 @@ if __name__ == "__main__":
     opt = parser.parse_args()
     soundfont_path = hf_hub_download_retry(repo_id="skytnt/midi-model", filename="soundfont.sf2")
     models_info = {"generic pretrain model": ["skytnt/midi-model", ""],
-                   "j-pop finetune model": ["skytnt/midi-model-ft", "jpop/"],
-                   "touhou finetune model": ["skytnt/midi-model-ft", "touhou/"],
+                   # "j-pop finetune model": ["skytnt/midi-model-ft", "jpop/"],
+                   # "touhou finetune model": ["skytnt/midi-model-ft", "touhou/"],
                    }
     models = {}
     tokenizer = MIDITokenizer()
@@ -301,7 +309,10 @@ if __name__ == "__main__":
 
         tab1.select(lambda: 0, None, tab_select, queue=False)
         tab2.select(lambda: 1, None, tab_select, queue=False)
-        input_gen_events = gr.Slider(label="generate n midi events", minimum=1, maximum=opt.max_gen,
+        input_seed = gr.Slider(label="seed", minimum=0, maximum=2 ** 31 - 1,
+                               step=1, value=0)
+        input_seed_rand = gr.Checkbox(label="random seed", value=True)
+        input_gen_events = gr.Slider(label="generate max n midi events", minimum=1, maximum=opt.max_gen,
                                      step=1, value=opt.max_gen // 2)
         with gr.Accordion("options", open=False):
             input_temp = gr.Slider(label="temperature", minimum=0.1, maximum=1.2, step=0.01, value=1)
@@ -316,9 +327,9 @@ if __name__ == "__main__":
         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, input_instruments, input_drum_kit, input_bpm,
-                                        input_midi, input_midi_events, input_gen_events, input_temp,
-                                        input_top_p, input_top_k, input_allow_cc],
-                                  [output_midi_seq, output_midi, output_audio, js_msg],
+                                        input_midi, input_midi_events, 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],
                                   concurrency_limit=3)
         stop_btn.click(cancel_run, [output_midi_seq], [output_midi, output_audio, js_msg], cancels=run_event, queue=False)
     app.launch(server_port=opt.port, share=opt.share, inbrowser=True)