Update app.py

app.py

@@ -4,124 +4,135 @@ import torch
 import torchaudio
 from transformers import (
     WhisperProcessor, WhisperForConditionalGeneration,
-    SpeechT5Processor, SpeechT5ForTextToSpeech,
-    MarianMTModel, MarianTokenizer
+    MarianMTModel, MarianTokenizer,
+    SpeechT5Processor, SpeechT5ForTextToSpeech
 )
 import ffmpeg
 import soundfile as sf
+import numpy as np
 
 # Configurações
 UPLOAD_FOLDER = "uploads"
 OUTPUT_FOLDER = "outputs"
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(OUTPUT_FOLDER, exist_ok=True)
 
-if not os.path.exists(UPLOAD_FOLDER):
-    os.makedirs(UPLOAD_FOLDER)
-if not os.path.exists(OUTPUT_FOLDER):
-    os.makedirs(OUTPUT_FOLDER)
-
-# Verificar se a GPU está disponível
+# Configurar dispositivo (GPU se disponível)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Usando dispositivo: {device}")
+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-en-pt"
+TTS_MODEL = "microsoft/speecht5_tts"
 
-# Inicializar modelos
-whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-small")
-whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to(device)
+# 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)
 
-tts_processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
-tts_model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts").to(device)
+translation_tokenizer = MarianTokenizer.from_pretrained(TRANSLATION_MODEL)
+translation_model = MarianMTModel.from_pretrained(TRANSLATION_MODEL).to(device)
 
-# Usar um modelo alternativo de tradução
-translation_model_name = "Helsinki-NLP/opus-mt-tc-big-en-pt"
-translation_tokenizer = MarianTokenizer.from_pretrained(translation_model_name)
-translation_model = MarianMTModel.from_pretrained(translation_model_name).to(device)
+tts_processor = SpeechT5Processor.from_pretrained(TTS_MODEL)
+tts_model = SpeechT5ForTextToSpeech.from_pretrained(TTS_MODEL).to(device)
 
-# Funções auxiliares
+# Funções otimizadas
 def transcribe_audio(audio_path):
     waveform, sample_rate = torchaudio.load(audio_path)
     waveform = waveform.to(device)
-    inputs = whisper_processor(waveform.squeeze().cpu().numpy(), sampling_rate=sample_rate, return_tensors="pt").to(device)
-    with torch.no_grad():
-        predicted_ids = whisper_model.generate(inputs.input_features)
-    transcription = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)
-    return transcription[0]
+    
+    # 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.no_grad():
+    with torch.inference_mode():
         speech = tts_model.generate_speech(inputs["input_ids"], tts_model.speaker_embeddings)
-    sf.write(output_path, speech.cpu().numpy(), samplerate=22050)
-
-def translate_text(text, target_language="pt"):
-    inputs = translation_tokenizer(text, return_tensors="pt", truncation=True).to(device)
-    with torch.no_grad():
-        translated_ids = translation_model.generate(**inputs)
-    translated_text = translation_tokenizer.decode(translated_ids[0], skip_special_tokens=True)
-    return translated_text
-
-def extract_audio(video_path, audio_path):
-    ffmpeg.input(video_path).output(audio_path, ac=1, ar="16000").run(overwrite_output=True)
-
-def replace_audio_in_video(video_path, audio_path, output_path):
-    # Extrair a taxa de amostragem do áudio original
-    probe = ffmpeg.probe(video_path)
-    audio_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'audio'), None)
-    original_sample_rate = int(audio_stream['sample_rate'])
-
-    # Converter o áudio sintetizado para a taxa de amostragem do vídeo original
-    converted_audio_path = os.path.join(UPLOAD_FOLDER, "converted_audio.wav")
-    ffmpeg.input(audio_path).output(converted_audio_path, ar=original_sample_rate).run(overwrite_output=True)
-
-    # Substituir o áudio no vídeo
-    video = ffmpeg.input(video_path)
-    audio = ffmpeg.input(converted_audio_path)
-    ffmpeg.output(video.video, audio.audio, output_path, vcodec="copy", acodec="aac").run(overwrite_output=True)
-
-# Função principal para Gradio
+    
+    # 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()):
     try:
-        # O Gradio passa o caminho do arquivo como uma string
-        video_path = video
-
-        # Extrair áudio
-        progress(0.1, desc="Extraindo áudio do vídeo...")
+        # Etapa 1: Extrair áudio
+        progress(0.1, "Extraindo áudio...")
         audio_path = os.path.join(UPLOAD_FOLDER, "audio.wav")
-        extract_audio(video_path, audio_path)
-
-        # Transcrever áudio
-        progress(0.3, desc="Transcrevendo áudio...")
-        transcribed_text = transcribe_audio(audio_path)
-        print("Texto transcrito:", transcribed_text)
-
-        # Traduzir texto
-        progress(0.5, desc="Traduzindo texto...")
-        translated_text = translate_text(transcribed_text, target_language="pt")
-        print("Texto traduzido:", translated_text)
-
-        # Sintetizar áudio
-        progress(0.7, desc="Sintetizando áudio em português...")
-        synthesized_audio_path = os.path.join(UPLOAD_FOLDER, "synthesized_audio.wav")
-        synthesize_speech(translated_text, synthesized_audio_path)
-
-        # Substituir áudio no vídeo
-        progress(0.9, desc="Substituindo áudio no vídeo...")
-        output_video_path = os.path.join(OUTPUT_FOLDER, "translated_video.mp4")
-        replace_audio_in_video(video_path, synthesized_audio_path, output_video_path)
-
-        # Finalizar progresso
-        progress(1.0, desc="Concluído!")
-        return output_video_path
+        (
+            ffmpeg
+            .input(video)
+            .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)
+            .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 Gradio
+# Interface otimizada
 iface = gr.Interface(
     fn=translate_video,
-    inputs=gr.Video(),
-    outputs=gr.Video(),
-    title="Tradutor de Vídeo",
-    description="Carregue um vídeo em qualquer idioma e ele será traduzido para português."
+    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"
 )
 
-iface.launch()
-
+if __name__ == "__main__":
+    iface.launch(server_port=7860, show_error=True)
+