import torch import torch.nn as nn import librosa import numpy as np import gradio as gr from transformers import XLMRobertaModel, XLMRobertaTokenizer, WavLMModel from huggingface_hub import hf_hub_download import warnings warnings.filterwarnings('ignore') # ── CONFIGURATION ── REPO_ID = "anggars/neural-mathrock" DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") TARGET_COLS = ['mbti', 'emotion', 'vibe', 'intensity', 'tempo'] 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=DEVICE, weights_only=False) # ── 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'] 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 HybridMultimodalModel(nn.Module): def __init__(self): super().__init__() 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.Linear(1024, 512), nn.BatchNorm1d(512), nn.ReLU(), nn.Dropout(0.4), ) 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 = HybridMultimodalModel().to(DEVICE) model.load_state_dict(ckpt['model_state'], strict=False) model.eval() tokenizer = XLMRobertaTokenizer.from_pretrained('anggars/xlm-mbti') # ── HYBRID INFERENCE PIPELINE ── def predict_multimodal(audio_path, lyrics): if audio_path is None: return [{"Error": "Audio file cannot be empty."}] * 5 try: # 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) 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}") # 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} # 4. Block-by-block Prediction with torch.no_grad(): 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]) # 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 if col == 'tempo': 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))} 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": f"Processing failed: {str(e)}"}] * 5 # ── 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.") 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 HYBRID ANALYSIS", variant="primary") with gr.Column(): 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="Intensity Level") out_tempo = gr.Label(label="Tempo Classification (BPM Adjusted)") 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()