import gradio as gr
import torch
from diffusers import DiffusionPipeline

# --- 1. CONFIGURACIÓN DEL MODELO (LCM para velocidad) ---
print("🚀 Cargando modelo de difusión rápida (LCM)...")
model_id = "SimianLuo/LCM_Dreamshaper_v7"

try:
    pipe = DiffusionPipeline.from_pretrained(model_id)
    
    pipe.to("cpu")
    
    pipe.safety_checker = None 
    
    print("✅ Modelo LCM cargado. Listo para generar imágenes rápidas.")
except Exception as e:
    print(f"❌ Error al cargar modelo: {e}")
    pipe = None

# --- 2. LÓGICA DE GENERACIÓN ---
def generar_pictograma(texto_usuario, estilo):
    if pipe is None:
        return None
    
    prompt_base = f"{texto_usuario}, {estilo}, white background, centered, vector illustration, flat design, high contrast, minimal, simple shapes, no shadows"
    prompt_negativo = "complex, realistic, photo, details, text, watermark, blurry, distorted, 3d render"

    print(f"🎨 Generando: {texto_usuario}...")
    
    try:

        imagen = pipe(
            prompt=prompt_base,
            negative_prompt=prompt_negativo,
            num_inference_steps=4, 
            guidance_scale=8.0,
            width=512, 
            height=512
        ).images[0]
        return imagen
    except Exception as e:
        print(f"Error generando: {e}")
        return None

# --- 3. INTERFAZ GRÁFICA ---
css_accesible = """
.gradio-container {background-color: #f0f8ff !important;}
h1 {font-family: 'Arial', sans-serif; color: #2c3e50;}
.boton-grande {font-size: 20px !important; height: 60px !important;}
.texto-grande textarea {font-size: 24px !important; line-height: 1.5 !important;}
.info-box {background-color: #e8f4f8; padding: 15px; border-radius: 10px; border: 2px solid #3498db;}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=css_accesible) as demo:
    
    with gr.Row():
        gr.Markdown("# 🗣️ PictoGen: Comunicación Visual")
    
    gr.Markdown(
        """
        <div class='info-box'>
        <b>Herramienta de Accesibilidad:</b> Escribe lo que necesitas comunicar y la IA creará una tarjeta visual simple al instante.
        Ideal para crear tableros de comunicación, agendas visuales o expresar necesidades.
        </div>
        """
    )

    with gr.Row():
        # Columna Izquierda: Controles
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                label="¿Qué imagen necesitas?", 
                placeholder="Ej: Comer manzana, Ir al baño, Jugar pelota...",
                elem_classes="texto-grande"
            )
            
            radio_style = gr.Radio(
                ["icon style, thick lines", "cartoon style, colorful", "line art, black and white"],
                label="Estilo Visual",
                value="icon style, thick lines",
                info="Selecciona el tipo de dibujo que mejor se entienda."
            )
            
            btn = gr.Button("🎨 Crear Pictograma", variant="primary", elem_classes="boton-grande")

        # Columna Derecha: Resultado
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Tu Pictograma", 
                type="pil", 
                interactive=False,
                height=512
            )


    btn.click(fn=generar_pictograma, inputs=[input_text, radio_style], outputs=output_image)

    gr.Examples(
        examples=[
            ["Un vaso de agua", "icon style, thick lines"],
            ["Sentirse feliz", "cartoon style, colorful"],
            ["Autobús escolar", "icon style, thick lines"]
        ],
        inputs=[input_text, radio_style]
    )

if __name__ == "__main__":
    demo.launch()