import gradio as gr
import requests
import torch
from PIL import Image
from io import BytesIO
import os

# ---------- 配置 ----------
HF_TOKEN = os.getenv("HF_TOKEN")  # 从 Secrets 中读取
API_URLS = {
    "flux": "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev",
    "sd3": "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large",
    "sdxl": "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0",
}
DEFAULT_API = "flux"

# 轻量 CPU 模型（仅用于演示，质量较低）
LIGHTWEIGHT_MODEL_ID = "OFA-Sys/small-stable-diffusion-v0"  # 极小，CPU 可跑

# GPU 模型（需要 GPU 硬件）
GPU_MODEL_ID = "runwayml/stable-diffusion-v1-5"

# 检测是否有 GPU
has_gpu = torch.cuda.is_available()

# ---------- 推理函数 ----------
def generate_api(prompt, api_choice):
    """调用 Hugging Face Inference API"""
    if not HF_TOKEN:
        return None, "❌ 未设置 HF_TOKEN，请在 Space Secrets 中添加。"
    url = API_URLS.get(api_choice, API_URLS[DEFAULT_API])
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    payload = {"inputs": prompt, "parameters": {"negative_prompt": "ugly, blurry"}}
    try:
        response = requests.post(url, headers=headers, json=payload, timeout=60)
        if response.status_code == 200:
            img = Image.open(BytesIO(response.content))
            return img, "✅ API 生成成功"
        else:
            return None, f"❌ API 错误：{response.status_code} - {response.text}"
    except Exception as e:
        return None, f"❌ 请求异常：{e}"

def generate_cpu_lightweight(prompt):
    """CPU 轻量模型（使用 diffusers，强制 CPU）"""
    try:
        from diffusers import StableDiffusionPipeline
        import torch
        # 加载极轻量模型（首次会下载，约 500MB）
        pipe = StableDiffusionPipeline.from_pretrained(
            LIGHTWEIGHT_MODEL_ID,
            torch_dtype=torch.float32
        )
        pipe = pipe.to("cpu")
        # 生成
        image = pipe(prompt, num_inference_steps=20).images[0]
        return image, "✅ CPU 轻量模型生成成功（速度较慢）"
    except Exception as e:
        return None, f"❌ CPU 模型错误：{e}"

def generate_gpu(prompt):
    """本地 GPU 模型（需要 T4 或更高）"""
    if not has_gpu:
        return None, "❌ 未检测到 GPU，请为 Space 分配 T4 small 硬件。"
    try:
        from diffusers import StableDiffusionPipeline
        pipe = StableDiffusionPipeline.from_pretrained(
            GPU_MODEL_ID,
            torch_dtype=torch.float16
        ).to("cuda")
        image = pipe(prompt, num_inference_steps=25).images[0]
        return image, "✅ GPU 模型生成成功（快速）"
    except Exception as e:
        return None, f"❌ GPU 模型错误：{e}"

# ---------- Gradio 界面 ----------
def inference(prompt, mode, api_choice):
    if not prompt.strip():
        return None, "请输入提示词"
    if mode == "官方推理 API":
        return generate_api(prompt, api_choice)
    elif mode == "CPU 轻量模型":
        return generate_cpu_lightweight(prompt)
    elif mode == "本地 GPU 模型":
        return generate_gpu(prompt)
    else:
        return None, "未知模式"

# 界面布局
with gr.Blocks(title="三合一文生图 AI", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎨 三合一文生图 AI")
    gr.Markdown("支持三种模式：官方推理 API（快，有额度）、CPU 轻量模型（慢，无限）、本地 GPU 模型（快，需 T4）")
    
    with gr.Row():
        with gr.Column(scale=3):
            prompt_input = gr.Textbox(label="提示词 (Prompt)", placeholder="A beautiful sunset over a mountain range", lines=3)
            with gr.Row():
                mode_radio = gr.Radio(
                    choices=["官方推理 API", "CPU 轻量模型", "本地 GPU 模型"],
                    label="选择生成模式",
                    value="官方推理 API"
                )
                api_choice = gr.Dropdown(
                    choices=list(API_URLS.keys()),
                    label="API 模型选择（仅官方API模式）",
                    value=DEFAULT_API,
                    visible=True
                )
            generate_btn = gr.Button("生成图片", variant="primary")
        with gr.Column(scale=2):
            output_image = gr.Image(label="生成结果", type="pil")
            output_status = gr.Textbox(label="状态信息", lines=2)
    
    # 动态显示/隐藏 API 模型下拉框
    def update_visibility(mode):
        if mode == "官方推理 API":
            return gr.update(visible=True)
        else:
            return gr.update(visible=False)
    mode_radio.change(fn=update_visibility, inputs=mode_radio, outputs=api_choice)
    
    generate_btn.click(
        fn=inference,
        inputs=[prompt_input, mode_radio, api_choice],
        outputs=[output_image, output_status]
    )
    
    gr.Markdown("""
    ### 📌 注意事项
    - **官方推理 API**：需要设置 `HF_TOKEN`（在 Space Secrets 中添加），免费但每日有调用限制。
    - **CPU 轻量模型**：不需要 GPU，但生成一张图约 3-10 分钟，质量较低。
    - **本地 GPU 模型**：需要为 Space 分配 **T4 small** 硬件（免费，但会休眠），速度快，质量高。
    - 建议首选 **官方推理 API** 或 **本地 GPU 模型**。
    """)

demo.launch(server_port=7860)