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}
}
```
"""


with gr.Blocks() as demo:
    with gr.Tab("Text to Image Generator"):
        with gr.Row():
            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)
            with gr.Column():
                control_image = gr.Image(label="Enter an image of a shoe, that you want to use as a reference", type='numpy')
                # neg_prompt = gr.Textbox(label="Enter a negative prompt", value="low quality, watermark, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, signature, cut off, draft, closed eyes, text, logo")
        with gr.Row():
            with gr.Column():
                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")
        with gr.Row(): 
            with gr.Column():
                image_nobg = gr.Image(label="Generated Image", 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=[image_nobg])

    with gr.Row(variant="panel"):
        with gr.Column():
            with gr.Row():
                input_image = gr.Image(
                    label="Input Image",
                    image_mode="RGBA",
                    sources="upload",
                    #width=256,
                    #height=256,
                    type="pil",
                    elem_id="content_image",
                )
                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.Row(variant="panel"):
                gr.Examples(
                    examples=[
                        os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
                    ],
                    inputs=[input_image],
                    label="Examples",
                    cache_examples=False,
                )

        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():
                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=[input_image]).success(
        fn=preprocess,
        inputs=[input_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]
    )

demo.launch()