app.py

"""MuScriptor Gradio Space — multi-instrument AMT -> MIDI + MusicXML + piano roll."""
from __future__ import annotations

import json
import os
import tempfile
import time
import warnings
import uuid
from datetime import datetime, timezone
from pathlib import Path

import gradio as gr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import partitura as pt
import pretty_midi
import soundfile as sf
from muscriptor import TranscriptionModel

# ---- Friendly UI labels -> model vocabulary ----
# The model expects MT3-style program group names. Passing the wrong name
# silently masks the channel. Map the UI's friendly buckets to the model's
# actual group names. Run `muscriptor list-instruments` to see the full list.
INSTRUMENT_MAP: dict[str, list[str]] = {
    "piano":  ["acoustic_piano", "electric_piano", "grand_piano", "harpsichord", "upright_piano"],
    "drums":  ["drum_kit"],
    "bass":   ["bass", "acoustic_bass", "electric_bass", "synth_bass", "double_bass", "upright_bass", "fretless_bass", "slap_bass"],
    "guitar": ["acoustic_guitar", "clean_electric_guitar", "distorted_electric_guitar", "electric_guitar", "synth_guitar", "classical_guitar", "acoustic_nylon_guitar", "acoustic_steel_guitar"],
    "vocals": ["lead_vocal", "vocal_harmonies", "background_vocal", "synth_choir", "choir"],
    "other":  [
        "accordion", "banjo", "bassoon", "brass", "celesta", "cello", "clarinet",
        "cymbals", "fiddle", "flute", "french_horn", "glockensharp", "glockenspiel",
        "harp", "marimba", "melodic_percussion", "oboe", "organ", "pad_synth",
        "percussion", "pizzicato_strings", "saxophone", "sitar", "steel_drum",
        "string_section", "synth", "synth_brass", "synth_drum", "synth_lead",
        "synth_pad", "synth_strings", "timpani", "trombone", "trumpet",
        "tubular_bells", "ukulele", "viola", "violin", "woodwind", "xylophone",
    ],
}
FRIENDLY_LABELS = list(INSTRUMENT_MAP.keys()) + ["all"]


def _resolve_instruments(selection: list[str] | None) -> list[str] | None:
    """Map the UI's friendly buckets to the model's group names.

    - None or ["all"] -> None (model decides)
    - empty list -> None (treat as "all")
    - any "all" in the list -> None
    - otherwise, flatten the mapped model names, dedup, preserve order
    """
    if not selection:
        return None
    sel = [s.lower() for s in selection]
    if "all" in sel:
        return None
    out: list[str] = []
    for label in sel:
        for name in INSTRUMENT_MAP.get(label, []):
            if name not in out:
                out.append(name)
    return out or None


# ---- Model cache (one process, multiple sizes) ----
_MODELS: dict[str, TranscriptionModel] = {}


def get_model(size: str) -> TranscriptionModel:
    size = (size or "medium").lower()
    if size not in _MODELS:
        _MODELS[size] = TranscriptionModel.load_model(size)
    return _MODELS[size]


# ---- Piano roll rendering (one row per instrument) ----
def render_piano_roll(midi_path: str) -> str:
    pm = pretty_midi.PrettyMIDI(midi_path)
    instruments = [inst for inst in pm.instruments if inst.notes]
    if not instruments:
        # Empty roll, return a placeholder.
        fig, ax = plt.subplots(figsize=(10, 2.5))
        ax.text(0.5, 0.5, "No notes detected", ha="center", va="center",
                transform=ax.transAxes, fontsize=14, color="#888")
        ax.set_axis_off()
        out = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        plt.savefig(out.name, dpi=110, bbox_inches="tight")
        plt.close(fig)
        return out.name

    cmap = plt.get_cmap("tab10")
    fig, ax = plt.subplots(figsize=(12, max(2.5, 0.45 * len(instruments) + 1.2)))
    end = pm.get_end_time() or 1.0
    for i, inst in enumerate(instruments):
        y = i
        for note in inst.notes:
            ax.add_patch(plt.Rectangle(
                (note.start, y - 0.4),
                max(note.end - note.start, 0.01),
                0.8,
                facecolor=cmap(i % 10),
                edgecolor="none",
                alpha=0.85,
            ))
    ax.set_xlim(0, max(end, 1.0))
    ax.set_ylim(-0.5, len(instruments) - 0.5)
    ax.set_yticks(range(len(instruments)))
    ax.set_yticklabels([f"{inst.name or inst.program}" for inst in instruments], fontsize=9)
    ax.set_xlabel("time (s)")
    ax.set_title(f"Piano roll \u2014 {len(instruments)} instruments, {pm.get_end_time():.1f}s")
    ax.grid(True, axis="x", alpha=0.2)
    plt.tight_layout()
    out = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    plt.savefig(out.name, dpi=110)
    plt.close(fig)
    return out.name


# ---- MIDI -> MusicXML ----
def render_musicxml(midi_path: str) -> str:
    """Convert MIDI to MusicXML using partitura.

    Returns the path to the .musicxml file (or empty string on failure).
    Multi-instrument MIDI -> a part per instrument, concatenated.
    """
    try:
        score = pt.load_score_midi(midi_path)
        # If multi-instrument, partitura's load_score_midi will return a Score
        # with one part per instrument; we serialize each to MusicXML and
        # combine into one file with a single <part-list>.
        out = tempfile.NamedTemporaryFile(suffix=".musicxml", delete=False)
        out.close()
        pt.save_musicxml(out.name, score)
        return out.name
    except Exception as e:  # noqa: BLE001 — we want a soft failure, not a 500
        warnings.warn(f"MusicXML export failed: {e}")
        return ""


