import os
import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image
import traceback

# --- 核心函数：调用HuggingFace推理API ---
def run_image_to_image(img: Image.Image, prompt: str, model: str, strength: float, guidance_scale: float) -> Image.Image:
    """
    使用HuggingFace InferenceClient运行image-to-image任务。

    参数:
    - img: 输入的PIL Image对象。
    - prompt: 描述图像变化的文本提示词。
    - model: 要使用的HuggingFace模型库ID。
    - strength: 控制对原始图像的修改程度 (0.0到1.0)。
    - guidance_scale: 控制提示词与输出图像的匹配程度。
    """
    token = os.getenv("HF_TOKEN")
    if not token:
        raise gr.Error("错误：环境变量 HF_TOKEN 未设置。请在HuggingFace Space的Secrets中设置您的Hugging Face Token。")

    # 在每次调用时，根据用户选择的模型ID创建客户端
    client = InferenceClient(model=model, token=token)

    try:
        # 直接将PIL图像和参数传递给客户端
        # 新版库支持直接传入PIL Image对象
        result_img = client.image_to_image(
            image=img,
            prompt=prompt,
            strength=strength,
            guidance_scale=guidance_scale,
        )
        return result_img
    except Exception as e:
        traceback.print_exc()
        # 将详细错误反馈给用户
        raise gr.Error(f"模型推理失败，请检查模型是否支持Image-to-Image任务或更换模型尝试。错误详情: {e}")


# --- Gradio Blocks 界面 ---
with gr.Blocks(title="图像风格转换", css=".gradio-container {max-width: 1080px}") as demo:
    gr.Markdown("# 图像风格转换工具 🎨")
    gr.Markdown("上传一张图片，输入提示词，选择一个模型，然后点击“一键转换”来生成新的图片。")

    with gr.Row(variant="panel"):
        with gr.Column(scale=1):
            in_img = gr.Image(type="pil", label="🖼️ 上传图片", height=512)
            prompt = gr.Textbox(
                label="📝 提示词",
                value="A beautiful painting of a castle in a forest, fantasy, detailed, epic",
                lines=3,
                placeholder="输入你想要的画面效果，例如：'A cute cat wearing a wizard hat'"
            )
            # 为了保证成功率，我们选择一个已知可靠的免费模型
            # 您的原始模型列表可以作为备选，但需确认它们在HF上部署了免费的推理API
            model = gr.Dropdown(
                choices=[
                    "stabilityai/stable-diffusion-xl-refiner-1.0",
                    "runwayml/stable-diffusion-v1-5",
                    "SG161222/Realistic_Vision_V5.1_noVAE"
                ],
                value="stabilityai/stable-diffusion-xl-refiner-1.0",
                label="🤖 选择模型",
                info="这些是经过验证的免费模型，主要支持英文提示词。"
            )
            with gr.Accordion("高级参数 (选填)", open=False):
                strength = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, step=0.1, label="修改强度", info="值越高，图片变化越大")
                guidance_scale = gr.Slider(minimum=0.0, maximum=20.0, value=7.5, step=0.5, label="引导强度", info="值越高，图片越接近提示词")

            with gr.Row():
                run_btn = gr.Button("🚀 一键转换", variant="primary", scale=2)

        with gr.Column(scale=1):
            out_img = gr.Image(type="pil", label="✨ 输出结果", height=512, interactive=False)
            gr.Markdown("### 💡 提示词示例 (Examples)")
            gr.Examples(
                examples=[
                    ["A beautiful oil painting of a landscape."],
                    ["photo of a cute white cat, cinematic, 8k"],
                    ["Make it in anime style, detailed, beautiful."],
                    ["cyberpunk city, neon lights, rainy night"]
                ],
                inputs=prompt,
                label="点击示例以填充提示词"
            )

    # --- 按钮点击事件 ---
    run_btn.click(
        fn=run_image_to_image,
        inputs=[in_img, prompt, model, strength, guidance_scale],
        outputs=out_img
    )

# --- 启动应用 ---
if __name__ == "__main__":
    demo.launch(share=True)