import streamlit as st
from transformers import MarianMTModel, MarianTokenizer, pipeline
import torch
from io import BytesIO

# Define the translation model
model_name = "Helsinki-NLP/opus-mt-en-es"  # Default to Spanish for the example
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)

# Define the TTS model
tts_model_name = "microsoft/speecht5_tts"
tts_pipeline = pipeline("text-to-speech", model=tts_model_name)

# Define the language options
languages = {
    'es': 'Spanish',
    'fr': 'French',
    'de': 'German',
    'ur': 'Urdu',
    'hi': 'Hindi',
    'bn': 'Bengali'
}

# Streamlit app
st.title("Real-Time Language Translator and TTS")

# Text input
text_to_translate = st.text_area("Enter text to translate:", "Hello, how are you?")

# Language selection
target_language = st.selectbox("Select target language:", list(languages.keys()))

def translate_text(text, target_lang):
    # Update the model and tokenizer based on the selected language
    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, forced_bos_token_id=tokenizer.lang_code_to_id[target_lang])
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    return translated_text

def synthesize_speech(text, lang):
    # Synthesize speech
    try:
        audio = tts_pipeline(text, lang=lang)
        audio_bytes = BytesIO(audio["audio"])
        return audio_bytes, None
    except Exception as e:
        return None, str(e)

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
    audio_bytes, error = synthesize_speech(translated_text, target_language)
    if error:
        st.error(f"Error: {error}")
    else:
        st.audio(audio_bytes, format="audio/wav")