Update utils.py

utils.py

@@ -3,16 +3,26 @@ from pydantic import BaseModel, ValidationError
 from typing import List, Literal
 import os
 import tiktoken
-from gtts import gTTS
 import tempfile
 import json
 import re
+from transformers import pipeline
+import torch
+import soundfile as sf
 
 groq_client = Groq(api_key=os.environ["GROQ_API_KEY"])
 tokenizer = tiktoken.get_encoding("cl100k_base")
 
+# Initialize TTS pipelines
+tts_male = pipeline("text-to-speech", model="microsoft/speecht5_tts", device="cpu")
+tts_female = pipeline("text-to-speech", model="microsoft/speecht5_tts", device="cpu")
+
+# Load speaker embeddings
+male_embedding = torch.load("https://huggingface.co/microsoft/speecht5_tts/resolve/main/en_speaker_1.pt")
+female_embedding = torch.load("https://huggingface.co/microsoft/speecht5_tts/resolve/main/en_speaker_9.pt")
+
 class DialogueItem(BaseModel):
-    speaker: Literal["Host", "Guest"]
+    speaker: Literal["John", "Sarah"]
     text: str
 
 class Dialogue(BaseModel):
@@ -37,16 +47,13 @@ def generate_script(system_prompt: str, input_text: str, tone: str):
         temperature=0.7
     )
     
-    # Extract content and remove any markdown code block syntax
     content = response.choices[0].message.content
     content = re.sub(r'```json\s*|\s*```', '', content)
     
     try:
-        # First, try to parse as JSON
         json_data = json.loads(content)
         dialogue = Dialogue.model_validate(json_data)
     except json.JSONDecodeError as json_error:
-        # If JSON parsing fails, try to extract JSON from the text
         match = re.search(r'\{.*\}', content, re.DOTALL)
         if match:
             try:
@@ -62,7 +69,11 @@ def generate_script(system_prompt: str, input_text: str, tone: str):
     return dialogue
 
 def generate_audio(text: str, speaker: str) -> str:
-    tts = gTTS(text, lang='en', tld='com' if speaker == "Host" else 'co.uk')
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_audio:
-        tts.save(temp_audio.name)
+    if speaker == "John":
+        speech = tts_male(text, speaker_embeddings=male_embedding)
+    else:  # Sarah
+        speech = tts_female(text, speaker_embeddings=female_embedding)
+    
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
+        sf.write(temp_audio.name, speech["audio"], speech["sampling_rate"])
         return temp_audio.name