import streamlit as st
from transformers import pipeline
import soundfile as sf

# Define translation models
translation_models = {
    "French": "Helsinki-NLP/opus-mt-en-fr",
    "Chinese": "Helsinki-NLP/opus-mt-en-zh",
    "Italian": "Helsinki-NLP/opus-mt-en-it",
    "Urdu": "Helsinki-NLP/opus-mt-en-ur",
    "Hindi": "Helsinki-NLP/opus-mt-en-hi",
    "Punjabi": "Helsinki-NLP/opus-mt-en-pa",
    "Pashto": "Helsinki-NLP/opus-mt-en-ps"
}

# Define TTS model
tts_model_name = "facebook/fastspeech2-en-ljspeech"
tts_pipeline = pipeline("text-to-speech", model=tts_model_name)

# Function to translate text
def translate_text(text, target_lang):
    model_name = translation_models.get(target_lang)
    if model_name:
        translation_pipeline = pipeline("translation", model=model_name)
        translated = translation_pipeline(text)[0]['translation_text']
        return translated
    else:
        return "Translation model not found."

# Function to generate speech
def text_to_speech(text):
    # Generate speech
    audio = tts_pipeline(text)
    return audio['array']

# Function to save audio to file
def save_audio(audio_numpy, file_name):
    sf.write(file_name, audio_numpy, 22050)  # Adjust sample rate if needed
    return file_name

# Streamlit UI layout
st.title("TextLangAudioGenerator")

# Text input
text_input = st.text_area("Enter your text in English:")

# Language selection dropdown
target_lang = st.selectbox("Select target language:", list(translation_models.keys()))

if st.button("Translate and Generate Audio"):
    if text_input:
        # Translate text
        translated_text = translate_text(text_input, target_lang)
        st.write(f"Translated Text ({target_lang}): {translated_text}")

        # Generate speech from translated text
        speech_audio = text_to_speech(translated_text)

        # Save and play audio
        audio_file = save_audio(speech_audio, 'output.wav')
        st.audio(audio_file)

        # Clear the input for new text
        text_input = ""

# Footer
st.write("Powered by Hugging Face Transformers and TTS")