Upload modeling_apex.py with huggingface_hub

modeling_apex.py

@@ -0,0 +1,219 @@
+import os
+import json
+import torch
+import torch.nn as nn
+import numpy as np
+import soundfile as sf
+import torchaudio.functional as TAF
+from transformers import PreTrainedModel, AutoProcessor, AutoModel
+from .configuration_apex import APEXConfig
+
+
+# Building blocks
+class SharedBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Linear(in_dim, out_dim),
+            nn.BatchNorm1d(out_dim),
+            nn.GELU(),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class BranchBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout, use_bn=True):
+        super().__init__()
+        layers = [nn.Linear(in_dim, out_dim)]
+        if use_bn:
+            layers.append(nn.BatchNorm1d(out_dim))
+        layers += [nn.GELU(), nn.Dropout(dropout)]
+        self.block = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class TaskBranch(nn.Module):
+    def __init__(self, in_dim, branch_dims, dropout, scale, shift):
+        super().__init__()
+        layers = []
+        prev   = in_dim
+        for dim in branch_dims:
+            layers.append(BranchBlock(prev, dim, dropout=dropout, use_bn=True))
+            prev = dim
+        layers.append(nn.Linear(prev, 1))
+        self.branch = nn.Sequential(*layers)
+        self.scale  = scale
+        self.shift  = shift
+
+    def forward(self, x):
+        return torch.sigmoid(self.branch(x)) * self.scale + self.shift
+
+
+# APEX MODEL
+class APEXModel(PreTrainedModel):
+    config_class = APEXConfig
+
+    def __init__(self, config: APEXConfig):
+        super().__init__(config)
+
+        # MERT encoder + processor
+        self.mert_processor = AutoProcessor.from_pretrained(
+            config.mert_model_name, trust_remote_code=True
+        )
+        self.mert = AutoModel.from_pretrained(
+            config.mert_model_name, trust_remote_code=True
+        )
+        self.mert.eval()
+        for param in self.mert.parameters():
+            param.requires_grad = False
+
+        self.target_sr = self.mert_processor.sampling_rate
+
+        # Conv1d aggregator with fixed seed
+        torch.manual_seed(config.seed)
+        self.aggregator = nn.Conv1d(
+            in_channels  = len(config.layer_indices),
+            out_channels = 1,
+            kernel_size  = 1
+        )
+
+        # Shared layers: 768 → 512 → 256
+        shared_layers = []
+        prev_dim      = config.input_dim
+        for dim in config.shared_dims:
+            shared_layers.append(SharedBlock(prev_dim, dim, dropout=config.dropout_shared))
+            prev_dim = dim
+        self.shared = nn.Sequential(*shared_layers)
+
+        out_dim = config.shared_dims[-1]  # 256
+
+        # Task branches: 256 → 128 → 64 → 1
+        self.branch_score_streams = TaskBranch(out_dim, config.branch_dims, config.dropout_branch, scale=100, shift=0)
+        self.branch_score_likes   = TaskBranch(out_dim, config.branch_dims, config.dropout_branch, scale=100, shift=0)
+        self.branch_coherence     = TaskBranch(out_dim, config.branch_dims, config.dropout_branch, scale=4,   shift=1)
+        self.branch_musicality    = TaskBranch(out_dim, config.branch_dims, config.dropout_branch, scale=4,   shift=1)
+        self.branch_memorability  = TaskBranch(out_dim, config.branch_dims, config.dropout_branch, scale=4,   shift=1)
+        self.branch_clarity       = TaskBranch(out_dim, config.branch_dims, config.dropout_branch, scale=4,   shift=1)
+        self.branch_naturalness   = TaskBranch(out_dim, config.branch_dims, config.dropout_branch, scale=4,   shift=1)
+
+    def forward(self, embedding):
+        shared = self.shared(embedding)
+        return {
+            "score_streams": self.branch_score_streams(shared).squeeze(1),
+            "score_likes"  : self.branch_score_likes(shared).squeeze(1),
+            "coherence"    : self.branch_coherence(shared).squeeze(1),
+            "musicality"   : self.branch_musicality(shared).squeeze(1),
+            "memorability" : self.branch_memorability(shared).squeeze(1),
+            "clarity"      : self.branch_clarity(shared).squeeze(1),
+            "naturalness"  : self.branch_naturalness(shared).squeeze(1),
+        }
+
+    def _load_audio(self, audio_path):
+        waveform, sr = sf.read(audio_path, dtype="float32")
+        waveform     = torch.from_numpy(waveform)
+
+        # Stereo to mono
+        if len(waveform.shape) > 1 and waveform.shape[1] > 1:
+            waveform = waveform.mean(dim=1)
+
+        waveform = waveform.to(self.device)
+
+        # Resample if needed
+        if sr != self.target_sr:
+            waveform = TAF.resample(waveform, sr, self.target_sr)
+
+        return waveform
+
+    def _extract_embedding(self, waveform):
+        segment_len        = self.config.segment_sec * self.target_sr
+        segment_embeddings = []
+
+        for start in range(0, waveform.shape[0], segment_len):
+            segment = waveform[start:start + segment_len]
+            if segment.numel() == 0:
+                break
+
+            # Zero-pad if needed
+            if segment.shape[0] < segment_len:
+                pad_len = segment_len - segment.shape[0]
+                segment = torch.nn.functional.pad(segment, (0, pad_len))
+
+            # MERT forward
+            inputs = self.mert_processor(
+                segment.cpu().numpy(),
+                sampling_rate  = self.target_sr,
+                return_tensors = "pt"
+            )
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+            with torch.no_grad():
+                outputs = self.mert(**inputs, output_hidden_states=True)
+
+            # Extract layers and aggregate
+            all_hidden = torch.stack([
+                outputs.hidden_states[i].mean(dim=1)
+                for i in self.config.layer_indices
+            ])  # [4, 1, 768]
+            all_hidden = all_hidden.squeeze(1)  # [4, 768]
+
+            # Conv1d aggregation
+            pooled = self.aggregator(
+                all_hidden.unsqueeze(0)
+            ).squeeze()  # [768]
+
+            segment_embeddings.append(pooled)
+
+            del segment, inputs, outputs, all_hidden, pooled
+
+        # Average across segments to song-level embedding
+        song_embedding = torch.stack(segment_embeddings).mean(dim=0)
+        return song_embedding
+
+    @torch.no_grad()
+    def predict(self, audio_path, save_json=None):
+        self.eval()
+        print(f"\nProcessing: {audio_path}")
+
+        waveform = self._load_audio(audio_path)
+        duration = waveform.shape[0] / self.target_sr
+        n_segs   = int(np.ceil(duration / self.config.segment_sec))
+        print(f"Duration: {duration:.1f}s | Segments: {n_segs}")
+
+        print("Extracting MERT embeddings...")
+        embedding = self._extract_embedding(waveform)
+
+        print("Running APEX model...")
+        preds = self.forward(embedding.unsqueeze(0))
+
+        results = {
+            task: float(preds[task].squeeze().cpu())
+            for task in preds
+        }
+
+        print(f"\n{'='*50}")
+        print(f"  APEX Predictions")
+        print(f"{'='*50}")
+        print(f"\n  Popularity:")
+        print(f"  {'─'*40}")
+        print(f"  {'Streams Score':<20} {results['score_streams']:>8.2f} / 100")
+        print(f"  {'Likes Score':<20} {results['score_likes']:>8.2f} / 100")
+        print(f"\n  Aesthetic Quality:")
+        print(f"  {'─'*40}")
+        for dim in ["coherence", "musicality", "memorability", "clarity", "naturalness"]:
+            print(f"  {dim.capitalize():<20} {results[dim]:>8.2f} / 5.00")
+        print(f"{'='*50}\n")
+
+        if save_json:
+            with open(save_json, "w") as f:
+                json.dump({
+                    "audio_path" : audio_path,
+                    "predictions": results
+                }, f, indent=2)
+            print(f"Results saved to {save_json}")
+
+        return results