import gradio as gr
from transformers import pipeline, SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
from datasets import load_dataset
import torch




replacements = [
    ("á", "a"),
    ("í", "i"),
    ("ñ", "n"),
    ("ó", "o"),
    ("ú", "u"),
    ("ü", "u"),
]


def cleanup_text(text):
    for src, dst in replacements:
        text = text.replace(src, dst)
    return text


def modelo1(image):
    imageToText = pipeline(
        "image-to-text", model="Salesforce/blip-image-captioning-large")

    resultado = imageToText(image)
    resultado = resultado[0]["generated_text"].replace("araffe ", "")
    return resultado


def modelo2(text):
    enToEs = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")

    resultado = enToEs(text)
    return resultado[0]["translation_text"]


def modelo3En(text):
    processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
    model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
    vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
    embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
    speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)

    inputs = processor(text=text, return_tensors="pt")

    speech = model.generate_speech(inputs["input_ids"], speaker_embeddings, vocoder=vocoder)

    return (16000, speech.cpu().numpy())


def modelo3Es(text):
    model = SpeechT5ForTextToSpeech.from_pretrained("Sandiago21/speecht5_finetuned_facebook_voxpopuli_spanish")
    vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
    embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
    processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
    speaker_embeddings = torch.tensor(embeddings_dataset[7440]["xvector"]).unsqueeze(0)

    text = cleanup_text(text)
    inputs = processor(text=text, return_tensors="pt")

    speech = model.generate_speech(inputs["input_ids"], speaker_embeddings, vocoder=vocoder)

    return (16000, speech.cpu().numpy())


def executionIMG(image, lan):
    if lan == 'english':
        model1res = modelo1(image)
        model3res = modelo3En(model1res)
        return model3res
    elif lan == 'spanish':
        model1res = modelo1(image)
        model2res = modelo2(model1res)
        model3res = modelo3Es(model2res)
        return model3res


def executionTEXT(text, lan):
    if lan == 'english':
        model3res = modelo3En(text)
        return model3res
    elif lan == 'spanish':
        model3res = modelo3Es(text)
        return model3res


with gr.Blocks() as demo:
    gr.Markdown(
        """
        Se quiere hacer un programa que saque un audio de una imagen o de un texto, el cual tiene que ser introducido por el usuario. Para resolver este problema se realiza el siguiente programa. Se van a usar
        tres modelos (Salesforce/blip-image-captioning-large, Helsinki-NLP/opus-mt-en-es, microsoft/speecht5_tts), los cuales se describen a continuación cuál es la función de cada uno:  \n
            - Primero necesitaremos dos Tabs, uno con un input tipo Image (IMAGE) en el que pasaremos una imagen y otro con un input tipo Textbox (TEXT) en el que pasaremos un texto.  \n
            - En el caso de la imagen, la pasaremos a texto usando un modelo con esta función (Salesforce/blip-image-captioning-large). Este modelo está entrenado para sacar texto describiendo qué hay en la
            foto. El modelo nos sacará un texto en inglés.  \n
            - En caso del texto, no hace falta usar el modelo anterior ya que directamente tenemos el texto que queremos pasar a audio.  \n
            - Luego, tenemos un evento de tipo Radio, con el cual podemos elegir el idioma en el que vamos a sacar el audio. En el caso de la imagen, dado que el modelo saca el texto de esta imagen en inglés,
            si hemos seleccionado que queremos sacar el audio en español tendremos que traducir este texto de inglés a español. En el caso del texto se da por hecho que el texto va a ser introducido en el
            mismo idioma que se quiere sacar el audio.  \n
            - Para traducir el texto usaremos un modelo que está entrenado para pasar texto de inglés a español (Helsinki-NLP/opus-mt-en-es), por lo que nos devolverá un texto casi perfectamente traducido al
            español.  \n
            - Una vez tenemos el texto que queremos pasar a audio en el idioma deseado, con el último modelo pasaremos este texto a audio (microsoft/speecht5_tts). Este modelo está entrenado para sacar audio a
            raíz de un texto, en el que
            se escucha justo lo que pone en el texto que le mandamos.  \n
            - Por último tendremos un output de tipo Audio que nos mostrará el audio que hemos conseguido con el último modelo.
      """)

    with gr.Tab("IMAGE"):
        inp = gr.Image(type="pil")
        language = gr.Radio(["english", "spanish"], label="Language", info="Choose the language in which you want the audio to appear", value='english', interactive=True)
        out = gr.Audio()
        btn = gr.Button("RUN")
        btn.click(fn=executionIMG, inputs=[inp, language], outputs=out)

    with gr.Tab("TEXT"):
        inp = gr.Textbox()
        language = gr.Radio(["english", "spanish"], label="Language", info="Choose the language in which you want the audio to appear", value='english', interactive=True)
        out = gr.Audio()
        btn = gr.Button("RUN")
        btn.click(fn=executionTEXT, inputs=[inp, language], outputs=out)

    demo.launch()