Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 from transformers import MarianMTModel, MarianTokenizer, SpeechT5Processor, SpeechT5ForTextToSpeech
 from datasets import load_dataset
 import torch
+from scipy.io.wavfile import write as write_wav
 from io import BytesIO
 
 # Define the language options
@@ -36,7 +37,7 @@ def translate_text(text, target_lang):
     translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
     return translated_text
 
-def synthesize_speech(text, lang):
+def synthesize_speech(text):
     # Load the TTS model and processor
     processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
     model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
@@ -49,8 +50,14 @@ def synthesize_speech(text, lang):
     inputs = processor(text=text, return_tensors="pt")
     speech = model.generate_speech(inputs["input_ids"], speaker_embeddings)
     
-    audio_bytes = BytesIO(speech.numpy())
-    return audio_bytes, None
+    # 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
 
 if st.button("Translate and Synthesize Speech"):
     # Perform translation
@@ -58,11 +65,9 @@ if st.button("Translate and Synthesize Speech"):
     st.write(f"Translated text ({languages[target_language]}): {translated_text}")
     
     # Perform text-to-speech
-    audio_bytes, error = synthesize_speech(translated_text, target_language)
-    if error:
-        st.error(f"Error: {error}")
-    else:
-        st.audio(audio_bytes, format="audio/wav")
+    audio_bytes = synthesize_speech(translated_text)
+    st.audio(audio_bytes, format="audio/wav")
+