Update app.py

app.py

@@ -6,22 +6,40 @@ import os, stat
 import uuid
 from googletrans import Translator
 from TTS.api import TTS
+import ffmpeg
 from faster_whisper import WhisperModel
+from scipy.signal import wiener
 import soundfile as sf
+from pydub import AudioSegment
 import numpy as np
+import librosa
+from zipfile import ZipFile
+import shlex
 import cv2
+import torch
+import torchvision
+from tqdm import tqdm
+from numba import jit
 from huggingface_hub import HfApi
-import shlex
 
+# Environment setup
 HF_TOKEN = os.environ.get("HF_TOKEN")
 os.environ["COQUI_TOS_AGREED"] = "1"
 api = HfApi(token=HF_TOKEN)
 repo_id = "artificialguybr/video-dubbing"
 
-# Whisper
+# Extract ffmpeg
+ZipFile("ffmpeg.zip").extractall()
+st = os.stat('ffmpeg')
+os.chmod('ffmpeg', st.st_mode | stat.S_IEXEC)
+
+# Initialize Whisper model
 model_size = "small"
 model = WhisperModel(model_size, device="cpu", compute_type="int8")
 
+# Initialize TTS model
+tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2")
+
 def check_for_faces(video_path):
     face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
     cap = cv2.VideoCapture(video_path)
@@ -39,6 +57,17 @@ def check_for_faces(video_path):
 
     return False
 
+@spaces.GPU
+def transcribe_audio(audio_path):
+    segments, info = model.transcribe(audio_path, beam_size=5)
+    whisper_text = " ".join(segment.text for segment in segments)
+    whisper_language = info.language
+    return whisper_text, whisper_language
+
+@spaces.GPU
+def generate_tts(text, speaker_wav, language_code):
+    tts.tts_to_file(text, speaker_wav=speaker_wav, file_path="output_synth.wav", language=language_code)
+
 @spaces.GPU
 def process_video(radio, video, target_language, has_closeup_face):
     if target_language is None:
@@ -46,37 +75,33 @@ def process_video(radio, video, target_language, has_closeup_face):
         
     run_uuid = uuid.uuid4().hex[:6]
     output_filename = f"{run_uuid}_resized_video.mp4"
-    
-    # Use subprocess for ffmpeg operations
-    subprocess.run(["ffmpeg", "-i", video, "-vf", "scale=-2:720", output_filename])
+    ffmpeg.input(video).output(output_filename, vf='scale=-2:720').run()
 
     video_path = output_filename
     
     if not os.path.exists(video_path):
         return f"Error: {video_path} does not exist."
 
-    # Check video duration
-    video_info = subprocess.check_output(["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", video_path])
-    video_duration = float(video_info)
+    video_info = ffmpeg.probe(video_path)
+    video_duration = float(video_info['streams'][0]['duration'])
 
     if video_duration > 60:
         os.remove(video_path)
         return gr.Error("Video duration exceeds 1 minute. Please upload a shorter video.")
 
-    subprocess.run(["ffmpeg", "-i", video_path, "-acodec", "pcm_s24le", "-ar", "48000", "-map", "a", f"{run_uuid}_output_audio.wav"])
+    ffmpeg.input(video_path).output(f"{run_uuid}_output_audio.wav", acodec='pcm_s24le', ar=48000, map='a').run()
 
-    subprocess.run(["ffmpeg", "-y", "-i", f"{run_uuid}_output_audio.wav", "-af", "lowpass=3000,highpass=100", f"{run_uuid}_output_audio_final.wav"])
+    shell_command = f"ffmpeg -y -i {run_uuid}_output_audio.wav -af lowpass=3000,highpass=100 {run_uuid}_output_audio_final.wav".split(" ")
+    subprocess.run([item for item in shell_command], capture_output=False, text=True, check=True)
     
     print("Attempting to transcribe with Whisper...")
     try:
-        segments, info = model.transcribe(f"{run_uuid}_output_audio_final.wav", beam_size=5)
-        whisper_text = " ".join(segment.text for segment in segments)
-        whisper_language = info.language
+        whisper_text, whisper_language = transcribe_audio(f"{run_uuid}_output_audio_final.wav")
         print(f"Transcription successful: {whisper_text}")
     except RuntimeError as e:
         print(f"RuntimeError encountered: {str(e)}")
         if "CUDA failed with error device-side assert triggered" in str(e):
-            gr.Warning("Error. Space need to restart. Please retry in a minute")
+            gr.Warning("Error. Space needs to restart. Please retry in a minute")
             api.restart_space(repo_id=repo_id)
             
     language_mapping = {'English': 'en', 'Spanish': 'es', 'French': 'fr', 'German': 'de', 'Italian': 'it', 'Portuguese': 'pt', 'Polish': 'pl', 'Turkish': 'tr', 'Russian': 'ru', 'Dutch': 'nl', 'Czech': 'cs', 'Arabic': 'ar', 'Chinese (Simplified)': 'zh-cn'}
