Update app.py

app.py

@@ -11,7 +11,8 @@ import shutil
 from pathlib import Path
 import time
 from tqdm import tqdm
-from gtts import gTTS
+import torch
+from TTS.api import TTS
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, 
@@ -30,20 +31,35 @@ LANGUAGES = {
     "Hindi": "hi"
 }
 
-# TTS voice mapping for different languages
-TTS_VOICES = {
-    "en": "en-US",
-    "es": "es-ES",
-    "fr": "fr-FR",
-    "de": "de-DE",
-    "ja": "ja-JP",
-    "hi": "hi-IN"
+# TTS model mapping for different languages
+TTS_MODELS = {
+    "en": "tts_models/en/ljspeech/tacotron2-DDC_ph",
+    "es": "tts_models/es/css10/vits",
+    "fr": "tts_models/fr/css10/vits",
+    "de": "tts_models/de/thorsten/tacotron2-DDC",
+    "ja": "tts_models/ja/kokoro/tacotron2-DDC",
+    "hi": "tts_models/hi/kb/tacotron2-DDC"
 }
 
 # Create a permanent output directory
 OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "outputs")
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
+# Initialize TTS
+def init_tts():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    tts_models = {}
+    for lang_code, model_name in TTS_MODELS.items():
+        try:
+            tts = TTS(model_name=model_name, progress_bar=False).to(device)
+            tts_models[lang_code] = tts
+            logger.info(f"Loaded TTS model for {lang_code}: {model_name}")
+        except Exception as e:
+            logger.warning(f"Failed to load TTS model for {lang_code}: {str(e)}")
+    return tts_models
+
+tts_models = init_tts()
+
 def extract_audio(video_path):
     """Extract audio from video file using ffmpeg"""
     try:
@@ -124,7 +140,7 @@ def translate_subtitles(srt_path, target_langs):
         raise Exception(f"Translation failed: {str(e)}")
 
 def generate_translated_audio(srt_path, target_lang):
-    """Generate translated audio using text-to-speech"""
+    """Generate translated audio using Coqui TTS"""
     try:
         logger.info(f"Generating translated audio for {target_lang}")
         subs = pysrt.open(srt_path, encoding="utf-8")
@@ -138,6 +154,11 @@ def generate_translated_audio(srt_path, target_lang):
         audio_files = []
         timings = []
         
+        # Get the appropriate TTS model
+        tts = tts_models.get(target_lang)
+        if tts is None:
+            raise Exception(f"No TTS model available for language: {target_lang}")
+        
         for i, sub in enumerate(tqdm(subs, desc=f"Generating {target_lang} speech")):
             text = sub.text.strip()
             if not text:
@@ -157,31 +178,11 @@ def generate_translated_audio(srt_path, target_lang):
             duration = end_time - start_time
             
             # Generate TTS audio
-            tts_lang = TTS_VOICES.get(target_lang, target_lang)
-            audio_file = os.path.join(temp_dir, f"chunk_{i:04d}.mp3")
+            audio_file = os.path.join(temp_dir, f"chunk_{i:04d}.wav")
             
             try:
-                # Add a retry mechanism for Hindi and other potentially problematic languages
-                retry_count = 0
-                max_retries = 3
-                while retry_count < max_retries:
-                    try:
-                        # For Hindi, use slower speed which might improve reliability
-                        slow_option = target_lang == "hi" 
-                        tts = gTTS(text=text, lang=target_lang, slow=slow_option)
-                        tts.save(audio_file)
-                        break
-                    except Exception as e:
-                        retry_count += 1
-                        logger.warning(f"TTS attempt {retry_count} failed for {target_lang}: {str(e)}")
-                        time.sleep(1)  # Wait before retrying
-                        
-                        # If still failing after retries, try with shorter text
-                        if retry_count == max_retries and len(text) > 100:
-                            logger.warning(f"Trying with shortened text for {target_lang}")
-                            shortened_text = text[:100] + "..."
-                            tts = gTTS(text=shortened_text, lang=target_lang, slow=True)
-                            tts.save(audio_file)
+                # For multi-speaker models, we might need to specify speaker
+                tts.tts_to_file(text=text, file_path=audio_file)
                 
                 if os.path.exists(audio_file) and os.path.getsize(audio_file) > 0:
                     audio_files.append(audio_file)
@@ -587,11 +588,11 @@ if __name__ == "__main__":
         missing_deps.append("assemblyai")
         
     try:
-        import gtts
-        logger.info("gTTS package found")
+        import TTS
+        logger.info("Coqui TTS package found")
     except ImportError:
-        logger.warning("gTTS package not found - required for text-to-speech")
-        missing_deps.append("gtts")
+        logger.warning("Coqui TTS package not found - required for text-to-speech")
+        missing_deps.append("TTS")
     
     try:
         import deep_translator