import gradio as gr
import spaces
from ominicontrol import generate_image
import os


USE_ZERO_GPU = os.environ.get("USE_ZERO_GPU", "0") == "1"

css = """
.inputPanel {
    width: 320px;
    display: flex;
    align-items: center;
}
.outputPanel {
    display: flex;
    align-items: center;
}
.hint {
    font-size: 14px;
    color: #777;
    # border: 1px solid #ccc;
    padding: 4px;
    border-radius: 5px;
    # background-color: #efefef;
}
"""

header = """
# 🎨 OminiControl Art
<div style="text-align: center; display: flex; justify-content: left; gap: 5px;">
<a href="https://arxiv.org/abs/2411.15098"><img src="https://img.shields.io/badge/ariXv-Paper-A42C25.svg" alt="arXiv"></a>
<a href="https://huggingface.co/spaces/Yuanshi/OminiControl"><img src="https://img.shields.io/badge/🤗OminiControl-Demo-ffbd45.svg" alt="HuggingFace"></a>
<a href="https://github.com/Yuanshi9815/OminiControl"><img src="https://img.shields.io/badge/GitHub-Code-blue.svg?logo=github&" alt="GitHub"></a>
</div>
"""


def style_transfer(image, style):
    return image


styles = [
    "Studio Ghibli",
    "Irasutoya Illustration",
    "The Simpsons",
    "Snoopy",
]


def gradio_interface():
    with gr.Blocks(css=css) as demo:
        gr.Markdown(header)
        with gr.Row(equal_height=False):
            with gr.Column(variant="panel", elem_classes="inputPanel"):
                original_image = gr.Image(
                    type="pil",
                    label="Condition Image",
                    width=400,
                    height=400,
                )
                style = gr.Radio(
                    styles,
                    label="🎨 Select Style",
                    value=styles[0],
                )
                # Advanced settings
                with gr.Accordion(
                    "⚙️ Advanced Settings", open=False
                ) as advanced_settings:
                    inference_mode = gr.Radio(
                        ["High Quality", "Fast"],
                        value="High Quality",
                        label="Generating Mode",
                    )
                    image_ratio = gr.Radio(
                        ["Auto", "Square(1:1)", "Portrait(2:3)", "Landscape(3:2)"],
                        label="Image Ratio",
                        value="Auto",
                    )
                    use_random_seed = gr.Checkbox(label="Use Random Seed", value=True)
                    seed = gr.Number(
                        label="Seed",
                        value=42,
                        visible=(not use_random_seed.value),
                    )
                    use_random_seed.change(
                        lambda x: gr.update(visible=(not x)),
                        use_random_seed,
                        seed,
                        show_progress="hidden",
                    )
                    image_guidance = gr.Slider(
                        label="Image Guidance",
                        minimum=1.1,
                        maximum=5,
                        value=1.5,
                        step=0.1,
                    )
                    steps = gr.Slider(
                        label="Steps",
                        minimum=10,
                        maximum=50,
                        value=20,
                        step=1,
                    )
                    inference_mode.change(
                        lambda x: gr.update(interactive=(x == "High Quality")),
                        inference_mode,
                        image_guidance,
                        show_progress="hidden",
                    )

                btn = gr.Button("Generate Image")

            with gr.Column(elem_classes="outputPanel"):
                output_images = gr.Image(
                    type="pil",
                    width=640,
                    height=640,
                    label="Output Image",
                )
                hint = gr.Markdown(
                    """
                    <div style="text-align: center; width: 100%;">
                    <b>Note: The selected style is in beta testing.</b> Feel free to try a few more times to get a the better result.
                    </div>
                    """,
                    visible=False,
                )
                style.change(
                    lambda x: gr.update(visible=x in styles[1:]),
                    style,
                    hint,
                )

        with gr.Row():
            examples = gr.Examples(
                examples=[
                    ["examples/DistractedBoyfriend.webp", styles[0]],
                    ["examples/steve.webp", styles[0]],
                    ["examples/oiiai.png", styles[1]],
                    ["examples/doge.jpg", styles[1]],
                    ["examples/breakingbad.jpg", styles[2]],
                    ["examples/PulpFiction.jpg", styles[3]],
                ],
                inputs=[original_image, style],
            )

        btn.click(
            fn=infer,
            inputs=[
                style,
                original_image,
                inference_mode,
                image_guidance,
                image_ratio,
                use_random_seed,
                seed,
                steps,
            ],
            outputs=output_images,
        )

    return demo


def infer(
    style,
    original_image,
    inference_mode,
    image_guidance,
    image_ratio,
    use_random_seed,
    seed,
    steps,
):
    print(
        f"Style: {style}, Inference Mode: {inference_mode}, Image Guidance: {image_guidance}, Image Ratio: {image_ratio}, Use Random Seed: {use_random_seed}, Seed: {seed}"
    )
    result_image = generate_image(
        image=original_image,
        style=style,
        inference_mode=inference_mode,
        image_guidance=image_guidance,
        image_ratio=image_ratio,
        use_random_seed=use_random_seed,
        seed=seed,
        steps=steps,
    )
    return result_image


if USE_ZERO_GPU:
    infer = spaces.GPU(infer)

if __name__ == "__main__":
    demo = gradio_interface()
    demo.launch(ssr_mode=False)