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, WavLMModel, AutoFeatureExtractor
+from transformers import XLMRobertaModel, XLMRobertaTokenizer, WavLMModel
 from huggingface_hub import hf_hub_download
 import warnings
 
@@ -11,131 +11,149 @@ warnings.filterwarnings('ignore')
 
 # ── CONFIGURATION ──
 REPO_ID = "anggars/neural-mathrock"
-SR = 16000
-DURATION = 10
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 TARGET_COLS = ['mbti', 'emotion', 'vibe', 'intensity', 'tempo']
 
-print("Downloading model from HF Hub...")
+print("Downloading model from Hugging Face Hub...")
 model_path = hf_hub_download(repo_id=REPO_ID, filename="model.pt")
-ckpt = torch.load(model_path, map_location=torch.device('cpu'), weights_only=False)
+ckpt = torch.load(model_path, map_location=DEVICE, weights_only=False)
 
-class DummyEncoder:
-    def __init__(self, classes_list):
-        self.classes_ = np.array(classes_list)
+# ── LOAD LABEL ENCODERS ──
+le_mbti = ckpt['le_mbti']
+le_emotion = ckpt['le_emotion']
+le_vibe = ckpt['le_vibe']
+le_intensity = ckpt['le_intensity']
+le_tempo = ckpt['le_tempo']
 
-# Label Encoders
-label_encoders = {
-    'mbti': ckpt.get('le_mbti', DummyEncoder(['ENFJ', 'ENFP', 'ENTJ', 'ENTP', 'ESFJ', 'ESFP', 'ESTJ', 'ESTP', 'INFJ', 'INFP', 'INTJ', 'INTP', 'ISFJ', 'ISFP', 'ISTJ', 'ISTP'])),
-    'emotion': ckpt.get('le_emotion', DummyEncoder(['amusement', 'anger', 'annoyance', 'approval', 'caring', 'confusion', 'curiosity', 'desire', 'disappointment', 'disapproval', 'disgust', 'embarrassment', 'excitement', 'fear', 'gratitude', 'grief', 'love', 'nervousness', 'neutral', 'pride', 'realization', 'relief', 'remorse', 'sadness'])),
-    'vibe': DummyEncoder(['Aggressive', 'Atmospheric', 'Melancholic', 'Technical']),
-    'intensity': DummyEncoder(['High', 'Low', 'Medium']),
-    'tempo': DummyEncoder(['Fast', 'Moderate', 'Slow'])
-}
-
-num_classes = {col: len(le.classes_) for col, le in label_encoders.items()}
+NUM_MBTI = len(le_mbti.classes_)
+NUM_EMOTION = len(le_emotion.classes_)
+NUM_VIBE = len(le_vibe.classes_)
+NUM_INTENSITY = len(le_intensity.classes_)
+NUM_TEMPO = len(le_tempo.classes_)
 
 # ── ARCHITECTURE ──
-class MultimodalMathRock(nn.Module):
+class HybridMultimodalModel(nn.Module):
     def __init__(self):
         super().__init__()
-        self.audio_encoder = WavLMModel.from_pretrained("microsoft/wavlm-base")
-        self.text_encoder  = AutoModel.from_pretrained('xlm-roberta-base')
+        self.text_model = XLMRobertaModel.from_pretrained('anggars/xlm-mbti')
+        self.audio_model = WavLMModel.from_pretrained('microsoft/wavlm-base')
+        
+        self.audio_proj = nn.Linear(768, 256)
+        
         self.fusion = nn.Sequential(
-            nn.Dropout(0.3),
-            nn.Linear(768 + 768, 512),
-            nn.ReLU()
+            nn.Linear(1024, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.4),
         )
