Spaces:
Runtime error
Runtime error
File size: 4,882 Bytes
738e954 33e5b8e 0faa3aa 9b103ac 0faa3aa 87b7f71 0faa3aa 87b7f71 0faa3aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | import gradio as gr
import numpy as np
import torch
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"
logger.info(f"Device set to use {device}")
# Supported languages
SUPPORTED_LANGUAGES = {
"English": "en",
"Hindi": "hi",
}
# Load ASR (Speech Recognition) Model
asr_model_name = "openai/whisper-base"
asr_processor_name = "openai/whisper-base"
# Load TTS (Text-to-Speech) Model
tts_model_name = "microsoft/speecht5_tts"
tts_processor_name = "microsoft/speecht5_tts"
if 'SPACE_ID' in os.environ: # Check if running on Hugging Face Spaces
cache_dir = "./models"
asr_processor = AutoProcessor.from_pretrained(asr_processor_name, cache_dir=cache_dir)
asr_model = AutoModelForSpeechSeq2Seq.from_pretrained(asr_model_name, cache_dir=cache_dir).to(device)
tts_processor = AutoProcessor.from_pretrained(tts_processor_name, cache_dir=cache_dir)
tts_model = AutoModelForTextToSpectrogram.from_pretrained(tts_model_name, cache_dir=cache_dir).to(device)
else:
asr_processor = AutoProcessor.from_pretrained(asr_processor_name)
asr_model = AutoModelForSpeechSeq2Seq.from_pretrained(asr_model_name).to(device)
tts_processor = AutoProcessor.from_pretrained(tts_processor_name)
tts_model = AutoModelForTextToSpectrogram.from_pretrained(tts_model_name).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() |