app.py

import os
import gc
import torch
import torch.nn as nn
import torch.nn.functional as F
import soundfile as sf
import torchaudio
import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download
from transformers import (
    AutoModel, 
    AutoFeatureExtractor, 
    AutoTokenizer,
    XLMRobertaForSequenceClassification, 
    XLMRobertaTokenizer
)
import warnings

warnings.filterwarnings('ignore')

# ------------------------------------------------------------------------------
# CONFIGURATION & LABEL MAPPING
# ------------------------------------------------------------------------------
REPO_MAIN      = "anggars/neural-mathrock"
REPO_TEXT_MBTI = "anggars/xlm-mbti"
REPO_TEXT_EMO  = "anggars/xlm-emotion"
MODEL_FILE     = "model.pth"
DEVICE         = torch.device("cuda" if torch.cuda.is_available() else "cpu")
CROP_SEC       = 5
SR_TARGET      = 24000

MBTI_LABELS = sorted(["INTJ", "INTP", "ENTJ", "ENTP", "INFJ", "INFP", "ENFJ", "ENFP", "ISTJ", "ISFJ", "ESTJ", "ESFJ", "ISTP", "ISFP", "ESTP", "ESFP"])
EMOTION_CLASSES_AUDIO = ["anger", "disgust", "fear", "joy", "neutral", "sadness", "surprise"]
EMOTION_CLASSES_TEXT  = ['admiration', 'amusement', 'anger', 'annoyance', 'approval', 'caring', 'confusion', 'curiosity', 'desire', 'disappointment', 'disapproval', 'disgust', 'embarrassment', 'excitement', 'fear', 'gratitude', 'grief', 'joy', 'love', 'nervousness', 'optimism', 'pride', 'realization', 'relief', 'remorse', 'sadness', 'surprise', 'neutral']

text_emo2id = {e: i for i, e in enumerate(EMOTION_CLASSES_TEXT)}

# Calibration Weights derived from SOTA Report 3
RAW_CE_WEIGHTS = [0.563, 16.326, 4.525, 0.968, 2.704, 0.473, 0.702]
SOFT_WEIGHTS   = torch.sqrt(torch.tensor(RAW_CE_WEIGHTS, dtype=torch.float32)).to(DEVICE)
TEMPERATURE    = 2.0

# ------------------------------------------------------------------------------
# NEURAL ARCHITECTURE (7-CLASS AUDIO FUSION)
# ------------------------------------------------------------------------------
class MultimodalFusionClassifier(nn.Module):
    def __init__(self, num_classes=7):
        super().__init__()
        self.audio_model = AutoModel.from_pretrained("m-a-p/MERT-v1-330M", trust_remote_code=True)
        self.text_model = AutoModel.from_pretrained("FacebookAI/roberta-base")
        self.fusion_head = nn.Sequential(
            nn.Linear(1024 + 768, 512),
            nn.LayerNorm(512),
            nn.GELU(),
            nn.Dropout(0.5),
            nn.Linear(512, 256),
            nn.LayerNorm(256),
            nn.GELU(),
            nn.Dropout(0.4),
            nn.Linear(256, num_classes)
        )

    def forward(self, audio_vals, text_ids, text_mask):
        audio_feats = self.audio_model(audio_vals, output_hidden_states=True).last_hidden_state.mean(dim=1)
        text_feats = self.text_model(input_ids=text_ids, attention_mask=text_mask).last_hidden_state.mean(dim=1)
        return self.fusion_head(torch.cat((audio_feats, text_feats), dim=-1))

# ------------------------------------------------------------------------------
# MODEL INITIALIZATION
# ------------------------------------------------------------------------------
print("[SYSTEM] Booting Models into Memory...")
mert_ext = AutoFeatureExtractor.from_pretrained("m-a-p/MERT-v1-330M", trust_remote_code=True)
roberta_tok = AutoTokenizer.from_pretrained("FacebookAI/roberta-base")

fusion_model = MultimodalFusionClassifier(num_classes=len(EMOTION_CLASSES_AUDIO)).to(DEVICE)
ckpt_path = hf_hub_download(repo_id=REPO_MAIN, filename=MODEL_FILE)
fusion_model.load_state_dict(torch.load(ckpt_path, map_location=DEVICE)['model_state_dict'])
fusion_model.eval()

# Load Text-Only Models (Isolated)
mbti_tok = XLMRobertaTokenizer.from_pretrained(REPO_TEXT_MBTI)
mbti_model = XLMRobertaForSequenceClassification.from_pretrained(REPO_TEXT_MBTI).to(DEVICE)
mbti_model.eval()

# Load High-Resolution 28-Class Text Emotion Model
emo28_model = XLMRobertaForSequenceClassification.from_pretrained(REPO_TEXT_EMO).to(DEVICE)
emo28_model.eval()

# ------------------------------------------------------------------------------
# INFERENCE ENGINE
# ------------------------------------------------------------------------------
def process_audio_chunks(audio_path):
    waveform_np, orig_sr = sf.read(audio_path)
    if len(waveform_np.shape) > 1: 
        waveform_np = waveform_np.mean(axis=1) 
    
    waveform = torch.tensor(waveform_np, dtype=torch.float32).unsqueeze(0)
    if orig_sr != SR_TARGET:
        waveform = torchaudio.functional.resample(waveform, orig_sr, SR_TARGET)
    
    waveform = waveform.squeeze(0)
    chunk_samples = SR_TARGET * CROP_SEC
    
    chunks = [waveform[i:i+chunk_samples].numpy() for i in range(0, len(waveform), chunk_samples) if len(waveform[i:i+chunk_samples]) > SR_TARGET * 2]
    if not chunks:
        chunks = [F.pad(waveform, (0, chunk_samples - waveform.shape[0])).numpy()]
        
    return chunks[:10]

