app.py

import os
import gradio as gr
import whisperx
import numpy as np
import moviepy.editor as mp
from moviepy.audio.AudioClip import AudioArrayClip
from pytube import YouTube
import deepl
import torch
import pyrubberband as pyrb
import soundfile as sf
import librosa
from TTS.api import TTS

os.environ["COQUI_TOS_AGREED"] = "1"
HF_TOKEN = os.environ["HF_TOKEN"]
DEEPL_TOKEN = os.environ["DEEPL_TOKEN"]

# Download video from Youtube
def download_youtube_video(url):
 yt = YouTube(url)
 stream = yt.streams.filter(file_extension='mp4').first()
 output_path = stream.download()
 return output_path


# Extract audio from video
def extract_audio(video_path):
 clip = mp.VideoFileClip(video_path)
 audio_path = os.path.splitext(video_path)[0] + ".wav"
 clip.audio.write_audiofile(audio_path)
 return audio_path


# Perform speech diarization
def speech_diarization(audio_path, hf_token):
 device = "cuda"
 batch_size = 16
 compute_type = "float16"
 model = whisperx.load_model("large-v2", device, compute_type=compute_type)

 # 1. Transcribe audio
 audio = whisperx.load_audio(audio_path)
 result = model.transcribe(audio, batch_size=batch_size)

 # delete model if low on GPU resources
 import gc; gc.collect(); torch.cuda.empty_cache(); del model

 # 2. Align whisper output
 model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
 result = whisperx.align(result["segments"], model_a, metadata, audio, device, return_char_alignments=False)

 # delete model if low on GPU resources
 import gc; gc.collect(); torch.cuda.empty_cache(); del model_a

 # 3. Assign speaker labels
 diarize_model = whisperx.DiarizationPipeline(model_name='pyannote/[email protected]', use_auth_token=hf_token, device=device)

 # add min/max number of speakers if known
 diarize_segments = diarize_model(audio)
 # diarize_model(audio, min_speakers=min_speakers, max_speakers=max_speakers)

 result = whisperx.assign_word_speakers(diarize_segments, result)
 print(f'\n[Original transcript]:\n{result["segments"]}\n')

 return result["segments"]


# Create per speaker voice clips for tts voice cloning
def speaker_voice_clips(transcription, audio_path):
 # Create 3 uninterrupted per speaker timecodes
 snippets_timecodes = {}
 for segment in transcription:
 speaker = segment['speaker']

 if speaker not in snippets_timecodes:
 snippets_timecodes[speaker] = []

 if len(snippets_timecodes[speaker]) < 3:
 snippet = {
 'start': segment['start'],
 'end': segment['end']
 }
 snippets_timecodes[speaker].append(snippet)

 # Cut voice clips and stitch them together
 original_audio = mp.AudioFileClip(audio_path)
 audio_file_directory = os.path.dirname(audio_path)

 voice_clips = {}
 for speaker, speaker_snippets in snippets_timecodes.items():
 subclips = []
 for snippet in speaker_snippets:
 start, end = snippet['start'], snippet['end']
 subclip = original_audio.subclip(start, end)
 subclips.append(subclip)

 concatenated_clip = mp.concatenate_audioclips(subclips)

 output_filename = os.path.join(audio_file_directory, f"{speaker}_voice_clips.wav")
 concatenated_clip.write_audiofile(output_filename)
 voice_clips[speaker] = output_filename

 return voice_clips


# Perform text translation
def translate_transcript(transcript, target_language, deepl_token):
 translator = deepl.Translator(deepl_token)

 translated_transcript = []
 for segment in transcript:
 text_to_translate = segment['text']
 translated_text = translator.translate_text(text_to_translate, target_lang=target_language)

 translated_segment = {
 'start': segment['start'],
 'end': segment['end'],
 'text': translated_text.text,
 'speaker': segment['speaker']
 }

 translated_transcript.append(translated_segment)

 print(f'\n[Translated transcript]:\n{translated_transcript}\n')

 return translated_transcript


# Adjust voice pace
def adjust_voice_pace(sound_array, sample_rate, target_duration):
 duration = len(sound_array) / sample_rate
 tempo_change = duration / target_duration
 sound_array_stretched = pyrb.time_stretch(sound_array, sample_rate, tempo_change)
 return sound_array_stretched


