use javascript

app.py

@@ -2,12 +2,11 @@ import argparse
 import glob
 import os.path
 
-import PIL
-import PIL.ImageColor
 import gradio as gr
 import numpy as np
 import onnxruntime as rt
 import tqdm
+import json
 from huggingface_hub import hf_hub_download
 
 import MIDI
@@ -107,44 +106,14 @@ def generate(model, prompt=None, max_len=512, temp=1.0, top_p=0.98, top_k=20,
                 break
 
 
-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
-    img = np.full((128 * 2, img_len, 3), 255, dtype=np.uint8)
-    state = {"t1": 0, "t": 0, "cur_pos": 0}
-    colors = ['navy', 'blue', 'deepskyblue', 'teal', 'green', 'lightgreen', 'lime', 'orange',
-              'brown', 'grey', 'red', 'pink', 'aqua', 'orchid', 'bisque', 'coral']
-    colors = [PIL.ImageColor.getrgb(color) for color in colors]
+def create_msg(name, data):
+    return {"name": name, "data": data}
 
-    def draw_event(tokens):
-        if tokens[0] in tokenizer.id_events:
-            name = tokenizer.id_events[tokens[0]]
-            if len(tokens) <= len(tokenizer.events[name]):
-                return
-            params = tokens[1:]
-            params = [params[i] - tokenizer.parameter_ids[p][0] for i, p in enumerate(tokenizer.events[name])]
-            if not all([0 <= params[i] < tokenizer.event_parameters[p] for i, p in enumerate(tokenizer.events[name])]):
-                return
-            event = [name] + params
-            state["t1"] += event[1]
-            t = state["t1"] * 16 + event[2]
-            state["t"] = t
-            if name == "note":
-                tr, d, c, p = event[3:7]
-                shift = t + d - (state["cur_pos"] + img_len)
-                if shift > 0:
-                    img[:, :-shift] = img[:, shift:]
-                    img[:, -shift:] = 255
-                    state["cur_pos"] += shift
-                t = t - state["cur_pos"]
-                img[p * 2:(p + 1) * 2, t: t + d] = colors[c]
 
-    def get_img():
-        t = state["t"] - state["cur_pos"]
-        img_new = img.copy()
-        img_new[:, t: t + 2] = 0
-        return PIL.Image.fromarray(np.flip(img_new, 0))
+def run(model_name, tab, instruments, drum_kit, mid, midi_events, gen_events, temp, top_p, top_k, allow_cc):
+    mid_seq = []
+    gen_events = int(gen_events)
+    max_len = gen_events
 
     disable_patch_change = False
     disable_channels = None
@@ -170,25 +139,25 @@ def run(model_name, tab, instruments, drum_kit, mid, midi_events, gen_events, te
         mid = mid[:int(midi_events)]
         max_len += len(mid)
         for token_seq in mid:
-            mid_seq.append(token_seq)
-            draw_event(token_seq)
+            mid_seq.append(token_seq.tolist())
+    init_msgs = [create_msg("visualizer_clear", None)]
+    for tokens in mid_seq:
+        init_msgs.append(create_msg("visualizer_append", tokenizer.tokens2event(tokens)))
+    yield mid_seq, None, None, init_msgs
     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)
-    try:
-        for token_seq in generator:
-            mid_seq.append(token_seq)
-            draw_event(token_seq)
-            yield mid_seq, get_img(), None, None
-    except Exception as e:
-        print(e)
-
+    for i, token_seq in enumerate(generator):
+        token_seq = token_seq.tolist()
+        mid_seq.append(token_seq)
+        event = tokenizer.tokens2event(token_seq)
+        yield mid_seq, None, None, [create_msg("visualizer_append", event), create_msg("progress", [i + 1, gen_events])]
     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)
-    yield mid_seq, get_img(), "output.mid", (44100, audio)
+    audio = synthesis(MIDI.score2opus(mid), opt.soundfont_path)
+    yield mid_seq, "output.mid", (44100, audio), [create_msg("visualizer_end", None)]
 
 
 def cancel_run(mid_seq):
