Update app.py

app.py

@@ -1,104 +1,140 @@
-from transformers import pipeline, MarianMTModel, MarianTokenizer, SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
+import gradio as gr
+import os
 import torch
+import torchaudio
+from transformers import (
+    WhisperProcessor, WhisperForConditionalGeneration,
+    MarianMTModel, MarianTokenizer,
+    SpeechT5Processor, SpeechT5ForTextToSpeech
+)
 import ffmpeg
-import os
-from pydub import AudioSegment
+import soundfile as sf
+import numpy as np
 
-# 1. Reconhecimento de Fala com Whisper (Inglês)
-def transcribe_audio(audio_path):
-    print("Transcrevendo áudio...")
-    # Carregar o modelo Whisper para inglês
-    asr = pipeline("automatic-speech-recognition", model="openai/whisper-large-v2")
-    
-    # Processar o áudio
-    result = asr(audio_path, chunk_length_s=30)
-    text = result['text']
-    print(f"Texto transcrito: {text}")
-    return text
+# 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
 
-# 2. Tradução Automática com MarianMT (Inglês → Português)
-def translate_text(text):
-    print("Traduzindo texto para o português...")
-    # Carregar o modelo MarianMT para inglês → português
-    model_name = "Helsinki-NLP/opus-mt-en-pt"
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
+# 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)
     
-    # Tokenizar e traduzir
-    inputs = tokenizer(text, return_tensors="pt")
-    translated_ids = model.generate(**inputs)
-    translated_text = tokenizer.decode(translated_ids[0], skip_special_tokens=True)
-    print(f"Texto traduzido: {translated_text}")
-    return translated_text
-
-
-# 3. Síntese de Voz com VITS/FastSpeech (Português)
-def synthesize_speech(text, output_path="output_speech.wav"):
-    print("Sintetizando voz em português...")
-    # Carregar o processador e o modelo VITS
-    processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
-    model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
-    vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
+    # 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)
     
-    # Obter embeddings de voz (opcional: usar embeddings pré-carregados)
-    speaker_embeddings = torch.randn((1, 512))  # Gerar embeddings aleatórios
+    with torch.inference_mode():
+        predicted_ids = whisper_model.generate(**inputs)
     
-    # Processar o texto
-    inputs = processor(text=text, return_tensors="pt")
-    speech = model.generate_speech(inputs["input_ids"], speaker_embeddings, vocoder=vocoder)
+    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)
     
-    # Salvar o áudio gerado
-    speech_np = speech.numpy()
-    AudioSegment(speech_np.tobytes(), frame_rate=16000, sample_width=2, channels=1).export(output_path, format="wav")
-    print(f"Áudio sintetizado salvo em: {output_path}")
-    return output_path
+    # 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')
+    )
 
-# 4. Substituição de Áudio no Vídeo com FFmpeg
-def replace_audio_in_video(input_video, new_audio, output_video):
-    print("Substituindo áudio no vídeo...")
+# Fluxo principal otimizado
+def translate_video(video, progress=gr.Progress()):
     try:
-        # Usar FFmpeg para substituir o áudio
+        # 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(input_video)
-            .output(new_audio, output_video, map_video=0, map_audio=1)
+            .input(video_path)
+            .output(audio_path, ac=1, ar=16000)
+            .global_args('-loglevel', 'error')
             .run(overwrite_output=True)
         )
-        print(f"Vídeo com áudio substituído salvo em: {output_video}")
-    except Exception as e:
-        print(f"Erro ao substituir áudio: {e}")
 
+        # Etapa 2: Transcrição paralela
+        progress(0.3, "Transcrevendo...")
+        transcription = transcribe_audio(audio_path)
 
-# Função Principal
-def main():
-    # Caminhos dos arquivos
-    input_video_path = "input_video.mp4"  # Substitua pelo caminho do seu vídeo
-    temp_audio_path = "temp_audio.wav"
-    output_video_path = "output_video.mp4"
-    
-    # Passo 1: Extrair áudio do vídeo
-    print("Extraindo áudio do vídeo...")
-    (
-        ffmpeg
-        .input(input_video_path)
-        .output(temp_audio_path)
-        .run(overwrite_output=True)
-    )
-    
-    # Passo 2: Transcrever o áudio
-    transcribed_text = transcribe_audio(temp_audio_path)
-    
-    # Passo 3: Traduzir o texto para português
-    translated_text = translate_text(transcribed_text)
-    
-    # Passo 4: Sintetizar a voz traduzida
-    synthesized_audio_path = synthesize_speech(translated_text, output_path="synthesized_audio.wav")
-    
-    # Passo 5: Substituir o áudio no vídeo
-    replace_audio_in_video(input_video_path, synthesized_audio_path, output_video_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
+
+    except Exception as e:
+        return f"Erro: {str(e)}"
+
+# Interface otimizada
+iface = gr.Interface(
+    fn=translate_video,
+    inputs=gr.Video(label="Vídeo de Entrada"),
+    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"
+)
 
 if __name__ == "__main__":
-    main()
-    
+    iface.launch(server_port=7860, show_error=True)