import gradio as gr

from service import Service

prompt = ""
scale = 10
n_samples = 4
rows = 2
cols = 2

service = Service()


def generate_pkm(prompt):
    return service.generate(prompt, n_samples=n_samples, rows=rows, cols=cols)


with gr.Blocks() as demo:
    gr.Markdown("Write a prompt to generate a Pokemon!")
    prompt = gr.Textbox(label="Prompt")
    image_output = gr.Image()
    generate = gr.Button("Generate")
    generate.click(generate_pkm, inputs=prompt, outputs=image_output)
gr.Markdown("Models finetuned by: lambdalabs")

if __name__ == "__main__":
    demo.launch()