import gradio as gr
from transformers import pipeline
import time
from fairseq.checkpoint_utils import load_model_ensemble_and_task_from_hf_hub
from fairseq.models.text_to_speech.hub_interface import TTSHubInterface
import torch
from fairseq.utils import move_to_cuda

"""
Notes:
pip install sentencepiece
pip install phonemizer

install ffmpeg and add to path

Must install e-speak to path 

solution
1.Download and install the Windows version of espeak: http://espeak.sourceforge.net/download.html

2. set PATH=%PATH%;"C:\Program Files (x86)\eSpeak\command_line"_

3. Install .msi from https://github.com/espeak-ng/espeak-ng/releases

4.Enter environment variable

    1.PHONEMIZER_ESPEAK_LIBRARY="c:\Program Files\eSpeak NG\libespeak-ng.dll"
    2.PHONEMIZER_ESPEAK_PATH =“c:\Program Files\eSpeak NG”

and Restart your Computer. Run the same command again from the Command Prompt (cmd.exe):


"""
asr = pipeline("automatic-speech-recognition",model="facebook/wav2vec2-base-960h" )
translation_pipeline = pipeline('translation_en_to_es',model = "Helsinki-NLP/opus-mt-en-es" ) #This model version is built for en- to -fr , less mistakes

models, cfg, task = load_model_ensemble_and_task_from_hf_hub(
    "facebook/tts_transformer-es-css10",
    arg_overrides={"vocoder": "hifigan", "fp16": False}
)

model = models[0]
model = model.to(torch.device("cuda:0")) if torch.cuda.is_available() else model

TTSHubInterface.update_cfg_with_data_cfg(cfg, task.data_cfg)
generator = task.build_generator(models, cfg)


def transcribe_translate(audio):
    time.sleep(3)
    text_en = asr(audio)["text"]
    
    text_fr = translation_pipeline(text_en.lower()) # for some reason all audio converted to all caps and it translates differently???
    text_fr = text_fr[0]['translation_text']       # good evening = bonsoir  but GOOD EVENING = BONNES SÉANCES . WEIRD

    sample = TTSHubInterface.get_model_input(task, text_fr)
    sample = move_to_cuda(sample) if torch.cuda.is_available() else sample

    wav, rate = TTSHubInterface.get_prediction(task, model, generator, sample)
    wav = wav.to('cpu')
    wav = wav.numpy()
    print(wav.dtype)

    return text_en,text_fr , (rate,wav)


gr.Interface(
    fn=transcribe_translate, 
    inputs=[
        gr.Audio(source="microphone", type="filepath")
    ],
    outputs=[
        gr.Textbox(label= "English Transcription"),
        gr.Textbox(label= "Spanish Translation"),
        gr.Audio(label = "Spanish Audio")
    ],
    live=True).launch(share=True)