# app.py
import streamlit as st
import torch
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC, MarianMTModel, MarianTokenizer
import soundfile as sf
import tempfile

# Load models and tokenizers
@st.cache_resource
def load_models():
    # Load ASR model (Wav2Vec2 for Urdu)
    asr_processor = Wav2Vec2Processor.from_pretrained("m3hrdadfi/wav2vec2-large-xlsr-ur")
    asr_model = Wav2Vec2ForCTC.from_pretrained("m3hrdadfi/wav2vec2-large-xlsr-ur")

    # Load translation model (Urdu to German)
    translation_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ur-de")
    translation_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-ur-de")

    return asr_processor, asr_model, translation_tokenizer, translation_model

asr_processor, asr_model, translation_tokenizer, translation_model = load_models()

# Streamlit App UI
st.title("Real-Time Urdu to German Voice Translator")
st.markdown("Upload an Urdu audio file, and the app will translate it to German.")

uploaded_file = st.file_uploader("Upload an audio file (in .wav format)", type=["wav"])

if uploaded_file is not None:
    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
        temp_file.write(uploaded_file.read())
        temp_file_path = temp_file.name

    # Load audio file
    audio_input, sample_rate = sf.read(temp_file_path)

    # Ensure proper sampling rate
    if sample_rate != 16000:
        st.error("Please upload a .wav file with a sampling rate of 16kHz.")
    else:
        st.info("Processing the audio...")

        # Convert speech to text (ASR)
        input_values = asr_processor(audio_input, return_tensors="pt", sampling_rate=16000).input_values
        with torch.no_grad():
            logits = asr_model(input_values).logits
            predicted_ids = torch.argmax(logits, dim=-1)
            transcription = asr_processor.batch_decode(predicted_ids)[0]

        st.text(f"Transcribed Urdu Text: {transcription}")

        # Translate Urdu text to German
        translated = translation_model.generate(**translation_tokenizer(transcription, return_tensors="pt", padding=True))
        german_translation = translation_tokenizer.decode(translated[0], skip_special_tokens=True)

        st.success(f"Translated German Text: {german_translation}")