import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# Configurar el dispositivo (usar GPU si está disponible)
device = "cuda" if torch.cuda.is_available() else "cpu"

# Cargar el modelo Stable Diffusion
pipeline = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",  # El nombre del modelo de Stable Diffusion
    torch_dtype=torch.float16  # Cambia a float32 si no estás usando una GPU
).to(device)

# Función para generar imágenes
def generar_imagen(descripcion):
    imagen = pipeline(descripcion).images[0]
    return imagen

# Crear la interfaz de Gradio
interfaz = gr.Interface(
    fn=generar_imagen,
    inputs="text",
    outputs="image",
    title="Generador de Imágenes con IA",
    description="Introduce una descripción y genera una imagen usando Stable Diffusion."
)

# Lanzar la interfaz
interfaz.launch()