Update app.py

app.py

@@ -5,6 +5,11 @@ import numpy as np
 import gradio as gr
 from transformers import XLMRobertaModel, XLMRobertaTokenizer, WavLMModel
 from huggingface_hub import hf_hub_download
+import yt_dlp
+import os
+import re
+import lyricsgenius
+import syncedlyrics
 import warnings
 
 warnings.filterwarnings('ignore')
@@ -14,6 +19,9 @@ REPO_ID = "anggars/neural-mathrock"
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 TARGET_COLS = ['mbti', 'emotion', 'vibe', 'intensity', 'tempo']
 
+# Set your Genius token here or use HF Secrets/Environment Variables
+GENIUS_TOKEN = os.environ.get("GENIUS_TOKEN", "z2XGBWXalGUtAdC1qxxXBxUnK1ZuoHPkCu5eP9q-fed-DW1uCJ3NSFpHemk3Unmg")
+
 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)
@@ -56,10 +64,13 @@ class HybridMultimodalModel(nn.Module):
         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):
+    def forward(self, input_ids, attention_mask, audio_values, text_missing=False):
         text_out = self.text_model(input_ids=input_ids, attention_mask=attention_mask)
         text_feat = text_out.pooler_output
         
+        if text_missing:
+            text_feat = torch.zeros_like(text_feat)
+            
         audio_out = self.audio_model(audio_values).last_hidden_state
         audio_feat = self.audio_proj(audio_out.mean(dim=1))
         
@@ -83,24 +94,86 @@ model.eval()
 
 tokenizer = XLMRobertaTokenizer.from_pretrained('anggars/xlm-mbti')
 
+# ── UTILITIES ──
+def get_audio_from_youtube(query):
+    temp_filename = "temp_downloaded_audio"
+    
+    ydl_opts = {
+        'format': 'bestaudio/best',
+        'outtmpl': f'{temp_filename}.%(ext)s',
+        'postprocessors': [{
+            'key': 'FFmpegExtractAudio',
+            'preferredcodec': 'wav',
+            'preferredquality': '192',
+        }],
+        'noplaylist': True,
+        'quiet': True,
+        'default_search': 'ytsearch1',
+    }
+    
+    try:
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            ydl.download([query])
+            return f"{temp_filename}.wav"
+    except Exception as e:
+        print(f"Error yt-dlp: {str(e)}")
+        return None
+
+def get_lyrics_fallback(query):
+    try:
+        genius = lyricsgenius.Genius(GENIUS_TOKEN, verbose=False)
+        song = genius.search_song(query)
+        if song and song.lyrics:
+            clean_text = re.sub(r'\[.*?\]', '', song.lyrics)
+            return clean_text.strip()
+    except Exception as e:
+        print(f"Genius API Error: {str(e)}")
+        
+    try:
+        lrc_lyrics = syncedlyrics.search(query)
+        if lrc_lyrics:
+            clean_text = re.sub(r'\[\d{2}:\d{2}\.\d{2}\]', '', lrc_lyrics)
+            return clean_text.strip()
+    except Exception as e:
+        print(f"SyncedLyrics Error: {str(e)}")
+        
+    return None
+
 # ── HYBRID INFERENCE PIPELINE ──
-def predict_multimodal(audio_path, lyrics):
-    if audio_path is None:
-        return [{"Error": "Audio file cannot be empty."}] * 5
+def predict_multimodal(audio_path, yt_query, lyrics_input):
+    target_audio = audio_path
+    
+    if target_audio is None and yt_query:
+        print(f"Searching audio on YouTube for: {yt_query}")
+        target_audio = get_audio_from_youtube(yt_query)
+        
+    if target_audio is None:
+        return [{"Error": "Audio file or YouTube query is required."}] * 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)
+        final_lyrics = lyrics_input
+        if not final_lyrics or str(final_lyrics).strip() == "":
+            if yt_query:
+                print(f"Fetching lyrics for: {yt_query}")
+                fetched_lyrics = get_lyrics_fallback(yt_query)
+                if fetched_lyrics:
+                    final_lyrics = fetched_lyrics
+
+        is_instrumental = False
+        if not final_lyrics or str(final_lyrics).strip() == "":
+            is_instrumental = True
+            text_to_encode = "[INSTRUMENTAL]"
+        else:
+            text_to_encode = str(final_lyrics).strip()
+            
+        enc = tokenizer(text_to_encode, 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)
+        wav, sr = librosa.load(target_audio, 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):
@@ -114,15 +187,13 @@ def predict_multimodal(audio_path, lyrics):
 
         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)
+                out = model(enc['input_ids'], enc['attention_mask'], audio_tensor, text_missing=is_instrumental)
                 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, 
@@ -133,7 +204,6 @@ def predict_multimodal(audio_path, lyrics):
             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:
@@ -153,9 +223,16 @@ def predict_multimodal(audio_path, lyrics):
             sorted_preds = dict(sorted(pred_dict.items(), key=lambda item: item[1], reverse=True)[:3])
             final_results.append(sorted_preds)
             
+        if yt_query and target_audio == "temp_downloaded_audio.wav":
+            if os.path.exists(target_audio):
+                os.remove(target_audio)
+                
         return final_results
         
     except Exception as e:
+        if yt_query and target_audio == "temp_downloaded_audio.wav":
+            if os.path.exists(target_audio):
+                os.remove(target_audio)
         return [{"Error": f"Processing failed: {str(e)}"}] * 5
 
 # ── GRADIO UI ──
@@ -165,8 +242,9 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
     
     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...")
+            audio_input = gr.Audio(type="filepath", label="1. Upload Audio File (Optional, overrides YouTube search)")
+            yt_input = gr.Textbox(lines=1, label="2. Search YouTube (Format: Artist - Song Title)", placeholder="e.g. eleventwelfth - front-and-centre")
+            lyrics_input = gr.Textbox(lines=5, label="3. Manual Lyrics Input (Optional, auto-fetched if empty)", placeholder="Paste lyrics here or leave empty for instrumental / auto-fetch...")
             btn = gr.Button("START HYBRID ANALYSIS", variant="primary")
         
         with gr.Column():
@@ -178,7 +256,7 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
 
     btn.click(
         fn=predict_multimodal,
-        inputs=[audio_input, lyrics_input],
+        inputs=[audio_input, yt_input, lyrics_input],
         outputs=[out_mbti, out_emotion, out_vibe, out_intensity, out_tempo]
     )