import gradio as gr
from diffusers import DiffusionPipeline,StableDiffusionPipeline
#from googletrans import Translator
from transformers import pipeline

es_en_translator = pipeline("translation",model = "Helsinki-NLP/opus-mt-es-en")


def translate_text(text):
    text = es_en_translator(text)[0].get("translation_text")
    return text

# pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",negative_prompt="bad quality")
import base64

with open("Iso_Logotipo_Ceibal.png", "rb") as image_file:
    encoded_image = base64.b64encode(image_file.read()).decode()


pipeline_image = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5"
)


def get_completion(prompt,steps,guidance,width,height,negative_prompt):
#    translator = Translator()
#    prompt=translator.translate(prompt, src='es',dest="en").text
#    print(prompt)
#    negative_prompt = translator.translate(negative_prompt, src='es',dest="en").text
    print(prompt)
    prompt = translate_text(prompt)
    negative_prompt = translate_text(negative_prompt)

    print(prompt)
    return pipeline_image(prompt,
                    num_inference_steps=steps,
                    height=height,
                    width=width,
                    guidance_scale = guidance,
                    negative_prompt = negative_prompt).images[0]    


with gr.Blocks() as demo:
    gr.Markdown("""
    <center>
    <h1>
    Uso de AI para la generación de imagenes a partir de texto.
    </h1>
    <img src='data:image/jpg;base64,{}' width=200px>
    <h3>
    Con este espacio podrás generar imagenes a partir de texto. Sepuede usar directamente o podés cambiar ajustes, que impacto tiene cada uno está detallado en su descripción. Cambiá valores y mirá los resultados!
    </h3>
    <h4>El texto se traduce del español al inglés para alimentar al modelo, también se puede escribir el texto de entrada en inglés.</h4>
    </center>
    """.format(encoded_image))
    with gr.Row():
        with gr.Column():
            gr.Markdown("Primero debes ingresar el texto para generar la imagen:")
            with gr.Row():
                with gr.Column(scale=4):
                    prompt = gr.Textbox(label="Texo base para generar la imagen") #Give prompt some real estate
                with gr.Column(scale=1, min_width=50):
                    btn = gr.Button("Generar") #Submit button side by side!
            with gr.Row():
                with gr.Accordion("Opciones avanzadas", open=False): #Let's hide the advanced options!
                        negative_prompt = gr.Textbox(label="Texto negativo para la generación", info='Al ingresar texto en este campo el modelo intentará alejarse lo mas posible del mismo, este puede ser "baja calidad"')
                        with gr.Row():
                            with gr.Column():
                                steps = gr.Slider(label="Pasos de inferencia", minimum=1, maximum=100, value=10, step = 1,
                                info="Cuánto mayor sea la cantidad de pasos de inferencia el resultado será mejor y mas nitido pero el tiempo de procesamiento será mayor (usar números enteros).")
                                guidance = gr.Slider(label="Guidance Scale", minimum=1, maximum=20, value=7,
                                info="Controla cuánto el texto de entrada influye en la generación de la imagen.")
                            with gr.Column():
                                width = gr.Slider(label="Ancho", minimum=64, maximum=256, step=64, value=512)
                                height = gr.Slider(label="Alto", minimum=64, maximum=256, step=64, value=512)
        with gr.Column():
            output = gr.Image(label="Resultado",height=512, width=512) #Move the output up too
            examples = gr.Examples(inputs=[prompt,negative_prompt],examples=[["Un perro en el parque","low quality"]])
            
    btn.click(fn=get_completion, inputs=[prompt,steps,guidance,width,height,negative_prompt], outputs=[output])  #steps,guidance,width,height]

gr.close_all()
demo.launch()