import gradio as gr
import os

from util.instantmesh import generate_mvs, make3d, preprocess, check_input_image
from util.text_img import generate_image, check_prompt

_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:
    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")
                negative_prompt = gr.Textbox(label="Negative Prompt", value="low quality, bad quality, sketches, legs")
                scale = gr.Slider(label="Control Image Scale", minimum=0.1, maximum=1.0, step=0.1, value=0.5)
                control_image = gr.Image(label="Enter an image of a shoe, that you want to use as a reference", type='numpy')
                gr.Examples(
                       examples=[
                           os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
                       ],
                       inputs=[control_image],
                       label="Examples",
                       cache_examples=False,
                   )
            with gr.Column():
                button_gen = 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_gen.click(check_prompt, inputs=[prompt]).success(generate_image, inputs=[prompt, negative_prompt, control_image, scale], outputs=[gen_image])

    with gr.Tab("Image to 3D Model Generator"):
        with gr.Row(variant="panel"):
            with gr.Column():
                with gr.Row():
                    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("glb"):
                        output_model_glb = gr.Model3D(
                            label="Output Model (GLB Format)",
                            interactive=False,
                        )
                    with gr.Tab("obj"):
                        output_model_obj = gr.Model3D(
                            label="Output Model (OBJ 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]).success(
        fn=preprocess,
        inputs=[gen_image, 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()