import gradio as gr

def ui(controller):#get_first_result,get_second_result): #controller):
    with gr.Blocks() as ui:
        with gr.Row():
            with gr.Column():
                sketch=gr.Image(sources = 'upload', label='Model image', type = 'pil')
                
            with gr.Column():
                improved_sketch_view = gr.Image(type="pil", label="Improved Sketch")        

        #with gr.Row():
            with gr.Column():
                result=gr.Image(type="pil", label="Colored Image")
                

        with gr.Row():
            with gr.Column():
                first_prompt = gr.Textbox(label="Prompt", lines=3)
                first_negative_prompt = gr.Textbox(label="Negative prompt", lines=3, value="sketch, lowres, error, extra digit, fewer digits, cropped, worst quality,low quality, normal quality, jpeg artifacts, blurry")
                firts_prompt_helper = gr.Button(value="Prompt Helper", variant="primary")
                #controlnet_scale = gr.Slider(minimum=0.5, maximum=1.25, value=0.5, step=0.01, label="Contr")
                improve_sketch = gr.Button(value="Improve Sketch", variant="primary")

            with gr.Column():
                second_prompt = gr.Textbox(label="Prompt", lines=3)
                second_negative_prompt = gr.Textbox(label="Negative prompt", lines=3, value="disfigured, extra digit, fewer digits, cropped, worst quality, low quality")
                second_prompt_helper = gr.Button(value="Prompt Helper", variant="primary")
                result_button = gr.Button(value="Paint It", variant="primary")

        with gr.Row():
            gr.Examples([["./examples/sketch.png"], 
                         ["solo, open mouth, white background, standing, tail, full body,bmonochrome, greyscale, from side, no humans, animal, dog, animal focus, realistic, walk on the street, add city on the backstage, city street"],
                         ["solo, open mouth, simple background, grey background, standing, full body, monochrome, greyscale, from side, no humans, animal, dog, animal focus, lineart, wolf"]],
                        [sketch, first_prompt, second_prompt],
                        label='Examples. Be concrete as much as possible.')
                

        improve_sketch.click(fn=controller.get_first_result,
                             inputs=[sketch, first_prompt, first_negative_prompt],
                             outputs=improved_sketch_view)

        result_button.click(fn=controller.get_second_result,
                             inputs=[improved_sketch_view, second_prompt, second_negative_prompt],
                             outputs=result)

        firts_prompt_helper.click(fn=controller.get_help_w_prompt,
                                 inputs=[sketch],
                                 outputs=first_prompt)

        second_prompt_helper.click(fn=controller.get_help_w_prompt,
                                 inputs=[improved_sketch_view],
                                 outputs=second_prompt)
        


    ui.queue().launch(debug=True)