Update app.py

app.py

@@ -1,73 +1,56 @@
 import streamlit as st
-from transformers import MarianMTModel, MarianTokenizer, SpeechT5Processor, SpeechT5ForTextToSpeech
-from datasets import load_dataset
+from transformers import MarianMTModel, MarianTokenizer, pipeline
+from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech
+import soundfile as sf
 import torch
-from scipy.io.wavfile import write as write_wav
-from io import BytesIO
 
-# Define the language options
-languages = {
-    'ur': 'Urdu',
-    'hi': 'Hindi',
-    'bn': 'Bengali'
-}
+# 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)
 
-# Streamlit app
-st.title("Real-Time Language Translator and TTS")
-
-# Text input
-text_to_translate = st.text_area("Enter text to translate:", "My Name is Meer Hassan Farhad, I live in Islamabad")
-
-# Language selection
-target_language = st.selectbox("Select target language:", list(languages.keys()))
+# Load the text-to-speech model
+tts_model_name = "microsoft/speecht5_tts"
+processor = SpeechT5Processor.from_pretrained(tts_model_name)
+tts_model = SpeechT5ForTextToSpeech.from_pretrained(tts_model_name)
 
+# Function to translate text
 def translate_text(text, target_lang):
-    # Use a more specific model for translation
-    model_name = f"Helsinki-NLP/opus-mt-en-{target_lang}"
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-    
-    # Translate text
-    encoded_text = tokenizer.encode(text, return_tensors="pt")
-    translated = model.generate(encoded_text)
-    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
+    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)
     return translated_text
 
-def synthesize_speech(text):
-    # Load the TTS model and processor
-    processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
-    model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
-    
-    # Load speaker embeddings
-    embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
-    speaker_embeddings = torch.tensor(embeddings_dataset[0]["xvector"]).unsqueeze(0)
-
-    # Synthesize speech
-    inputs = processor(text=text, return_tensors="pt")
-    speech = model.generate_speech(inputs["input_ids"], speaker_embeddings)
-    
-    # Convert speech tensor to numpy array and save as wav
-    audio_np = speech.squeeze().cpu().numpy()
-    sample_rate = 16000  # Define a sample rate
-    audio_buffer = BytesIO()
-    write_wav(audio_buffer, sample_rate, audio_np)  # Write the numpy array as a WAV file to the buffer
-    audio_buffer.seek(0)
-    
-    return audio_buffer
+# 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)
+    return "output.wav"
 
-if st.button("Translate and Synthesize Speech"):
-    # Perform translation
-    translated_text = translate_text(text_to_translate, target_language)
-    st.write(f"Translated text ({languages[target_language]}): {translated_text}")
-    
-    # Perform text-to-speech
-    try:
-        audio_bytes = synthesize_speech(translated_text)
-        st.audio(audio_bytes, format="audio/wav")
-    except Exception as e:
-        st.error(f"Error generating audio: {e}")
+# Streamlit app
+st.title("Language Translator with TTS")
+st.write("Enter the text you want to translate and hear the translation.")
 
+# Select target language
+target_language = st.selectbox("Select Target Language", ["Urdu", "Hindi", "Bengali"])
 
+# Text input
+text_to_translate = st.text_input("Enter text here")
 
+if st.button("Translate and Generate Audio"):
+    # Clear input for new text
+    st.session_state.text_to_translate = ""
 
+    # 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)