import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# Cargar el modelo de Stable Diffusion desde Hugging Face
model_id = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cuda")  # Cambiar a CPU si no hay GPU disponible

# Función para generar la imagen
def generar_imagen(texto):
    with torch.autocast("cuda"):
        image = pipe(texto, guidance_scale=7.5).images[0]
    return image

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

# Ejecuta la aplicación
if __name__ == "__main__":
    interfaz.launch()