@@ -198,7 +167,39 @@ def cancel_run(mid_seq):
     with open(f"output.mid", 'wb') as f:
         f.write(MIDI.score2midi(mid))
     audio = synthesis(MIDI.score2opus(mid), soundfont_path)
-    return "output.mid", (44100, audio)
+    return "output.mid", (44100, audio), [create_msg("visualizer_end", None)]
+
+
+def load_javascript(dir="javascript"):
+    scripts_list = glob.glob(f"{dir}/*.js")
+    javascript = ""
+    for path in scripts_list:
+        with open(path, "r", encoding="utf8") as jsfile:
+            javascript += f"\n<!-- {path} --><script>{jsfile.read()}</script>"
+    template_response_ori = gr.routes.templates.TemplateResponse
+
+    def template_response(*args, **kwargs):
+        res = template_response_ori(*args, **kwargs)
+        res.body = res.body.replace(
+            b'</head>', f'{javascript}</head>'.encode("utf8"))
+        res.init_headers()
+        return res
+
+    gr.routes.templates.TemplateResponse = template_response
+
+
+class JSMsgReceiver(gr.HTML):
+
+    def __init__(self, **kwargs):
+        super().__init__(elem_id="msg_receiver", visible=False, **kwargs)
+
+    def postprocess(self, y):
+        if y:
+            y = f"<p>{json.dumps(y)}</p>"
+        return super().postprocess(y)
+
+    def get_block_name(self) -> str:
+        return "html"
 
 
 number2drum_kits = {-1: "None", 0: "Standard", 8: "Room", 16: "Power", 24: "Electric", 25: "TR-808", 32: "Jazz",
@@ -226,6 +227,7 @@ if __name__ == "__main__":
         model_token = rt.InferenceSession(model_token_path, providers=providers)
         models[name] = [model_base, model_token]
 
+    load_javascript()
     app = gr.Blocks()
     with app:
         gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Midi Composer</h1>")
@@ -236,6 +238,7 @@ if __name__ == "__main__":
                     "(https://colab.research.google.com/github/SkyTNT/midi-model/blob/main/demo.ipynb)"
                     " for faster running and longer generation"
                     )
+        js_msg = JSMsgReceiver()
         input_model = gr.Dropdown(label="select model", choices=list(models.keys()),
                                   type="value", value=list(models.keys())[0])
         tab_select = gr.Variable(value=0)
@@ -277,12 +280,12 @@ if __name__ == "__main__":
         run_btn = gr.Button("generate", variant="primary")
         stop_btn = gr.Button("stop and output")
         output_midi_seq = gr.Variable()
-        output_midi_img = gr.Image(label="output image")
-        output_midi = gr.File(label="output midi", file_types=[".mid"])
+        output_midi_visualizer = gr.HTML(elem_id="midi_visualizer_container")
         output_audio = gr.Audio(label="output audio", format="mp3")
+        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_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)
+                                  [output_midi_seq, output_midi, output_audio, js_msg])
+        stop_btn.click(cancel_run, output_midi_seq, [output_midi, output_audio, js_msg], cancels=run_event, queue=False)
     app.queue(2).launch(server_port=opt.port, share=opt.share, inbrowser=True)

javascript/app.js

