Update app.py

app.py

@@ -33,6 +33,8 @@ TEMPO_LABELS = ['slow', 'moderate', 'fast']
 WAVLM_SR = 16000
 CLAP_SR = 48000
 
+emo2id = {e: i for i, e in enumerate(EMO_LABELS)}
+
 # -- HYBRID ARCHITECTURE ALIGNED WITH TRAINING EXECUTION --
 class AudioMathRockModel(nn.Module):
     def __init__(self):
@@ -75,6 +77,7 @@ class AudioMathRockModel(nn.Module):
         self.tmp_head  = nn.Linear(512, len(TEMPO_LABELS))
 
     def forward(self, wavlm_values, clap_values):
+        # Extract features from frozen/unfrozen backbones
         wavlm_feats = self.wavlm(wavlm_values).last_hidden_state.mean(dim=1)
         clap_feats  = self.clap(clap_values).pooler_output
         
@@ -83,10 +86,12 @@ class AudioMathRockModel(nn.Module):
             mel = self.mel_spectrogram(wavlm_values.float())
             mel_db = self.amplitude_to_db(mel)
         
+        # Align precision and pass through CNN branch
         cnn_feats = F.gelu(self.cnn_extractor(mel_db.unsqueeze(1))).to(wavlm_feats.dtype)
         wlm_p = F.gelu(self.wlm_proj(wavlm_feats))
         clp_p = F.gelu(self.clp_proj(clap_feats))
         
+        # Late fusion sequence
         fused = self.fusion(torch.cat([wlm_p, clp_p, cnn_feats], dim=-1))
         return self.emo_head(fused), self.vibe_head(fused), self.int_head(fused), self.tmp_head(fused)
 
@@ -115,6 +120,7 @@ def analyze_track(audio_path, lyrics_input):
 
     res_mbti, res_emo, res_vibe, res_int, res_tmp = {}, {}, {}, {}, {}
 
+    # --- TEXT ONLY PROCESSING ---
     if has_lyrics and not has_audio:
         t_in = tokenizer(str(lyrics_input), truncation=True, padding=True, max_length=256, return_tensors="pt").to(DEVICE)
         with torch.no_grad():
@@ -128,16 +134,19 @@ def analyze_track(audio_path, lyrics_input):
         res_emo = dict(sorted(emo_dict.items(), key=lambda x: x[1], reverse=True)[:3])
         return res_mbti, res_emo, {}, {}, {}
 
+    # --- AUDIO PROCESSING ---
     try:
         wav_orig, orig_sr = librosa.load(audio_path, sr=None, mono=True)
         waveform = torch.tensor(wav_orig).unsqueeze(0)
         
+        # Resampling constraints based on model requirements
         wlm_wave = torchaudio.functional.resample(waveform, orig_sr, WAVLM_SR).squeeze(0)
         clp_wave = torchaudio.functional.resample(waveform, orig_sr, CLAP_SR).squeeze(0)
 
         wavlm_samples = WAVLM_SR * 15
         clap_samples = CLAP_SR * 15
 
+        # Chunk generation for robust global representation
         wlm_chunks = [wlm_wave[i:i+wavlm_samples].numpy() for i in range(0, len(wlm_wave), wavlm_samples) if len(wlm_wave[i:i+wavlm_samples]) > WAVLM_SR]
         clp_chunks = [clp_wave[i:i+clap_samples].numpy() for i in range(0, len(clp_wave), clap_samples) if len(clp_wave[i:i+clap_samples]) > CLAP_SR]
         
@@ -158,26 +167,31 @@ def analyze_track(audio_path, lyrics_input):
                 il_list.append(il.cpu())
                 tl_list.append(tl.cpu())
 
+        # Extract pure raw logits
         raw_el = torch.cat(el_list, dim=0).mean(dim=0, keepdim=True)
         raw_vl = torch.cat(vl_list, dim=0).mean(dim=0, keepdim=True)
         raw_il = torch.cat(il_list, dim=0).mean(dim=0, keepdim=True)
         raw_tl = torch.cat(tl_list, dim=0).mean(dim=0, keepdim=True)
 
-        # FIXED LOGIC UNIFIED TO SIGMOID INFERENCE AS SPECIFIED IN MULTI-LABEL TRAINING
+        # Apply correct activation logic (Pure Model Inference)
+        # Emotion remains Sigmoid (Multi-label distribution)
         probs_el = torch.sigmoid(raw_el).squeeze().numpy()
-        probs_vl = torch.sigmoid(raw_vl).squeeze().numpy()
-        probs_il = torch.sigmoid(raw_il).squeeze().numpy()
-        probs_tl = torch.sigmoid(raw_tl).squeeze().numpy()
+        # Vibe, Intensity, Tempo strictly evaluated as Softmax (Single-label exclusivity)
+        probs_vl = F.softmax(raw_vl, dim=-1).squeeze().numpy()
+        probs_il = F.softmax(raw_il, dim=-1).squeeze().numpy()
+        probs_tl = F.softmax(raw_tl, dim=-1).squeeze().numpy()
 
         res_emo_raw = {EMO_LABELS[i]: float(probs_el[i]) for i in range(len(EMO_LABELS))}
         res_vibe = {VIBE_LABELS[i]: float(probs_vl[i]) for i in range(len(VIBE_LABELS))}
         res_int = {INTENSITY_LABELS[i]: float(probs_il[i]) for i in range(len(INTENSITY_LABELS))}
         res_tmp = {TEMPO_LABELS[i]: float(probs_tl[i]) for i in range(len(TEMPO_LABELS))}
 
+        # Sort and truncate dicts for UI presentation
         res_vibe = dict(sorted(res_vibe.items(), key=lambda x: x[1], reverse=True)[:3])
         res_int = dict(sorted(res_int.items(), key=lambda x: x[1], reverse=True)[:3])
         res_tmp = dict(sorted(res_tmp.items(), key=lambda x: x[1], reverse=True)[:3])
 
+        # --- LATE FUSION PIPELINE WITH TEXT MODEL ---
         if has_lyrics:
             t_in = tokenizer(str(lyrics_input), truncation=True, padding=True, max_length=256, return_tensors="pt").to(DEVICE)
             with torch.no_grad():
@@ -187,8 +201,9 @@ def analyze_track(audio_path, lyrics_input):
             m_dict = {MBTI_LABELS[i]: float(t_mbti_probs[i]) for i in range(len(MBTI_LABELS))}
             res_mbti = dict(sorted(m_dict.items(), key=lambda x: x[1], reverse=True)[:3])
             
+            # Weighted ensemble calculation
             for k in res_emo_raw.keys():
-                idx = EMO_LABELS.index(k)
+                idx = emo2id[k]
                 res_emo_raw[k] = (res_emo_raw[k] * 0.3) + (float(t_emo_probs[idx]) * 0.7)
             res_emo = dict(sorted(res_emo_raw.items(), key=lambda x: x[1], reverse=True)[:3])
         else: