import gradio as gr

def generateImage(prompt, n_prompt, modelName):

    return models[modelName].process(prompt, n_prompt)




def create_demo():

    with gr.Blocks() as demo:
        with gr.Column():
            prompt = gr.Textbox(label='Prompt')
            n_prompt = gr.Textbox(
                label='Negative Prompt',
                value=
                'ugly, disfigured, deformed'
            )

            modelName = gr.Dropdown(choices = list(models.keys()), 
                                        label = "Model", 
                                        value=list(models.keys())[0])

            run_button = gr.Button('Run')
        gr.Markdown("### [Stable Diffusion Art](https://stable-diffusion-art.com/) -- tutorials and resources. Read [Model license](https://huggingface.co/spaces/CompVis/stable-diffusion-license).")
        result = gr.Gallery(label='Output',
                            show_label=False,
                            elem_id='gallery').style(columns=1, rows=1, preview=True)        


        inputs = [
            prompt,
            n_prompt,
            modelName,
        ]

        prompt.submit(
            fn=generateImage,
            inputs=inputs,
            outputs=result
        )
        n_prompt.submit(
            fn=generateImage,
            inputs=inputs,
            outputs=result
        )

        run_button.click(
            fn=generateImage,
            inputs=inputs,
            outputs=result
        )
    return demo


if __name__ == '__main__':
    from model import Model
    models = {
        "Stable Diffusion v1.5": Model("runwayml/stable-diffusion-v1-5"),
        "Realistic Vision v2.0": Model("SG161222/Realistic_Vision_V2.0"),
        "Anything v3.0": Model("Linaqruf/anything-v3.0")
              }
    demo = create_demo()
    demo.queue().launch()