import torch import torch.nn as nn import torch.nn.functional as F import numpy as np import soundfile as sf import librosa from transformers import PretrainedConfig, PreTrainedModel from huggingface_hub import PyTorchModelHubMixin # Suppress warnings import warnings warnings.filterwarnings('ignore') class MusicNNConfig(PretrainedConfig): model_type = 'musicnn' def __init__( self, num_classes=50, mid_filt=64, backend_units=200, dataset='MTT', **kwargs ): self.num_classes = num_classes self.mid_filt = mid_filt self.backend_units = backend_units self.dataset = dataset super().__init__(**kwargs) # ------------------------- # Building blocks # ------------------------- class ConvReLUBN(nn.Module): def __init__(self, in_ch, out_ch, kernel_size, padding=0): super().__init__() self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, padding=padding) self.bn = nn.BatchNorm2d(out_ch, eps=0.001, momentum=0.01) def forward(self, x): return self.bn(F.relu(self.conv(x))) class TimbralBlock(nn.Module): def __init__(self, mel_bins, out_ch): super().__init__() self.conv_block = ConvReLUBN(1, out_ch, kernel_size=(7, mel_bins), padding=0) def forward(self, x): x = F.pad(x, (0, 0, 3, 3)) x = self.conv_block(x) return torch.max(x, dim=3).values class TemporalBlock(nn.Module): def __init__(self, kernel_size, out_ch): super().__init__() self.conv_block = ConvReLUBN(1, out_ch, kernel_size=(kernel_size, 1), padding='same') def forward(self, x): x = self.conv_block(x) return torch.max(x, dim=3).values class MidEnd(nn.Module): def __init__(self, in_ch, num_filt): super().__init__() self.c1_conv = nn.Conv2d(1, num_filt, kernel_size=(7, in_ch), padding=0) self.c1_bn = nn.BatchNorm2d(num_filt, eps=0.001, momentum=0.01) self.c2_conv = nn.Conv2d(1, num_filt, kernel_size=(7, num_filt), padding=0) self.c2_bn = nn.BatchNorm2d(num_filt, eps=0.001, momentum=0.01) self.c3_conv = nn.Conv2d(1, num_filt, kernel_size=(7, num_filt), padding=0) self.c3_bn = nn.BatchNorm2d(num_filt, eps=0.001, momentum=0.01) def forward(self, x): x = x.transpose(1, 2).unsqueeze(3) x_perm = x.permute(0, 2, 3, 1) x1_pad = F.pad(x_perm, (3, 3, 0, 0)) x1 = x1_pad.permute(0, 2, 3, 1) x1 = self.c1_bn(F.relu(self.c1_conv(x1))) x1_t = x1.permute(0, 2, 1, 3) x2_perm = x1_t.permute(0, 2, 3, 1) x2_pad = F.pad(x2_perm, (3, 3, 0, 0)) x2 = x2_pad.permute(0, 2, 3, 1) x2 = self.c2_bn(F.relu(self.c2_conv(x2))) x2_t = x2.permute(0, 2, 1, 3) res_conv2 = x2_t + x1_t x3_perm = res_conv2.permute(0, 2, 3, 1) x3_pad = F.pad(x3_perm, (3, 3, 0, 0)) x3 = x3_pad.permute(0, 2, 3, 1) x3 = self.c3_bn(F.relu(self.c3_conv(x3))) x3_t = x3.permute(0, 2, 1, 3) res_conv3 = x3_t + res_conv2 return [x.squeeze(3), x1_t.squeeze(3), res_conv2.squeeze(3), res_conv3.squeeze(3)] class Backend(nn.Module): def __init__(self, in_ch, num_classes, hidden): super().__init__() self.bn_in = nn.BatchNorm1d(in_ch * 2, eps=0.001, momentum=0.01) self.fc1 = nn.Linear(in_ch * 2, hidden) self.bn_fc1 = nn.BatchNorm1d(hidden, eps=0.001, momentum=0.01) self.fc2 = nn.Linear(hidden, num_classes) def forward(self, x): max_pool = torch.max(x, dim=1).values mean_pool = torch.mean(x, dim=1) z = torch.stack([max_pool, mean_pool], dim=2) z = z.view(z.size(0), -1) z = self.bn_in(z) z = F.dropout(z, p=0.5, training=self.training) z = self.bn_fc1(F.relu(self.fc1(z))) z = F.dropout(z, p=0.5, training=self.training) logits = self.fc2(z) return logits, mean_pool, max_pool class MusicNN(PreTrainedModel, PyTorchModelHubMixin): config_class = MusicNNConfig def __init__(self, config): super().__init__(config) self.bn_input = nn.BatchNorm2d(1, eps=0.001, momentum=0.01) self.timbral_1 = TimbralBlock(int(0.4 * 96), int(1.6 * 128)) self.timbral_2 = TimbralBlock(int(0.7 * 96), int(1.6 * 128)) self.temp_1 = TemporalBlock(128, int(1.6 * 32)) self.temp_2 = TemporalBlock(64, int(1.6 * 32)) self.temp_3 = TemporalBlock(32, int(1.6 * 32)) self.midend = MidEnd(in_ch=561, num_filt=config.mid_filt) self.backend = Backend(in_ch=config.mid_filt * 3 + 561, num_classes=config.num_classes, hidden=config.backend_units) def forward(self, x): x = x.unsqueeze(1) x = self.bn_input(x) f74 = self.timbral_1(x).transpose(1, 2) f77 = self.timbral_2(x).transpose(1, 2) s1 = self.temp_1(x).transpose(1, 2) s2 = self.temp_2(x).transpose(1, 2) s3 = self.temp_3(x).transpose(1, 2) frontend_features = torch.cat([f74, f77, s1, s2, s3], dim=2) mid_feats = self.midend(frontend_features.transpose(1, 2)) z = torch.cat(mid_feats, dim=2) logits, mean_pool, max_pool = self.backend(z) return logits, mean_pool, max_pool @staticmethod def preprocess_audio(audio_file, sr=16000): # Try librosa first (works well for many formats) try: audio, file_sr = librosa.load(audio_file, sr=None) if len(audio) == 0: raise ValueError("Empty audio from librosa") except Exception: # Fallback to soundfile (better for some MP3s) try: audio, file_sr = sf.read(audio_file) # Convert to mono if stereo if len(audio.shape) > 1: audio = np.mean(audio, axis=1) except Exception as e: raise ValueError(f'Could not load audio file {audio_file}: {e}') # Resample to target sample rate if necessary if file_sr != sr: audio = librosa.resample(audio, orig_sr=file_sr, target_sr=sr) if len(audio) == 0: raise ValueError(f'Audio file {audio_file} is empty or could not be loaded.') # Create mel spectrogram audio_rep = librosa.feature.melspectrogram( y=audio, sr=sr, hop_length=256, n_fft=512, n_mels=96 ).T audio_rep = audio_rep.astype(np.float32) audio_rep = np.log10(10000 * audio_rep + 1) return audio_rep def predict_tags(self, audio_file, top_k=5): # Auto-detect device and move model to it device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.to(device) # Use the same batching approach as the original implementation # This matches musicnn_torch.py extractor function # Load and preprocess audio (similar to batch_data in musicnn_torch.py) audio, file_sr = sf.read(audio_file) # Convert to mono if stereo if len(audio.shape) > 1: audio = np.mean(audio, axis=1) # Resample to 16000 if necessary if file_sr != 16000: audio = librosa.resample(audio, orig_sr=file_sr, target_sr=16000) if len(audio) == 0: raise ValueError(f'Audio file {audio_file} is empty or could not be loaded.') # Create mel spectrogram audio_rep = librosa.feature.melspectrogram( y=audio, sr=16000, hop_length=256, n_fft=512, n_mels=96 ).T audio_rep = audio_rep.astype(np.float32) audio_rep = np.log10(10000 * audio_rep + 1) # Batch the data (same as original implementation) n_frames = 187 # librosa.time_to_frames(3, sr=16000, n_fft=512, hop_length=256) + 1 overlap = n_frames # No overlap for simplicity last_frame = audio_rep.shape[0] - n_frames + 1 batches = [] if last_frame <= 0: # Pad with zeros if audio is too short patch = np.zeros((n_frames, 96), dtype=np.float32) patch[:audio_rep.shape[0], :] = audio_rep batches.append(patch) else: # Create overlapping windows for time_stamp in range(0, last_frame, overlap): patch = audio_rep[time_stamp : time_stamp + n_frames, :] batches.append(patch) # Convert to tensor and run inference batch_tensor = torch.from_numpy(np.stack(batches)).to(device) all_probs = [] with torch.no_grad(): self.eval() for i in range(0, len(batches), 1): # Process in batches if needed batch_subset = batch_tensor[i:i+1] logits, _, _ = self(batch_subset) probs = torch.sigmoid(logits).squeeze(0).cpu().numpy() all_probs.append(probs) # Average probabilities across all windows avg_probs = np.mean(all_probs, axis=0) # Get labels based on config if self.config.dataset == 'MTT': labels = [ 'guitar', 'classical', 'slow', 'techno', 'strings', 'drums', 'electronic', 'rock', 'fast', 'piano', 'ambient', 'beat', 'violin', 'vocal', 'synth', 'female', 'indian', 'opera', 'male', 'singing', 'vocals', 'no vocals', 'harpsichord', 'loud', 'quiet', 'flute', 'woman', 'male vocal', 'no vocal', 'pop', 'soft', 'sitar', 'solo', 'man', 'classic', 'choir', 'voice', 'new age', 'dance', 'male voice', 'female vocal', 'beats', 'harp', 'cello', 'no voice', 'weird', 'country', 'metal', 'female voice', 'choral' ] elif self.config.dataset == 'MSD': labels = [ 'rock', 'pop', 'alternative', 'indie', 'electronic', 'female vocalists', 'dance', '00s', 'alternative rock', 'jazz', 'beautiful', 'metal', 'chillout', 'male vocalists', 'classic rock', 'soul', 'indie rock', 'Mellow', 'electronica', '80s', 'folk', '90s', 'chill', 'instrumental', 'punk', 'oldies', 'blues', 'hard rock', 'ambient', 'acoustic', 'experimental', 'female vocalist', 'guitar', 'Hip-Hop', '70s', 'party', 'country', 'easy listening', 'sexy', 'catchy', 'funk', 'electro', 'heavy metal', 'Progressive rock', '60s', 'rnb', 'indie pop', 'sad', 'House', 'happy' ] else: raise ValueError(f"Unknown dataset: {self.config.dataset}") # Get top k tags top_indices = np.argsort(avg_probs)[-top_k:][::-1] return [labels[i] for i in top_indices] def extract_embeddings(self, audio_file, layer=None, pool='mean'): """ Extract embeddings from audio file. Args: audio_file: path to audio file layer: which layer to extract from (ignored for simplicity, uses final embeddings) pool: pooling method ('mean', 'max', or 'both') Returns: embeddings as numpy array """ # Auto-detect device and move model to it device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.to(device) # Load and preprocess audio audio, file_sr = sf.read(audio_file) # Convert to mono if stereo if len(audio.shape) > 1: audio = np.mean(audio, axis=1) # Resample to 16000 if necessary if file_sr != 16000: audio = librosa.resample(audio, orig_sr=file_sr, target_sr=16000) if len(audio) == 0: raise ValueError(f'Audio file {audio_file} is empty or could not be loaded.') # Create mel spectrogram audio_rep = librosa.feature.melspectrogram( y=audio, sr=16000, hop_length=256, n_fft=512, n_mels=96 ).T audio_rep = audio_rep.astype(np.float32) audio_rep = np.log10(10000 * audio_rep + 1) # Batch the data n_frames = 187 # librosa.time_to_frames(3, sr=16000, n_fft=512, hop_length=256) + 1 overlap = n_frames # No overlap last_frame = audio_rep.shape[0] - n_frames + 1 batches = [] if last_frame <= 0: # Pad with zeros if audio is too short patch = np.zeros((n_frames, 96), dtype=np.float32) patch[:audio_rep.shape[0], :] = audio_rep batches.append(patch) else: # Create windows for time_stamp in range(0, last_frame, overlap): patch = audio_rep[time_stamp : time_stamp + n_frames, :] batches.append(patch) # Convert to tensor and run inference batch_tensor = torch.from_numpy(np.stack(batches)).to(device) all_embeddings = [] with torch.no_grad(): self.eval() for i in range(0, len(batches), 1): batch_subset = batch_tensor[i:i+1] logits, mean_pool, max_pool = self(batch_subset) if pool == 'mean': embeddings = mean_pool.squeeze(0).cpu().numpy() elif pool == 'max': embeddings = max_pool.squeeze(0).cpu().numpy() elif pool == 'both': embeddings = torch.cat([mean_pool, max_pool], dim=1).squeeze(0).cpu().numpy() else: embeddings = mean_pool.squeeze(0).cpu().numpy() # default to mean all_embeddings.append(embeddings) # Average embeddings across all windows avg_embeddings = np.mean(all_embeddings, axis=0) return avg_embeddings # For uploading to Hugging Face Hub if __name__ == '__main__': import json import os from huggingface_hub import HfApi import shutil # Create the model with MTT config config = MusicNNConfig( num_classes=50, mid_filt=64, backend_units=200, dataset='MTT' ) model = MusicNN(config) # Load the weights state_dict = torch.load('weights/MTT_musicnn.pt') model.load_state_dict(state_dict) # Save and push to Hugging Face save_dir = 'musicnn-pytorch' os.makedirs(save_dir, exist_ok=True) model.save_pretrained(save_dir) shutil.copy('musicnn.py', save_dir) # Create config.json config_dict = config.to_dict() config_dict.update({ '_name_or_path': 'oriyonay/musicnn-pytorch', 'architectures': ['MusicNN'], 'auto_map': { 'AutoConfig': 'musicnn.MusicNNConfig', 'AutoModel': 'musicnn.MusicNN' }, 'model_type': 'musicnn' }) with open(os.path.join(save_dir, 'config.json'), 'w') as f: json.dump(config_dict, f, indent=4) # Push to Hugging Face api = HfApi() api.upload_folder( folder_path=save_dir, repo_id='oriyonay/musicnn-pytorch', repo_type='model' ) print("✅ Model uploaded to Hugging Face!") print("Usage: model = MusicNN.from_pretrained('oriyonay/musicnn-pytorch')")