import gradio as gr
import os

from model import VirtualStagingToolV2


def predict(init_image, style, backyard_style, color_preference, additional_info):
    vs_tool = VirtualStagingToolV2(diffusion_version="stabilityai/stable-diffusion-2-inpainting")
    if backyard_style:
        style = backyard_style

    output_images, transparent_mask_image = vs_tool.virtual_stage(
        image=init_image, style=style, color_preference=color_preference,
        additional_info=additional_info, number_images=3)

    return output_images[0], output_images[1], output_images[2], transparent_mask_image


image_blocks = gr.Blocks()
with image_blocks as demo:
    gr.Markdown(
        """
        # Virtual Home Staging
        """)
    with gr.Group():
        with gr.Box():
            with gr.Row():
                with gr.Column():
                    image = gr.Image(source='upload', elem_id="image_upload",
                                     type="pil", label="Upload",
                                     ).style(height=400)
                    with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
                        style = gr.Dropdown(
                            ["Bohemian", "Coastal", "Contemporary", "Farmhouse", "French country",
                             "Glam", "Industrial", "Japandi", "Mid-century modern", "Minimal", "Modern",
                             "Rustic", "Scandinavian", "Traditional", "Transitional", ],
                            label="Design theme", elem_id="input-style",
                            info="only select if the image is NOT a backyard"
                        )

                        backyard_style = gr.Dropdown(
                            ["Beautiful garden", "Charming playhouse garden landscaping",
                             "Cottage garden's colorful planting palette", "Cozy corner with fire pit and seating",
                             "Garden landscaping with gravel landscaping", "Hip california garden landscaping",
                             "Lush green lawn", "Mediterranean garden landscaping",
                             "Moss garden", "Outdoor dining and sitting area",
                             "Party-ready outdoor space with pool, spa, and fire feature",
                             "Resort-style landscaping and pool", "Round swimming pool with lawn and pool house"],
                            label="Backyard design theme", elem_id="input-backyard-style",
                            info="only select if the image is a backyard"
                        )

                    color_preference = gr.Textbox(placeholder='Enter color preference',
                                                  label="Color preference", elem_id="input-color")

                    additional_info = gr.Textbox(placeholder='Enter additional information',
                                                  label="Additional information", elem_id="input-add-info")

                    btn = gr.Button("Inpaint!").style(
                        margin=False,
                        rounded=(False, True, True, False),
                        full_width=False,
                    )

                    gr.Markdown("## Image Examples")
                    gr.Examples(
                        examples=[os.path.join(os.path.dirname(__file__),
                                               "examples/exciting-small-kitchen-ideas-1821197-hero-d00f516e2fbb4dcabb076ee9685e877a.jpg"),
                                  os.path.join(os.path.dirname(__file__),
                                               "examples/oct-2019-idh-bathroom-reno-ideas-new-gallery-2.jpg"),
                                  os.path.join(os.path.dirname(__file__),
                                               "examples/tips-for-decorating-a-beautiful-bedroom-1976169-hero-e960fbb8311c4b9b875a1813962d34eb.jpg"),
                                  os.path.join(os.path.dirname(__file__),
                                               "examples/backyard_additions.jpg"),
                                  os.path.join(os.path.dirname(__file__),
                                               "examples/living-room-gallery-shelves-l-shaped-couch-ELeyNpyyqpZ8hosOG3EG1X-b5a39646574544e8a75f2961332cd89a.jpg"),
                                  os.path.join(os.path.dirname(__file__),
                                               "examples/modern-dining-room-ideas-4147451-hero-d6333998f8b34620adfd4d99ac732586.jpg"),
                                  ],
                        inputs=image
                    )

                with gr.Column():
                    mask_image = gr.Image(label="Mask image", elem_id="mask-img", type="pil").style(height=400)
                    image_out_1 = gr.Image(label="Output 1", elem_id="output-img-1", type="pil").style(height=400)
                    image_out_2 = gr.Image(label="Output 2", elem_id="output-img-2", type="pil").style(height=400)
                    image_out_3 = gr.Image(label="Output 3", elem_id="output-img-3", type="pil").style(height=400)

                btn.click(fn=predict, inputs=[image, style, backyard_style, color_preference, additional_info],
                          outputs=[image_out_1, image_out_2, image_out_3, mask_image])



image_blocks.queue(concurrency_count=3)
image_blocks.launch()