Update app.py

app.py

@@ -41,7 +41,7 @@ def sample_top_p_k(probs, p, k):
     return next_token
 
 
-def generate(prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
+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):
     if disable_channels is not None:
         disable_channels = [tokenizer.parameter_ids["channel"][c] for c in disable_channels]
@@ -63,7 +63,7 @@ def generate(prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
     with bar:
         while cur_len < max_len:
             end = False
-            hidden = model_base.run(None, {'x': input_tensor})[0][:, -1]
+            hidden = model[0].run(None, {'x': input_tensor})[0][:, -1]
             next_token_seq = np.empty((1, 0), dtype=np.int64)
             event_name = ""
             for i in range(max_token_seq):
@@ -81,7 +81,7 @@ def generate(prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
                     if param_name == "channel":
                         mask_ids = [i for i in mask_ids if i not in disable_channels]
                     mask[mask_ids] = 1
-                logits = model_token.run(None, {'x': next_token_seq, "hidden": hidden})[0][:, -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)
                 if i == 0:
@@ -107,7 +107,7 @@ def generate(prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
                 break
 
 
-def run(tab, instruments, drum_kit, mid, midi_events, gen_events, temp, top_p, top_k, allow_cc):
+def run(model_name, tab, instruments, drum_kit, mid, midi_events, gen_events, temp, top_p, top_k, allow_cc):
     mid_seq = []
     max_len = int(gen_events)
     img_len = 1024
@@ -172,7 +172,8 @@ def run(tab, instruments, drum_kit, mid, midi_events, gen_events, temp, top_p, t
         for token_seq in mid:
             mid_seq.append(token_seq)
             draw_event(token_seq)
-    generator = generate(mid, max_len=max_len, temp=temp, top_p=top_p, top_k=top_k,
+    model = models[model_name]
+    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 token_seq in generator:
@@ -208,13 +209,18 @@ if __name__ == "__main__":
     parser.add_argument("--max-gen", type=int, default=1024, help="max")
     opt = parser.parse_args()
     soundfont_path = hf_hub_download(repo_id="skytnt/midi-model", filename="soundfont.sf2")
-    model_base_path = hf_hub_download(repo_id="skytnt/midi-model", filename="onnx/model_base.onnx")
-    model_token_path = hf_hub_download(repo_id="skytnt/midi-model", filename="onnx/model_token.onnx")
-
+    models_info = {"generic pretrain model": ["skytnt/midi-model", ""],
+                   "symphony finetune model": ["skytnt/midi-model-ft", "symphony/"],
+                   "touhou finetune model": ["skytnt/midi-model-ft", "touhou/"]}
+    models = {}
     tokenizer = MIDITokenizer()
     providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
-    model_base = rt.InferenceSession(model_base_path, providers=providers)
-    model_token = rt.InferenceSession(model_token_path, providers=providers)
+    for name, (repo_id, path) in models_info.items():
+        model_base_path = hf_hub_download(repo_id=repo_id, filename=f"{path}onnx/model_base.onnx")
+        model_token_path = hf_hub_download(repo_id=repo_id, filename=f"{path}onnx/model_token.onnx")
+        model_base = rt.InferenceSession(model_base_path, providers=providers)
+        model_token = rt.InferenceSession(model_token_path, providers=providers)
+        models[name] = [model_base, model_token]
 
     app = gr.Blocks()
     with app:
@@ -229,6 +235,8 @@ if __name__ == "__main__":
 
         tab_select = gr.Variable(value=0)
         with gr.Tabs():
+            input_model = gr.Dropdown(label="select model", choices=list(models.keys()),
+                                      type="value", value=list(models.keys())[0])
             with gr.TabItem("instrument prompt") as tab1:
                 input_instruments = gr.Dropdown(label="instruments (auto if empty)", choices=list(patch2number.keys()),
                                                 multiselect=True, max_choices=15, type="value")
@@ -260,7 +268,7 @@ if __name__ == "__main__":
         with gr.Accordion("options", open=False):
             input_temp = gr.Slider(label="temperature", minimum=0.1, maximum=1.2, step=0.01, value=1)
             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=20, step=1, value=12)
+            input_top_k = gr.Slider(label="top k", minimum=1, maximum=20, step=1, value=20)
             input_allow_cc = gr.Checkbox(label="allow midi cc event", value=True)
             example3 = gr.Examples([[1, 0.98, 12], [1.2, 0.95, 8]], [input_temp, input_top_p, input_top_k])
         run_btn = gr.Button("generate", variant="primary")
@@ -269,8 +277,8 @@ if __name__ == "__main__":
         output_midi_img = gr.Image(label="output image")
         output_midi = gr.File(label="output midi", file_types=[".mid"])
         output_audio = gr.Audio(label="output audio", format="mp3")
-        run_event = run_btn.click(run, [tab_select, input_instruments, input_drum_kit, input_midi, input_midi_events,
-                                        input_gen_events, input_temp, input_top_p, input_top_k,
+        run_event = run_btn.click(run, [input_model, tab_select, input_instruments, input_drum_kit, input_midi,
+                                        input_midi_events, input_gen_events, input_temp, input_top_p, input_top_k,
                                         input_allow_cc],
                                   [output_midi_seq, output_midi_img, output_midi, output_audio])
         stop_btn.click(cancel_run, output_midi_seq, [output_midi, output_audio], cancels=run_event, queue=False)