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=[["./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"]],
                        inputs=[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)