Update app.py

app.py

@@ -1,72 +1,108 @@
 import gradio as gr
 import numpy as np
 import torch
-from datasets import load_dataset
-
-from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor, pipeline
-
-
+import os
+import tempfile
+import logging
+import asyncio
+from typing import Optional
+from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq, AutoModelForTextToSpectrogram
+from deep_translator import GoogleTranslator
+import time
+import threading
+import torchaudio
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+# Device setup
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
-
-# load speech translation checkpoint
-asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device)
-
-# load text-to-speech checkpoint and speaker embeddings
-processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
-
-model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts").to(device)
-vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan").to(device)
-
-embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
-speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
-
-
-def translate(audio):
-    outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate"})
-    return outputs["text"]
-
-
-def synthesise(text):
-    inputs = processor(text=text, return_tensors="pt")
-    speech = model.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=vocoder)
-    return speech.cpu()
-
-
-def speech_to_speech_translation(audio):
-    translated_text = translate(audio)
-    synthesised_speech = synthesise(translated_text)
-    synthesised_speech = (synthesised_speech.numpy() * 32767).astype(np.int16)
-    return 16000, synthesised_speech
-
-
-title = "Cascaded STST"
-description = """
-Demo for cascaded speech-to-speech translation (STST), mapping from source speech in any language to target speech in English. Demo uses OpenAI's [Whisper Base](https://huggingface.co/openai/whisper-base) model for speech translation, and Microsoft's
-[SpeechT5 TTS](https://huggingface.co/microsoft/speecht5_tts) model for text-to-speech:
-
-![Cascaded STST](https://huggingface.co/datasets/huggingface-course/audio-course-images/resolve/main/s2st_cascaded.png "Diagram of cascaded speech to speech translation")
-"""
-
-demo = gr.Blocks()
-
-mic_translate = gr.Interface(
-    fn=speech_to_speech_translation,
-    inputs=gr.Audio(source="microphone", type="filepath"),
-    outputs=gr.Audio(label="Generated Speech", type="numpy"),
-    title=title,
-    description=description,
-)
-
-file_translate = gr.Interface(
-    fn=speech_to_speech_translation,
-    inputs=gr.Audio(source="upload", type="filepath"),
-    outputs=gr.Audio(label="Generated Speech", type="numpy"),
-    examples=[["./example.wav"]],
-    title=title,
-    description=description,
-)
-
-with demo:
-    gr.TabbedInterface([mic_translate, file_translate], ["Microphone", "Audio File"])
-
-demo.launch()
+logger.info(f"Device set to use {device}")
+
+# Supported languages
+SUPPORTED_LANGUAGES = {
+    "English": "en",
+    "Hindi": "hi",
+}
+
+# Load ASR (Speech Recognition) Model
+asr_processor = AutoProcessor.from_pretrained("openai/whisper-base")
+asr_model = AutoModelForSpeechSeq2Seq.from_pretrained("openai/whisper-base").to(device)
+
+# Load TTS (Text-to-Speech) Model
+tts_processor = AutoProcessor.from_pretrained("microsoft/speecht5_tts")
+tts_model = AutoModelForTextToSpectrogram.from_pretrained("microsoft/speecht5_tts").to(device)
+
+def speech_to_text(audio_data: np.ndarray, sample_rate: int) -> str:
+    try:
+        if len(audio_data.shape) > 1:
+            audio_data = np.mean(audio_data, axis=1)
+        audio_data = torch.tensor(audio_data, dtype=torch.float32).to(device)
+        
+        inputs = asr_processor(audio_data, sampling_rate=sample_rate, return_tensors="pt").to(device)
+        outputs = asr_model.generate(**inputs)
+        text = asr_processor.batch_decode(outputs, skip_special_tokens=True)[0]
+        logger.info(f"Transcribed text: {text}")
+        return text if text else ""
+    except Exception as e:
+        logger.error(f"STT failed: {str(e)}")
+        return ""
+
+def text_to_speech(text: str, lang: str) -> Optional[str]:
+    try:
+        if not text or not text.strip():
+            logger.warning("Empty text received for TTS.")
+            return None
+        target_lang_code = SUPPORTED_LANGUAGES.get(lang, "en")
+        translated_text = GoogleTranslator(source="auto", target=target_lang_code).translate(text)
+        
+        inputs = tts_processor(translated_text, return_tensors="pt").to(device)
+        spectrogram = tts_model.generate(**inputs)
+        
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
+            torchaudio.save(temp_file.name, spectrogram.cpu(), 16000)
+            return temp_file.name
+    except Exception as e:
+        logger.error(f"TTS failed: {str(e)}")
+        return None
+
+async def process_audio_stream(audio_stream, target_lang: str):
+    async for sample_rate, audio_chunk in audio_stream:
+        try:
+            if audio_chunk is None or not isinstance(audio_chunk, np.ndarray):
+                yield "", None
+                continue
+            text = speech_to_text(audio_chunk, sample_rate)
+            dubbed_audio = text_to_speech(text, target_lang) if text and text.strip() else None
+            yield text, dubbed_audio
+        except Exception as e:
+            logger.error(f"Stream processing failed: {str(e)}")
+            yield "", None
+
+def run_cleanup():
+    while True:
+        for root, dirs, files in os.walk(tempfile.gettempdir()):
+            for file in files:
+                if file.endswith(".wav"):
+                    try:
+                        os.remove(os.path.join(root, file))
+                    except:
+                        pass
+        time.sleep(300)
+
+with gr.Blocks(title="Real-Time Multilingual Dubbing") as demo:
+    gr.Markdown("<h1>Real-Time Multilingual Dubbing</h1>")
+    with gr.Row():
+        audio_input = gr.Audio(sources=["microphone"], type="numpy", streaming=True, label="Speak")
+        lang_dropdown = gr.Dropdown(choices=list(SUPPORTED_LANGUAGES.keys()), label="Target Language", value="English")
+    with gr.Row():
+        stt_output = gr.Textbox(label="Transcription")
+        dub_output = gr.Audio(label="Dubbed Audio", autoplay=True)
+    
+    audio_input.stream(fn=process_audio_stream, inputs=[audio_input, lang_dropdown], outputs=[stt_output, dub_output])
+
+if __name__ == "__main__":
+    cleanup_thread = threading.Thread(target=run_cleanup, daemon=True)
+    cleanup_thread.start()
+    demo.launch()