app.py

import torch
import torch.nn as nn
import librosa
import numpy as np
import gradio as gr
from transformers import AutoTokenizer, AutoModel, WavLMModel, AutoFeatureExtractor
from huggingface_hub import hf_hub_download
import warnings

warnings.filterwarnings('ignore')

REPO_ID = "anggars/neural-mathrock"
SR = 16000
DURATION = 10
TARGET_COLS = ['mbti', 'emotion', 'vibe', 'intensity', 'tempo']

print("Downloading model from HF 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)

class DummyEncoder:
    def __init__(self, classes_list):
        self.classes_ = np.array(classes_list)

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()}

class MultimodalMathRock(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.fusion = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(768 + 768, 512),
            nn.ReLU()
        )
        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, 70 if 'tempo' not in num_classes else 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}

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.eval()

tokenizer = AutoTokenizer.from_pretrained('xlm-roberta-base')
audio_processor = AutoFeatureExtractor.from_pretrained("microsoft/wavlm-base")

def predict_multimodal(audio_path, lyrics):
    if audio_path is None:
        return [None]*5
        
    try:
        y_full, _ = librosa.load(audio_path, sr=SR, mono=True)
        total_dur = librosa.get_duration(y=y_full, sr=SR)
        
        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')
        all_logits = {col: [] for col in TARGET_COLS}
        
        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 col in TARGET_COLS:
                    all_logits[col].append(out[col][0])
        
        final_results = []
        for col in TARGET_COLS:
            avg_logits = torch.stack(all_logits[col]).mean(dim=0)
            probs = torch.nn.functional.softmax(avg_logits, dim=0).numpy()
            classes = label_encoders[col].classes_
            
            pred_dict = {str(classes[i]): float(probs[i]) for i in range(len(classes)) if i < len(probs)}
            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

# ── GRADIO UI (Glass Theme) ──
with gr.Blocks(theme=gr.themes.Glass()) as demo:
    gr.Markdown("# Neural Math Rock & Midwest Emo Analysis")
    gr.Markdown("Multimodal Affective Computing System for Math Rock & Midwest Emo Genres.")
    
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(type="filepath", label="Upload Song (Full Analysis)")
            lyrics_input = gr.Textbox(lines=5, label="Lyrics Content", placeholder="Paste lyrics here...")
            btn = gr.Button("START MULTIMODAL ANALYSIS", variant="primary")
        
        with gr.Column():
            out_mbti = gr.Label(label="Personality Type (MBTI)")
            out_emotion = gr.Label(label="Emotional State")
            out_vibe = gr.Label(label="Acoustic Vibe")
            out_intensity = gr.Label(label="Energy Level")
            out_tempo = gr.Label(label="Tempo Classification")

    btn.click(
        fn=predict_multimodal,
        inputs=[audio_input, lyrics_input],
        outputs=[out_mbti, out_emotion, out_vibe, out_intensity, out_tempo]
    )

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