| from __future__ import annotations |
|
|
| import importlib |
| import json |
| import sys |
| import threading |
| from pathlib import Path |
|
|
| import librosa |
| import mir_eval |
| import numpy as np |
| import torch |
| import torchaudio |
| from huggingface_hub import snapshot_download |
| from music21 import note, stream |
| from torch import nn |
| from transformers import AutoModel, Wav2Vec2FeatureExtractor |
|
|
| SOURCE_REPO = "amaai-lab/music2emo" |
| SOURCE_REVISION = "b036e59471583c3d5b30c69e63e8c7323cc36c4a" |
| MERT_REPO = "m-a-p/MERT-v1-95M" |
| MERT_REVISION = "12af15fef9d0ac838c3f475bfbbf26d2060dd4f5" |
| SAMPLE_RATE = 24000 |
| WINDOW_SECONDS = 30 |
| MOOD_CLASSES = 56 |
| _LOCK = threading.Lock() |
| _RUNTIME = None |
|
|
|
|
| class PositionalEncoding(nn.Module): |
| def __init__(self, width: int, max_length: int = 100): |
| super().__init__() |
| encoding = torch.zeros(max_length, width) |
| position = torch.arange(max_length, dtype=torch.float32).unsqueeze(1) |
| scale = torch.exp( |
| torch.arange(0, width, 2).float() * (-np.log(10000.0) / width) |
| ) |
| encoding[:, 0::2] = torch.sin(position * scale) |
| encoding[:, 1::2] = torch.cos(position * scale) |
| self.register_buffer("encoding", encoding.unsqueeze(0), persistent=False) |
|
|
| def forward(self, values: torch.Tensor) -> torch.Tensor: |
| return values + self.encoding[:, : values.size(1)] |
|
|
|
|
| class EmotionHead(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.root_embedding = nn.Embedding(14, 4) |
| self.attribute_embedding = nn.Embedding(14, 4) |
| self.position = PositionalEncoding(8) |
| layer = nn.TransformerEncoderLayer( |
| d_model=8, |
| nhead=8, |
| dim_feedforward=64, |
| dropout=0.1, |
| batch_first=True, |
| ) |
| self.chord_transformer = nn.TransformerEncoder(layer, num_layers=2) |
| self.input_projection = nn.Sequential(nn.Linear(1545, 512), nn.ReLU()) |
| self.classifier = nn.Sequential( |
| nn.Linear(512, 256), |
| nn.ReLU(), |
| nn.Linear(256, MOOD_CLASSES), |
| ) |
| self.regressor = nn.Sequential( |
| nn.Linear(512, 256), |
| nn.ReLU(), |
| nn.Linear(256, 2), |
| ) |
|
|
| def forward( |
| self, |
| mert: torch.Tensor, |
| chord_roots: torch.Tensor, |
| chord_attributes: torch.Tensor, |
| mode: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| chord_values = torch.cat( |
| ( |
| self.root_embedding(chord_roots), |
| self.attribute_embedding(chord_attributes), |
| ), |
| dim=-1, |
| ) |
| chord_values = self.position(chord_values) |
| cls_token = torch.zeros_like(chord_values[:, :1]) |
| chord_values = self.chord_transformer( |
| torch.cat((cls_token, chord_values), dim=1) |
| )[:, 0] |
| combined = torch.cat((mert, chord_values, mode.float()), dim=1) |
| hidden = self.input_projection(combined) |
| return self.classifier(hidden), self.regressor(hidden) |
|
|
|
|
| class Music2EmoRuntime: |
| def __init__(self): |
| self.source_dir = Path( |
| snapshot_download( |
| repo_id=SOURCE_REPO, |
| revision=SOURCE_REVISION, |
| allow_patterns=[ |
| "inference/data/*", |
| "saved_models/J_all.ckpt", |
| "utils/*.py", |
| ], |
| ) |
| ) |
| sys.path.insert(0, str(self.source_dir)) |
| self._load_source_modules() |
|
|
| self.mert = AutoModel.from_pretrained( |
| MERT_REPO, |
| revision=MERT_REVISION, |
| trust_remote_code=True, |
| ) |
| self.processor = Wav2Vec2FeatureExtractor.from_pretrained( |
| MERT_REPO, |
| revision=MERT_REVISION, |
| trust_remote_code=True, |
| ) |
| self.head = EmotionHead() |
| self._load_emotion_checkpoint() |
| self.chord_model = self.BTCModel(config=self.config.model) |
| self._load_chord_checkpoint() |
|
|
| tags = np.load(self.data_dir / "tag_list.npy", allow_pickle=True) |
| self.mood_labels = [ |
| str(tag).replace("mood/theme---", "") for tag in tags[-MOOD_CLASSES:] |
| ] |
| self.root_map = self._read_json("chord_root.json") |
| self.attribute_map = self._read_json("chord_attr.json") |
|
|
| @property |
| def data_dir(self) -> Path: |
| return self.source_dir / "inference" / "data" |
|
|
| def _load_source_modules(self) -> None: |
| hparams = importlib.import_module("utils.hparams") |
| btc_model = importlib.import_module("utils.btc_model") |
| chords = importlib.import_module("utils.mir_eval_modules") |
| self.config = hparams.HParams.load(self.data_dir / "run_config.yaml") |
| self.config.feature["large_voca"] = True |
| self.config.model["num_chords"] = 170 |
| self.BTCModel = btc_model.BTC_model |
| self.chord_vocabulary = chords.idx2voca_chord() |
|
|
| def _read_json(self, name: str) -> dict[str, int]: |
| return json.loads((self.data_dir / name).read_text(encoding="utf-8")) |
|
|
| def _load_emotion_checkpoint(self) -> None: |
| checkpoint = torch.load( |
| self.source_dir / "saved_models" / "J_all.ckpt", |
| map_location="cpu", |
| weights_only=False, |
| ) |
| state = { |
| key.removeprefix("model."): value |
| for key, value in checkpoint["state_dict"].items() |
| } |
| rename = { |
| "chord_root_embedding.": "root_embedding.", |
| "chord_attr_embedding.": "attribute_embedding.", |
| "positional_encoding.": "position.", |
| "input_proj.": "input_projection.", |
| "classification_branch.": "classifier.", |
| "regression_branch.": "regressor.", |
| } |
| converted = {} |
| for key, value in state.items(): |
| for source, target in rename.items(): |
| if key.startswith(source): |
| key = target + key[len(source) :] |
| break |
| converted[key] = value |
| expected = self.head.state_dict() |
| converted = {key: value for key, value in converted.items() if key in expected} |
| self.head.load_state_dict(converted, strict=True) |
| self.head.eval() |
|
|
| def _load_chord_checkpoint(self) -> None: |
| checkpoint = torch.load( |
| self.data_dir / "btc_model_large_voca.pt", |
| map_location="cpu", |
| weights_only=False, |
| ) |
| self.chord_mean = checkpoint["mean"] |
| self.chord_std = checkpoint["std"] |
| self.chord_model.load_state_dict(checkpoint["model"]) |
| self.chord_model.eval() |
|
|
| @staticmethod |
| def _audio(path: str) -> tuple[torch.Tensor, int]: |
| waveform, sample_rate = torchaudio.load(path) |
| waveform = waveform.mean(dim=0) |
| if sample_rate != SAMPLE_RATE: |
| waveform = torchaudio.functional.resample( |
| waveform, |
| sample_rate, |
| SAMPLE_RATE, |
| ) |
| return waveform, SAMPLE_RATE |
|
|
| def _mert_embedding( |
| self, |
| waveform: torch.Tensor, |
| device: torch.device, |
| ) -> torch.Tensor: |
| window = WINDOW_SECONDS * SAMPLE_RATE |
| chunks = waveform.split(window) |
| embeddings = [] |
| for chunk in chunks: |
| inputs = self.processor( |
| chunk, |
| sampling_rate=SAMPLE_RATE, |
| return_tensors="pt", |
| ) |
| inputs = {key: value.to(device) for key, value in inputs.items()} |
| outputs = self.mert(**inputs, output_hidden_states=True) |
| layer_means = torch.stack(outputs.hidden_states[1:]).mean(dim=2) |
| embeddings.append(torch.cat((layer_means[5], layer_means[6]), dim=1)) |
| return torch.stack(embeddings).mean(dim=0) |
|
|
| def _chord_intervals( |
| self, |
| audio_path: str, |
| device: torch.device, |
| ) -> list[tuple[float, float, str]]: |
| config = self.config |
| audio, sample_rate = librosa.load( |
| audio_path, |
| sr=config.mp3["song_hz"], |
| mono=True, |
| ) |
| feature = librosa.cqt( |
| audio, |
| sr=sample_rate, |
| n_bins=config.feature["n_bins"], |
| bins_per_octave=config.feature["bins_per_octave"], |
| hop_length=config.feature["hop_length"], |
| ) |
| feature = np.log(np.abs(feature) + 1e-6).T |
| feature = (feature - self.chord_mean) / self.chord_std |
|
|
| timestep = config.model["timestep"] |
| pad = timestep - (feature.shape[0] % timestep) |
| feature = np.pad(feature, ((0, pad), (0, 0))) |
| blocks = feature.shape[0] // timestep |
| frame_seconds = config.mp3["inst_len"] / timestep |
| changes: list[tuple[float, float, str]] = [] |
| start = 0.0 |
| previous = None |
|
|
| tensor = torch.tensor(feature, dtype=torch.float32).unsqueeze(0).to(device) |
| for block in range(blocks): |
| section = tensor[:, block * timestep : (block + 1) * timestep] |
| encoded, _ = self.chord_model.self_attn_layers(section) |
| prediction, _ = self.chord_model.output_layer(encoded) |
| for offset, chord_index in enumerate(prediction.squeeze().tolist()): |
| frame = block * timestep + offset |
| if frame >= feature.shape[0] - pad: |
| break |
| if previous is None: |
| previous = chord_index |
| elif chord_index != previous: |
| end = frame * frame_seconds |
| changes.append((start, end, self.chord_vocabulary[previous])) |
| start = end |
| previous = chord_index |
|
|
| duration = len(audio) / sample_rate |
| if previous is not None and duration > start: |
| changes.append((start, duration, self.chord_vocabulary[previous])) |
| return changes |
|
|
| @staticmethod |
| def _key(intervals: list[tuple[float, float, str]]) -> tuple[str, str]: |
| score = stream.Stream() |
| note_count = 0 |
| for start, end, chord in intervals: |
| root, bitmap, _ = mir_eval.chord.encode(chord) |
| if root < 0: |
| continue |
| chroma = mir_eval.chord.rotate_bitmap_to_root(bitmap, root) |
| for pitch_class, active in enumerate(chroma): |
| if active: |
| value = note.Note(48 + pitch_class) |
| value.duration.quarterLength = max(end - start, 0.01) |
| score.insert(start, value) |
| note_count += 1 |
| if note_count == 0: |
| return "C", "major" |
| key = score.analyze("key") |
| tonic = str(key.tonic).replace("-", "b") |
| return tonic, str(key.mode) |
|
|
| def _encode_chords( |
| self, |
| intervals: list[tuple[float, float, str]], |
| tonic: str, |
| mode: str, |
| device: torch.device, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| pitch_classes = [ |
| "C", |
| "C#", |
| "D", |
| "D#", |
| "E", |
| "F", |
| "F#", |
| "G", |
| "G#", |
| "A", |
| "A#", |
| "B", |
| ] |
| flat_to_sharp = { |
| "Cb": "B", |
| "Db": "C#", |
| "Eb": "D#", |
| "Fb": "E", |
| "Gb": "F#", |
| "Ab": "G#", |
| "Bb": "A#", |
| } |
| tonic = flat_to_sharp.get(tonic, tonic) |
| reference = "A" if mode == "minor" else "C" |
| shift = (pitch_classes.index(tonic) - pitch_classes.index(reference)) % 12 |
| roots = [] |
| attributes = [] |
| for _, _, chord in intervals[:100]: |
| if chord in {"N", "X"}: |
| root, attribute = chord, 0 |
| else: |
| parts = chord.split(":", 1) |
| source_root = flat_to_sharp.get(parts[0], parts[0]) |
| root = pitch_classes[ |
| (pitch_classes.index(source_root) - shift) % 12 |
| ] |
| attribute_name = parts[1] if len(parts) == 2 else "maj" |
| attribute = self.attribute_map.get(attribute_name, 0) |
| roots.append(self.root_map.get(root, 0)) |
| attributes.append(attribute) |
| roots.extend([0] * (100 - len(roots))) |
| attributes.extend([0] * (100 - len(attributes))) |
| mode_value = 1 if mode == "minor" else 0 |
| return ( |
| torch.tensor(roots, dtype=torch.long, device=device).unsqueeze(0), |
| torch.tensor(attributes, dtype=torch.long, device=device).unsqueeze(0), |
| torch.tensor([[mode_value]], dtype=torch.long, device=device), |
| ) |
|
|
| def predict(self, audio_path: str, threshold: float) -> dict: |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self.mert.to(device).eval() |
| self.head.to(device).eval() |
| self.chord_model.to(device).eval() |
|
|
| waveform, _ = self._audio(audio_path) |
| with torch.inference_mode(): |
| mert = self._mert_embedding(waveform, device) |
| intervals = self._chord_intervals(audio_path, device) |
| tonic, mode = self._key(intervals) |
| roots, attributes, mode_tensor = self._encode_chords( |
| intervals, |
| tonic, |
| mode, |
| device, |
| ) |
| logits, dimensions = self.head( |
| mert, |
| roots, |
| attributes, |
| mode_tensor, |
| ) |
| probabilities = torch.sigmoid(logits).squeeze().cpu().tolist() |
| valence, arousal = dimensions.squeeze().cpu().tolist() |
|
|
| ranked = sorted( |
| ( |
| {"label": label, "probability": round(float(score), 4)} |
| for label, score in zip(self.mood_labels, probabilities) |
| if score >= threshold |
| ), |
| key=lambda item: item["probability"], |
| reverse=True, |
| ) |
| return { |
| "model": "Music2Emo", |
| "moods": ranked, |
| "valence": round(float(valence), 4), |
| "arousal": round(float(arousal), 4), |
| "scale": {"valence": [1, 9], "arousal": [1, 9]}, |
| "threshold": float(threshold), |
| "estimated_key": f"{tonic} {mode}", |
| } |
|
|
|
|
| def analyze_music(audio_path: str, threshold: float = 0.5) -> dict: |
| global _RUNTIME |
| with _LOCK: |
| if _RUNTIME is None: |
| _RUNTIME = Music2EmoRuntime() |
| return _RUNTIME.predict(audio_path, threshold) |
|
|
