Update app.py

app.py

@@ -1,142 +1,127 @@
-import gradio as gr
 import os
-import torch
-import torchaudio
-from transformers import (
-    WhisperProcessor, WhisperForConditionalGeneration,
-    MarianMTModel, MarianTokenizer,
-    SpeechT5Processor, SpeechT5ForTextToSpeech
-)
-import ffmpeg
+import subprocess
 import soundfile as sf
-import numpy as np
+from transformers import WhisperProcessor, WhisperForConditionalGeneration, MarianMTModel, MarianTokenizer
+from gtts import gTTS
+import gradio as gr
 
 # Configurações
-UPLOAD_FOLDER = "uploads"
-OUTPUT_FOLDER = "outputs"
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
-os.makedirs(OUTPUT_FOLDER, exist_ok=True)
-
-# Configurar dispositivo (GPU se disponível)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-torch.backends.cudnn.benchmark = True  # Acelera GPU NVIDIA
-
-# Carregar modelos uma única vez (cache)
-WHISPER_MODEL = "openai/whisper-tiny"  # Modelo mais rápido
-TRANSLATION_MODEL = "Helsinki-NLP/opus-mt-tc-big-en-pt"  # Modelo alternativo
-TTS_MODEL = "microsoft/speecht5_tts"
-
-# Inicialização rápida dos modelos
-print("Carregando modelos...")
-whisper_processor = WhisperProcessor.from_pretrained(WHISPER_MODEL)
-whisper_model = WhisperForConditionalGeneration.from_pretrained(WHISPER_MODEL).to(device)
-
-translation_tokenizer = MarianTokenizer.from_pretrained(TRANSLATION_MODEL)
-translation_model = MarianMTModel.from_pretrained(TRANSLATION_MODEL).to(device)
-
-tts_processor = SpeechT5Processor.from_pretrained(TTS_MODEL)
-tts_model = SpeechT5ForTextToSpeech.from_pretrained(TTS_MODEL).to(device)
-
-# Funções otimizadas
-def transcribe_audio(audio_path):
-    waveform, sample_rate = torchaudio.load(audio_path)
-    waveform = waveform.to(device)
-    
-    # Processamento em chunks para áudios longos
-    inputs = whisper_processor(
-        waveform.squeeze().cpu().numpy(), 
-        sampling_rate=sample_rate, 
-        return_tensors="pt",
-        chunk_length_s=30  # Processar em chunks de 30 segundos
-    ).to(device)
-    
-    with torch.inference_mode():
-        predicted_ids = whisper_model.generate(**inputs)
-    
-    return whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
-
-def translate_text(text):
-    inputs = translation_tokenizer(text, return_tensors="pt", truncation=True).to(device)
-    with torch.inference_mode():
-        translated_ids = translation_model.generate(**inputs)
-    return translation_tokenizer.decode(translated_ids[0], skip_special_tokens=True)
-
-def synthesize_speech(text, output_path):
-    inputs = tts_processor(text, return_tensors="pt").to(device)
-    with torch.inference_mode():
-        speech = tts_model.generate_speech(inputs["input_ids"], tts_model.speaker_embeddings)
-    
-    # Converter para formato compatível com vídeo (44100 Hz, stereo)
-    sf.write(output_path, np.tile(speech.cpu().numpy(), (2, 1)).T, 44100, subtype='PCM_16')
-
-def process_video(video_path, output_path):
-    # Processamento paralelizado com FFmpeg
-    (
-        ffmpeg
-        .input(video_path)
-        .output(output_path, vcodec='copy', acodec='aac', strict='experimental')
-        .global_args('-loglevel', 'error')  # Reduzir logs
-        .run(overwrite_output=True, cmd='ffmpeg')
-    )
-
-# Fluxo principal otimizado
-def translate_video(video, progress=gr.Progress()):
+WHISPER_MODEL = "openai/whisper-medium"
+TRANSLATION_MODEL = "Helsinki-NLP/opus-mt-en-pt"
+TEMP_AUDIO_FILE = "temp_audio.wav"
+TEMP_TRANSLATED_AUDIO = "translated_audio.mp3"
+OUTPUT_VIDEO = "output_video.mp4"
+
+def extract_audio(video_path: str, output_audio_path: str) -> None:
+    """
+    Extrai o áudio de um vídeo usando FFmpeg.
+    """
+    try:
+        command = [
+            "ffmpeg", "-i", video_path, "-q:a", "0", "-map", "a", output_audio_path, "-y"
+        ]
+        subprocess.run(command, check=True)
+    except subprocess.CalledProcessError as e:
+        raise Exception(f"Erro ao extrair áudio: {e}")
+
+def transcribe_audio(audio_path: str) -> str:
+    """
+    Transcreve o áudio para texto usando o modelo Whisper.
+    """
+    try:
+        processor = WhisperProcessor.from_pretrained(WHISPER_MODEL)
+        model = WhisperForConditionalGeneration.from_pretrained(WHISPER_MODEL)
+        audio, _ = sf.read(audio_path)
+        input_features = processor(audio, sampling_rate=16000, return_tensors="pt").input_features
+        predicted_ids = model.generate(input_features)
+        transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+        return transcription
+    except Exception as e:
+        raise Exception(f"Erro ao transcrever áudio: {e}")
+
+def translate_text(text: str, source_lang: str = "en", target_lang: str = "pt") -> str:
+    """
+    Traduz o texto para o idioma desejado usando o modelo MarianMT.
+    """
+    try:
+        model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        model = MarianMTModel.from_pretrained(model_name)
+        translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True))
+        translation = tokenizer.decode(translated[0], skip_special_tokens=True)
+        return translation
+    except Exception as e:
+        raise Exception(f"Erro ao traduzir texto: {e}")
+
+def text_to_speech(text: str, output_audio_path: str, lang: str = "pt") -> None:
+    """
+    Converte texto em áudio usando gTTS.
+    """
+    try:
+        tts = gTTS(text, lang=lang)
+        tts.save(output_audio_path)
+    except Exception as e:
+        raise Exception(f"Erro ao gerar áudio: {e}")
+
+def merge_audio_video(video_path: str, audio_path: str, output_path: str) -> None:
+    """
+    Combina o áudio traduzido com o vídeo original usando FFmpeg.
+    """
+    try:
+        command = [
+            "ffmpeg", "-i", video_path, "-i", audio_path, "-c:v", "copy", "-map", "0:v:0", "-map", "1:a:0", "-shortest", output_path, "-y"
+        ]
+        subprocess.run(command, check=True)
+    except subprocess.CalledProcessError as e:
+        raise Exception(f"Erro ao combinar áudio e vídeo: {e}")
+
+def translate_video(video_path: str) -> str:
+    """
+    Função principal que orquestra a tradução do vídeo.
+    """
     try:
