import gradio as gr
import spaces
from ominicontrol import generate_image, vote_feedback
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>

***OminiControl Art*** distills the artistic style of [GPT-4o](https://openai.com/index/introducing-4o-image-generation/) into the [FLUX.1](https://blackforestlabs.ai/) model, building on the foundation of [OminiControl](https://github.com/Yuanshi9815/OminiControl)✨.
Enjoy playing around! 🌈
"""


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", variant="primary")

                with gr.Accordion("🏞️ Examples", open=True) as advanced_settings:
                    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],
                    )

            with gr.Column(elem_classes="outputPanel"):
                output_image = gr.Image(
                    type="pil",
                    width=600,
                    height=600,
                    label="Output Image",
                    interactive=False,
                    sources=None,
                )
                inference_id = gr.Textbox(
                    visible=False,
                    interactive=False,
                )

                # Feedback buttons
                with gr.Column(visible=False) as feedback:
                    gr.Markdown(
                        """
                        Your feedback improves the model! Please let us know how you feel about the generated image.
                        """,
                    )
                    with gr.Row() as feedback_buttons:
                        upvote = gr.Button("👍 I like it", variant="primary")
                        downvote = gr.Button("👎 It looks bad")

                        def feedback_func(feedback):
                            def func(inputs):
                                print(f"Feedback: {feedback}, Inference ID: {inputs}")
                                vote_feedback(log_id=inputs, feedback=feedback)
                                # Here you can add your feedback logging logic
                                return gr.update(visible=False)

                            return func

                        upvote.click(feedback_func("1"), inference_id, feedback)
                        downvote.click(feedback_func("0"), inference_id, feedback)

                inference_id.change(
                    lambda x: gr.update(visible=True), output_image, feedback
                )

                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():
        btn.click(
            fn=infer,
            inputs=[
                style,
                original_image,
                inference_mode,
                image_guidance,
                image_ratio,
                use_random_seed,
                seed,
                steps,
            ],
            outputs=[
                output_image,
                inference_id,
            ],
        )

    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, inference_id = 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, inference_id


if USE_ZERO_GPU:
    infer = spaces.GPU(infer)

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