Update app.py

app.py

@@ -2,45 +2,43 @@ 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
 import json
+import tempfile
 
-
+# Carregar mapeamento de idiomas
 with open('language_codes.json', 'r') as f:
     lang_codes = json.load(f)
-# 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")
+# Inicializar modelos
+tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
+model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
+whisper_model = WhisperModel("large-v2", 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"])
+    # 1. Extrair áudio
+    audio_file = tempfile.NamedTemporaryFile(suffix=".wav").name
+    run(["ffmpeg", "-i", video.name, audio_file])
 
-    # 2. Transcrição usando Whisper
-    segments, _ = whisper_model.transcribe("audio.wav")
+    # 2. Transcrição
+    segments, _ = whisper_model.transcribe(audio_file)
     transcript = " ".join([segment.text for segment in segments])
 
-    # 3. Tradução usando NLLB
-    flores_code = lang_codes.get(target_language, "eng_Latn")  # Default para inglês se não encontrado
+    # 3. Tradução
+    flores_code = lang_codes.get(target_language, "eng_Latn")
     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_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[flores_code], max_length=100)
     translated_text = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
-    
-    return "Tradução realizada para o idioma com código FLORES-200: " + flores_code
-    # 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"
+    # 4. Criar arquivo de legenda
+    subtitle_file = tempfile.NamedTemporaryFile(suffix=".srt", delete=False).name
+    with open(subtitle_file, "w") as f:
+        f.write("1\n00:00:00,000 --> 00:00:10,000\n" + translated_text)
+
+    # 5. Incorporar legenda
+    output_video = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    run(["ffmpeg", "-i", video.name, "-vf", f"subtitles={subtitle_file}", output_video])
+
+    return output_video
 
 # Interface Gradio
 iface = gr.Interface(
@@ -49,8 +47,9 @@ iface = gr.Interface(
         gr.Radio(["Upload", "Record"], value="Upload", show_label=False),
         gr.Video(),
         gr.Dropdown(choices=list(lang_codes.keys()), label="Target Language for Dubbing", value="English"),
+        gr.Checkbox(label="Video has a close-up face. Use Wav2lip.", value=False)
     ],
-    outputs=gr.Textbox(),
+    outputs=gr.Video(),
     live=False,
     title="AI Video Dubbing"
 )