import gradio as gr

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

def create_demo():
    with gr.Blocks() as demo:
        with gr.Column():
            model_name = gr.Dropdown(choices = list(models.keys()), label = 'Modelo', value=list(models.keys())[0])
            prompt = gr.Textbox(label='Prompt', value='La luna en el espacio')
            negative_prompt = gr.Textbox( label='Negative Prompt (prompt negativo)', value='mala calidad')
            generate_button = gr.Button('Generar')

        gr.Markdown('''
        🌖 lunarnaut - intro a stable diffusion 
        ---
        1. Selecciona un Modelo
        2. Agrega un prompt y prompt negativo
        3. Haz clic en 'Generar'
        ''')
        result = gr.Gallery(label='Resultado', show_label=False, elem_id='gallery').style(columns=1, rows=1, preview=True)

        inputs = [ prompt, negative_prompt, model_name ]

        prompt.submit(fn=generate, inputs=inputs, outputs=result)
        negative_prompt.submit(fn=generate, inputs=inputs, outputs=result)
        generate_button.click(fn=generate, inputs=inputs, outputs=result)
    return demo


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