import gradio as gr

import spaces

from generation import generate_material


@spaces.GPU
def generate(prompts, seed, resolution, refinement):
    image = generate_material(
        prompts, seed=seed, resolution=int(resolution), refinement=refinement
    )
    return image.basecolor, image.normal, image.height, image.metallic, image.roughness
    pass


def interface_function(
    prompt_type, text_prompt, image_prompt, seed, resolution, refinement
):
    if prompt_type == "Text":
        return generate(text_prompt, seed, resolution, refinement)
    elif prompt_type == "Image":
        return generate(image_prompt, seed, resolution, refinement)


def update_visibility(prompt_type):
    if prompt_type == "Text":
        return gr.update(visible=True), gr.update(visible=False)
    elif prompt_type == "Image":
        return gr.update(visible=False), gr.update(visible=True)


with gr.Blocks() as demo:
    title = """
    <center> 

    <h1> StableMaterials </h1>
    <b> Try the <a href="https://huggingface.co/gvecchio/StableMaterials">StableMaterials</a> model. Start generating using a text o image prompt, or simply try one of the examples below. </b>
    <br>
    <b> NB: </b> The demo currently uses only the base model and does not include the refinement step. The refiner model will be released in a future update.

    </center>
    """
    gr.HTML(title)

    with gr.Row():
        with gr.Column():
            prompt_type = gr.Radio(
                choices=["Text", "Image"], label="Prompt Type", value="Text"
            )
            text_prompt = gr.Textbox(
                label="Text Prompt", visible=True, lines=3, placeholder="A brick wall"
            )
            image_prompt = gr.Image(type="pil", label="Image Prompt", visible=False)

        with gr.Column():
            seed = gr.Number(value=-1, label="Seed (-1 for random)")
            resolution = gr.Dropdown(
                ["512", "1024", "2048"],
                value="512",
                label="Resolution",
                interactive=False,
            )
            refinement = gr.Checkbox(label="Refinement", interactive=False)
            generate_button = gr.Button("Generate")

    prompt_type.change(
        fn=update_visibility, inputs=prompt_type, outputs=[text_prompt, image_prompt]
    )

    with gr.Row():
        output_basecolor = gr.Image(label="Base Color", format="png", image_mode="RGB")
        output_normal = gr.Image(label="Normal Map", format="png", image_mode="RGB")
        output_height = gr.Image(label="Height Map", format="png", image_mode="L")
        output_metallic = gr.Image(label="Metallic Map", format="png", image_mode="L")
        output_roughness = gr.Image(label="Roughness Map", format="png", image_mode="L")

    generate_button.click(
        fn=interface_function,
        inputs=[prompt_type, text_prompt, image_prompt, seed, resolution, refinement],
        outputs=[
            output_basecolor,
            output_normal,
            output_height,
            output_metallic,
            output_roughness,
        ],
    )

    gr.Markdown("### Examples")
    gr.Examples(
        examples=[
            ["Text", "A brick wall", None],
            ["Text", "A wooden floor", None],
            ["Image", None, "./samples/bricks.jpg"],
            ["Image", None, "./samples/metal.jpg"],
        ],
        inputs=[prompt_type, text_prompt, image_prompt],
    )
    gr.Markdown(
        "Try the examples above to see how the Material Generator works. Adjust the parameters and click 'Generate' to create new material maps."
    )


if __name__ == "__main__":
    demo.launch()