from diffusers import StableDiffusionPipeline,  EulerDiscreteScheduler
from transformers import pipeline
import torch
from IPython.display import clear_output
import gradio as gr

model_name = "stabilityai/stable-diffusion-2"
scheduler = EulerDiscreteScheduler.from_pretrained(model_name, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_name, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16)

# Movemos el pipeline a GPU para tener una inferencia más rápida.
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipe.to(device)
pipe.enable_attention_slicing() 

# Cargamos el modelo traductor y creamos el pipeline para traducir al ingles
spanish_model_name = "Helsinki-NLP/opus-mt-es-en"
translator_es_en = pipeline("translation", model=spanish_model_name)

def predict(text):
  prompt_es = text
  english_text = translator_es_en(prompt_es)
  prompt_en = english_text[0]['translation_text']
  image = pipe(prompt_en).images[0]
  return image

description = """
<h2 style="text-align:center">Programa para generar imágenes a partir de texto con Stable Diffusion 2</h2>
<h2 style="text-align:center">Prompt en español!!</h2>
"""

gr.Interface(fn=predict,
              title="Texto a Imagen en Español",
              inputs= gr.Textbox("", max_lines = 2, label = "Inserte su texto aqui"),
              outputs = "image",
              description = description).launch(debug=True)