-        # O Gradio passa o caminho do arquivo como uma string
-        video_path = video
-
-        # Etapa 1: Extrair áudio
-        progress(0.1, "Extraindo áudio...")
-        audio_path = os.path.join(UPLOAD_FOLDER, "audio.wav")
-        (
-            ffmpeg
-            .input(video_path)
-            .output(audio_path, ac=1, ar=16000)
-            .global_args('-loglevel', 'error')
-            .run(overwrite_output=True)
-        )
-
-        # Etapa 2: Transcrição paralela
-        progress(0.3, "Transcrevendo...")
-        transcription = transcribe_audio(audio_path)
-
-        # Etapa 3: Tradução em lote
-        progress(0.5, "Traduzindo...")
-        translated_text = translate_text(transcription)
-
-        # Etapa 4: Síntese de voz acelerada
-        progress(0.7, "Sintetizando voz...")
-        synthesized_audio = os.path.join(UPLOAD_FOLDER, "synthesized_audio.wav")
-        synthesize_speech(translated_text, synthesized_audio)
-
-        # Etapa 5: Processamento final do vídeo
-        progress(0.9, "Montando vídeo...")
-        output_path = os.path.join(OUTPUT_FOLDER, "video_traduzido.mp4")
-        (
-            ffmpeg
-            .input(video_path)
-            .output(output_path, vcodec='copy', acodec='copy', map='0:v:0')
-            .input(synthesized_audio)
-            .global_args('-loglevel', 'error')
-            .run(overwrite_output=True)
-        )
-
-        return output_path
+        # 1. Extrair áudio do vídeo
+        extract_audio(video_path, TEMP_AUDIO_FILE)
+
+        # 2. Transcrever áudio para texto
+        transcription = transcribe_audio(TEMP_AUDIO_FILE)
+        print(f"Transcrição: {transcription}")
+
+        # 3. Traduzir texto para português
+        translation = translate_text(transcription)
+        print(f"Tradução: {translation}")
+
+        # 4. Converter texto traduzido em áudio
+        text_to_speech(translation, TEMP_TRANSLATED_AUDIO)
 
+        # 5. Combinar áudio traduzido com vídeo original
+        merge_audio_video(video_path, TEMP_TRANSLATED_AUDIO, OUTPUT_VIDEO)
+
+        # 6. Limpar arquivos temporários
+        os.remove(TEMP_AUDIO_FILE)
+        os.remove(TEMP_TRANSLATED_AUDIO)
+
+        return OUTPUT_VIDEO
+    except Exception as e:
+        print(f"Erro durante o processamento: {e}")
+        raise
+
+# Interface Gradio
+def gradio_interface(video):
+    try:
+        output_video = translate_video(video)
+        return output_video
     except Exception as e:
-        return f"Erro: {str(e)}"
+        return f"Erro: {e}"
 
-# Interface otimizada
+# Configuração da interface
 iface = gr.Interface(
-    fn=translate_video,
-    inputs=gr.Video(label="Vídeo de Entrada"),
+    fn=gradio_interface,
+    inputs=gr.Video(label="Upload do Vídeo"),
     outputs=gr.Video(label="Vídeo Traduzido"),
-    title="🚀 Tradutor de Vídeo Ultra-Rápido",
-    description="Carregue um vídeo e receba a versão em português com áudio traduzido!",
-    allow_flagging="never"
+    title="Tradutor de Vídeos para Português",
+    description="Faça upload de um vídeo em qualquer idioma e receba o vídeo com áudio traduzido para português."
 )
 
-if __name__ == "__main__":
-    iface.launch(server_port=7860, show_error=True)
-    
-    
+# Iniciar a interface
+iface.launch()