import os
import gradio as gr
from huggingface_hub import InferenceClient

# 1. 讀取 Token 並檢查
HF_TOKEN = os.getenv("HF_TOKEN")
client = InferenceClient(api_key=HF_TOKEN)

def generate_image(prompt, model_id):
    if not prompt:
        return None, "⚠️ 請輸入描述文字！"
    if not model_id:
        return None, "⚠️ 請輸入模型 ID！"
    if not HF_TOKEN:
        return None, "❌ 錯誤：請在 Space Settings > Secrets 設定 HF_TOKEN。"

    try:
        image = client.text_to_image(
            prompt, 
            model=model_id.strip()
        )
        return image, f"✅ 成功！使用模型：{model_id}"
    except Exception as e:
        return None, f"❌ 發生錯誤：{str(e)}"

# 2. 建立介面
with gr.Blocks() as demo:
    gr.Markdown("# 🎨 AI 多模型生圖工具 (穩定版)")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="1. 描述你想看到的畫面", 
                placeholder="例如：一個穿著太空衣的貓...",
                lines=3
            )
            
            model_input = gr.Textbox(
                label="2. 模型 ID (可手動修改)", 
                value="black-forest-labs/FLUX.1-schnell"
            )
            
            model_presets = gr.Radio(
                choices=[
                    "black-forest-labs/FLUX.1-schnell", 
                    "stabilityai/stable-diffusion-3.5-large",
                    "digiplay/2K"
                ],
                label="常用模型快捷選擇",
                value="black-forest-labs/FLUX.1-schnell"
            )
            model_presets.change(fn=lambda x: x, inputs=model_presets, outputs=model_input)

            run_btn = gr.Button("立即生成 ✨", variant="primary")
            status_text = gr.Markdown("準備就緒")
        
        with gr.Column():
            # 💡 核心修改：固定高度且不顯示標籤，防止畫面閃爍跳動
            output_img = gr.Image(
                label="生成結果", 
                type="pil", 
                height=512,       # 固定高度，防止版面撐開
                show_label=False, # 減少渲染零件
                interactive=False
            )

    # 💡 核心修改：show_progress="minimal" 讓載入時不變暗
    run_btn.click(
        fn=generate_image, 
        inputs=[input_text, model_input], 
        outputs=[output_img, status_text],
        show_progress="minimal" 
    )

# 3. 啟動介面
demo.launch(theme=gr.themes.Soft())