Update app.py

app.py

@@ -3,22 +3,21 @@ import google.generativeai as genai
 import numpy as np
 import re
 import torch
+import torchaudio
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import snapshot_download, login
 import logging
 import os
 import spaces
 import warnings
-import librosa
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 # Suppress specific warnings
-warnings.filterwarnings("ignore", category=UserWarning, message="Trying to convert audio automatically from float32 to 16-bit int format.")
-warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value encountered in divide")
-warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value encountered in cast")
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore", category=RuntimeWarning)
 
 def get_device():
     if torch.cuda.is_available():
@@ -115,26 +114,32 @@ def text_to_speech(text, voice):
         if model is None or tokenizer is None:
             load_model()
         
-        inputs = tokenizer(text, return_tensors="pt").to(device)
+        # Remove emotion tags for TTS processing
+        clean_text = re.sub(r'<[^>]+>', '', text)
+        
+        inputs = tokenizer(clean_text, return_tensors="pt").to(device)
         with torch.no_grad():
             output = model.generate(**inputs, max_new_tokens=256)
-        mel = output[0].cpu().numpy()
-        audio = mel_to_audio(mel)
-        return audio
+        
+        # Convert output tensor to mel spectrogram
+        mel = output[0].cpu()
+        
+        # Normalize the mel spectrogram
+        mel = (mel - mel.min()) / (mel.max() - mel.min())
+        
+        # Convert mel spectrogram to audio using torchaudio
+        griffin_lim = torchaudio.transforms.GriffinLim(n_fft=2048, n_iter=10)
+        audio = griffin_lim(mel.unsqueeze(0))
+        
+        # Convert to numpy array and ensure it's in the correct format
+        audio_np = audio.squeeze().numpy()
+        audio_np = np.clip(audio_np, -1, 1)
+        
+        return (24000, audio_np.astype(np.float32))  # Assuming 24kHz sample rate
     except Exception as e:
         logger.error(f"Error in text_to_speech: {str(e)}")
         raise
 
-def mel_to_audio(mel):
-    try:
-        # Convert mel spectrogram to audio using librosa
-        audio = librosa.feature.inverse.mel_to_audio(mel, sr=24000, n_iter=10)
-        return audio
-    except Exception as e:
-        logger.error(f"Error in mel_to_audio conversion: {str(e)}")
-        # Return silence if conversion fails
-        return np.zeros(24000, dtype=np.float32)
-
 @spaces.GPU()
 def render_podcast(api_key, script, voice1, voice2, num_hosts):
     try:
@@ -144,7 +149,7 @@ def render_podcast(api_key, script, voice1, voice2, num_hosts):
         for i, line in enumerate(lines):
             voice = voice1 if num_hosts == 1 or i % 2 == 0 else voice2
             try:
-                audio = text_to_speech(line, voice)
+                _, audio = text_to_speech(line, voice)
                 audio_segments.append(audio)
             except Exception as e:
                 logger.error(f"Error processing audio segment: {str(e)}")