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)