import gradio as gr
from huggingface_hub import InferenceClient
import os

# Usa tu token de Hugging Face (guárdalo como variable de entorno HF_TOKEN)
HF_TOKEN = os.getenv("HF_TOKEN")

# Cliente de inferencia
client = InferenceClient(api_key=HF_TOKEN)

def generate_image(prompt, steps, guidance):
    """
    Genera una imagen con Flux.1 Schnell usando Hugging Face Inference API.
    """
    image = client.text_to_image(
        model="black-forest-labs/flux-1-schnell",  # Modelo Flux
        inputs=prompt,
        parameters={
            "num_inference_steps": steps,
            "guidance_scale": guidance
        }
    )
    return image

# Interfaz con Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 Generador de imágenes con Flux.1 Schnell")
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", value="Astronauta montando un caballo")
    with gr.Row():
        steps = gr.Slider(1, 50, value=5, step=1, label="Pasos de inferencia")
        guidance = gr.Slider(1, 20, value=7, step=1, label="Guidance scale")
    output = gr.Image(type="pil")
    btn = gr.Button("Generar imagen")
    btn.click(generate_image, [prompt, steps, guidance], output)

demo.launch()