import gradio as gr
from huggingface_hub import InferenceClient

# Text-to-image models currently served by Hugging Face Inference Providers.
# (The original list pointed at the retired api-inference.huggingface.co endpoint.)
DRAW_MODELS = {
    "FLUX.1-schnell (fast)": "black-forest-labs/FLUX.1-schnell",
    "FLUX.1-dev": "black-forest-labs/FLUX.1-dev",
    "Stable Diffusion XL 1.0": "stabilityai/stable-diffusion-xl-base-1.0",
    "Stable Diffusion 3.5 Large": "stabilityai/stable-diffusion-3.5-large",
    "Stable Diffusion 3.5 Medium": "stabilityai/stable-diffusion-3.5-medium",
    "Qwen-Image": "Qwen/Qwen-Image",
    "Z-Image-Turbo": "Tongyi-MAI/Z-Image-Turbo",
    "HiDream-I1-Fast": "HiDream-ai/HiDream-I1-Fast",
}
DEFAULT_MODEL = "FLUX.1-schnell (fast)"


def draw(prompt, model_label, oauth_token: gr.OAuthToken | None = None):
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt.")
    if oauth_token is None:
        raise gr.Error("Sign in with Hugging Face (button in the sidebar) to generate images.")
    client = InferenceClient(model=DRAW_MODELS[model_label], token=oauth_token.token)
    try:
        return client.text_to_image(prompt)
    except Exception as e:
        raise gr.Error(f"{model_label} failed: {e}")


with gr.Blocks(title="DrawApp") as demo:
    with gr.Sidebar():
        gr.LoginButton()
        model = gr.Dropdown(list(DRAW_MODELS), value=DEFAULT_MODEL, label="Draw Models")
    gr.Markdown("# 🎨 DrawApp")
    image = gr.Image(label="Result", format="png")
    prompt = gr.Textbox(placeholder="Send your prompt", show_label=False, submit_btn=True)
    prompt.submit(draw, [prompt, model], image, api_name="draw")

demo.launch()