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

# API client for the external Space
space_client = Client("prithivMLmods/Qwen-Image-Edit-2511-LoRAs-Fast")

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**32-1

def safe_image_path(path):
    """Ensure the given path points to an existing, readable file."""
    if not path or not isinstance(path, str):
        return None
    if os.path.isfile(path):
        return path
    # Try to handle Gradio's nested path structure if needed
    # (Could be adjusted depending how Gradio saves upload images)
    return None

def infer(
    image,
    prompt,
    lora_adapter,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    progress=gr.Progress(track_tqdm=True),
):
    # Prepare images input as per API (expects Gallery [list of dicts])
    images = []
    if image is not None:
        if isinstance(image, list):  # Gradio Gallery
            for im in image:
                file_path = safe_image_path(im)
                if file_path:
                    images.append({"image": {"path": file_path}, "caption": None})
                else:
                    print(f"警告: 路径无效或找不到文件: {im}")
        else:
            file_path = safe_image_path(image)
            if file_path:
                images.append({"image": {"path": file_path}, "caption": None})
            else:
                print(f"警告: 路径无效或找不到文件: {image}")

    if len(images) == 0:
        print("未检测到有效图片路径，未能上传图片。")
        return None, seed

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

    print("[调用API] space_client.predict 输入参数：")
    print(f"  images: {images}")
    print(f"  prompt: {prompt}")
    print(f"  lora_adapter: {lora_adapter}")
    print(f"  seed: {seed}")
    print(f"  randomize_seed: {randomize_seed}")
    print(f"  guidance_scale: {guidance_scale}")
    print(f"  steps: {steps}")

    try:
        result = space_client.predict(
            images=images,
            prompt=prompt,
            lora_adapter=lora_adapter,
            seed=float(seed),
            randomize_seed=bool(randomize_seed),
            guidance_scale=float(guidance_scale),
            steps=float(steps),
            api_name="/infer",
        )
        print(f"[调用API] space_client.predict 返回值: {result}")
        # result is a tuple: (image_dict, seed)
        image_info, seed_used = result
        # The API may return image at .url or .path, we use .url if available
        img_url = image_info.get("url") or image_info.get("path")
        return img_url, seed_used
    except Exception as e:
        import traceback
        traceback.print_exc()
        print(f"[调用API] 调用接口异常: {e}")
        # More verbose error message for image-processing errors
        if hasattr(e, 'message') and "Could not process uploaded images" in str(e):
            print(
                "\n[错误] 上传图片处理失败。请确保图片文件有效且未损坏。\n"
                "可以尝试重新选择图片或修改图片格式。"
            )
        return None, seed

# For examples, use None for image input to avoid Gradio directory/File errors on startup.
examples = [
    [None, "Astronaut in jungle, anime style", "Photo-to-Anime", 0, True, 1.0, 4],
    [None, "A delicious ceviche cheesecake slice", "Style-Transfer", 0, True, 1.0, 4],
]

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

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

        with gr.Row():
            image = gr.Image(
                label="上传图片",
                sources=["upload"],
                type="filepath",
                elem_id="input-image"
            )
        with gr.Row():
            prompt = gr.Text(
                label="编辑描述（Prompt）",
                placeholder="请输入图片编辑描述...",
            )

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

        run_button = gr.Button("执行编辑", scale=1, 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,
            )

        # Only show example text/inputs, but avoid file path errors (set image to None)
        gr.Examples(
            examples=examples,
            inputs=[image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
            label="示例",
        )

    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__":
    # See: https://prithivmlmods-qwen-image-edit-2511-loras-fast.hf.space
    # To use SSR, remove 'ssr_mode=False' (can help with speed in some settings)
    # To create a public link, set share=True
    demo.launch(css=css, ssr_mode=True, share=True)