Views
No views yet
(1, 1500, 768).Mean, Max, and Min pooling.librosa, transformers, torch, and huggingface_hub.1import os
2import torch
3import numpy as np
4import librosa
5from huggingface_hub import hf_hub_download
6from transformers import WhisperModel, WhisperProcessor
7from model_architecture import MusicAestheticsModel # Downloaded from this repo
8
9# Configuration
10# 1. The Audio Encoder (The Music Whisper Model)
11WHISPER_REPO = "laion/music-whisper"
12# 2. This Aesthetics Model
13AESTHETICS_REPO = "laion/music-aesthetics"
14
15DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
16
17def load_models():
18 print("Loading Whisper Encoder...")
19 processor = WhisperProcessor.from_pretrained(WHISPER_REPO)
20 # We only need the encoder part of Whisper
21 whisper = WhisperModel.from_pretrained(WHISPER_REPO).encoder.to(DEVICE)
22 whisper.eval()
23
24 print("Loading Aesthetics Experts...")
25 # Initialize the architecture
26 model = MusicAestheticsModel().to(DEVICE)
27
28 # Download and load weights
29 # 1. Load Shared Bottleneck
30 bt_path = hf_hub_download(repo_id=AESTHETICS_REPO, filename="stage1_bottleneck.pt")
31 model.bottleneck.load_state_dict(torch.load(bt_path, map_location=DEVICE))
32
33 # 2. Load Expert Heads
34 for metric in model.metrics:
35 head_path = hf_hub_download(repo_id=AESTHETICS_REPO, filename=f"expert_{metric}.pt")
36 model.heads[metric].load_state_dict(torch.load(head_path, map_location=DEVICE))
37
38 model.eval()
39 return processor, whisper, model
40
41def predict_score(audio_path, processor, whisper, aesthetic_model):
42 # 1. Load and Preprocess Audio
43 # Resample to 16kHz and pad/crop to exactly 30s
44 audio, sr = librosa.load(audio_path, sr=16000)
45 target_len = 16000 * 30
46 if len(audio) > target_len:
47 start = (len(audio) - target_len) // 2
48 audio = audio[start : start + target_len]
49 else:
50 audio = np.pad(audio, (0, target_len - len(audio)))
51
52 # 2. Extract Whisper Features
53 inputs = processor(audio, sampling_rate=16000, return_tensors="pt")
54 with torch.no_grad():
55 # Get last hidden state from encoder
56 outputs = whisper(inputs.input_features.to(DEVICE))
57 last_hidden = outputs.last_hidden_state # (1, 1500, 768)
58
59 # 3. Apply Feature Pooling (Expert Model Logic)
60 # Reshape to (1, 10 segments, 150 frames, 768 dim)
61 feats = last_hidden.view(1, 10, 150, 768)
62
63 mean_pool = torch.mean(feats, dim=2)
64 max_pool = torch.max(feats, dim=2).values
65 min_pool = torch.min(feats, dim=2).values
66
67 # Concat -> Flatten -> (23040,)
68 concat = torch.cat([mean_pool, max_pool, min_pool], dim=2)
69 embedding = concat.view(-1).unsqueeze(0) # Add batch dim
70
71 # 4. Predict Scores
72 with torch.no_grad():
73 outputs = aesthetic_model(embedding)
74
75 results = {k: v.item() for k, v in outputs.items()}
76
77 # Calculate Average Global Score
78 avg_score = sum(results.values()) / len(results)
79 results["Overall_Aesthetics"] = avg_score
80
81 return results
82
83# Example Usage
84if __name__ == "__main__":
85 processor, whisper, model = load_models()
86
87 # Replace with your audio file
88 audio_file = "test_song.mp3"
89
90 if os.path.exists(audio_file):
91 scores = predict_score(audio_file, processor, whisper, model)
92
93 print("-" * 30)
94 print(f"Aesthetics Analysis for {audio_file}")
95 print("-" * 30)
96 for metric, score in scores.items():
97 print(f"{metric:<20}: {score:.2f} / 5.0")
98 print("-" * 30)
99 else:
100 print("Please provide a valid audio file path.")