Update app.py

app.py

@@ -3,7 +3,7 @@ import torch.nn as nn
 import librosa
 import numpy as np
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModel, WavLMModel, AutoFeatureExtractor
 from huggingface_hub import hf_hub_download
 import warnings
 
@@ -12,14 +12,7 @@ warnings.filterwarnings('ignore')
 # ── CONFIGURATION ──
 REPO_ID = "anggars/neural-mathrock"
 SR = 16000
-DURATION = 10
-N_MELS = 32
-N_MFCC = 13
-HOP_LENGTH = 512
-MAX_LENGTH = 128
-AUDIO_EMB = 256
-TEXT_EMB = 768
-FUSED_EMB = 1024
+DURATION = 10  # Potong 10 detik aja biar RAM HF Space gratisan kaga jebol
 TARGET_COLS = ['mbti', 'emotion', 'vibe', 'intensity', 'tempo']
 
 print("Downloading model and extracting label encoders...")
@@ -29,28 +22,21 @@ ckpt = torch.load(model_path, map_location=torch.device('cpu'), weights_only=Fal
 label_encoders = ckpt['label_encoders']
 num_classes = {col: len(le.classes_) for col, le in label_encoders.items()}
 
-# ── ARCHITECTURE ──
-class AudioCNNTransformer(nn.Module):
+# ============================================================
+# ── ARCHITECTURE (WAJIB SESUAIKAN SAMA IPYNB TRAINING LU) ──
+# ============================================================
+# HAPUS DAN GANTI blok arsitektur ini pake class asli lu dari jupyter notebook.
+
+class AudioTransformer(nn.Module):
     def __init__(self):
         super().__init__()
-        self.cnn = nn.Sequential(
-            nn.Conv2d(3, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(), nn.MaxPool2d(2),
-            nn.Conv2d(16, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2),
-            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2),
-            nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
-            nn.AdaptiveAvgPool2d((8, 8)) 
-        )
-        encoder_layer = nn.TransformerEncoderLayer(d_model=128, nhead=4, dim_feedforward=512, batch_first=True)
-        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2)
-        self.fc_out = nn.Linear(128, AUDIO_EMB)
+        self.wavlm = WavLMModel.from_pretrained("microsoft/wavlm-base")
+        for p in self.wavlm.parameters():
+            p.requires_grad = False
 
     def forward(self, x):
-        features = self.cnn(x)
-        B, C, H, W = features.shape
-        features = features.view(B, C, H * W).transpose(1, 2) 
-        out = self.transformer(features)
-        out = out.mean(dim=1) 
-        return self.fc_out(out)
+        out = self.wavlm(x).last_hidden_state
+        return out.mean(dim=1) 
 
 class TextTransformer(nn.Module):
     def __init__(self):
@@ -66,10 +52,10 @@ class TextTransformer(nn.Module):
 class MultimodalMathRock(nn.Module):
     def __init__(self):
         super().__init__()
-        self.audio  = AudioCNNTransformer()
+        self.audio  = AudioTransformer()
         self.text   = TextTransformer()
         self.fusion = nn.Sequential(
-            nn.Linear(FUSED_EMB, 512),
+            nn.Linear(768 + 768, 512), 
             nn.ReLU(),
             nn.Dropout(0.3)
         )
@@ -77,17 +63,21 @@ class MultimodalMathRock(nn.Module):
             col: nn.Linear(512, num_classes[col]) for col in TARGET_COLS
         })
 
-    def forward(self, audio, input_ids, attention_mask):
-        x_a = self.audio(audio)
+    def forward(self, audio_values, input_ids, attention_mask):
+        x_a = self.audio(audio_values)
         x_t = self.text(input_ids, attention_mask)
         fused = self.fusion(torch.cat([x_a, x_t], dim=1))
         return {col: self.heads[col](fused) for col in TARGET_COLS}
 
+# ============================================================
+
 print("Initializing architecture and loading weights...")
 model = MultimodalMathRock()
 model.load_state_dict(ckpt['model_state'])
 model.eval()
+
 tokenizer = AutoTokenizer.from_pretrained('xlm-roberta-base')
+audio_processor = AutoFeatureExtractor.from_pretrained("microsoft/wavlm-base")
 
 # ── INFERENCE PIPELINE ──
 def predict_multimodal(audio_path, lyrics):
@@ -95,39 +85,23 @@ def predict_multimodal(audio_path, lyrics):
         return [None]*5
         
     try:
+        # 1. Load audio mentah (otomatis jadi 16kHz mono)
         y, _ = librosa.load(audio_path, sr=SR, duration=DURATION, mono=True)
-        target_len = SR * DURATION
         
-        if len(y) < target_len:
-            y = np.pad(y, (0, target_len - len(y)))
-        else:
-            y = y[:target_len]
-
-        mel = librosa.power_to_db(librosa.feature.melspectrogram(y=y, sr=SR, n_mels=N_MELS, hop_length=HOP_LENGTH), ref=np.max)
-        mfcc = librosa.feature.mfcc(y=y, sr=SR, n_mfcc=N_MFCC, hop_length=HOP_LENGTH)
-        chroma = librosa.feature.chroma_stft(y=y, sr=SR, hop_length=HOP_LENGTH)
-
-        T = mel.shape[1]
-        feat = np.stack([
-            mel,
-            np.resize(mfcc, (N_MELS, T)),
-            np.resize(chroma, (N_MELS, T))
-        ], axis=0).astype(np.float32)
-
-        for i in range(3):
-            std = feat[i].std()
-            feat[i] = (feat[i] - feat[i].mean()) / (std + 1e-8)
-            
-        audio_tensor = torch.from_numpy(feat).unsqueeze(0)
+        # 2. Ekstraksi fitur pake WavLM processor
+        inputs = audio_processor(y, sampling_rate=SR, return_tensors="pt")
+        audio_tensor = inputs.input_values
 
+        # 3. Proses Lirik
         if not lyrics or lyrics.strip() == "":
             lyrics = "instrumental"
             
         enc = tokenizer(
-            lyrics, max_length=MAX_LENGTH, padding='max_length',
+            lyrics, max_length=128, padding='max_length',
             truncation=True, return_tensors='pt'
         )
         
+        # 4. Forward pass
         with torch.no_grad():
             out = model(audio_tensor, enc['input_ids'], enc['attention_mask'])
             
@@ -148,7 +122,7 @@ def predict_multimodal(audio_path, lyrics):
 # ── UI ──
 with gr.Blocks() as demo:
     gr.Markdown("# Neural Math Rock Classifier")
-    gr.Markdown("Audio feature-based MBTI and emotion classification system. Upload an MP3 or WAV file!")
+    gr.Markdown("Multimodal WavLM & XLM-RoBERTa based MBTI and emotion classification system.")
     
     with gr.Row():
         with gr.Column():