-        self.heads = nn.ModuleDict({
-            'mbti': nn.Linear(512, num_classes['mbti']),
-            'emotion': nn.Linear(512, num_classes['emotion']),
-            'vibe': nn.Linear(512, num_classes['vibe']),
-            'intensity': nn.Linear(512, num_classes['intensity']),
-            'tempo': nn.Linear(512, num_classes['tempo'])
-        })
-
-    def forward(self, audio_values, input_ids, attention_mask):
-        x_a = self.audio_encoder(audio_values).last_hidden_state.mean(dim=1)
-        x_t = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0, :]
-        fused = self.fusion(torch.cat([x_a, x_t], dim=1))
-        return {col: self.heads[col](fused) for col in TARGET_COLS}
+        
+        self.head_mbti = nn.Linear(512, NUM_MBTI)
+        self.head_emotion = nn.Linear(512, NUM_EMOTION)
+        self.head_vibe = nn.Linear(512, NUM_VIBE)
+        self.head_intensity = nn.Linear(512, NUM_INTENSITY)
+        self.head_tempo = nn.Linear(512, NUM_TEMPO)
+
+    def forward(self, input_ids, attention_mask, audio_values):
+        text_out = self.text_model(input_ids=input_ids, attention_mask=attention_mask)
+        text_feat = text_out.pooler_output
+        
+        audio_out = self.audio_model(audio_values).last_hidden_state
+        audio_feat = self.audio_proj(audio_out.mean(dim=1))
+        
+        fused = self.fusion(torch.cat([text_feat, audio_feat], dim=-1))
+        
+        return {
+            'mbti': self.head_mbti(fused),
+            'emotion': self.head_emotion(fused),
+            'vibe': self.head_vibe(fused),
+            'intensity': self.head_intensity(fused),
+            'tempo': self.head_tempo(fused),
+        }
 
 print("Initializing architecture and loading weights...")
-model = MultimodalMathRock()
-state_dict = ckpt['model_state'] if 'model_state' in ckpt else ckpt
-model.load_state_dict(state_dict, strict=False)
+model = HybridMultimodalModel().to(DEVICE)
+model.load_state_dict(ckpt['model_state'], strict=False)
 model.eval()
-tokenizer = AutoTokenizer.from_pretrained('xlm-roberta-base')
-audio_processor = AutoFeatureExtractor.from_pretrained("microsoft/wavlm-base")
+
+tokenizer = XLMRobertaTokenizer.from_pretrained('anggars/xlm-mbti')
 
 # ── HYBRID INFERENCE PIPELINE ──
 def predict_multimodal(audio_path, lyrics):
     if audio_path is None:
-        return [None]*5
+        return [{"Error": "Audio file cannot be empty."}] * 5
         
     try:
-        # 1. Load Audio & Hitung Real Tempo (BPM) pake Librosa
-        y_full, _ = librosa.load(audio_path, sr=SR, mono=True)
-        total_dur = librosa.get_duration(y=y_full, sr=SR)
+        # 1. Text Processing
+        text = str(lyrics).strip() if lyrics else "[INSTRUMENTAL]"
+        enc = tokenizer(text, truncation=True, padding='max_length', max_length=128, return_tensors='pt').to(DEVICE)
+
+        # 2. Audio Processing & DSP
+        wav, sr = librosa.load(audio_path, sr=16000)
         
-        # Beat Tracking Logic
-        onset_env = librosa.onset.onset_strength(y=y_full, sr=SR)
-        tempo_bpm, _ = librosa.beat.beat_track(onset_envelope=onset_env, sr=SR)
+        onset_env = librosa.onset.onset_strength(y=wav, sr=sr)
+        tempo_bpm, _ = librosa.beat.beat_track(onset_envelope=onset_env, sr=sr)
         print(f"Calculated BPM: {tempo_bpm}")
 
