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

device = "cuda" if torch.cuda.is_available() else "cpu"

model_id = "OFA-Sys/small-stable-diffusion-v0"

try:
    pipe = StableDiffusionPipeline.from_pretrained(model_id)
    pipe = pipe.to(device)
    print(f"Modelo cargado y movido a: {device.upper()}")
except Exception as e:
    print(f"Error al cargar el modelo: {e}")
    pipe = None


def generar_imagen(prompt):
    """Genera una imagen a partir de un prompt usando Stable Diffusion."""
    if not prompt:
        return None
    if pipe is None:
        return Image.new('RGB', (512, 512), color = 'red') 
    
    try:
        image = pipe(prompt).images[0] 
        return image
    except Exception as e:
        print(f"Error al generar la imagen: {e}")
        return Image.new('RGB', (512, 512), color = 'red')


with gr.Blocks(title="Generador de Imágenes Ligero") as demo:
    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Prompt (Describe la imagen que quieres)", 
                placeholder="Un astronauta montando a caballo, estilo fotorealista"
            )
            generate_button = gr.Button("🖼️ Generar Imagen")
        
        with gr.Column(scale=1):
            image_output = gr.Image(
                type="pil",
                label="Imagen Generada",
                height=512,
                width=512
            )
    
    generate_button.click(
        fn=generar_imagen, 
        inputs=prompt_input, 
        outputs=image_output
    )


demo.launch(inbrowser=True)