# ---- Duration sniff (for the >3 min warning) ----
def audio_duration_sec(path: str) -> float | None:
    try:
        info = sf.info(path)
        return float(info.frames) / float(info.samplerate)
    except Exception:  # noqa: BLE001
        return None


# ---- Transcribe ----
def transcribe(audio_path, model_size, instruments, _progress=gr.Progress()):
    if audio_path is None:
        raise gr.Error("Upload an audio file or record from the mic first.")

    instr = _resolve_instruments(instruments)
    duration = audio_duration_sec(audio_path)
    if duration is not None and duration > 180:
        # Surface the warning but don't refuse — the user may want to try.
        warnings.warn(
            f"Input is {duration:.0f}s; the free CPU tier can take 10\u201330 min "
            "per transcribe on the medium model. Consider using the 'small' "
            "model or upgrading to a GPU Space for longer files.",
            stacklevel=2,
        )

    t0 = time.time()
    _progress(0.05, desc=f"Loading {model_size} model\u2026")
    model = get_model(model_size)

    _progress(0.20, desc="Transcribing\u2026")
    midi_bytes = model.transcribe_to_midi(audio_path, instruments=instr)

    out_midi = tempfile.NamedTemporaryFile(suffix=".mid", delete=False)
    out_midi.write(midi_bytes)
    out_midi.close()

    _progress(0.85, desc="Rendering piano roll + MusicXML\u2026")
    png = render_piano_roll(out_midi.name)
    musicxml = render_musicxml(out_midi.name)

    elapsed = time.time() - t0
    summary = {
        "model": model_size,
        "instruments": instruments or "all",
        "instruments_resolved": instr or "all",
        "elapsed_sec": round(elapsed, 2),
        "input_duration_sec": round(duration, 2) if duration else None,
        "midi_path": out_midi.name,
        "musicxml_path": musicxml or None,
        "piano_roll_path": png,
    }
    # Stringify once: used both for the analytics event's model_size log and the
    # user-facing Run summary panel — keeping a single source of truth.
    summary_json: str = json.dumps(summary, indent=2)

    # Append-only analytics event. JSONL file at /tmp — survives container lifetime,
    # resets on HF restart. Upgrade path: swap path for `EVENT_LOG_PATH` env var and
    # configure a durable backend once we know what we actually want to track.
    try:
        with open("/tmp/muscriptor-events.jsonl", "a") as f:
            f.write(json.dumps({
                "event": "transcribe",
                "ts": datetime.now(timezone.utc).isoformat(),
                "session_id": str(uuid.uuid4())[:8],  # request-level, not user-level
                "model_size": model_size,
                "instruments": instruments or ["all"],
                "input_duration_sec": duration,
                "elapsed_sec": round(elapsed, 2),
                "success": True,
            }) + "\n")
    except OSError:
        # Analytics must never break the request.
        pass

    # Return midi + musicxml as gr.File list (single-element lists are valid).
    musicxml_files = [musicxml] if musicxml else []
    return png, out_midi.name, musicxml_files, summary_json


# ---- UI ----
with gr.Blocks(title="MuScriptor \u2014 Multi-Instrument AMT", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """# MuScriptor \u2014 Multi-Instrument Music Transcription

Drag in an audio file (WAV/MP3/FLAC) or record from the mic, pick a model size,
and get back a piano roll, downloadable MIDI, and (when possible) MusicXML.

Model: [muscriptor/muscriptor-medium](https://huggingface.co/MuScriptor/muscriptor-medium) \u00b7
code: [github.com/muscriptor/muscriptor](https://github.com/muscriptor/muscriptor) \u00b7
licence: weights CC BY-NC 4.0, code MIT."""
    )
    with gr.Row():
        with gr.Column():
            audio = gr.Audio(type="filepath", label="Audio input")
            model_size = gr.Radio(
                ["small", "medium", "large"],
                value="small",
                label="Model size (small is fast on free CPU; medium/large want a GPU)",
            )
            instruments = gr.CheckboxGroup(
                FRIENDLY_LABELS, value=["all"], label="Instruments (uncheck to filter)"
            )
            run = gr.Button("Transcribe", variant="primary")
        with gr.Column():
            roll = gr.Image(label="Piano roll", type="filepath")
            midi_dl = gr.File(label="Download MIDI")
            musicxml_dl = gr.File(label="Download MusicXML (engraving)", file_count="single")
            summary = gr.Code(label="Run summary", language="json")
    run.click(
        transcribe,
        inputs=[audio, model_size, instruments],
        outputs=[roll, midi_dl, musicxml_dl, summary],
    )
    gr.Markdown(
        """### Notes

- First run downloads model weights (~600 MB for `medium`). Subsequent runs are cached.
- Free CPU tier is fine for files up to 3 minutes. For longer files, pick `small` or upgrade to a GPU Space.
- The first user must accept the [CC BY-NC 4.0 licence](https://huggingface.co/MuScriptor/muscriptor-medium) on the model page once per HF account, and the Space must have an `HF_TOKEN` secret set in Settings \u2192 Variables.
- Weights are CC BY-NC 4.0 (gated); inference code is MIT. Non-commercial use only."""
    )


if __name__ == "__main__":
    # HF Spaces injects HF_TOKEN at runtime if it's set as a Space secret.
    # The muscriptor package will pick it up via huggingface_hub.
    port = int(os.environ.get("PORT", "7860"))
    demo.queue(max_size=8).launch(server_name="0.0.0.0", server_port=port)