-        # 2. Windowing Strategy (8 Windows)
-        num_windows = min(8, int(total_dur // DURATION))
-        if num_windows == 0: num_windows = 1
-        offsets = np.linspace(0, total_dur - DURATION, num_windows)
-        
-        if not lyrics or lyrics.strip() == "":
-            lyrics = "instrumental math rock"
-            
-        enc = tokenizer(lyrics, max_length=128, padding='max_length', truncation=True, return_tensors='pt')
+        # 3. Chunking (20-second sequential blocks)
+        chunk_size = 16000 * 20
+        chunks = []
+        for i in range(0, len(wav), chunk_size):
+            chunk = wav[i:i + chunk_size]
+            if len(chunk) < chunk_size:
+                chunk = np.pad(chunk, (0, chunk_size - len(chunk)))
+            chunks.append(chunk)
+
+        if not chunks:
+            chunks = [np.zeros(chunk_size)]
+
         all_logits = {col: [] for col in TARGET_COLS}
         
-        # 3. Model Prediction
+        # 4. Block-by-block Prediction
         with torch.no_grad():
-            for offset in offsets:
-                start = int(offset * SR)
-                end = start + int(DURATION * SR)
-                y_chunk = y_full[start:end]
-                if len(y_chunk) < int(DURATION * SR):
-                    y_chunk = np.pad(y_chunk, (0, int(DURATION * SR) - len(y_chunk)))
-                
-                audio_inputs = audio_processor(y_chunk, sampling_rate=SR, return_tensors="pt")
-                out = model(audio_inputs.input_values, enc['input_ids'], enc['attention_mask'])
+            for chunk in chunks:
+                audio_tensor = torch.tensor(chunk, dtype=torch.float32).unsqueeze(0).to(DEVICE)
+                out = model(enc['input_ids'], enc['attention_mask'], audio_tensor)
                 for col in TARGET_COLS:
                     all_logits[col].append(out[col][0])
         
-        # 4. Result Processing with Tempo Override
+        # 5. Averaging and Formatting Results
         final_results = []
+        encoders_map = {
+            'mbti': le_mbti, 'emotion': le_emotion, 'vibe': le_vibe, 
+            'intensity': le_intensity, 'tempo': le_tempo
+        }
+        
         for col in TARGET_COLS:
             avg_logits = torch.stack(all_logits[col]).mean(dim=0)
+            le = encoders_map[col]
             
-            # ── TEMPO LOGIC OVERRIDE ──
+            # TEMPO LOGIC OVERRIDE
             if col == 'tempo':
-                tempo_classes = list(label_encoders['tempo'].classes_)
-                # Indexing: Fast (0), Moderate (1), Slow (2) berdasarkan dummy encoder lo
-                if tempo_bpm > 125: # Jika BPM di atas 125, paksa ke arah Fast
-                    avg_logits[0] += 5.0 
-                elif tempo_bpm > 90: # Jika BPM di atas 90, paksa ke arah Moderate
-                    avg_logits[1] += 3.0
-                # Jika di bawah itu, biarin model mutusin (biasanya Slow)
-
-            probs = torch.nn.functional.softmax(avg_logits, dim=0).numpy()
-            classes = label_encoders[col].classes_
+                tempo_classes = list(le.classes_)
+                try:
+                    fast_idx = tempo_classes.index('Fast')
+                    mod_idx = tempo_classes.index('Moderate')
+                    if tempo_bpm > 125:
+                        avg_logits[fast_idx] += 5.0 
+                    elif tempo_bpm > 90:
+                        avg_logits[mod_idx] += 3.0
+                except ValueError:
+                    pass
+            
+            probs = torch.nn.functional.softmax(avg_logits, dim=0).cpu().numpy()
+            classes = le.classes_
             
-            pred_dict = {str(classes[i]): float(probs[i]) for i in range(len(classes)) if i < len(probs)}
+            pred_dict = {str(classes[i]): float(probs[i]) for i in range(len(classes))}
             sorted_preds = dict(sorted(pred_dict.items(), key=lambda item: item[1], reverse=True)[:3])
             final_results.append(sorted_preds)
             
         return final_results
         
     except Exception as e:
-        return [{"Error": str(e)}]*5
+        return [{"Error": f"Processing failed: {str(e)}"}] * 5
 
-# ── GRADIO UI (Glass Theme) ──
-with gr.Blocks(theme=gr.themes.Glass()) as demo:
+# ── GRADIO UI ──
+with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
     gr.Markdown("# Neural Math Rock Multimodal Analysis")
     gr.Markdown("Analyzing MBTI, Emotion, and Vibes using Hybrid Deep Learning & DSP.")
     
@@ -149,7 +167,7 @@ with gr.Blocks(theme=gr.themes.Glass()) as demo:
             out_mbti = gr.Label(label="Personality (MBTI)")
             out_emotion = gr.Label(label="Emotional State")
             out_vibe = gr.Label(label="Acoustic Vibe")
-            out_intensity = gr.Label(label="Energy Level")
+            out_intensity = gr.Label(label="Intensity Level")
             out_tempo = gr.Label(label="Tempo Classification (BPM Adjusted)")
 
     btn.click(