Add per-instrument piano-roll canvas; fix note-count summary

app.py

@@ -11,6 +11,7 @@ import re
 import base64
 from pathlib import Path
 
+import pretty_midi
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
 
@@ -20,7 +21,11 @@ from muscriptor.modules.conditioners import (
     ClassConditioner,
     ConditioningProvider,
 )
-from muscriptor.tokenizer.mt3 import MT3Tokenizer, MT3_FULL_PLUS_GROUP_NAMES
+from muscriptor.tokenizer.mt3 import (
+    MT3Tokenizer,
+    MT3_FULL_PLUS_GROUP_NAMES,
+    get_group_program_map,
+)
 from muscriptor.transcription_model import (
     TranscriptionModel,
     _resolve_source,
@@ -87,11 +92,33 @@ print(f"[muscriptor-space] Model loaded in {time.perf_counter() - t0:.1f}s", fil
 # Build the instrument choices list (sorted by group ID for stable ordering)
 INSTRUMENT_CHOICES = sorted(MT3_FULL_PLUS_GROUP_NAMES.keys(), key=lambda k: MT3_FULL_PLUS_GROUP_NAMES[k])
 
+# MIDI program -> human-readable label using the model's own MT3_FULL_PLUS
+# instrument groups (the model emits the representative program of a group,
+# so this names lanes with the taxonomy the model actually decodes into).
+# misc_programs="OMIT" keeps only the named groups 0-35; drums are handled
+# via is_drum. Falls back to General MIDI names for unmapped programs.
+_PROGRAM_TO_GROUP_LABEL: dict[int, str] = {}
+try:
+    _gid_to_name = {v: k for k, v in MT3_FULL_PLUS_GROUP_NAMES.items()}
+    for _gid, _programs in get_group_program_map(
+        "MT3_FULL_PLUS", misc_programs="OMIT", is_mt3=True
+    ).items():
+        _name = _gid_to_name.get(_gid)
+        if _name is None:
+            continue
+        _label = _name.replace("_", " ").title()
+        for _p in _programs:
+            _PROGRAM_TO_GROUP_LABEL[_p] = _label
+except Exception:
+    pass
+
 CSS = """
 #col-container { max-width: 1100px; margin: 0 auto; }
 .dark .gradio-container { color: var(--body-text-color); }
 midi-player { width: 100%; display: block; margin-bottom: 8px; }
 midi-visualizer { width: 100%; display: block; overflow: auto; }
+.midi-roll-container { margin-top: 10px; }
+.midi-roll-canvas { width: 100%; display: block; }
 """
 
 # The html-midi-player web components (<midi-player> / <midi-visualizer>) rely
@@ -139,6 +166,226 @@ watch('value', syncPlayer);
 """
 
 
+# --- Piano-roll canvas (visual design ported from the space-demo-kit MIDI roll) ---
+# One lane per instrument track: rounded lane background, label chip (colored
+# dot + instrument name + note count), gridlines, note bars placed by
+# onset/offset on x and pitch on y (per-lane pitch range), and a playhead
+# synced to the <midi-player> above. Notes glow while they are being played.
+_MIDI_ROLL_TEMPLATE = (
+    '<div class="midi-roll-container" style="display:none">'
+    '<canvas class="midi-roll-canvas"></canvas>'
+    "</div>"
+)
+
+_MIDI_ROLL_JS = r"""
+const LANE_COLORS = ['#f97316', '#3b82f6', '#22c55e', '#a855f7', '#ec4899', '#eab308'];
+const container = element.querySelector('.midi-roll-container');
+const canvas = element.querySelector('.midi-roll-canvas');
+let mt = null;       // {span, tracks: [{label, drum, notes: [[t0, t1, pitch], ...]}]}
+let lastT = 0;       // last known playback time (seconds)
+let live = false;    // player currently playing
+let rafId = null;
+let wired = null;    // the <midi-player> element we attached listeners to
+
+function isDark() {
+  return document.body.classList.contains('dark') ||
+         document.documentElement.classList.contains('dark');
+}
+
+const PAD_L = 20, PAD_R = 20, PAD_T = 14, PAD_B = 18, LANE_H = 84, LANE_GAP = 12;
+
+function layout() {
+  if (!mt) return;
+  const dpr = window.devicePixelRatio || 1;
+  const w = container.clientWidth || 600;
+  const nT = mt.tracks.length;
+  const h = PAD_T + PAD_B + nT * LANE_H + (nT - 1) * LANE_GAP;
+  canvas.width = Math.round(w * dpr);
+  canvas.height = Math.round(h * dpr);
+  canvas.style.width = w + 'px';
+  canvas.style.height = h + 'px';
+}
+
+function draw() {
+  if (!mt) return;
+  const dpr = window.devicePixelRatio || 1;
+  const cx = canvas.getContext('2d');
+  const W = canvas.width / dpr, H = canvas.height / dpr;
+  cx.setTransform(dpr, 0, 0, dpr, 0, 0);
+  cx.clearRect(0, 0, W, H);
+  const dark = isDark();
+  const span = mt.span;
+  const tNow = Math.min(Math.max(lastT, 0), span);
+  const innerW = W - PAD_L - PAD_R;
+  const tx = s => PAD_L + (s / span) * innerW;
+
+  // gridlines across the full roll (step adapts to the clip length)
+  const step = span > 240 ? 30 : span > 120 ? 15 : span > 60 ? 10 : span > 30 ? 5 : 1;
+  cx.strokeStyle = dark ? 'rgba(255,255,255,.07)' : 'rgba(17,24,39,.07)';
+  cx.lineWidth = 1;
+  for (let s = step; s < span; s += step) {
+    cx.beginPath(); cx.moveTo(tx(s), PAD_T); cx.lineTo(tx(s), H - PAD_B); cx.stroke();
+  }
+
+  mt.tracks.forEach((tr, i) => {
+    const y0 = PAD_T + i * (LANE_H + LANE_GAP);
+    const color = LANE_COLORS[i % LANE_COLORS.length];
+    // lane background
+    cx.fillStyle = dark ? 'rgba(255,255,255,.035)' : 'rgba(17,24,39,.04)';
+    cx.beginPath(); cx.roundRect(PAD_L - 10, y0, innerW + 20, LANE_H, 12); cx.fill();
+
+    // note area (below the label band)
+    const labelBand = 34;
+    const nTop = y0 + labelBand, nBot = y0 + LANE_H - 8;
+    let pMin = 127, pMax = 0;
+    for (const nn of tr.notes) { if (nn[2] < pMin) pMin = nn[2]; if (nn[2] > pMax) pMax = nn[2]; }
+    const pSpan = Math.max(1, pMax - pMin);
+    const nH = Math.max(4, Math.min(12, (nBot - nTop) / (pSpan + 1)));
+
+    for (const [t0, t1, p] of tr.notes) {
+      const x0 = tx(t0);
+      const w = Math.max(4, tx(t1) - x0);
+      const y = pMax === pMin
+        ? (nTop + nBot - nH) / 2
+        : nTop + ((pMax - p) / pSpan) * (nBot - nTop - nH);
+      const active = live && tNow >= t0 && tNow <= t1;
+      const past = tNow > t1;
+      cx.save();
+      if (active) { cx.shadowColor = color; cx.shadowBlur = 12; cx.globalAlpha = 1; }
+      else cx.globalAlpha = past ? .85 : .3;
+      cx.fillStyle = color;
+      cx.beginPath();
+      cx.roundRect(x0, active ? y - 1 : y, w, active ? nH + 2 : nH, 2);
+      cx.fill();
+      cx.restore();
+    }
+
+    // label chip: colored dot + instrument name + note count
+    const label = tr.label.toUpperCase();
+    const count = tr.notes.length + (tr.notes.length === 1 ? ' NOTE' : ' NOTES');
+    cx.font = '700 13px system-ui, -apple-system, sans-serif';
+    const lw = cx.measureText(label).width;
+    cx.font = '500 12px ui-monospace, SFMono-Regular, Menlo, monospace';
+    const cw = cx.measureText(count).width;
+    const chipW = 12 + 9 + 7 + lw + 10 + cw + 12, chipH = 24;
+    cx.fillStyle = dark ? 'rgba(22,26,34,.92)' : 'rgba(255,255,255,.94)';
+    cx.shadowColor = 'rgba(0,0,0,.10)'; cx.shadowBlur = 8;
+    cx.beginPath(); cx.roundRect(PAD_L, y0 + 5, chipW, chipH, 999); cx.fill();
+    cx.shadowBlur = 0;
+    cx.fillStyle = color;
+    cx.beginPath(); cx.arc(PAD_L + 12 + 4, y0 + 5 + chipH / 2, 4, 0, 7); cx.fill();
+    cx.textBaseline = 'middle'; cx.textAlign = 'left';
+    cx.fillStyle = dark ? '#e5e7eb' : '#1f2430';
+    cx.font = '700 13px system-ui, -apple-system, sans-serif';
+    cx.fillText(label, PAD_L + 12 + 9 + 7, y0 + 6 + chipH / 2);
+    cx.fillStyle = dark ? '#9aa2ae' : '#6b7280';
+    cx.font = '500 12px ui-monospace, SFMono-Regular, Menlo, monospace';
+    cx.fillText(count, PAD_L + 12 + 9 + 7 + lw + 10, y0 + 6 + chipH / 2);
+  });
+
+  // playhead
+  const x = tx(tNow);
+  cx.fillStyle = dark ? '#e5e7eb' : '#111827';
+  cx.globalAlpha = live ? .85 : (lastT > 0 ? .3 : 0);
+  cx.beginPath(); cx.roundRect(x - 1.5, PAD_T - 4, 3, H - PAD_T - PAD_B + 8, 2); cx.fill();
+  cx.globalAlpha = 1;
+}
+
+// --- playhead sync with the <midi-player> web component (lives in the
+// sibling gr.HTML component, so we look it up on document) ---
+function tick() {
+  if (!live) { rafId = null; return; }
+  if (wired && typeof wired.currentTime === 'number') lastT = wired.currentTime;
+  draw();
+  rafId = requestAnimationFrame(tick);
+}
+
+function wirePlayer() {
+  const player = document.querySelector('midi-player');
+  if (!player || wired === player) return;
+  wired = player;
+  player.addEventListener('start', () => {
+    live = true;
+    if (rafId === null) rafId = requestAnimationFrame(tick);
+  });
+  player.addEventListener('stop', () => {   // fires on both pause and finish
+    live = false;
+    if (rafId !== null) { cancelAnimationFrame(rafId); rafId = null; }
+    if (wired && typeof wired.currentTime === 'number') lastT = wired.currentTime;
+    draw();
+  });
+}
+
+// Low-rate poll: wires the player once it mounts, and catches scrubbing /
+// theme flips while paused. Deliberately NOT a ResizeObserver.
+setInterval(() => {
+  wirePlayer();
+  if (!live && wired && mt) {
+    const t = wired.currentTime;
+    if (typeof t === 'number' && Math.abs(t - lastT) > 0.05) { lastT = t; draw(); }
+  }
+}, 400);
+
+let resizeTimer = null;
+window.addEventListener('resize', () => {
+  clearTimeout(resizeTimer);
+  resizeTimer = setTimeout(() => { layout(); draw(); }, 150);
+});
+
+function syncRoll() {
+  let data = null;
+  try { data = props.value ? JSON.parse(props.value) : null; } catch (e) { data = null; }
+  if (data && data.tracks && data.tracks.length) {
+    mt = data; lastT = 0; live = false;
+    container.style.display = 'block';
+    layout(); draw();
+  } else {
+    mt = null;
+    container.style.display = 'none';
+  }
+}
+syncRoll();
+watch('value', syncRoll);
+wirePlayer();
+"""
+
+
+def _midi_roll_data(midi_bytes: bytes | None) -> dict:
+    """Per-track note data for the piano-roll canvas (and the summary stats).
+
+    Returns {"span": seconds, "tracks": [{"label", "drum", "notes":
+    [[onset, offset, pitch], ...]}, ...]} computed from the real returned
+    MIDI bytes with pretty_midi. All tracks are included (vocals too).
+    """
+    if not midi_bytes:
+        return {}
+    try:
+        pm = pretty_midi.PrettyMIDI(io.BytesIO(midi_bytes))
+    except Exception:
+        return {}
+    tracks = []
+    for inst in pm.instruments:
+        if not inst.notes:
+            continue
+        if inst.is_drum:
+            label = "Drums"
+        elif inst.name and inst.name.strip():
+            label = inst.name.strip()
+        else:
+            label = _PROGRAM_TO_GROUP_LABEL.get(
+                inst.program
+            ) or pretty_midi.program_to_instrument_name(inst.program)
+        notes = sorted(
+            [round(float(n.start), 4), round(float(n.end), 4), int(n.pitch)]
+            for n in inst.notes
+        )
+        tracks.append({"label": label, "drum": bool(inst.is_drum), "notes": notes})
+    if not tracks:
+        return {}
+    span = max(float(pm.get_end_time()), 0.001)
+    return {"span": round(span, 4), "tracks": tracks}
+
+
 def _midi_player_src(midi_bytes: bytes | None) -> str:
     """Return the transcribed MIDI as a base64 data URI for the player.
 
@@ -175,7 +422,7 @@ def transcribe_audio(
         temperature: Sampling temperature (only used when use_sampling is True).
     """
     if audio_path is None:
-        return None, "Please upload an audio file first.", ""
+        return None, "Please upload an audio file first.", "", ""
 
     t0 = time.perf_counter()
 
@@ -188,7 +435,7 @@ def transcribe_audio(
             temperature=temperature,
         )
     except Exception as e:
-        return None, f"Transcription failed: {e}", ""
+        return None, f"Transcription failed: {e}", "", ""
 
     elapsed = time.perf_counter() - t0
 
@@ -197,29 +444,30 @@ def transcribe_audio(
     tmp.write(midi_bytes)
     tmp.close()
 
-    # Also collect note statistics for the summary
-    note_count = 0
-    instrument_set = set()
-    try:
-        from mido import MidiFile
-        midi = MidiFile(tmp.name)
-        for track in midi.tracks:
-            for msg in track:
-                if msg.type == "note_on" and msg.velocity > 0:
-                    note_count += 1
-                    instrument_set.add(msg.program if not msg.is_meta else 0)
-    except Exception:
-        pass
-
-    summary = (
-        f"Transcription complete in {elapsed:.1f}s. "
-        f"Found {note_count} notes across {len(instrument_set)} instrument program(s). "
-        f"Play it back below or download the MIDI file."
-    )
+    # Per-track note data drives both the piano roll and the summary stats.
+    # (The previous mido-based stats loop hit an AttributeError on the first
+    # note — note_on messages have no .program — so it always reported 1 note.)
+    roll = _midi_roll_data(midi_bytes)
+    tracks = roll.get("tracks", [])
+    total_notes = sum(len(t["notes"]) for t in tracks)
+    if tracks:
+        per_track = ", ".join(f"{t['label']} {len(t['notes'])}" for t in tracks)
+        summary = (
+            f"Transcription complete in {elapsed:.1f}s. "
+            f"Found {total_notes} note{'s' if total_notes != 1 else ''} "
+            f"across {len(tracks)} track(s): {per_track}. "
+            f"Play it back below or download the MIDI file."
+        )
+    else:
+        summary = (
+            f"Transcription complete in {elapsed:.1f}s, but no notes were "
+            f"detected in the result."
+        )
 
     player_src = _midi_player_src(midi_bytes)
+    roll_json = json.dumps(roll) if tracks else ""
 
-    return tmp.name, summary, player_src
+    return tmp.name, summary, player_src, roll_json
 
 
 with gr.Blocks() as demo:
@@ -275,6 +523,12 @@ with gr.Blocks() as demo:
                     html_template=_MIDI_PLAYER_TEMPLATE,
                     js_on_load=_MIDI_PLAYER_JS,
                 )
+                midi_roll = gr.HTML(
+                    value="",
+                    label="Piano roll (per instrument)",
+                    html_template=_MIDI_ROLL_TEMPLATE,
+                    js_on_load=_MIDI_ROLL_JS,
+                )
                 midi_output = gr.File(label="Download MIDI file")
 
         summary_output = gr.Textbox(label="Summary", interactive=False, visible=False)
@@ -282,7 +536,7 @@ with gr.Blocks() as demo:
     transcribe_btn.click(
         fn=transcribe_audio,
         inputs=[audio_input, instrument_checkbox, use_sampling, temperature],
-        outputs=[midi_output, summary_output, midi_player],
+        outputs=[midi_output, summary_output, midi_player, midi_roll],
         api_name="transcribe",
     )
 
@@ -293,7 +547,7 @@ with gr.Blocks() as demo:
             ["example_water_afro_pop.mp3", [], False, 1.0],
         ],
         inputs=[audio_input, instrument_checkbox, use_sampling, temperature],
-        outputs=[midi_output, summary_output, midi_player],
+        outputs=[midi_output, summary_output, midi_player, midi_roll],
         fn=transcribe_audio,
         cache_examples=True,
         cache_mode="lazy",

requirements.txt

@@ -2,5 +2,6 @@ git+https://github.com/muscriptor/muscriptor.git
 numpy>=1.24
 einops>=0.4
 mido>=1.3
+pretty_midi>=0.2.10
 safetensors>=0.4
 soundfile>=0.14.0