Added app.py

app.py

@@ -0,0 +1,223 @@
+import gradio as gr
+import os
+import torch
+from openvoice import se_extractor
+from openvoice.api import ToneColorConverter
+import whisper
+from moviepy.editor import VideoFileClip
+from pydub import AudioSegment
+from df.enhance import enhance, init_df, load_audio, save_audio
+import translators as ts
+from melo.api import TTS
+from concurrent.futures import ThreadPoolExecutor
+import ffmpeg
+
+def process_video(video_file, language_choice):
+    if video_file == None or language_choice == None:
+        return None
+    
+    # Initialize paths and devices
+    ckpt_converter = 'checkpoints_v2/converter'
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    output_dir = 'outputs_v2'
+    os.makedirs(output_dir, exist_ok=True)
+
+    tone_color_converter = ToneColorConverter(f'{ckpt_converter}/config.json', device=device)
+    tone_color_converter.load_ckpt(f'{ckpt_converter}/checkpoint.pth')
+
+    # Process the reference video
+    reference_video = VideoFileClip(video_file)
+    reference_audio = os.path.join(output_dir, "reference_audio.wav")
+    reference_video.audio.write_audiofile(reference_audio)
+    audio = AudioSegment.from_file(reference_audio)
+    resampled_audio = audio.set_frame_rate(48000)
+    resampled_audio.export(reference_audio, format="wav")
+
+    # Enhance the audio
+    model, df_state, _ = init_df()
+    audio, _ = load_audio(reference_audio, sr=df_state.sr())
+    enhanced = enhance(model, df_state, audio)
+    save_audio(reference_audio, enhanced, df_state.sr())
+    reference_speaker = reference_audio  # This is the voice you want to clone
+    target_se, audio_name = se_extractor.get_se(reference_speaker, tone_color_converter, vad=False)
+
+    src_path = os.path.join(output_dir, "tmp.wav")
+
+    # Speed is adjustable
+    speed = 1.0
+
+    # Transcribe the original audio with timestamps
+    sttmodel = whisper.load_model("base")
+    sttresult = sttmodel.transcribe(reference_speaker, verbose=True)
+
+    # Print the original transcription
+    print(sttresult["text"])
+    print(sttresult["language"])
+
+    # Get the segments with start and end times
+    segments = sttresult['segments']
+
+    # Choose the target language for translation
+    language = 'EN_NEWEST'
+    match language_choice:
+        case 'en':
+            language = 'EN_NEWEST'
+        case 'es':
+            language = 'ES'
+        case 'fr':
+            language = 'FR'
+        case 'zh':
+            language = 'ZH'
+        case 'ja':
+            language = 'JP'
+        case 'ko':
+            language = 'KR'
+        case _:
+            language = 'EN_NEWEST'
+
+    # Translate the transcription segment by segment
+    def translate_segment(segment):
+        return segment["start"], segment["end"], ts.translate_text(query_text=segment["text"], translator="google", to_language=language_choice)
+
+    # Batch translation to reduce memory load
+    batch_size = 2
+    translation_segments = []
+    for i in range(0, len(segments), batch_size):
+        batch = segments[i:i + batch_size]
+        with ThreadPoolExecutor(max_workers=5) as executor:
+            batch_translations = list(executor.map(translate_segment, batch))
+        translation_segments.extend(batch_translations)
+
+    # Generate the translated audio for each segment
+    model = TTS(language=language, device=device)
+    speaker_ids = model.hps.data.spk2id
+
+    def generate_segment_audio(segment, speaker_id):
+        start, end, translated_text = segment
+        segment_path = os.path.join(output_dir, f'segment_{start}_{end}.wav')
+        model.tts_to_file(translated_text, speaker_id, segment_path, speed=speed)
+        return segment_path, start, end, translated_text
+
+    for speaker_key in speaker_ids.keys():
+        speaker_id = speaker_ids[speaker_key]
+        speaker_key = speaker_key.lower().replace('_', '-')
+
+        source_se = torch.load(f'checkpoints_v2/base_speakers/ses/{speaker_key}.pth', map_location=device)
+
+        segment_files = []
+        subtitle_entries = []
+        for segment in translation_segments:
+            segment_file, start, end, translated_text = generate_segment_audio(segment, speaker_id)
+
+            # Run the tone color converter
+            encode_message = "@MyShell"
+            tone_color_converter.convert(
+            audio_src_path=segment_file,
+            src_se=source_se,
+            tgt_se=target_se,
+            output_path=segment_file,
+            message=encode_message)
+            
+            segment_files.append((segment_file, start, end, translated_text))
+
+        # Combine the audio segments
+        combined_audio = AudioSegment.empty()
+        video_segments = []
+        previous_end = 0
+        subtitle_counter = 1
+        for segment_file, start, end, translated_text in segment_files:
+            segment_audio = AudioSegment.from_file(segment_file)
+            combined_audio += segment_audio
+            
+            # Calculate the duration of the audio segment
+            audio_duration = len(segment_audio) / 1000.0
+
+            # Add the subtitle entry for this segment
+            subtitle_entries.append((subtitle_counter, previous_end, previous_end + audio_duration, translated_text))
+            subtitle_counter += 1
+
+            # Get the corresponding video segment and adjust its speed to match the audio duration
+            video_segment = (
+                ffmpeg
+                .input(reference_video.filename, ss=start, to=end)
+                .filter('setpts', f'PTS / {(end - start) / audio_duration}')
+            )
+            video_segments.append((video_segment, ffmpeg.input(segment_file)))
+            previous_end += audio_duration
+
+        save_path = os.path.join(output_dir, f'output_v2_{speaker_key}.wav')
+        combined_audio.export(save_path, format="wav")
+
+        # Combine video and audio segments using ffmpeg
+        video_and_audio_files = [item for sublist in video_segments for item in sublist]
+        joined = (
+            ffmpeg
+            .concat(*video_and_audio_files, v=1, a=1)
+            .node
+        )
+
+        final_video_path = os.path.join(output_dir, f'final_video_{speaker_key}.mp4')
+        try:
+            (
+                ffmpeg
+                .output(joined[0], joined[1], final_video_path, vcodec='libx264', acodec='aac')
+                .run(overwrite_output=True)
+            )
+        except ffmpeg.Error as e:
+            print('ffmpeg error:', e)
+            print(e.stderr.decode('utf-8'))
+
+        print(f"Final video without subtitles saved to: {final_video_path}")
+
+        # Generate subtitles file in SRT format
+        srt_path = os.path.join(output_dir, 'subtitles.srt')
+        with open(srt_path, 'w', encoding='utf-8') as srt_file:
+            for entry in subtitle_entries:
+                index, start, end, text = entry
+                start_hours, start_minutes = divmod(int(start), 3600)
+                start_minutes, start_seconds = divmod(start_minutes, 60)
+                start_milliseconds = int((start * 1000) % 1000)
+
+                end_hours, end_minutes = divmod(int(end), 3600)
+                end_minutes, end_seconds = divmod(end_minutes, 60)
+                end_milliseconds = int((end * 1000) % 1000)
+
+                srt_file.write(f"{index}\n")
+                srt_file.write(f"{start_hours:02}:{start_minutes:02}:{start_seconds:02},{start_milliseconds:03} --> "
+                               f"{end_hours:02}:{end_minutes:02}:{end_seconds:02},{end_milliseconds:03}\n")
+                srt_file.write(f"{text}\n\n")
+
+        # Add subtitles to the video
+        final_video_with_subs_path = os.path.join(output_dir, f'final_video_with_subs_{speaker_key}.mp4')
+        try:
+            (
+                ffmpeg
+                .input(final_video_path)
+                .output(final_video_with_subs_path, vf=f"subtitles={srt_path}")
+                .run(overwrite_output=True)
+            )
+        except ffmpeg.Error as e:
+            print('ffmpeg error:', e)
+            print(e.stderr.decode('utf-8'))
+
+        print(f"Final video with subtitles saved to: {final_video_with_subs_path}")
+
+        return final_video_with_subs_path
+
+
+# Define Gradio interface
+def gradio_interface(video_file, language_choice):
+    return process_video(video_file, language_choice)
+
+language_choices = ts.get_languages("google")["en"]
+
+gr.Interface(
+    fn=gradio_interface,
+    inputs=[
+        gr.Video(label="Upload Video"),
+        gr.Dropdown(choices=language_choices, label="Choose Language for Translation")
+    ],
+    outputs=gr.Video(label="Translated Video"),
+    title="Video Translation and Voice Cloning",
+    description="Upload a video, choose a language to translate the audio, and download the processed video with translated audio."
+).launch()