add kv cache for onnx

README.md

@@ -5,7 +5,7 @@ colorFrom: red
 colorTo: indigo
 sdk: gradio
 sdk_version: 4.43.0
-app_file: app.py
+app_file: app_onnx.py
 pinned: true
 license: apache-2.0
 ---

app.py

@@ -415,7 +415,8 @@ if __name__ == "__main__":
                     "(https://colab.research.google.com/github/SkyTNT/midi-model/blob/main/demo.ipynb)"
                     " or [download windows app](https://github.com/SkyTNT/midi-model/releases)"
                     " for unlimited generation\n\n"
-                    "**Update v1.3**: MIDITokenizerV2 and new MidiVisualizer"
+                    "**Update v1.3**: MIDITokenizerV2 and new MidiVisualizer\n\n"
+                    "The current **best** model: generic pretrain model (tv2o-medium) by skytnt"
                     )
         js_msg = gr.Textbox(elem_id="msg_receiver", visible=False)
         js_msg.change(None, [js_msg], [], js="""

app_onnx.py

@@ -47,6 +47,37 @@ def sample_top_p_k(probs, p, k, generator=None):
     return next_token
 
 
+def apply_io_binding(model: rt.InferenceSession, inputs, outputs, batch_size, past_len, cur_len):
+    io_binding = model.io_binding()
+    for input_ in  model.get_inputs():
+        name = input_.name
+        if name.startswith("past_key_values"):
+            present_name = name.replace("past_key_values", "present")
+            if present_name in outputs:
+                v = outputs[present_name]
+            else:
+                v = rt.OrtValue.ortvalue_from_shape_and_type(
+                    (batch_size, input_.shape[1], past_len, input_.shape[3]),
+                    element_type=np.float32,
+                    device_type=device)
+            inputs[name] = v
+        else:
+            v = inputs[name]
+        io_binding.bind_ortvalue_input(name, v)
+
+    for output in model.get_outputs():
+        name = output.name
+        if name.startswith("present"):
+            v = rt.OrtValue.ortvalue_from_shape_and_type(
+                (batch_size, output.shape[1], cur_len, output.shape[3]),
+                element_type=np.float32,
+                device_type=device)
+            outputs[name] = v
+        else:
+            v = outputs[name]
+        io_binding.bind_ortvalue_output(name, v)
+    return io_binding
+
 def generate(model, 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[2]
@@ -77,12 +108,31 @@ def generate(model, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98
         input_tensor = prompt
     cur_len = input_tensor.shape[1]
     bar = tqdm.tqdm(desc="generating", total=max_len - cur_len)
+    model0_inputs = {}
+    model0_outputs = {}
+    emb_size = 1024
+    for output in model[0].get_outputs():
+        if output.name == "hidden":
+            emb_size = output.shape[2]
+    past_len = 0
     with bar:
         while cur_len < max_len:
             end = [False] * batch_size
-            hidden = model[0].run(None, {'x': input_tensor})[0][:, -1]
-            next_token_seq = np.empty((batch_size, 0), dtype=np.int64)
+            model0_inputs["x"] = rt.OrtValue.ortvalue_from_numpy(input_tensor[:, past_len:], device_type=device)
+            model0_outputs["hidden"] = rt.OrtValue.ortvalue_from_shape_and_type(
+                (batch_size, cur_len - past_len, emb_size),
+                element_type=np.float32,
+                device_type=device)
+            io_binding = apply_io_binding(model[0], model0_inputs, model0_outputs, batch_size, past_len, cur_len)
+            io_binding.synchronize_inputs()
+            model[0].run_with_iobinding(io_binding)
+            io_binding.synchronize_outputs()
+
+            hidden = model0_outputs["hidden"].numpy()[:, -1:]
+            next_token_seq = np.zeros((batch_size, 0), dtype=np.int64)
             event_names = [""] * batch_size
+            model1_inputs = {"hidden": rt.OrtValue.ortvalue_from_numpy(hidden, device_type=device)}
+            model1_outputs = {}
             for i in range(max_token_seq):
                 mask = np.zeros((batch_size, tokenizer.vocab_size), dtype=np.int64)
                 for b in range(batch_size):
@@ -107,7 +157,24 @@ def generate(model, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98
                             mask_ids = [i for i in mask_ids if i not in disable_channels]
                         mask[b, mask_ids] = 1
                 mask = mask[:, None, :]
-                logits = model[1].run(None, {'x': next_token_seq, "hidden": hidden})[0][:, -1:]
+                x = next_token_seq
+                if i != 0:
+                    # cached
+                    if i == 1:
+                        hidden = np.zeros((batch_size, 0, emb_size), dtype=np.float32)
+                        model1_inputs["hidden"] = rt.OrtValue.ortvalue_from_numpy(hidden, device_type=device)
+                    x = x[:, -1:]
+                model1_inputs["x"] = rt.OrtValue.ortvalue_from_numpy(x, device_type=device)
+                model1_outputs["y"] = rt.OrtValue.ortvalue_from_shape_and_type(
+                    (batch_size, 1, tokenizer.vocab_size),
+                    element_type=np.float32,
+                    device_type=device
+                )
+                io_binding = apply_io_binding(model[1], model1_inputs, model1_outputs, batch_size, i, i+1)
+                io_binding.synchronize_inputs()
+                model[1].run_with_iobinding(io_binding)
+                io_binding.synchronize_outputs()
+                logits = model1_outputs["y"].numpy()
                 scores = softmax(logits / temp, -1) * mask
                 samples = sample_top_p_k(scores, top_p, top_k, generator)
                 if i == 0:
@@ -130,6 +197,7 @@ def generate(model, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98
                                         mode="constant", constant_values=tokenizer.pad_id)
             next_token_seq = next_token_seq[:, None, :]
             input_tensor = np.concatenate([input_tensor, next_token_seq], axis=1)
+            past_len = cur_len
             cur_len += 1
             bar.update(1)
             yield next_token_seq[:, 0]
@@ -145,24 +213,13 @@ def send_msgs(msgs):
     return json.dumps(msgs)
 
 
-def calc_time(x):
-    return 5.849e-5*x**2 + 0.04781*x + 0.1168
-
 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 tab == 0:
-        start_events = 1
-    elif tab == 1 and mid is not None:
-        start_events = midi_events
-    elif tab == 2 and mid_seq is not None:
-        start_events = len(mid_seq[0])
-    else:
-        start_events = 1
-    t = calc_time(start_events + gen_events) - calc_time(start_events) + 5
+    t = gen_events // 23
     if "large" in model_name:
-        t *= 2
-    return t
+        t = gen_events // 14
+    return t + 5
 
 
 @spaces.GPU(duration=get_duration)
@@ -428,6 +485,7 @@ if __name__ == "__main__":
     }
     models = {}
     providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+    device = "cuda"
 
     for name, (repo_id, path, config, loras) in models_info.items():
         model_base_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}onnx/model_base.onnx")
@@ -451,7 +509,8 @@ if __name__ == "__main__":
                     "(https://colab.research.google.com/github/SkyTNT/midi-model/blob/main/demo.ipynb)"
                     " or [download windows app](https://github.com/SkyTNT/midi-model/releases)"
                     " for unlimited generation\n\n"
-                    "**Update v1.3**: MIDITokenizerV2 and new MidiVisualizer"
+                    "**Update v1.3**: MIDITokenizerV2 and new MidiVisualizer\n\n"
+                    "The current **best** model: generic pretrain model (tv2o-medium) by skytnt"
                     )
         js_msg = gr.Textbox(elem_id="msg_receiver", visible=False)
         js_msg.change(None, [js_msg], [], js="""