import os
from huggingface_hub import InferenceClient
import gradio as gr
from dotenv import load_dotenv

# Cargar variables del .env
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")

if HF_TOKEN is None:
    raise ValueError("❌ No se encontró la variable HF_TOKEN en el .env")

# Inicializar cliente de inferencia HF
client = InferenceClient(provider="hf-inference", api_key=HF_TOKEN)

def generar_imagen(prompt):
    if not prompt.strip():
        return None

    # Llamada a la API
    imagen = client.text_to_image(prompt, model="black-forest-labs/FLUX.1-schnell")
    
    return imagen  # es un objeto PIL.Image

# Interfaz Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🎨 Generador de Imágenes con FLUX.1-schnell")
    gr.Markdown("Escribe un prompt y el modelo generará la imagen usando HF Inference API gratuita.")

    prompt_input = gr.Textbox(label="Escribe tu prompt aquí", lines=3)
    boton = gr.Button("Generar Imagen")
    imagen_output = gr.Image(label="Imagen Generada")

    boton.click(fn=generar_imagen, inputs=prompt_input, outputs=imagen_output)

demo.launch()