import gradio as gr
import requests
from PIL import Image
from io import BytesIO
import os, time

HF_TOKEN = os.environ.get("HF_TOKEN", "")

# Only models available on HF free serverless inference tier
MODELS = {
    "FLUX.1 Schnell (Fast, Free)":   "black-forest-labs/FLUX.1-schnell",
    "Stable Diffusion XL (Free)":    "stabilityai/stable-diffusion-xl-base-1.0",
    "Stable Diffusion 2.1 (Free)":   "stabilityai/stable-diffusion-2-1",
    "OpenJourney (Free)":            "prompthero/openjourney",
    "Deliberate v2 (Free)":          "XpucT/Deliberate",
}

def generate(prompt, model_name, steps=20):
    if not prompt.strip():
        return None, "Enter a prompt"
    if not HF_TOKEN:
        return None, "HF_TOKEN secret not set on this Space"

    model_id = MODELS.get(model_name, "black-forest-labs/FLUX.1-schnell")
    t0 = time.time()

    payload = {"inputs": prompt}
    if "FLUX" not in model_name:
        payload["parameters"] = {"num_inference_steps": steps}

    for attempt in range(3):
        try:
            r = requests.post(
                f"https://api-inference.huggingface.co/models/{model_id}",
                headers={"Authorization": f"Bearer {HF_TOKEN}"},
                json=payload, timeout=120
            )
            if r.status_code == 503:
                est = r.json().get("estimated_time", 20)
                time.sleep(min(est, 30))
                continue
            if r.status_code == 200 and "image" in r.headers.get("Content-Type", ""):
                img = Image.open(BytesIO(r.content))
                return img, f"✅ {model_name} | {int(time.time()-t0)}s"
            err = ""
            try: err = r.json().get("error", r.text[:100])
            except: err = r.text[:100]
            return None, f"❌ HTTP {r.status_code}: {err}"
        except Exception as e:
            if attempt == 2:
                return None, f"❌ {str(e)[:100]}"
            time.sleep(5)

    return None, "❌ All attempts failed"

with gr.Blocks(title="HF Inference - Free Models") as demo:
    gr.Markdown("""# 🖼️ HuggingFace Free Inference
These models are available on the **free** HF Inference API tier.
*First call may take 30s while the model loads.*""")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Describe your image...")
            model = gr.Dropdown(list(MODELS.keys()),
                               value="FLUX.1 Schnell (Fast, Free)", label="Model")
            steps = gr.Slider(10, 50, 20, label="Steps (ignored for FLUX)")
            btn = gr.Button("🎨 Generate", variant="primary")
            status = gr.Textbox(label="Status", interactive=False)
        with gr.Column():
            out = gr.Image(label="Result", type="pil", height=500)
    btn.click(generate, inputs=[prompt, model, steps], outputs=[out, status])
    prompt.submit(generate, inputs=[prompt, model, steps], outputs=[out, status])

demo.launch()