Update app.py

app.py

@@ -10,58 +10,7 @@ import os
 import requests
 from tqdm import tqdm
 
-# Function to download the model file
-def download_model(url, filename):
-    response = requests.get(url, stream=True)
-    total_size = int(response.headers.get('content-length', 0))
-    block_size = 1024  # 1 KB
-    progress_bar = tqdm(total=total_size, unit='iB', unit_scale=True)
-    
-    os.makedirs(os.path.dirname(filename), exist_ok=True)
-    with open(filename, 'wb') as file:
-        for data in response.iter_content(block_size):
-            size = file.write(data)
-            progress_bar.update(size)
-    progress_bar.close()
-
-# Check if model file exists, if not, download it
-model_path = "ckpts/E2TTS_Base/model_1200000.pt"
-if not os.path.exists(model_path):
-    print("Downloading model file...")
-    model_url = "https://huggingface.co/SWivid/E2-TTS/resolve/main/E2TTS_Base/model_1200000.pt"
-    download_model(model_url, model_path)
-    print("Model file downloaded successfully.")
-
-# Initialize E2-TTS model
-duration_predictor = DurationPredictor(
-    transformer=dict(
-        dim=512,
-        depth=8,
-    )
-)
-
-e2tts = E2TTS(
-    duration_predictor=duration_predictor,
-    transformer=dict(
-        dim=512,
-        depth=8
-    ),
-)
-
-# Load the pre-trained model
-checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
-if 'model_state_dict' in checkpoint:
-    state_dict = checkpoint['model_state_dict']
-elif 'ema_model_state_dict' in checkpoint:
-    state_dict = checkpoint['ema_model_state_dict']
-else:
-    state_dict = checkpoint  # Assume the checkpoint is the state dict itself
-
-# Filter out unexpected keys
-model_dict = e2tts.state_dict()
-filtered_state_dict = {k: v for k, v in state_dict.items() if k in model_dict}
-e2tts.load_state_dict(filtered_state_dict, strict=False)
-e2tts.eval()
+# (Keep the model loading and initialization code as before)
 
 def generate_podcast_script(api_key, content, duration):
     genai.configure(api_key=api_key)
@@ -93,24 +42,11 @@ def text_to_speech(text, speaker_id):
     with torch.no_grad():
         sampled = e2tts.sample(mel[:, :5], text=[text])
     
-    return sampled.cpu().numpy()
+    return sampled.cpu().numpy().squeeze()
 
 def create_podcast(api_key, content, duration, voice1, voice2):
     script = generate_podcast_script(api_key, content, duration)
-    lines = script.split('\n')
-    audio_segments = []
-    
-    for line in lines:
-        if line.startswith("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)
-            audio_segments.append(audio)
-    
-    # Concatenate audio segments
-    podcast_audio = np.concatenate(audio_segments)
-    return (22050, podcast_audio)  # Assuming 22050 Hz sample rate
+    return render_podcast(api_key, script, voice1, voice2)
 
 def gradio_interface(api_key, content, duration, voice1, voice2):
     script = generate_podcast_script(api_key, content, duration)
@@ -128,8 +64,12 @@ def render_podcast(api_key, script, voice1, voice2):
             audio = text_to_speech(line[7:], speaker_id=1)
             audio_segments.append(audio)
     
+    if not audio_segments:
+        return (22050, np.zeros(22050))  # Return silence if no audio was generated
+    
+    # Concatenate audio segments
     podcast_audio = np.concatenate(audio_segments)
-    return (22050, podcast_audio)
+    return (22050, podcast_audio)  # Assuming 22050 Hz sample rate
 
 # Gradio Interface
 with gr.Blocks() as demo: