import gradio as gr
import os
import whisper
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from gtts import gTTS
import numpy as np

# Load models
model_stt = whisper.load_model("base")
model_translation = AutoModelForSeq2SeqLM.from_pretrained("alirezamsh/small100")
tokenizer_translation = AutoTokenizer.from_pretrained("alirezamsh/small100")

def speech_to_speech(input_audio, to_lang):
    # Save the uploaded audio file
    input_file = "input_audio" + os.path.splitext(input_audio.name)[1]
    input_audio.save(input_file)

    # Speech-to-Text (STT)
    audio = whisper.load_audio(input_file)
    audio = whisper.pad_or_trim(audio)
    mel = whisper.log_mel_spectrogram(audio).to(model_stt.device)
    _, probs = model_stt.detect_language(mel)
    options = whisper.DecodingOptions()
    result = whisper.decode(model_stt, mel, options)
    text = result.text
    lang = max(probs, key=probs.get)

    # Translate
    tokenizer_translation.src_lang = lang
    tokenizer_translation.tgt_lang = to_lang
    encoded_bg = tokenizer_translation(text, return_tensors="pt")
    generated_tokens = model_translation.generate(**encoded_bg)
    translated_text = tokenizer_translation.batch_decode(generated_tokens, skip_special_tokens=True)[0]

    # Text-to-Speech (TTS)
    tts = gTTS(text=translated_text, lang=to_lang)
    output_file = "output_audio.mp3"
    tts.save(output_file)

    # Load output audio as numpy array
    audio_np = np.array(output_file)

    return audio_np

languages = ["ru", "fr", "es", "de"]  # Example languages: Russian, French, Spanish, German
file_input = gr.inputs.File(label="Upload Audio", accept="audio/*")
dropdown = gr.inputs.Dropdown(languages, label="Translation Language")
audio_output = gr.outputs.Audio(label="Translated Voice", type="numpy")

gr.Interface(
    fn=speech_to_speech,
    inputs=[file_input, dropdown],
    outputs=audio_output,
    title="Speech-to-Speech Translator",
    description="Upload an audio file (MP3, WAV, or FLAC) and choose the target language for translation.",
    theme="default"
).launch()