# Perform voice cloning
def voice_cloning_translation(translated_transcription, speakers_voice_clips, target_language, speaker_model, audio_path):
 device = "cuda"

 vits_language_map = {
 'en':'eng',
 'ru':'rus',
 'uk':'ukr',
 'pl':'pol'
 }

 # Select model
 selected_model = None

 if 'vits' in speaker_model.lower() or target_language is 'uk':
 selected_model = f'tts_models/{vits_language_map[target_language]}/fairseq/vits'
 else:
 selected_model = 'tts_models/multilingual/multi-dataset/xtts_v2'

 print(selected_model)


 tts = None
 final_audio_track = None

 try:
 # TODO uncomment when https://github.com/coqui-ai/TTS/issues/3224 is resolved
 # tts = TTS(selected_model).to(device)

 # Generate and concatenate voice clips per speaker

 last_end_time = 0
 clips = []

 # Generate sentences
 for speech_item in translated_transcription:

 speech_item_duration = speech_item['end'] - speech_item['start']

 # Silence
 gap_duration = speech_item['start'] - last_end_time
 if gap_duration > 0:
 silent_audio = np.zeros((int(44100 * gap_duration), 2))
 silent_clip = AudioArrayClip(silent_audio, fps=44100)
 clips.append(silent_clip)
 print(f"\nAdded silence: Start={last_end_time}, Duration={gap_duration}")

 # Generate speech
 print(f"[{speech_item['speaker']}]")
 tts = TTS(selected_model).to(device)
 audio = tts.tts_with_vc(text=speech_item['text'], speaker_wav=speakers_voice_clips[speech_item['speaker']], language=target_language)
 sample_rate = tts.voice_converter.vc_config.audio.output_sample_rate

 # Adjust pace to fit the speech timeframe if translated audio is longer than phrase
 audio_duration = len(audio) / sample_rate
 if speech_item_duration < audio_duration:
 audio = adjust_voice_pace(audio, sample_rate, speech_item_duration)

 # Resample to higher rate
 new_sample_rate = 44100
 audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=new_sample_rate)

 # Transform to AudioArrayClip object
 audio = np.expand_dims(audio, axis=1)
 audio_stereo = np.repeat(audio, 2, axis=1)
 audio_clip = AudioArrayClip(audio_stereo, fps=44100)

 # Cut out possible glitch from AudioArrayClip end
 audio_clip = audio_clip.subclip(0, audio_clip.duration - 0.2)
 clips.append(audio_clip)
 print(f"Added speech: Start={speech_item['start']}, Final duration={audio_clip.duration}, Original duration={speech_item_duration}")

 last_end_time = speech_item['start'] + audio_clip.duration

 del tts; import gc; gc.collect(); torch.cuda.empty_cache()

 # Merge sentences
 final_audio_track = mp.concatenate_audioclips(clips)

 audio_files_directory = os.path.dirname(audio_path)
 final_audio_track.write_audiofile(os.path.join(audio_files_directory, "translated_voice_track.wav"), fps=44100)

 except Exception as e:
 if tts is not None:
 import gc; gc.collect(); torch.cuda.empty_cache(); del tts
 raise e

 return final_audio_track


def dub_video(video_path, translated_audio_track, target_language):
 video = mp.VideoFileClip(video_path)
 video = video.subclip(0, translated_audio_track.duration)
 original_audio = video.audio.volumex(0.2)
 dubbed_audio = mp.CompositeAudioClip([original_audio, translated_audio_track.set_start(0)])
 video_with_dubbing = video.set_audio(dubbed_audio)

 video_with_dubbing_path = os.path.splitext(video_path)[0] + "_" + target_language + ".mp4"
 video_with_dubbing.write_videofile(video_with_dubbing_path)

 return video_with_dubbing_path


# Perform video translation
def video_translation(video_path, target_language, speaker_model, hf_token, deepl_token):

 original_audio_path = extract_audio(video_path)

 transcription = speech_diarization(original_audio_path, hf_token)

 translated_transcription = translate_transcript(transcription, target_language, deepl_token)

 speakers_voice_clips = speaker_voice_clips(transcription, original_audio_path)

 translated_audio_track = voice_cloning_translation(translated_transcription, speakers_voice_clips, target_language, speaker_model, original_audio_path)

 video_with_dubbing = dub_video(video_path, translated_audio_track, target_language)

 return video_with_dubbing

def translate_video(_, video_path, __, youtube_link, ___, target_language, speaker_model):
 try:
 if not video_path and not youtube_link:
 gr.Warning("You should either upload video or input a YouTube link")
 return None
 if youtube_link:
 video_path = download_youtube_video(youtube_link)
 dubbed_video = video_translation(video_path, target_language, speaker_model, HF_TOKEN, DEEPL_TOKEN)
 except Exception as e:
 print(f"An error occurred: {e}")
 return gr.components.Video(dubbed_video)


inputs = [
 gr.Markdown("## Currently supported languages are: English, Polish, Ukrainian and Russian"),
 gr.Video(label="Upload a video file"),
 gr.Markdown("**OR**"),
 gr.Textbox(label="Paste YouTube link"),
 gr.Markdown("---"),
 gr.Dropdown(["en", "pl", "uk", "ru"], value="pl", label="Select translation target language"),
 gr.Dropdown(["(Recommended) XTTS_V2", "VITs (will be default for Ukrainian)"], value="(Recommended) XTTS_V2", label="Select text-to-speech generation model")
]

outputs = gr.Video(label="Translated video")

gr.Interface(fn=translate_video, 
 inputs=inputs, 
 outputs=outputs, 
 title="🌐AI Video Translation",
 theme=gr.themes.Base()
 ).launch(show_error=True, debug=True)