app.py

import os
import requests
import gradio as gr
import moviepy.editor as mp
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
import torch
import assemblyai as aai

# Download necessary models if not already present
model_files = {
    "wav2lip.pth": "https://github.com/justinjohn0306/Wav2Lip/releases/download/models/wav2lip.pth",
    "wav2lip_gan.pth": "https://github.com/justinjohn0306/Wav2Lip/releases/download/models/wav2lip_gan.pth",
    "resnet50.pth": "https://github.com/justinjohn0306/Wav2Lip/releases/download/models/resnet50.pth",
    "mobilenet.pth": "https://github.com/justinjohn0306/Wav2Lip/releases/download/models/mobilenet.pth",
    "s3fd.pth": "https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth"
}

# Download model files
for filename, url in model_files.items():
    file_path = os.path.join("checkpoints" if "pth" in filename else "face_detection", filename)
    if not os.path.exists(file_path):
        print(f"Downloading {filename}...")
        r = requests.get(url)
        with open(file_path, 'wb') as f:
            f.write(r.content)

# Initialize xtts model
def initialize_xtts_model():
    # Get the path to the xtts_v2 folder
    tts_dir = os.path.join(os.getcwd(), 'xtts_v2')
    
    # Load the configuration
    config_path = os.path.join(tts_dir, 'config.json')
    config = XttsConfig()
    config.load_json(config_path)
    
    # Initialize the model from the configuration
    model = Xtts.init_from_config(config)
    
    # Load the model checkpoint
    model.load_checkpoint(config, checkpoint_dir=tts_dir, eval=True)
    
    # Move the model to GPU (if available)
    if torch.cuda.is_available():
        model.cuda()
    
    return model

# Translation class
class Translation:
    def __init__(self, video_path, original_language, target_language):
        self.video_path = video_path
        self.original_language = original_language
        self.target_language = target_language
        self.model = initialize_xtts_model()  # Initialize TTS model

    def org_language_parameters(self, original_language):
        language_codes = {'English': 'en', 'German': 'de', 'Italian': 'it', 'Spanish': 'es'}
        self.lan_code = language_codes.get(original_language, '')

    def target_language_parameters(self, target_language):
        language_codes = {'English': 'en', 'German': 'de', 'Italian': 'it', 'Spanish': 'es'}
        self.tran_code = language_codes.get(target_language, '')

    def extract_audio(self):
        video = mp.VideoFileClip(self.video_path)
        audio = video.audio
        audio_path = "output_audio.wav"
        audio.write_audiofile(audio_path)
        return audio_path

    def transcribe_audio(self, audio_path):
        aai.settings.api_key = os.getenv("ASSEMBLYAI_API_KEY")
        config = aai.TranscriptionConfig(language_code=self.lan_code)
        transcriber = aai.Transcriber(config=config)
        transcript = transcriber.transcribe(audio_path)
        return transcript.text

    def translate_text(self, transcript_text):
        base_url = "https://api.cognitive.microsofttranslator.com/translate"
        headers = {
            "Ocp-Apim-Subscription-Key": os.getenv("MICROSOFT_TRANSLATOR_API_KEY"),
            "Content-Type": "application/json",
            "Ocp-Apim-Subscription-Region": "southeastasia"
        }
        params = {"api-version": "3.0", "from": self.lan_code, "to": self.tran_code}
        body = [{"text": transcript_text}]
        response = requests.post(base_url, headers=headers, params=params, json=body)
        translation = response.json()[0]["translations"][0]["text"]
        return translation

    def generate_audio(self, translated_text):
        # Generate audio using the xtts model
        config = XttsConfig()
        config.load_json(os.path.join(os.getcwd(), 'xtts_v2', 'config.json'))
        
        # Generate audio
        synthesized_audio_path = "output_synth.wav"
        outputs = self.model.synthesize(
            translated_text,
            config,
            speaker_wav='output_audio.wav',
            gpt_cond_len=3,
            language=self.tran_code,
        )
        
        # Save the output to file
        with open(synthesized_audio_path, 'wb') as f:
            f.write(outputs)
        
        return synthesized_audio_path

    def translate_video(self):
        audio_path = self.extract_audio()
        self.org_language_parameters(self.original_language)
        self.target_language_parameters(self.target_language)
        transcript_text = self.transcribe_audio(audio_path)
        translated_text = self.translate_text(transcript_text)
        translated_audio_path = self.generate_audio(translated_text)

        # Run Wav2Lip inference
        os.system(f"python inference.py --checkpoint_path 'checkpoints/wav2lip_gan.pth' --face {self.video_path} --audio {translated_audio_path} --outfile 'output_video.mp4'")
        return 'output_video.mp4'

# Gradio Interface
def app(video_path, original_language, target_language):
    translator = Translation(video_path, original_language, target_language)
    video_file = translator.translate_video()
    return video_file

interface = gr.Interface(
    fn=app,
    inputs=[
        gr.Video(label="Video Path"),
        gr.Dropdown(["English", "German", "Italian", "Spanish"], label="Original Language"),
        gr.Dropdown(["English", "German", "Italian", "Spanish"], label="Targeted Language"),
    ],
    outputs=gr.Video(label="Translated Video")
)

interface.launch()