def analyze_track(audio_path, lyrics_input):
    has_audio = audio_path is not None
    safe_lyrics = str(lyrics_input).strip() if lyrics_input else ""
    has_lyrics = len(safe_lyrics) > 10
    
    if not has_audio and not has_lyrics:
        return {"Error": 1.0}, {"Error": 1.0}

    res_mbti = {}
    res_emo_final = {}

    # 1. ISOLATED MBTI & 28-CLASS TEXT INFERENCE
    t_emo28_probs = None
    if has_lyrics:
        t_in = mbti_tok(safe_lyrics, truncation=True, padding=True, max_length=256, return_tensors="pt").to(DEVICE)
        with torch.no_grad():
            mbti_probs = F.softmax(mbti_model(**t_in).logits, dim=1).cpu().squeeze().numpy()
            t_emo28_probs = F.softmax(emo28_model(**t_in).logits, dim=1).cpu().squeeze().numpy()
            
        mbti_dict = {MBTI_LABELS[i]: float(mbti_probs[i]) for i in range(len(MBTI_LABELS))}
        res_mbti = dict(sorted(mbti_dict.items(), key=lambda x: x[1], reverse=True)[:3])
    else:
        res_mbti = {"Requires lyrics for MBTI": 1.0}

    # 2. AUDIO FUSION INFERENCE
    if has_audio:
        try:
            audio_chunks = process_audio_chunks(audio_path)
            
            t_inputs = roberta_tok([safe_lyrics] * len(audio_chunks), padding=True, truncation=True, max_length=128, return_tensors="pt")
            t_ids = t_inputs["input_ids"].to(DEVICE)
            t_mask = t_inputs["attention_mask"].to(DEVICE)
            
            a_inputs = mert_ext(audio_chunks, sampling_rate=SR_TARGET, return_tensors="pt", padding="max_length", truncation=True, max_length=SR_TARGET * CROP_SEC)
            a_vals = a_inputs["input_values"].to(DEVICE)
            
            with torch.no_grad():
                logits = fusion_model(a_vals, t_ids, t_mask)
                avg_logits = logits.mean(dim=0)
                
                scaled_logits = avg_logits / TEMPERATURE
                probs = F.softmax(scaled_logits, dim=-1)
                adjusted_probs = probs * SOFT_WEIGHTS
                audio_7_probs = (adjusted_probs / adjusted_probs.sum()).cpu().numpy()
                
            audio_emo_dict = {EMOTION_CLASSES_AUDIO[i]: float(audio_7_probs[i]) for i in range(len(EMOTION_CLASSES_AUDIO))}
            
            # --- TRUE MULTIMODAL LATE FUSION ENGAGEMENT ---
            if has_lyrics:
                # Bikin dictionary kosong untuk 28 emosi
                fusion_28_dict = {label: 0.0 for label in EMOTION_CLASSES_TEXT}
                
                # Masukin probabilitas murni dari Teks (Bobot 60%)
                for i, label in enumerate(EMOTION_CLASSES_TEXT):
                    fusion_28_dict[label] += float(t_emo28_probs[i]) * 0.60
                    
                # Suntikin probabilitas dari Audio 7 Kelas ke 28 Kelas Teks (Bobot 40%)
                for label in EMOTION_CLASSES_AUDIO:
                    if label in fusion_28_dict:
                        fusion_28_dict[label] += audio_emo_dict[label] * 0.40
                        
                # Normalisasi ulang biar totalnya 1.0
                total_prob = sum(fusion_28_dict.values())
                fusion_28_dict = {k: v / total_prob for k, v in fusion_28_dict.items()}
                
                res_emo_final = dict(sorted(fusion_28_dict.items(), key=lambda x: x[1], reverse=True)[:5])
            else:
                # Unimodal Audio (Hanya ngeluarin 7 Kelas)
                res_emo_final = dict(sorted(audio_emo_dict.items(), key=lambda x: x[1], reverse=True)[:4])
            
        except Exception as e:
            res_emo_final = {f"Audio Error: {str(e)}": 1.0}
    
    # 3. TEXT-ONLY FALLBACK (Kalau audio gak dimasukin)
    elif has_lyrics:
        emo_dict = {EMOTION_CLASSES_TEXT[i]: float(t_emo28_probs[i]) for i in range(len(EMOTION_CLASSES_TEXT))}
        res_emo_final = dict(sorted(emo_dict.items(), key=lambda x: x[1], reverse=True)[:5])

    return res_mbti, res_emo_final

# ------------------------------------------------------------------------------
# GRADIO INTERFACE
# ------------------------------------------------------------------------------
with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    gr.Markdown("# Neural Math Rock Multimodal Analysis")
    gr.Markdown("Identify personality (Text) and emotional states (MERT+RoBERTa Multimodal) from Math Rock & Midwest Emo tracks.")
    
    with gr.Row():
        with gr.Column():
            audio_box = gr.Audio(type="filepath", label="Audio Source (.wav / .mp3)")
            lyrics_box = gr.Textbox(lines=8, label="Lyrics Source", placeholder="Paste lyrics here... \n(If blank: Analyzes 7 Audio Emotions. If filled: Analyzes 28 High-Resolution Emotions.)")
            run_btn = gr.Button("RUN SOTA ANALYSIS", variant="primary")
        
        with gr.Column():
            res_mbti = gr.Label(label="Personality (MBTI - Text Only)")
            res_emo = gr.Label(label="Emotional State (Multimodal 28-Class Fusion)")

    run_btn.click(
        fn=analyze_track, 
        inputs=[audio_box, lyrics_box], 
        outputs=[res_mbti, res_emo]
    )

if __name__ == "__main__":
    demo.launch()