VaneM's picture
Upload 2 files
b058173
raw
history blame
1.44 kB
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)