import gradio as gr
from gradio_client import Client, handle_file
import random
import os

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

space_client = Client(
    "prithivMLmods/Qwen-Image-Edit-2511-LoRAs-Fast",
    token=HF_TOKEN  #不是hf_token
    )

LORA_STYLES = [
    'Multiple-Angles', 'Photo-to-Anime', 'Anime-V2', 'Light-Migration',
    'Upscaler', 'Style-Transfer', 'Manga-Tone', 'Anything2Real',
    'Fal-Multiple-Angles', 'Polaroid-Photo', 'Unblur-Anything',
    'Midnight-Noir-Eyes-Spotlight', 'Hyper-Realistic-Portrait',
    'Ultra-Realistic-Portrait', 'Pixar-Inspired-3D', 'Noir-Comic-Book',
    'Any-light', 'Studio-DeLight', 'Cinematic-FlatLog',
]
MAX_SEED = 2**31 - 1


def infer(
    image,
    prompt,
    lora_adapter,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    progress=gr.Progress(track_tqdm=True),
):
    if image is None:
        print("未上传图片")
        return None, seed

    if not os.path.exists(image):
        print(f"图片路径不存在: {image}")
        return None, seed

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # 关键：用 handle_file 上传到目标 Space，得到远端可访问的文件对象
    uploaded = handle_file(image)

    # 补全所有必要字段
    # uploaded["url"] = None
    # uploaded["size"] = os.path.getsize(image)
    # uploaded["mime_type"] = "image/jpeg"
    # uploaded["is_stream"] = False

    #images_input = [{"image": uploaded, "caption": None}]

    # Gallery 元素格式：{"image": <上传后的文件对象>, "caption": None}
    images_input = [{"image": uploaded, "caption": None}]

    print("[调用API] 输入参数：")
    print(f"  image path: {image}")
    print(f"  uploaded: {uploaded}")
    print(f"  prompt: {prompt}")
    print(f"  lora_adapter: {lora_adapter}")
    print(f"  seed: {seed}")
    print(f"  guidance_scale: {guidance_scale}")
    print(f"  steps: {steps}")

    try:
        result = space_client.predict(
            images=images_input,
            prompt=prompt,
            lora_adapter=lora_adapter,
            seed=int(seed),
            randomize_seed=bool(randomize_seed),
            guidance_scale=int(guidance_scale),
            steps=int(steps),
            api_name="/infer",
        )

        print(f"[调用API] 返回值: {result}")

        image_info, seed_used = result

        if isinstance(image_info, dict):
            img_out = image_info.get("path") or image_info.get("url")
        else:
            img_out = image_info

        return img_out, int(seed_used)

    except Exception as e:
        import traceback
        traceback.print_exc()
        print(f"[调用API] 异常: {e}")
        return None, seed


css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# 图像编辑 Demo\n基于 prithivMLmods/Qwen-Image-Edit-2511-LoRAs-Fast")

        image = gr.Image(
            label="上传图片",
            sources=["upload"],
            type="filepath",
        )

        prompt = gr.Text(
            label="编辑描述（Prompt）",
            placeholder="请输入图片编辑描述...",
        )

        lora_adapter = gr.Dropdown(
            label="编辑风格（Style）",
            choices=LORA_STYLES,
            value="Photo-to-Anime"
        )

        run_button = gr.Button("执行编辑", variant="primary")

        result = gr.Image(label="结果图片", show_label=True)

        with gr.Accordion("高级设置", open=False):
            seed = gr.Slider(
                label="随机种子",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            randomize_seed = gr.Checkbox(label="随机化种子", value=True)
            guidance_scale = gr.Slider(
                label="引导强度 (Guidance Scale)",
                minimum=0.1,
                maximum=10.0,
                step=0.1,
                value=1.0,
            )
            steps = gr.Slider(
                label="推理步数 (Steps)",
                minimum=1,
                maximum=50,
                step=1,
                value=4,
            )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.launch(ssr_mode=False, share=True)