import gradio as gr
#from transformers import pipeline

import torch
from diffusers import DiffusionPipeline

import platform

inCloud = True

#translator = pipeline("translation", model = "Helsinki-NLP/opus-mt-es-en")
pipe = ''

if inCloud:
    pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
    if platform.system() == "Darwin": #Apple OS
        pipe = pipe.to("mps")
    #pipe = pipe.to("cuda")

    # Recommended if your computer has < 64 GB of RAM
    pipe.enable_attention_slicing()

textos = [
    '''Érase una vez un hombre muy rico que vendió todo lo que tenía a cambio de varios lingotes de oro. Y para que nadie le robara, enterró el oro en un bosque. Todos los días acudía al lugar para comprobar que su oro seguía allí, sin saber que un ladrón lo vigilaba escondido.

Una noche, el ladrón desenterró el oro y se lo llevó. Cuando el rico descubrió el robo, dio tal grito que un vecino se acercó a ver qué pasaba. El hombre rico lloraba, desesperado. Entonces el vecino tomó unas piedras, las enterró en el mismo lugar y dijo:

—Aquí tiene su tesoro. Sabe que nunca habría gastado sus lingotes. ¿Qué más le da, entonces, que sean piedras? Así por lo menos dejará de sufrir.'''
]

#def translate(text):
#    return translator(text)[0]["translation_text"]

def generate_texts_array(text, max_length = 70):
    text_array = []
    start = 0

    while start < len(text):
        end = start + max_length
        part = text[start:end]
        text_array.append(part)
        start = end

    return text_array

def generation(text):
    images = []
    steps = 5
    #translate_text = translate(text)
    text_array = generate_texts_array(text)
    images.append(pipe(text_array[0],num_inference_steps=steps).images[0])
    images.append(pipe(text_array[1],num_inference_steps=steps).images[0])
    images.append(pipe(text_array[2],num_inference_steps=steps).images[0])
    images.append(pipe(text_array[3],num_inference_steps=steps).images[0])
    images.append(pipe(text_array[4],num_inference_steps=steps).images[0])
    return images

def generation_test(img):
    images = []
    images.append(img)
    images.append(img)
    images.append(img)
    images.append(img)
    images.append(img)
    return images

demo = gr.Blocks()

title = '# Generando historias graficas '
description = 'Escribe un texto y se generara una secuencia de imagenes basadas en el texto proporcionado'

with demo:
    gr.Markdown(title)
    gr.Markdown(description)
    
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox ( label="Cuento",
                                    info="Escribe una historia",
                                    lines=5, 
                                    value='Erase una vez..')
            #steps_input = gr.Number(value=20)
            #img_input = gr.Image (type='pil')
            button = gr.Button(value="Generate")

            gr.Examples(
                examples=textos,
                inputs=text_input,
                #outputs=im_2,
            )
            
        
    with gr.Row():
        imgs_output = []
        imgs_output.append(gr.Image())
        imgs_output.append(gr.Image())
        imgs_output.append(gr.Image())
        imgs_output.append(gr.Image())
        imgs_output.append(gr.Image())
    
    button.click(generation, inputs=text_input, outputs=imgs_output)
        


if __name__ == "__main__":
    demo.launch()