music-popularity-full-ft

import torch
import torch.nn as nn
import numpy as np
import librosa
import math
from transformers import WhisperProcessor, WhisperForConditionalGeneration
from huggingface_hub import hf_hub_download


# --- Define the MLP head ---

class PopularityMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.bottleneck = nn.Sequential(
            nn.Linear(23040, 1024), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(1024, 256), nn.ReLU(), nn.LayerNorm(256),
        )
        self.play_head = nn.Sequential(nn.Linear(256, 64), nn.ReLU(), nn.Linear(64, 1))
        self.upvote_head = nn.Sequential(nn.Linear(256, 64), nn.ReLU(), nn.Linear(64, 1))

    def forward(self, x):
        feat = self.bottleneck(x)
        return self.play_head(feat).squeeze(-1), self.upvote_head(feat).squeeze(-1)


# --- Load models ---

# Fine-tuned Whisper encoder (includes encoder weight updates from popularity training)
processor = WhisperProcessor.from_pretrained("laion/music-popularity-full-ft")
whisper = WhisperForConditionalGeneration.from_pretrained(
    "laion/music-popularity-full-ft", torch_dtype=torch.float16
).cuda().eval()
encoder = whisper.get_encoder()

# Popularity MLP head
head_path = hf_hub_download("laion/music-popularity-full-ft", "popularity_head.pt")
mlp = PopularityMLP().cuda()
mlp.load_state_dict(torch.load(head_path, map_location="cuda")["mlp_state_dict"])
mlp.eval()


# --- Run inference ---

audio, sr = librosa.load("song.mp3", sr=16000, mono=True)
audio = audio[:30 * 16000]  # first 30 seconds
if len(audio) < 30 * 16000:
    audio = np.pad(audio, (0, 30 * 16000 - len(audio)))

inputs = processor(audio, sampling_rate=16000, return_tensors="pt")

with torch.no_grad():
    enc_out = encoder(inputs.input_features.cuda().half()).last_hidden_state  # (1, 1500, 768)

    # Segment pooling: 10 segments, mean/max/min
    segments = enc_out.view(1, 10, 150, 768)
    pooled = torch.cat([segments.mean(2), segments.max(2).values, segments.min(2).values], dim=2)
    pooled = pooled.view(1, -1).float()  # (1, 23040)

    pred_play, pred_upvote = mlp(pooled)

print(f"Estimated plays:   {math.expm1(pred_play.item()):,.0f}")
print(f"Estimated upvotes: {math.expm1(pred_upvote.item()):,.0f}")