Spaces:
Running
Running
| import spaces # imported first, as ZeroGPU requires | |
| import os | |
| import sys | |
| import math | |
| import shutil | |
| import tempfile | |
| import pickle | |
| from pathlib import Path | |
| import torch | |
| import numpy as np | |
| import yaml | |
| import gradio as gr | |
| from pyharp import LabelList, MidiLabel, ModelCard, build_endpoint | |
| from pyharp.tags import Subcategory | |
| 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", | |
| revision="7fd45be21b2b917d0a3fc349b6f48e9486f6ad19", | |
| ) | |
| 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="cpu", 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 score_as_midi(score_path: Path) -> str: | |
| """The score as a MIDI file of its own, for the chord labels to be drawn over.""" | |
| handle, midi_path = tempfile.mkstemp(suffix=".mid") | |
| os.close(handle) | |
| if score_path.suffix.lower() in {".mid", ".midi"}: | |
| shutil.copyfile(score_path, midi_path) | |
| else: | |
| from music21 import converter | |
| converter.parse(str(score_path)).write("midi", fp=midi_path) | |
| return midi_path | |
| # The chord qualities of the vocabulary, as a symbol reads them | |
| QUALITY_SUFFIX = { | |
| "M": "", "m": "m", "M7": "maj7", "m7": "m7", "D7": "7", "+": "+", "+7": "+7", | |
| "o": "dim", "o7": "dim7", "/o7": "m7b5", "mM7": "mMaj7", "sus2": "sus2", "sus4": "sus4", | |
| "other": "?", | |
| } | |
| def chord_symbol(root: str, quality: str, bass: str) -> str: | |
| if root == "N" or quality == "N": | |
| return "N.C." | |
| symbol = root + QUALITY_SUFFIX.get(quality, quality) | |
| if bass not in (root, "N", "PAD"): | |
| symbol += "/" + bass | |
| return symbol | |
| def chord_labels(segments, midi_path: str) -> LabelList: | |
| """One label per chord segment, placed in seconds by the tempo map of the MIDI.""" | |
| import miditoolkit | |
| midi = miditoolkit.MidiFile(midi_path) | |
| ticks_per_beat = midi.ticks_per_beat or 480 | |
| mapping = midi.get_tick_to_time_mapping() | |
| final_tempo = midi.tempo_changes[-1].tempo if midi.tempo_changes else 120.0 | |
| def seconds(beat: float) -> float: | |
| tick = int(round(beat * ticks_per_beat)) | |
| if tick < len(mapping): | |
| return float(mapping[tick]) | |
| # Past the last event of the file, time runs on at the final tempo | |
| last = len(mapping) - 1 | |
| return float(mapping[last]) + (tick - last) / ticks_per_beat * 60.0 / final_tempo | |
| labels = LabelList() | |
| for start_beat, end_beat, root, quality, bass in segments: | |
| if root == "PAD": | |
| continue | |
| start, end = seconds(start_beat), seconds(end_beat) | |
| labels.append(MidiLabel( | |
| t=start, | |
| duration=max(0.0, end - start), | |
| label=chord_symbol(root, quality, bass), | |
| description=(f"Beats {start_beat:g} to {end_beat:g}: root {root}, quality {quality}, " | |
| f"bass {bass}"), | |
| )) | |
| return labels | |
| 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 [] | |
| 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() | |
| # Each label token covers this many beats, and a segment runs until the chord changes | |
| beats_per_token = 1.0 / max(1, config["model"]["label_resolution"]) | |
| segments = [] | |
| cur_r, cur_q, cur_b = r_seq[0], q_seq[0], b_seq[0] | |
| cur_start = 0 | |
| def close(end_token): | |
| segments.append((cur_start * beats_per_token, end_token * beats_per_token, | |
| inv_root.get(cur_r, str(cur_r)), inv_qual.get(cur_q, str(cur_q)), | |
| inv_bass.get(cur_b, str(cur_b)))) | |
| 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: | |
| close(t) | |
| cur_r, cur_q, cur_b = r_seq[t], q_seq[t], b_seq[t] | |
| cur_start = t | |
| close(valid_len) | |
| return segments | |
| model_card = ModelCard( | |
| name="BACHI", | |
| description=( | |
| "Automatic chord recognition from symbolic music scores (MIDI or " | |
| "MusicXML). Returns the score as MIDI, with a chord label over each " | |
| "segment, aligned to the beat. Repository: " | |
| "https://github.com/AndyWeasley2004/BACHI_Chord_Recognition Paper: " | |
| "https://arxiv.org/abs/2510.06528" | |
| ), | |
| author="Mingyang Yao, Ke Chen, Shlomo Dubnov, Taylor Berg-Kirkpatrick", | |
| tags=[Subcategory.MUSIC_ANALYSIS, 'domain:music'], | |
| ) | |
| def process_fn(input_file: str, model_type: str) -> tuple[str, LabelList]: | |
| # Checked before the GPU is reserved, so a request the model cannot take spends none of the | |
| # caller's ZeroGPU allowance | |
| if model_type not in CHECKPOINT_NAMES: | |
| raise gr.Error(f"Unknown model type: {model_type}") | |
| if not input_file: | |
| raise gr.Error("Provide the score to analyze.") | |
| config = get_model(model_type)[1] | |
| score_path = Path(input_file) | |
| try: | |
| pianoroll_np = extract_pianoroll(score_path, | |
| resolution=config["model"]["beat_resolution"]) | |
| midi_path = score_as_midi(score_path) | |
| except Exception as exc: | |
| raise gr.Error(f"Could not read the score: {exc}") from exc | |
| if pianoroll_np is None: | |
| raise gr.Error("Could not extract notes from the input file.") | |
| print(f"Processing with {model_type} model...", flush=True) | |
| segments = recognize(pianoroll_np, model_type) | |
| print("Done.", flush=True) | |
| if not segments: | |
| raise gr.Error("The model found no chords in this score.") | |
| return midi_path, chord_labels(segments, midi_path) | |
| def recognize(pianoroll_np, model_type: str) -> list: | |
| model, config, vocabs, use_key, device = get_model(model_type) | |
| pianoroll = torch.from_numpy(pianoroll_np.T).float() | |
| return predict_piece(pianoroll, model, config, vocabs, use_key, device) | |
| # pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the | |
| # server, so the model is loaded there, before the first request rather than during it. | |
| if __name__ != "__main__": | |
| for _model_type in CHECKPOINT_NAMES: | |
| get_model(_model_type) | |
| if __name__ == "__main__": | |
| with gr.Blocks() as demo: | |
| input_components = [ | |
| gr.File(type="filepath", | |
| label="Input Score (.mid, .midi, .musicxml, .mxl, .xml)", | |
| file_types=[".mid", ".midi", ".musicxml", ".mxl", ".xml"], | |
| ).set_info(( | |
| "A score to analyze. Chords are read from the notes, so a melody " | |
| "alone gives less to work with than a full texture." | |
| )), | |
| gr.Dropdown( | |
| choices=["Classical", "Pop"], | |
| value="Classical", | |
| label="Model Type", | |
| info=( | |
| "Which trained model to use. The default was trained on the widest " | |
| "range of material." | |
| ), | |
| ), | |
| ] | |
| output_components = [ | |
| gr.File( | |
| type="filepath", | |
| file_types=[".mid"], | |
| label="Score as MIDI", | |
| ).set_info("The score, as MIDI, with the chord labels over it."), | |
| gr.JSON(label="Chord Labels").set_info( | |
| "One label per chord, from its first beat to the next change." | |
| ), | |
| ] | |
| 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) | |