Update app.py

app.py

@@ -4,24 +4,34 @@ import gradio as gr
 import torch
 import torchaudio
 import google.generativeai as genai
-from transformers import AutoProcessor, AutoModel
+from e2_tts_pytorch import E2TTS, DurationPredictor
 import numpy as np
 import os
-import json
 
 # Initialize Gemini AI
 genai.configure(api_key='YOUR_GEMINI_API_KEY')
 model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
 
-# Initialize F5-TTS model
-processor = AutoProcessor.from_pretrained("SWivid/F5-TTS")
-f5_model = AutoModel.from_pretrained("SWivid/F5-TTS")
+# Initialize E2-TTS model
+duration_predictor = DurationPredictor(
+    transformer=dict(
+        dim=512,
+        depth=8,
+    )
+)
 
-def clone_voice(audio_file):
-    waveform, sample_rate = torchaudio.load(audio_file)
-    # Implement voice cloning logic here
-    # For demonstration, we'll just return a placeholder
-    return "Voice cloned successfully"
+e2tts = E2TTS(
+    duration_predictor=duration_predictor,
+    transformer=dict(
+        dim=512,
+        depth=8
+    ),
+)
+
+# Load the pre-trained model
+model_path = "ckpts/E2TTS_Base/model_1200000.safetensors"
+e2tts.load_state_dict(torch.load(model_path))
+e2tts.eval()
 
 def generate_podcast_script(content, duration):
     prompt = f"""
@@ -42,9 +52,15 @@ def generate_podcast_script(content, duration):
     return response.text
 
 def text_to_speech(text, speaker_id):
-    inputs = processor(text=text, return_tensors="pt")
-    speech = f5_model.generate_speech(inputs["input_ids"], speaker_id=speaker_id, vocoder_output=True)
-    return speech.cpu().numpy()
+    # For simplicity, we'll use a random mel spectrogram as input
+    # In a real scenario, you'd use the actual mel spectrogram from the cloned voice
+    mel = torch.randn(1, 80, 100)
+    
+    # Generate speech
+    with torch.no_grad():
+        sampled = e2tts.sample(mel[:, :5], text=[text])
+    
+    return sampled.cpu().numpy()
 
 def create_podcast(content, duration, voice1, voice2):
     script = generate_podcast_script(content, duration)
@@ -53,10 +69,10 @@ def create_podcast(content, duration, voice1, voice2):
     
     for line in lines:
         if line.startswith("Host 1:"):
-            audio = text_to_speech(line[7:], speaker_id=0)  # Assuming speaker_id 0 for Host 1
+            audio = text_to_speech(line[7:], speaker_id=0)
             audio_segments.append(audio)
         elif line.startswith("Host 2:"):
-            audio = text_to_speech(line[7:], speaker_id=1)  # Assuming speaker_id 1 for Host 2
+            audio = text_to_speech(line[7:], speaker_id=1)
             audio_segments.append(audio)
     
     # Concatenate audio segments