fix midi visualizer

app.py

@@ -1,6 +1,7 @@
 import argparse
 import glob
 import os.path
+import uuid
 
 import gradio as gr
 import numpy as np
@@ -107,7 +108,7 @@ def generate(model, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
 
 
 def create_msg(name, data):
-    return {"name": name, "data": data}
+    return {"name": name, "data": data, "uuid": uuid.uuid4().hex}
 
 
 def run(model_name, tab, instruments, drum_kit, mid, midi_events, gen_events, temp, top_p, top_k, allow_cc):
@@ -164,7 +165,7 @@ def run(model_name, tab, instruments, drum_kit, mid, midi_events, gen_events, te
 
 def cancel_run(mid_seq):
     if mid_seq is None:
-        return None, None
+        return None, None, []
     mid = tokenizer.detokenize(mid_seq)
     with open(f"output.mid", 'wb') as f:
         f.write(MIDI.score2midi(mid))
@@ -189,17 +190,25 @@ def load_javascript(dir="javascript"):
 
     gr.routes.templates.TemplateResponse = template_response
 
+
 # JSMsgReceiver
-HTML_postprocess_ori = gr.HTML.postprocess
+Textbox_postprocess_ori = gr.Textbox.postprocess
+
+msg_history = []
 
 
+# the change event may not trigger every time, so send msg history to avoid msg missing.
 def JSMsgReceiver_postprocess(self, y):
+    global msg_history
     if self.elem_id == "msg_receiver" and y:
-        y = f"<p>{json.dumps(y)}</p>"
-    return HTML_postprocess_ori(self, y)
+        msg_history.append(y)
+        if len(msg_history) > 50:
+            msg_history = msg_history[1:]
+        y = json.dumps(msg_history)
+    return Textbox_postprocess_ori(self, y)
 
 
-gr.HTML.postprocess = JSMsgReceiver_postprocess
+gr.Textbox.postprocess = JSMsgReceiver_postprocess
 
 number2drum_kits = {-1: "None", 0: "Standard", 8: "Room", 16: "Power", 24: "Electric", 25: "TR-808", 32: "Jazz",
                     40: "Blush", 48: "Orchestra"}
@@ -214,8 +223,8 @@ if __name__ == "__main__":
     opt = parser.parse_args()
     soundfont_path = hf_hub_download(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()
@@ -238,7 +247,14 @@ if __name__ == "__main__":
                     "(https://colab.research.google.com/github/SkyTNT/midi-model/blob/main/demo.ipynb)"
                     " for faster running and longer generation"
                     )
-        js_msg = gr.HTML(elem_id="msg_receiver", visible=False)
+        js_msg = gr.Textbox(elem_id="msg_receiver", visible=False)
+        js_msg.change(None, [js_msg], [], js="""
+        (msg_json) =>{
+            let msgs = JSON.parse(msg_json);
+            executeCallbacks(msgReceiveCallbacks, msgs);
+            return [];
+        }
+        """)
         input_model = gr.Dropdown(label="select model", choices=list(models.keys()),
                                   type="value", value=list(models.keys())[0])
         tab_select = gr.State(value=0)

javascript/app.js

@@ -76,33 +76,6 @@ document.addEventListener("DOMContentLoaded", function() {
     mutationObserver.observe( gradioApp(), { childList:true, subtree:true })
 });
 
-(()=>{
-    let mse_receiver_inited = null
-    onUiUpdate(()=>{
-        let app = gradioApp()
-        let msg_receiver = app.querySelector("#msg_receiver");
-        if(!!msg_receiver && mse_receiver_inited !== msg_receiver){
-            let mutationObserver = new MutationObserver(function(ms){
-                ms.forEach((m)=>{
-                    m.addedNodes.forEach((node)=>{
-                        if(node.nodeName === "P"){
-                            let obj = JSON.parse(node.innerText);
-                            if(obj instanceof Array){
-                                obj.forEach((o)=>{executeCallbacks(msgReceiveCallbacks, o);});
-                            }else{
-                                executeCallbacks(msgReceiveCallbacks, obj);
-                            }
-                        }
-                    })
-                })
-            });
-            mutationObserver.observe( msg_receiver, {childList:true, subtree:true, characterData:true})
-            console.log("receiver init");
-            mse_receiver_inited = msg_receiver;
-        }
-    })
-})()
-
 function HSVtoRGB(h, s, v) {
     let r, g, b, i, f, p, q, t;
     i = Math.floor(h * 6);
@@ -261,9 +234,11 @@ class MidiVisualizer extends HTMLElement{
                 this.midiTimes.push({ms:ms, t: t, tempo: tempo})
             }
             if(midiEvent[0]==="note"){
-                this.totalTimeMs = ms + (midiEvent[3]/ this.timePreBeat)*tempo
+                this.totalTimeMs = Math.max(this.totalTimeMs, ms + (midiEvent[3]/ this.timePreBeat)*tempo)
+            }else{
+                this.totalTimeMs = Math.max(this.totalTimeMs, ms);
             }
-            lastT = t
+            lastT = t;
         })
     }
 
@@ -431,7 +406,22 @@ customElements.define('midi-visualizer', MidiVisualizer);
         divInner.textContent = `${progress}/${total}`;
     }
 
-    onMsgReceive((msg)=>{
+    onMsgReceive((msgs)=>{
+        for(let msg of msgs){
+            if(msg instanceof Array){
+                msg.forEach((o)=>{handleMsg(o)});
+            }else{
+                handleMsg(msg);
+            }
+        }
+    })
+    let handled_msgs = [];
+    function handleMsg(msg){
+        if(handled_msgs.indexOf(msg.uuid)!== -1)
+            return;
+        handled_msgs.push(msg.uuid);
+        if(handled_msgs.length > 200)
+            handled_msgs = handled_msgs.slice(1);
         switch (msg.name) {
             case "visualizer_clear":
                 midi_visualizer.clearMidiEvents();
@@ -452,5 +442,5 @@ customElements.define('midi-visualizer', MidiVisualizer);
                 break;
             default:
         }
-    })
+    }
 })();