@@ -85,19 +110,20 @@ def process_video(radio, video, target_language, has_closeup_face):
     translated_text = translator.translate(whisper_text, src=whisper_language, dest=target_language_code).text
     print(translated_text)
 
-    tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2")
-    tts.tts_to_file(translated_text, speaker_wav=f"{run_uuid}_output_audio_final.wav", file_path=f"{run_uuid}_output_synth.wav", language=target_language_code)
+    generate_tts(translated_text, f"{run_uuid}_output_audio_final.wav", target_language_code)
     
     if has_closeup_face:
         try:
-            cmd = f"python Wav2Lip/inference.py --checkpoint_path 'Wav2Lip/checkpoints/wav2lip_gan.pth' --face {shlex.quote(video_path)} --audio '{run_uuid}_output_synth.wav' --pads 0 15 0 0 --resize_factor 1 --nosmooth --outfile '{run_uuid}_output_video.mp4'"
+            cmd = f"python Wav2Lip/inference.py --checkpoint_path 'Wav2Lip/checkpoints/wav2lip_gan.pth' --face {shlex.quote(video_path)} --audio 'output_synth.wav' --pads 0 15 0 0 --resize_factor 1 --nosmooth --outfile '{run_uuid}_output_video.mp4'"
             subprocess.run(cmd, shell=True, check=True)
         except subprocess.CalledProcessError as e:
             if "Face not detected! Ensure the video contains a face in all the frames." in str(e.stderr):
                 gr.Warning("Wav2lip didn't detect a face. Please try again with the option disabled.")
-                subprocess.run(["ffmpeg", "-i", video_path, "-i", f"{run_uuid}_output_synth.wav", "-c:v", "copy", "-c:a", "aac", "-strict", "experimental", "-map", "0:v:0", "-map", "1:a:0", f"{run_uuid}_output_video.mp4"])
+                cmd = f"ffmpeg -i {video_path} -i output_synth.wav -c:v copy -c:a aac -strict experimental -map 0:v:0 -map 1:a:0 {run_uuid}_output_video.mp4"
+                subprocess.run(cmd, shell=True)
     else:
-        subprocess.run(["ffmpeg", "-i", video_path, "-i", f"{run_uuid}_output_synth.wav", "-c:v", "copy", "-c:a", "aac", "-strict", "experimental", "-map", "0:v:0", "-map", "1:a:0", f"{run_uuid}_output_video.mp4"])
+        cmd = f"ffmpeg -i {video_path} -i output_synth.wav -c:v copy -c:a aac -strict experimental -map 0:v:0 -map 1:a:0 {run_uuid}_output_video.mp4"
+        subprocess.run(cmd, shell=True)
 
     if not os.path.exists(f"{run_uuid}_output_video.mp4"):
         raise FileNotFoundError(f"Error: {run_uuid}_output_video.mp4 was not generated.")
@@ -109,7 +135,7 @@ def process_video(radio, video, target_language, has_closeup_face):
         f"{run_uuid}_resized_video.mp4",
         f"{run_uuid}_output_audio.wav",
         f"{run_uuid}_output_audio_final.wav",
-        f"{run_uuid}_output_synth.wav"
+        "output_synth.wav"
     ]
     for file in files_to_delete:
         try:
@@ -120,11 +146,9 @@ def process_video(radio, video, target_language, has_closeup_face):
     return output_video_path
 
 def swap(radio):
-    if(radio == "Upload"):
-        return gr.update(source="upload")
-    else:
-        return gr.update(source="webcam")
-        
+    return gr.update(source="upload" if radio == "Upload" else "webcam")
+
+# Gradio interface setup
 video = gr.Video()
 radio = gr.Radio(["Upload", "Record"], value="Upload", show_label=False)
 iface = gr.Interface(
@@ -133,10 +157,7 @@ iface = gr.Interface(
         radio,
         video,
         gr.Dropdown(choices=["English", "Spanish", "French", "German", "Italian", "Portuguese", "Polish", "Turkish", "Russian", "Dutch", "Czech", "Arabic", "Chinese (Simplified)"], label="Target Language for Dubbing", value="Spanish"),
-        gr.Checkbox(
-                    label="Video has a close-up face. Use Wav2lip.",
-                    value=False,
-                    info="Say if video have close-up face. For Wav2lip. Will not work if checked wrongly.")
+        gr.Checkbox(label="Video has a close-up face. Use Wav2lip.", value=False, info="Say if video have close-up face. For Wav2lip. Will not work if checked wrongly.")
     ],
     outputs=gr.Video(),
     live=False,
@@ -159,5 +180,5 @@ with gr.Blocks() as demo:
     - If you incorrectly mark the 'Video has a close-up face' checkbox, the dubbing may not work as expected.
     """)
 
-demo.queue()
+demo.queue(concurrency_count=1, max_size=15)
 demo.launch()