Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
 from transformers import MarianMTModel, MarianTokenizer
 import soundfile as sf
+from datasets import load_dataset
 
 # Device setup
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -28,7 +29,6 @@ tts_pipe = pipeline("text-to-speech", "microsoft/speecht5_tts")
 
 # Load speaker embeddings
 def get_speaker_embedding():
-    from datasets import load_dataset
     dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
     # Use the first sample's embedding as an example
     speaker_embedding = torch.tensor(dataset[7306]["xvector"]).unsqueeze(0)
@@ -45,38 +45,54 @@ def load_translation_model(lang_code):
 
 st.title("TextLangAudioGenerator")
 
-# Text input box
-text_input = st.text_area("Enter text in English")
+# Text input
+text_input = st.text_area("Enter text in English", key="text_input")
 
-# Language selection dropdown
+# Select target language
 target_lang = st.selectbox(
     "Select target language",
     ["fr", "zh", "it", "ur", "hi"],  # Add more language codes as needed
-    format_func=lambda x: {"fr": "French", "zh": "Chinese", "it": "Italian", "ur": "Urdu", "hi": "Hindi"}.get(x, x)
+    format_func=lambda x: {"fr": "French", "zh": "Chinese", "it": "Italian", "ur": "Urdu", "hi": "Hindi"}.get(x, x),
+    key="target_lang"
 )
 
-# Buttons for actions
-if st.button("Translate and Generate Audio"):
+# Initialize session state for storing results
+if "translated_text" not in st.session_state:
+    st.session_state.translated_text = ""
+if "audio_path" not in st.session_state:
+    st.session_state.audio_path = ""
+
+# Submit button
+if st.button("Submit"):
     if text_input and target_lang:
         # Load translation model
         model, tokenizer = load_translation_model(target_lang)
         inputs = tokenizer(text_input, return_tensors="pt")
         translated = model.generate(**inputs)
-        translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
-        st.write(f"Translated text: {translated_text}")
+        st.session_state.translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
+        st.write(f"Translated text: {st.session_state.translated_text}")
 
+# Listen to Translated Audio button
+if st.button("Listen to Translated Audio"):
+    if st.session_state.translated_text:
         # Generate TTS
-        speech = tts_pipe(translated_text, forward_params={"speaker_embeddings": speaker_embedding})
-        audio_path = "translated_speech.wav"
-        sf.write(audio_path, speech["audio"], samplerate=speech["sampling_rate"])
-
-        st.audio(audio_path, format="audio/wav")
+        speech = tts_pipe(st.session_state.translated_text, forward_params={"speaker_embeddings": speaker_embedding})
+        st.session_state.audio_path = "translated_speech.wav"
+        sf.write(st.session_state.audio_path, speech["audio"], samplerate=speech["sampling_rate"])
+        st.audio(st.session_state.audio_path, format="audio/wav")
     else:
-        st.error("Please enter text and select a target language.")
+        st.error("Please submit the text first.")
 
-# Optional: Add a button for additional functionality if needed
+# Reset button
 if st.button("Reset"):
-    st.text_area("Enter text in English", value="")
-    st.selectbox("Select target language", ["fr", "zh", "it", "ur", "hi"])
+    st.session_state.translated_text = ""
+    st.session_state.audio_path = ""
+    st.experimental_rerun()  # Reload the app to reset the inputs
+
+# Display current state of translated text and audio
+if st.session_state.translated_text:
+    st.write(f"Translated text: {st.session_state.translated_text}")
+if st.session_state.audio_path:
+    st.audio(st.session_state.audio_path, format="audio/wav")