import gradio as gr
import os

from util.instantmesh import generate_mvs, make3d, preprocess, check_input_image
from util.text_img import generate_txttoimg, check_prompt, generate_imgtoimg, update_image

_CITE_ = r"""
```bibtex
@article{xu2024instantmesh,
  title={InstantMesh: Efficient 3D Mesh Generation from a Single Image with Sparse-view Large Reconstruction Models},
  author={Xu, Jiale and Cheng, Weihao and Gao, Yiming and Wang, Xintao and Gao, Shenghua and Shan, Ying},
  journal={arXiv preprint arXiv:2404.07191},
  year={2024}
}
```
"""

theme = gr.themes.Soft(
    primary_hue="orange",
    secondary_hue="gray",
    neutral_hue="slate",
    font=['Montserrat', gr.themes.GoogleFont('ui-sans-serif'), 'system-ui', 'sans-serif'],
)


with gr.Blocks(theme=theme) as GenDemo:
    gen_image_var = gr.State()
    # with gr.Tab("Image to Image Generator"):
    #     with gr.Row(variant="panel"):
    #         with gr.Column():
    #             prompt = gr.Textbox(label="Enter a discription of a shoe")
    #             image = gr.Image(label="Enter an image of a shoe, that you want to use as a reference", type='pil')
    #             strength = gr.Slider(label="Strength", minimum=0.1, maximum=1.0, value=0.5, step=0.1)
    #             gr.Examples(
    #                    examples=[
    #                        os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
    #                    ],
    #                    inputs=[image],
    #                    label="Examples",
    #                    cache_examples=False,
    #                )
    #         with gr.Column():
    #             button_img = gr.Button("Generate Image", elem_id="generateIm", variant="primary")
    #             gen_image = gr.Image(label="Generated Image", image_mode="RGBA", type='pil', show_download_button=True, show_label=False)
        
    #     button_img.click(check_prompt, inputs=[prompt]).success(generate_imgtoimg, inputs=[prompt, image, strength], outputs=[gen_image]).success(update_image, inputs=[gen_image], outputs=[gen_image_var])

    with gr.Tab("Text to Image Generator"):
        with gr.Row(variant="panel"):
            with gr.Column():
                prompt = gr.Textbox(label="Enter a discription of a shoe")
                select = gr.Dropdown(label="Select a controlnet model", choices=["Depth","Normal"])      
                controlNet_image = gr.Image(label="Enter an image of a shoe, that you want to use as a reference", type='pil')
                gr.Examples(
                       examples=[
                           os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
                       ],
                       inputs=[controlNet_image],
                       label="Examples",
                       cache_examples=False,
                   )
            with gr.Column():
                button_txt = gr.Button("Generate Image", elem_id="generateIm", variant="primary")
                gen_image = gr.Image(label="Generated Image", image_mode="RGBA", type='pil', show_download_button=True, show_label=False)
        
        button_txt.click(check_prompt, inputs=[prompt]).success(generate_txttoimg, inputs=[prompt, controlNet_image, select], outputs=[gen_image]).success(update_image, inputs=[gen_image], outputs=[gen_image_var])

    with gr.Tab("Image to 3D Model Generator"):
        with gr.Row(variant="panel"):
            with gr.Column():
                with gr.Row():
                    # input_image = gr.Image(
                    #     label="Input Image",
                    #     image_mode="RGBA",
                    #     type="pil",
                    #     interactive=True
                    # )
                    processed_image = gr.Image(
                        label="Processed Image", 
                        image_mode="RGBA", 
                        #width=256,
                        #height=256,
                        type="pil", 
                        interactive=False
                    )
                with gr.Row():
                    with gr.Group():
                        do_remove_background = gr.Checkbox(
                            label="Remove Background", value=True
                        )
                        sample_seed = gr.Number(value=42, label="Seed Value", precision=0)

                        sample_steps = gr.Slider(
                            label="Sample Steps",
                            minimum=30,
                            maximum=75,
                            value=75,
                            step=5
                        )

                with gr.Row():
                    submit = gr.Button("Generate", elem_id="generate", variant="primary")

            with gr.Column():

                with gr.Row():

                    with gr.Column():
                        mv_show_images = gr.Image(
                            label="Generated Multi-views",
                            type="pil",
                            width=379,
                            interactive=False
                        )

                with gr.Row():
                    with gr.Tab("obj"):
                        output_model_obj = gr.Model3D(
                            label="Output Model (OBJ Format)",
                            interactive=False,
                        )
                    with gr.Tab("glb"):
                        output_model_glb = gr.Model3D(
                            label="Output Model (GLB Format)",
                            interactive=False,
                        )

                with gr.Row():
                    gr.Markdown('''Try a different <b>seed value</b> if the result is unsatisfying (Default: 42).''')

        gr.Markdown(_CITE_)

    mv_images = gr.State()

    submit.click(fn=check_input_image, inputs=[gen_image_var]).success(
        fn=preprocess,
        inputs=[gen_image_var, do_remove_background],
        outputs=[processed_image],
    ).success(
        fn=generate_mvs,
        inputs=[processed_image, sample_steps, sample_seed],
        outputs=[mv_images, mv_show_images]   
    ).success(
        fn=make3d,
        inputs=[mv_images],
        outputs=[output_model_obj, output_model_glb]
    )

GenDemo.launch()