import gradio as gr
from multiprocessing import cpu_count
from utils.functions import generate, train_textual_inversion
from utils.shared import model_ids, scheduler_names, default_scheduler

default_img_size = 512

with open("html/header.html") as fp:
    header = fp.read()

with open("html/footer.html") as fp:
    footer = fp.read()

with gr.Blocks(css="html/style.css") as demo:

    pipe_state = gr.State(lambda: 1)

    gr.HTML(header)

    with gr.Row():

        with gr.Column(scale=70):

            # with gr.Row():
            prompt = gr.Textbox(
                label="Prompt", placeholder="<Shift+Enter> to generate", lines=2
            )
            neg_prompt = gr.Textbox(label="Negative Prompt", placeholder="", lines=2)

        with gr.Column(scale=30):
            model_name = gr.Dropdown(
                label="Model", choices=model_ids, value=model_ids[0]
            )
            scheduler_name = gr.Dropdown(
                label="Scheduler", choices=scheduler_names, value=default_scheduler
            )
            generate_button = gr.Button(value="Generate", elem_id="generate-button")

    with gr.Row():

        with gr.Column():

            with gr.Tab("Text to Image") as tab:
                tab.select(lambda: 1, [], pipe_state)

            with gr.Tab("Image to image") as tab:
                tab.select(lambda: 2, [], pipe_state)

                image = gr.Image(
                    label="Image to Image",
                    source="upload",
                    tool="editor",
                    type="pil",
                    elem_id="image_upload",
                ).style(height=default_img_size)
                strength = gr.Slider(
                    label="Denoising strength",
                    minimum=0,
                    maximum=1,
                    step=0.02,
                    value=0.8,
                )

            with gr.Tab("Inpainting") as tab:
                tab.select(lambda: 3, [], pipe_state)

                inpaint_image = gr.Image(
                    label="Inpainting",
                    source="upload",
                    tool="sketch",
                    type="pil",
                    elem_id="image_upload",
                ).style(height=default_img_size)
                inpaint_strength = gr.Slider(
                    label="Denoising strength",
                    minimum=0,
                    maximum=1,
                    step=0.02,
                    value=0.8,
                )
                inpaint_options = [
                    "preserve non-masked portions of image",
                    "output entire inpainted image",
                ]
                inpaint_radio = gr.Radio(
                    inpaint_options,
                    value=inpaint_options[0],
                    show_label=False,
                    interactive=True,
                )

            with gr.Tab("Textual Inversion") as tab:
                tab.select(lambda: 4, [], pipe_state)

                type_of_thing = gr.Dropdown(
                    label="What would you like to train?",
                    choices=["object", "person", "style"],
                    value="object",
                    interactive=True,
                )

                text_train_bsz = gr.Slider(
                    label="Training Batch Size",
                    minimum=1,
                    maximum=8,
                    step=1,
                    value=1,
                )

                files = gr.File(
                    label=f"""Upload the images for your concept""",
                    file_count="multiple",
                    interactive=True,
                    visible=True,
                )

                text_train_steps = gr.Number(label="How many steps", value=1000)

                text_learning_rate = gr.Number(label="Learning Rate", value=5.0e-4)

                concept_word = gr.Textbox(
                    label=f"""concept word - use a unique, made up word to avoid collisions"""
                )
                init_word = gr.Textbox(
                    label=f"""initial word - to init the concept embedding"""
                )

                textual_inversion_button = gr.Button(value="Train Textual Inversion")

                training_status = gr.Text(label="Training Status")

            with gr.Row():
                batch_size = gr.Slider(
                    label="Batch Size", value=1, minimum=1, maximum=8, step=1
                )
                seed = gr.Slider(-1, 2147483647, label="Seed", value=-1, step=1)

            with gr.Row():
                guidance = gr.Slider(
                    label="Guidance scale", value=7.5, minimum=0, maximum=20
                )
                steps = gr.Slider(
                    label="Steps", value=20, minimum=1, maximum=100, step=1
                )

            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    value=default_img_size,
                    minimum=64,
                    maximum=1024,
                    step=32,
                )
                height = gr.Slider(
                    label="Height",
                    value=default_img_size,
                    minimum=64,
                    maximum=1024,
                    step=32,
                )

        with gr.Column():
            gallery = gr.Gallery(
                label="Generated images", show_label=False, elem_id="gallery"
            ).style(height=default_img_size, grid=2)

            generation_details = gr.Markdown()

            pipe_kwargs = gr.Textbox(label="Pipe kwargs", value="{\n\t\n}")

            # if torch.cuda.is_available():
            #  giga = 2**30
            #  vram_guage = gr.Slider(0, torch.cuda.memory_reserved(0)/giga, label='VRAM Allocated to Reserved (GB)', value=0, step=1)
            #  demo.load(lambda : torch.cuda.memory_allocated(0)/giga, inputs=[], outputs=vram_guage, every=0.5, show_progress=False)

    gr.HTML(footer)

    inputs = [
        model_name,
        scheduler_name,
        prompt,
        guidance,
        steps,
        batch_size,
        width,
        height,
        seed,
        image,
        strength,
        inpaint_image,
        inpaint_strength,
        inpaint_radio,
        neg_prompt,
        pipe_state,
        pipe_kwargs,
    ]
    outputs = [gallery, generation_details]

    prompt.submit(generate, inputs=inputs, outputs=outputs)
    generate_button.click(generate, inputs=inputs, outputs=outputs)

    textual_inversion_inputs = [
        model_name,
        scheduler_name,
        type_of_thing,
        files,
        concept_word,
        init_word,
        text_train_steps,
        text_train_bsz,
        text_learning_rate,
    ]

    textual_inversion_button.click(
        train_textual_inversion,
        inputs=textual_inversion_inputs,
        outputs=[training_status],
    )


# demo = gr.TabbedInterface([demo, dreambooth_tab], ["Main", "Dreambooth"])

demo.queue(concurrency_count=cpu_count())

demo.launch()