import gradio as gr
from diffusers import StableDiffusionPipeline
from PIL import Image
import torch


#Modelo que genera imágenes, está poco probado ya que tarda muchísimo en generarlas
model_id = "OFA-Sys/small-stable-diffusion-v0"
pipe = StableDiffusionPipeline.from_pretrained(model_id)
pipe = pipe.to("cpu") 

# Función que Gradio llamará
def generar_imagen(prompt):
    if not prompt:
        return None
    image = pipe(prompt).images[0]
    return image 

# Interfaz Gradio
demo = gr.Interface(
    fn=generar_imagen,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.Image(type="pil"),
    title="Generador de imágenes ligero",
    description="Introduce un prompt y el modelo generará la imagen. Podrás descargarla con el botón de Gradio."
)

demo.launch()