app.py

import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from subprocess import run
from faster_whisper import WhisperModel
import soundfile as sf
import numpy as np

# Inicialize o modelo NLLB
tokenizer = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
model = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")

# Inicialize o modelo Whisper
model_size = "large-v2"
whisper_model = WhisperModel(model_size, device="cuda", compute_type="float16")

def process_video(radio, video, target_language, use_wav2lip):
    # 1. Extraia o áudio do vídeo usando FFMPEG
    run(["ffmpeg", "-i", video.name, "audio.wav"])

    # 2. Transcrição usando Whisper
    segments, _ = whisper_model.transcribe("audio.wav")
    transcript = " ".join([segment.text for segment in segments])

    # 3. Tradução usando NLLB
    inputs = tokenizer(transcript, return_tensors="pt")
    lang_code = {"English": "eng_Latn", "Spanish": "spa_Latn", "French": "fra_Latn"}  # Adicione mais idiomas conforme necessário
    translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[lang_code[target_language]], max_length=100)
    translated_text = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]

    # 4. Queimar a legenda traduzida no vídeo
    with open("subtitle.srt", "w") as f:
        f.write("1\n00:00:00,000 --> 00:00:10,000\n" + translated_text)  # Este é um exemplo simples. Você pode dividir o texto em várias partes e ajustar os tempos.
    
    run(["ffmpeg", "-i", video.name, "-vf", "subtitles=subtitle.srt", "output_video.mp4"])

    return "output_video.mp4"

# Interface Gradio
video = gr.Video()
radio = gr.Radio(["Upload", "Record"], value="Upload", show_label=False)
iface = gr.Interface(
    fn=process_video,
    inputs=[
        radio,
        video,
        gr.Dropdown(choices=["English", "Spanish", "French"], label="Target Language for Dubbing", value="Spanish"),
        gr.Checkbox(label="Video has a close-up face. Use Wav2lip.", value=False)
    ],
    outputs=gr.Video(),
    live=False
)

iface.launch()