import gradio as gr
from gradio_client import Client
import random

# 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 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:
                images.append({"image": {"path": im}, "caption": None})
        else:
            images.append({"image": {"path": image}, "caption": None})

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

    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",
        )
        # 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:
        return None, seed

examples = [
    ["", "Astronaut in jungle, anime style", "Photo-to-Anime", 0, True, 1.0, 4],
    ["", "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"],
                tool=None,
                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,
            )

        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__":
    demo.launch(css=css)