import gradio as gr
import os
from huggingface_hub import InferenceClient
from PIL import Image
from io import BytesIO

# Modelo a utilizar
MODEL_NAME = "black-forest-labs/FLUX.1-schnell"

# Función principal de generación de imágenes
def generate_image(hf_token: str, prompt: str, negative_prompt: str, size: str) -> Image.Image:
    """
    Genera una imagen usando el modelo FLUX.1-schnell a través de InferenceClient.
    """
    if not hf_token:
        raise gr.Error("Por favor, introduce tu HF Token para autenticarte.")
    
    if not prompt:
        raise gr.Error("El prompt no puede estar vacío.")

    try:
        # Inicializar el cliente con el token proporcionado por el usuario
        client = InferenceClient(
            provider="hf-inference",
            api_key=hf_token,
        )
        
        # Parsear el tamaño de la imagen
        # El formato es "Ancho x Alto (Descripción)", por lo que solo necesitamos la primera parte
        width, height = map(int, size.split(' ')[0].split('x'))

        # Parámetros de generación
        params = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "width": width,
            "height": height,
            "guidance_scale": 7.0,
            "num_inference_steps": 50,
        }

        # Llamada a la API de generación de imágenes
        image = client.text_to_image(
            model=MODEL_NAME,
            **params
        )
        
        return image

    except Exception as e:
        # Manejo de errores de la API o del cliente
        error_message = f"Error al generar la imagen: {e}"
        if "Invalid token" in str(e) or "Unauthorized" in str(e):
            error_message = "Error de autenticación: El HF Token proporcionado no es válido o no tiene permisos."
        elif "Model not found" in str(e):
            error_message = f"Error: El modelo {MODEL_NAME} no se encontró o no está disponible."
        
        raise gr.Error(error_message)

# Opciones de tamaño de imagen
SIZE_OPTIONS = [
    "1024x1024 (Cuadrado)",
    "1152x896 (Horizontal)",
    "896x1152 (Vertical)",
    "1216x832 (Horizontal)",
    "832x1216 (Vertical)",
    "1344x768 (Horizontal)",
    "768x1344 (Vertical)",
    "1536x640 (Horizontal)",
    "640x1536 (Vertical)",
]

# Interfaz Gradio
with gr.Blocks(title="FLUX.1-schnell Image Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"# ⚡ Generador de Imágenes con {MODEL_NAME}")
    gr.Markdown("Utiliza tu HF Token para generar imágenes de alta calidad con el modelo FLUX.1-schnell.")

    with gr.Row():
        with gr.Column(scale=1):
            # Entrada del HF Token
            hf_token_input = gr.Textbox(
                label="🔑 Tu HF Token (con permisos de lectura)",
                type="password",
                placeholder="Ingresa tu token de Hugging Face aquí...",
                info="Necesario para autenticar la llamada a la API de inferencia."
            )
            
            # Entrada del Prompt
            prompt_input = gr.Textbox(
                label="📝 Prompt (Descripción de la imagen)",
                lines=3,
                placeholder="Ej: Un astronauta montando a caballo en un planeta alienígena, estilo fotorrealista, 4k."
            )
            
            # Entrada del Prompt Negativo
            negative_prompt_input = gr.Textbox(
                label="🚫 Prompt Negativo (Cosas a evitar)",
                lines=2,
                placeholder="Ej: baja calidad, borroso, deforme, extra dedos, texto."
            )
            
            # Selector de Tamaño
            size_selector = gr.Dropdown(
                label="📐 Tamaño de la Imagen (Resolución)",
                choices=SIZE_OPTIONS,
                value=SIZE_OPTIONS[0],
                interactive=True
            )
            
            # Botones de acción
            with gr.Row():
                generate_btn = gr.Button("✨ Generar Imagen", variant="primary")
                # Botón para limpiar: Usa gr.ClearButton
                clear_btn = gr.ClearButton(
                    value="🗑️ Limpiar Entradas",
                    components=[hf_token_input, prompt_input, negative_prompt_input, size_selector]
                )

        with gr.Column(scale=2):
            # Salida de la Imagen
            image_output = gr.Image(
                label="🖼️ Imagen Generada",
                type="pil",
                show_download_button=True, # Habilita el botón de descarga
                height=512 # Altura inicial para la visualización
            )

    # Enlazar la función al botón de generación
    generate_btn.click(
        fn=generate_image,
        inputs=[hf_token_input, prompt_input, negative_prompt_input, size_selector],
        outputs=[image_output]
    )

# Ejecutar la aplicación
if __name__ == "__main__":
    demo.launch()