app.py

import os
import sys
import math
import tempfile
import pickle
from pathlib import Path

import torch
import numpy as np
import yaml
import gradio as gr
from pyharp import ModelCard, build_endpoint
from huggingface_hub import hf_hub_download

sys.path.insert(0, str(Path(__file__).parent))
from dataset import load_vocabs
from models.variants import build_model

os.environ["TOKENIZERS_PARALLELISM"] = "true"

REPO_ID = "Itsuki-music/BACHI_Chord_Recognition"
CHECKPOINT_NAMES = {
    "Classical": "classical_film_kdec",
    "Pop": "pop909_film_kdec",
}
CHECKPOINT_FILES = ["best_model.pt", "config.yaml", "vocab.pkl"]

loaded_models = {}

def get_model(model_type: str):
    if model_type not in loaded_models:
        folder = CHECKPOINT_NAMES[model_type]
        ckpt_dir = Path(__file__).parent / "ckpts" / folder
        ckpt_dir.mkdir(parents=True, exist_ok=True)

        for fname in CHECKPOINT_FILES:
            dest = ckpt_dir / fname
            if not dest.exists():
                print(f"Downloading {folder}/{fname}...", flush=True)
                hf_hub_download(
                    repo_id=REPO_ID,
                    filename=f"{folder}/{fname}",
                    repo_type="dataset",
                    local_dir=Path(__file__).parent / "ckpts",
                )
                print(f"Downloaded {fname}.", flush=True)

        with open(ckpt_dir / "config.yaml", "r") as f:
            config = yaml.safe_load(f)

        vocabs = load_vocabs(str(ckpt_dir / "vocab.pkl"))
        use_key = (
            bool(config.get("use_key", False))
            or bool(config["training"].get("use_key", False))
            or bool(config["model"].get("use_key", False))
        )
        experiment = config["experiment"]

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Loading {model_type} model...", flush=True)
        model = build_model(experiment, config["model"], vocabs, use_key=use_key).to(device)
        model.load_state_dict(
            torch.load(ckpt_dir / "best_model.pt", map_location=device, weights_only=True)
        )
        model.eval()
        print(f"{model_type} model loaded.", flush=True)

        loaded_models[model_type] = (model, config, vocabs, use_key, device)

    return loaded_models[model_type]


def extract_pianoroll(score_path: Path, resolution: int = 12):
    import miditoolkit
    from music21 import converter, note as m21_note, chord as m21_chord

    notes_data = []
    suffix = score_path.suffix.lower()

    if suffix in {".mid", ".midi"}:
        midi = miditoolkit.MidiFile(str(score_path))
        tpb = midi.ticks_per_beat or 480
        for inst in midi.instruments:
            if inst.is_drum:
                continue
            for n in inst.notes:
                notes_data.append((n.pitch, n.start / tpb, n.end / tpb))
    else:
        sc = converter.parse(str(score_path))
        parts = list(sc.parts) if sc.hasPartLikeStreams() else [sc]
        for part in parts:
            inst = part.getInstrument()
            if inst:
                classes = inst.classes if hasattr(inst, "classes") else []
                if "Percussion" in classes or "Unpitched" in classes:
                    continue
            for el in part.flat.notes:
                dur = float(el.quarterLength)
                start = float(el.offset)
                if isinstance(el, m21_note.Note):
                    notes_data.append((el.pitch.midi, start, start + dur))
                elif isinstance(el, m21_chord.Chord):
                    for p in el.pitches:
                        notes_data.append((p.midi, start, start + dur))

    if not notes_data:
        return None

    notes_data.sort(key=lambda x: x[1])
    last_end = max(nd[2] for nd in notes_data)
    total_frames = math.ceil(last_end * resolution)
    pianoroll = np.zeros((88, total_frames), dtype=np.int8)
    for midi_pitch, start_b, end_b in notes_data:
        row = midi_pitch - 21
        if not (0 <= row < 88):
            continue
        s_f = max(0, math.floor(start_b * resolution))
        e_f = min(total_frames, math.ceil(end_b * resolution))
        pianoroll[row, s_f:e_f] = 1

    return pianoroll


