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

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")

client = InferenceClient(provider="hf-inference", api_key=HF_TOKEN)

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

    imagen = client.text_to_image(prompt, model="black-forest-labs/FLUX.1-schnell")
    
    return imagen 

with gr.Blocks() as demo:
    gr.Markdown("#Generador de Imágenes")
    gr.Markdown("Escribe un prompt y el modelo generará la imagen usando Inference.")

    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()