Update app.py

app.py

@@ -30,6 +30,24 @@ import time
 import os
 import openai 
 from openai import OpenAI 
+import traceback
+from TTS.api import TTS
+import torch
+from TTS.tts.configs.xtts_config import XttsConfig
+
+# Accept license terms for Coqui XTTS
+os.environ["COQUI_TOS_AGREED"] = "1"
+torch.serialization.add_safe_globals([XttsConfig])
+
+# Load XTTS model
+try:
+    print("🔄 Loading XTTS model...")
+    tts = TTS(model_name="tts_models/multilingual/multi-dataset/xtts_v2")
+    print("✅ XTTS model loaded successfully.")
+except Exception as e:
+    print("❌ Error loading XTTS model:")
+    traceback.print_exc()
+    raise e
 
 client = OpenAI(
     api_key= os.environ.get("openAI_api_key"),  # This is the default and can be omitted
@@ -110,52 +128,80 @@ logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %
 logger = logging.getLogger(__name__)
 logger.info(f"MoviePy Version: {moviepy.__version__}")
 
-def transcribe_video(video_path):
-    # Load the video file and extract audio
+def transcribe_video_with_speakers(video_path):
+    # Extract audio from video
     video = VideoFileClip(video_path)
     audio_path = "audio.wav"
     video.audio.write_audiofile(audio_path)
+    logger.info(f"Audio extracted from video: {audio_path}")
     
-    # Load Whisper model
-    model = whisper.load_model("base")  # Options: tiny, base, small, medium, large
-
-    # Transcribe with Whisper
-    result = model.transcribe(audio_path, word_timestamps=True)
-    
-    # Extract timestamps, text, and compute word count
-    total_words = 0
-    total_duration = 0
-    transcript_with_timestamps = []
-    
+    # Set up device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    logger.info(f"Using device: {device}")
+
+    # Load WhisperX model
+    model = whisperx.load_model("large-v2", device)
+    logger.info("WhisperX model loaded")
+
+    # Transcribe with WhisperX
+    result = model.transcribe(audio_path)
+    logger.info("Audio transcription completed")
+
+    # Align transcription
+    model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
+    result = whisperx.align(result["segments"], model_a, metadata, audio_path, device)
+    logger.info("Transcription alignment completed")
+
+    # Perform speaker diarization
+    diarize_model = whisperx.DiarizationPipeline(use_auth_token=hf_api_key, device=device)
+    diarize_segments = diarize_model(audio_path)
+    logger.info("Speaker diarization completed")
+
+    # Assign speakers to transcribed segments
+    result = whisperx.assign_word_speakers(diarize_segments, result)
+    logger.info("Speakers assigned to transcribed segments")
+
+    # Extract timestamps, text, and speaker IDs
+    transcript_with_speakers = [
+        {
+            "start": segment["start"],
+            "end": segment["end"],
+            "text": segment["text"],
+            "speaker": segment["speaker"]
+        }
+        for segment in result["segments"]
+    ]
+
+    # Collect audio for each speaker
+    speaker_audio = {}
     for segment in result["segments"]:
-        start = segment["start"]
-        end = segment["end"]
-        text = segment["text"]
-
-        transcript_with_timestamps.append({
-            "start": start,
-            "end": end,
-            "text": text
-        })
-
-        word_count = count_words_or_characters(text)
-        total_words += word_count
-        total_duration += (end - start)
-    
-    # Compute average words per second
-    avg_words_per_second = total_words / total_duration if total_duration > 0 else 0
-    
-    # Add total statistics to the result
-    transcript_stats = {
-        "total_words": total_words,
-        "total_duration": total_duration,
-        "avg_words_per_second": avg_words_per_second
-    }
-    logger.debug(f"Transcription stats:\n{transcript_stats}")
+        speaker = segment["speaker"]
+        if speaker not in speaker_audio:
+            speaker_audio[speaker] = []
+        speaker_audio[speaker].append((segment["start"], segment["end"]))
+
+    # Collapse and truncate speaker audio
+    speaker_sample_paths = {}
+    audio_clip = AudioFileClip(audio_path)
+    for speaker, segments in speaker_audio.items():
+        speaker_clips = [audio_clip.subclip(start, end) for start, end in segments]
+        combined_clip = concatenate_audioclips(speaker_clips)
+        truncated_clip = combined_clip.subclip(0, min(30, combined_clip.duration))
+        sample_path = f"speaker_{speaker}_sample.wav"
+        truncated_clip.write_audiofile(sample_path)
+        speaker_sample_paths[speaker] = sample_path
+        logger.info(f"Created sample for {speaker}: {sample_path}")
+
     # Get the detected language
     detected_language = result["language"]
-    logger.debug(f"Detected language:\n{detected_language}")
-    return transcript_with_timestamps, detected_language
+    logger.debug(f"Detected language: {detected_language}")
+
+    # Clean up
+    video.close()
+    audio_clip.close()
+    os.remove(audio_path)
+
+    return transcript_with_speakers, detected_language
 
 # Function to get the appropriate translation model based on target language
 def get_translation_model(source_language, target_language):
@@ -259,7 +305,7 @@ def update_translations(file, edited_table, mode):
     except Exception as e:
         raise ValueError(f"Error updating translations: {e}")
 
-def process_entry(entry, i, video_width, video_height, add_voiceover, target_language):
+def process_entry(entry, i, video_width, video_height, add_voiceover, target_language, speaker_sample_paths=None):
     logger.debug(f"Processing entry {i}: {entry}")
 
     # Create text clip for subtitles
@@ -278,7 +324,9 @@ def process_entry(entry, i, video_width, video_height, add_voiceover, target_lan
     if add_voiceover:
         segment_audio_path = f"segment_{i}_voiceover.wav"
         desired_duration = entry["end"] - entry["start"]
-        generate_voiceover_OpenAI([entry], target_language, desired_duration, segment_audio_path)
+        speaker_wav_path = f"speaker_{entry["speaker"]}_sample.wav"
+        generate_voiceover_clone([entry], desired_duration, target_language, speaker_wav_path, segment_audio_path):
+        
         audio_clip = AudioFileClip(segment_audio_path)
         # Get and log all methods in AudioFileClip        
         logger.info("Methods in AudioFileClip:")
@@ -301,7 +349,7 @@ def process_entry(entry, i, video_width, video_height, add_voiceover, target_lan
 
     return i, txt_clip, audio_segment
 
-def add_transcript_voiceover(video_path, translated_json, output_path, add_voiceover=False, target_language="en"):
+def add_transcript_voiceover(video_path, translated_json, output_path, add_voiceover=False, target_language="en", speaker_sample_paths=None):
     """
     Add transcript and voiceover to a video, segment by segment.
     """
@@ -312,7 +360,7 @@ def add_transcript_voiceover(video_path, translated_json, output_path, add_voice
     audio_segments = []
 
     with concurrent.futures.ThreadPoolExecutor() as executor:
-        futures = [executor.submit(process_entry, entry, i, video.w, video.h, add_voiceover, target_language) 
+        futures = [executor.submit(process_entry, entry, i, video.w, video.h, add_voiceover, target_language, speaker_sample_paths)
                    for i, entry in enumerate(translated_json)]
 
         # Collect results with original index i
@@ -348,19 +396,35 @@ def add_transcript_voiceover(video_path, translated_json, output_path, add_voice
 
     logger.info("Video processing completed successfully.")
 
-def generate_voiceover(translated_json, language, output_audio_path):
-    """
-    Generate voiceover from translated text for a given language.
-    """
-    # Concatenate translated text into a single string
-    full_text = " ".join(entry["translated"] for entry in translated_json)
-
+# Voice cloning function with debug and error handling
+def generate_voiceover_clone(translated_json, desired_duration, target_language, speaker_wav_path, output_audio_path):
     try:
-        tts = gTTS(text=full_text, lang=language)
-        time.sleep(10)  # Add a delay of 10 seconds between requests
-        tts.save(output_audio_path)
+        full_text = " ".join(entry["translated"] for entry in translated_json)
+        speed_tts = calculate_speed(full_text, desired_duration)
+        if not speaker_wav_path or not os.path.exists(speaker_wav_path):
+            return None, "❌ Please upload a valid speaker audio file."
+
+        print(f"📥 Received text: {full_text}")
+        print(f"📁 Speaker audio path: {speaker_wav_path}")
+        print(f"🌐 Selected language: {target_language}")
+        print(f"⏱️ Target speed: {speed_tts}")
+
+        # Run TTS with speed control (if supported by model)
+        tts.tts_to_file(
+            text=full_text,
+            speaker_wav=speaker_wav_path,
+            language=language,
+            file_path=output_audio_path,
+            speed=speed_tts  # <- add speed control
+        )
+        print("✅ Voice cloning completed.")
+        return output_path, "✅ Voice cloning completed successfully."
+
     except Exception as e:
-        raise ValueError(f"Error generating voiceover: {e}")
+        print("❌ Error during voice cloning:")
+        traceback.print_exc()
+        error_msg = f"❌ An error occurred: {str(e)}"
+        return None, error_msg
 
 def truncated_linear(x):
     if x < 15:
@@ -381,47 +445,6 @@ def calculate_speed(text, desired_duration):
     
     return speed
 
-def generate_voiceover_OpenAI(translated_json, language, desired_duration, output_audio_path):
-    """
-    Generate voiceover from translated text for a given language using OpenAI TTS API.
-    """
-    # Concatenate translated text into a single string
-    full_text = " ".join(entry["translated"] for entry in translated_json)
-
-    # Define the voice based on the language (for now, use 'alloy' as default)
-    voice = "alloy"  # Adjust based on language if needed
-
-    # Define the model (use tts-1 for real-time applications)
-    model = "tts-1"
-
-    max_retries = 3
-    retry_count = 0
-
-    while retry_count < max_retries:
-        try:
-            speed_tts = calculate_speed(full_text, desired_duration)
-            # Create the speech using OpenAI TTS API
-            response = client.audio.speech.create(
-                model=model,
-                voice=voice,
-                input=full_text,
-                speed=speed_tts
-            )
-            # Save the audio to the specified path
-            with open(output_audio_path, 'wb') as f:
-                for chunk in response.iter_bytes():
-                    f.write(chunk)
-            logging.info(f"Voiceover generated successfully for {output_audio_path}")
-            break
-
-        except Exception as e:
-            retry_count += 1
-            logging.error(f"Error generating voiceover (retry {retry_count}/{max_retries}): {e}")
-            time.sleep(5)  # Wait 5 seconds before retrying
-
-    if retry_count == max_retries:
-        raise ValueError(f"Failed to generate voiceover after {max_retries} retries.")
-
 def upload_and_manage(file, target_language, mode="transcription"):
     if file is None:
         logger.info("No file uploaded. Please upload a video/audio file.")
@@ -439,7 +462,7 @@ def upload_and_manage(file, target_language, mode="transcription"):
 
         # Step 1: Transcribe audio from uploaded media file and get timestamps
         logger.info("Transcribing audio...")
-        transcription_json, source_language = transcribe_video(file.name)
+        transcription_json, source_language = transcribe_video_with_speakers(file.name)
         logger.info(f"Transcription completed. Detected source language: {source_language}")
 
         # Step 2: Translate the transcription
@@ -449,13 +472,13 @@ def upload_and_manage(file, target_language, mode="transcription"):
 
         # Step 3: Add transcript to video based on timestamps
         logger.info("Adding translated transcript to video...")
-        add_transcript_voiceover(file.name, translated_json, output_video_path, mode == "Transcription with Voiceover", target_language)
+        add_transcript_voiceover(file.name, translated_json, output_video_path, mode == "Transcription with Voiceover", target_language, speaker_sample_path)
         logger.info(f"Transcript added to video. Output video saved at {output_video_path}")
 
         # Convert translated JSON into a format for the editable table
         logger.info("Converting translated JSON into editable table format...")
         editable_table = [
-            [float(entry["start"]), entry["original"], entry["translated"], float(entry["end"])]
+            [float(entry["start"]), entry["original"], entry["translated"], float(entry["end"]), entry["speaker"]]
             for entry in translated_json
         ]
 
@@ -519,7 +542,7 @@ def build_interface():
 
             save_changes_button.click(
                 update_translations, 
-                inputs=[file_input, editable_table, process_mode], 
+                inputs=[file_input, editable_table, process_mode],
                 outputs=[processed_video_output, elapsed_time_display]
             )