def predict_piece(pianoroll, model, config, vocabs, use_key, device):
    beat_resolution = config["model"]["beat_resolution"]
    label_resolution = config["model"]["label_resolution"]
    segment_len = config["model"]["n_beats"] * beat_resolution
    pr_to_label_ratio = beat_resolution // label_resolution
    comps_eval = ["root", "quality", "bass"] + (["key"] if use_key else [])

    n_frames = pianoroll.shape[0]
    segments, masks = [], []
    for i in range(0, n_frames, segment_len):
        seg = pianoroll[i : i + segment_len]
        orig_len = seg.shape[0]
        if orig_len < segment_len:
            seg = torch.cat([seg, torch.zeros(segment_len - orig_len, seg.shape[1])], dim=0)
        mask = torch.ones(segment_len, dtype=torch.bool)
        if orig_len < segment_len:
            mask[orig_len:] = False
        segments.append(seg)
        masks.append(mask)

    piece_preds = {k: [] for k in comps_eval + ["boundary"]}
    for i in range(0, len(segments), 16):
        batch_segs = torch.stack(segments[i : i + 16]).to(device)
        batch_masks = torch.stack(masks[i : i + 16]).to(device)
        with torch.no_grad():
            out = model.forward_infer(batch_segs, src_key_padding_mask=~batch_masks)
        for k in comps_eval + ["boundary"]:
            if k in out:
                piece_preds[k].append(out[k].detach().cpu())

    n_target = math.ceil(n_frames / pr_to_label_ratio)
    piece_pred_ids = {}
    for k, parts in piece_preds.items():
        if not parts:
            continue
        cat = torch.cat([p.reshape(-1) for p in parts], dim=0)
        piece_pred_ids[k] = cat[:n_target] if k == "boundary" else cat[:n_target].long()

    inv_root = {v: k for k, v in vocabs["root"].items()}
    inv_qual = {v: k for k, v in vocabs["quality"].items()}
    inv_bass = {v: k for k, v in vocabs["bass"].items()}

    valid_len = len(piece_pred_ids.get("root", []))
    if valid_len == 0:
        return "No predictions generated."

    r_seq = piece_pred_ids["root"][:valid_len].tolist()
    q_seq = piece_pred_ids["quality"][:valid_len].tolist()
    b_seq = piece_pred_ids["bass"][:valid_len].tolist()

    time_per_token = 1.0 / max(1, config["model"]["label_resolution"])
    merged = []
    cur_r, cur_q, cur_b = r_seq[0], q_seq[0], b_seq[0]
    cur_start = 0
    for t in range(1, valid_len):
        if r_seq[t] != cur_r or q_seq[t] != cur_q or b_seq[t] != cur_b:
            label = f"{inv_root.get(cur_r, str(cur_r))}_{inv_qual.get(cur_q, str(cur_q))}_{inv_bass.get(cur_b, str(cur_b))}"
            merged.append(f"{cur_start * time_per_token:.2f} {label}")
            cur_r, cur_q, cur_b = r_seq[t], q_seq[t], b_seq[t]
            cur_start = t
    label = f"{inv_root.get(cur_r, str(cur_r))}_{inv_qual.get(cur_q, str(cur_q))}_{inv_bass.get(cur_b, str(cur_b))}"
    merged.append(f"{cur_start * time_per_token:.2f} {label}")

    return "\n".join(merged)


model_card = ModelCard(
    name="BACHI Chord Recognition",
    description="Automatic chord recognition from symbolic music scores (MIDI or MusicXML). Outputs beat-aligned chord labels.",
    author="Mingyang Yao, Ke Chen, Shlomo Dubnov, Taylor Berg-Kirkpatrick",
    tags=["chord-recognition", "symbolic-music", "midi", "musicxml"],
)


def process_fn(input_file: str, model_type: str) -> str:
    print(f"Processing with {model_type} model...", flush=True)
    model, config, vocabs, use_key, device = get_model(model_type)

    score_path = Path(input_file)
    pianoroll_np = extract_pianoroll(score_path, resolution=config["model"]["beat_resolution"])
    if pianoroll_np is None:
        return "Error: Could not extract notes from the input file."

    pianoroll = torch.from_numpy(pianoroll_np.T).float()
    result = predict_piece(pianoroll, model, config, vocabs, use_key, device)
    print("Done.", flush=True)
    return result


with gr.Blocks() as demo:
    input_components = [
        gr.File(
            label="Input Score (.mid, .midi, .musicxml, .mxl, .xml)",
            file_types=[".mid", ".midi", ".musicxml", ".mxl", ".xml"],
        ),
        gr.Dropdown(
            choices=["Classical", "Pop"],
            value="Classical",
            label="Model Type",
        ),
    ]
    output_components = [
        gr.Textbox(label="Chord Predictions", lines=20),
    ]
    app = build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

print("Launching Gradio...", flush=True)
demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True)