@@ -0,0 +1,252 @@
+function gradioApp() {
+    const elems = document.getElementsByTagName('gradio-app')
+    const gradioShadowRoot = elems.length == 0 ? null : elems[0].shadowRoot
+    return !!gradioShadowRoot ? gradioShadowRoot : document;
+}
+
+uiUpdateCallbacks = []
+msgReceiveCallbacks = []
+
+function onUiUpdate(callback){
+    uiUpdateCallbacks.push(callback)
+}
+
+function onMsgReceive(callback){
+    msgReceiveCallbacks.push(callback)
+}
+
+function runCallback(x, m){
+    try {
+        x(m)
+    } catch (e) {
+        (console.error || console.log).call(console, e.message, e);
+    }
+}
+function executeCallbacks(queue, m) {
+    queue.forEach(function(x){runCallback(x, m)})
+}
+
+document.addEventListener("DOMContentLoaded", function() {
+    var mutationObserver = new MutationObserver(function(m){
+        executeCallbacks(uiUpdateCallbacks, m);
+    });
+    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;
+        }
+    })
+})();
+
+class MidiVisualizer extends HTMLElement{
+    constructor() {
+        super();
+        this.midiEvents = [];
+        this.wrapper = null;
+        this.svg = null;
+        this.timeLine = null;
+        this.config = {
+            noteHeight : 4,
+            beatWidth: 32
+        }
+        this.svgWidth = 0;
+        this.t1 = 0;
+        this.playTime = 0
+        this.colorMap = new Map();
+        this.init();
+    }
+
+    init(){
+        this.innerHTML=''
+        const shadow = this.attachShadow({mode: 'open'});
+        const style = document.createElement("style");
+        const wrapper = document.createElement('div');
+        style.textContent = ".note.active {stroke: black;stroke-width: 0.75;stroke-opacity: 0.75;}";
+        wrapper.style.overflowX= "scroll"
+        const svg = document.createElementNS('http://www.w3.org/2000/svg', 'svg');
+        svg.style.height = `${this.config.noteHeight*128}px`;
+        svg.style.width = `${this.svgWidth}px`;
+        const timeLine = document.createElementNS('http://www.w3.org/2000/svg', 'line');
+        timeLine.style.stroke = "green"
+        timeLine.style.strokeWidth = 2;
+        shadow.appendChild(style)
+        shadow.appendChild(wrapper);
+        wrapper.appendChild(svg);
+        svg.appendChild(timeLine)
+        this.wrapper = wrapper;
+        this.svg = svg;
+        this.timeLine= timeLine;
+        this.setPlayTime(0);
+    }
+
+    setPlayTime(t){
+        this.playTime = t
+        let x = Math.round((t/16)*this.config.beatWidth)
+        this.timeLine.setAttribute('x1', `${x}`);
+        this.timeLine.setAttribute('y1', '0');
+        this.timeLine.setAttribute('x2', `${x}`);
+        this.timeLine.setAttribute('y2', `${this.config.noteHeight*128}`);
+    }
+
+    clearMidiEvents(){
+        this.midiEvents = [];
+        this.svgWidth = 0
+        this.svg.innerHTML = ''
+        this.t1 = 0
+        this.colorMap.clear()
+        this.svg.style.width = `${this.svgWidth}px`;
+        this.setPlayTime(0);
+        this.svg.appendChild(this.timeLine)
+    }
+
+    appendMidiEvent(midiEvent){
+        if(midiEvent instanceof Array && midiEvent.length > 0){
+            this.midiEvents.push(midiEvent);
+            this.t1 += midiEvent[1]
+            let t = this.t1*16 + midiEvent[2]
+            if(midiEvent[0] === "note"){
+                let track = midiEvent[3]
+                let duration = midiEvent[4]
+                let channel = midiEvent[5]
+                let pitch = midiEvent[6]
+                let velocity = midiEvent[7]
+                let x = (t/16)*this.config.beatWidth
+                let y = (127 - pitch)*this.config.noteHeight
+                let w = (duration/16)*this.config.beatWidth
+                let h = this.config.noteHeight
+                this.svgWidth = Math.ceil(Math.max(x + w, this.svgWidth))
+                let color = this.getColor(track, channel)
+                let opacity = Math.min(1, velocity/127 + 0.1).toFixed(2)
+                this.drawNote(x,y,w,h, `rgba(${color[0]}, ${color[1]}, ${color[2]}, ${opacity})`)
+                this.setPlayTime(t);
+                this.wrapper.scrollTo(this.svgWidth - this.wrapper.offsetWidth, 0)
+            }
+            this.svg.style.width = `${this.svgWidth}px`;
+        }
+
+    }
+
+    getColor(track, channel){
+        let key = `${track},${channel}`;
+        let color = this.colorMap.get(key);
+        if(!!color){
+            return color;
+        }
+        color = [Math.round(Math.random()*240) + 10, Math.round(Math.random()*240)+ 10, Math.round(Math.random()*240)+ 10];
+        this.colorMap.set(key, color);
+        return color;
+    }
+
+    drawNote(x, y, w, h, fill) {
+        if (!this.svg) {
+          return null;
+        }
+        const rect = document.createElementNS('http://www.w3.org/2000/svg', 'rect');
+        rect.classList.add('note');
+        rect.setAttribute('fill', fill);
+        // Round values to the nearest integer to avoid partially filled pixels.
+        rect.setAttribute('x', `${Math.round(x)}`);
+        rect.setAttribute('y', `${Math.round(y)}`);
+        rect.setAttribute('width', `${Math.round(w)}`);
+        rect.setAttribute('height', `${Math.round(h)}`);
+        this.svg.appendChild(rect);
+        return rect
+  }
+}
+
+customElements.define('midi-visualizer', MidiVisualizer);
+
+(()=>{
+    let midi_visualizer_container_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;
+        }
+    })
+
+    function createProgressBar(progressbarContainer){
+        let parentProgressbar = progressbarContainer.parentNode;
+        let divProgress = document.createElement('div');
+        divProgress.className='progressDiv';
+        let rect = progressbarContainer.getBoundingClientRect();
+        divProgress.style.width = rect.width + "px";
+        divProgress.style.background = "#b4c0cc";
+        divProgress.style.borderRadius = "8px";
+        let divInner = document.createElement('div');
+        divInner.className='progress';
+        divInner.style.color = "white";
+        divInner.style.background = "#0060df";
+        divInner.style.textAlign = "right";
+        divInner.style.fontWeight = "bold";
+        divInner.style.borderRadius = "8px";
+        divInner.style.height = "20px";
+        divInner.style.lineHeight = "20px";
+        divInner.style.paddingRight = "8px"
+        divProgress.appendChild(divInner);
+        parentProgressbar.insertBefore(divProgress, progressbarContainer);
+    }
+
+    function removeProgressBar(progressbarContainer){
+        let parentProgressbar = progressbarContainer.parentNode;
+        let divProgress = parentProgressbar.querySelector(".progressDiv");
+        parentProgressbar.removeChild(divProgress);
+    }
+
+    function setProgressBar(progressbarContainer, progress, total){
+        let parentProgressbar = progressbarContainer.parentNode;
+        let divProgress = parentProgressbar.querySelector(".progressDiv");
+        let divInner = parentProgressbar.querySelector(".progress");
+        if(total===0)
+            total = 1;
+        divInner.style.width = `${(progress/total)*100}%`;
+        divInner.textContent = `${progress}/${total}`;
+    }
+
+    onMsgReceive((msg)=>{
+        switch (msg.name) {
+            case "visualizer_clear":
+                midi_visualizer.clearMidiEvents();
+                createProgressBar(midi_visualizer_container_inited)
+                break;
+            case "visualizer_append":
+                midi_visualizer.appendMidiEvent(msg.data);
+                break;
+            case "progress":
+                let progress = msg.data[0]
+                let total = msg.data[1]
+                setProgressBar(midi_visualizer_container_inited, progress, total)
+                break;
+            case "visualizer_end":
+                midi_visualizer.setPlayTime(0)
+                removeProgressBar(midi_visualizer_container_inited)
+                break;
+            default:
+        }
+    })
+})();

midi_tokenizer.py

@@ -101,6 +101,19 @@ class MIDITokenizer:
         tokens += [self.pad_id] * (self.max_token_seq - len(tokens))
         return tokens
 
+    def tokens2event(self, tokens):
+        if tokens[0] in self.id_events:
+            name = self.id_events[tokens[0]]
+            if len(tokens) <= len(self.events[name]):
+                return []
+            params = tokens[1:]
+            params = [params[i] - self.parameter_ids[p][0] for i, p in enumerate(self.events[name])]
+            if not all([0 <= params[i] < self.event_parameters[p] for i, p in enumerate(self.events[name])]):
+                return []
+            event = [name] + params
+            return event
+        return []
+
     def detokenize(self, midi_seq):
         ticks_per_beat = 480
         tracks_dict = {}