Update app.py

app.py

@@ -1,56 +1,74 @@
 import streamlit as st
 from transformers import MarianMTModel, MarianTokenizer, pipeline
-from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech
-import soundfile as sf
+from datasets import load_dataset
 import torch
+import soundfile as sf
 
-# Define the language model and tokenizer
-translation_model_name = "Helsinki-NLP/opus-mt-en-ur"
-tokenizer = MarianTokenizer.from_pretrained(translation_model_name)
-translation_model = MarianMTModel.from_pretrained(translation_model_name)
+# Define the translation model and tokenizer
+model_name = "Helsinki-NLP/opus-mt-en-ur"
+model = MarianMTModel.from_pretrained(model_name)
+tokenizer = MarianTokenizer.from_pretrained(model_name)
 
-# Load the text-to-speech model
+# Define the TTS model
 tts_model_name = "microsoft/speecht5_tts"
-processor = SpeechT5Processor.from_pretrained(tts_model_name)
-tts_model = SpeechT5ForTextToSpeech.from_pretrained(tts_model_name)
+tts_model = pipeline("text-to-speech", model=tts_model_name)
+speaker_embeddings = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
 
 # Function to translate text
 def translate_text(text, target_lang):
-    if target_lang not in ["Urdu", "Hindi", "Bengali"]:
-        return "Error: Target language not supported."
-    tokens = tokenizer(text, return_tensors="pt", padding=True)
-    translated_tokens = translation_model.generate(**tokens)
-    translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+    inputs = tokenizer(text, return_tensors="pt", padding=True)
+    translated = model.generate(**inputs)
+    translated_text = tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
     return translated_text
 
-# Function to generate speech
-def synthesize_speech(text, speaker_embeddings=None):
-    inputs = processor(text, return_tensors="pt")
-    with torch.no_grad():
-        speech = tts_model.generate_speech(inputs, speaker_embeddings)
-    sf.write("output.wav", speech.numpy(), 16000)
+# Function to synthesize speech
+def synthesize_speech(text, target_lang):
+    if target_lang == 'ur':
+        lang = 'ur_PK'
+    elif target_lang == 'hi':
+        lang = 'hi_IN'
+    elif target_lang == 'bn':
+        lang = 'bn_BD'
+    else:
+        lang = 'en_US'
+    
+    speaker_embedding = torch.tensor(speaker_embeddings[0]["xvector"]).unsqueeze(0)
+    audio = tts_model(text, speaker_embeddings=speaker_embedding)
+    
+    with open("output.wav", "wb") as f:
+        sf.write(f, audio["waveform"].numpy(), samplerate=16000, format="WAV")
+    
     return "output.wav"
 
-# Streamlit app
-st.title("Language Translator with TTS")
-st.write("Enter the text you want to translate and hear the translation.")
+# Streamlit UI
+st.title("Language Translator with Speech Synthesis")
+
+# Input text
+text_input = st.text_input("Enter text in English:")
+if text_input:
+    st.session_state.text_input = text_input
+
+# Language selection
+target_lang = st.selectbox("Select Target Language:", [
+    "Urdu (ur)", "Hindi (hi)", "Bengali (bn)"
+])
 
-# Select target language
-target_language = st.selectbox("Select Target Language", ["Urdu", "Hindi", "Bengali"])
+# Translate button
+if st.button("Translate"):
+    if target_lang == "Urdu (ur)":
+        target_lang = "ur"
+    elif target_lang == "Hindi (hi)":
+        target_lang = "hi"
+    elif target_lang == "Bengali (bn)":
+        target_lang = "bn"
 
-# Text input
-text_to_translate = st.text_input("Enter text here")
+    translated_text = translate_text(st.session_state.text_input, target_lang)
+    st.text_area("Translated text:", value=translated_text, height=100)
+    
+    audio_file = synthesize_speech(translated_text, target_lang)
+    st.audio(audio_file)
 
-if st.button("Translate and Generate Audio"):
     # Clear input for new text
-    st.session_state.text_to_translate = ""
+    st.session_state.text_input = ""
 
-    # Perform translation
-    translated_text = translate_text(text_to_translate, target_language)
-    st.write(f"Translated text ({target_language}): {translated_text}")
 
-    # Generate speech
-    if translated_text and "Error" not in translated_text:
-        speaker_embeddings = None  # Placeholder, use actual embeddings if needed
-        audio_file = synthesize_speech(translated_text, speaker_embeddings)
-        st.audio(audio_file)