import streamlit as st

from utils import load_model, generate_images

st.title("butterflies generator")
st.write(
    "this is a simple app to generate butterflies using a pretrained lightweight GAN model"
)


def run():
    with st.spinner("loading model..."):
        ims = generate_images(load_model(), n_images=1, batch_size=1)
        st.session_state.ims = ims


if "ims" not in st.session_state:
    st.session_state.ims = None
    run()


ims = st.session_state.ims

run_button = st.button("generate new butterfly", on_click=run)

if ims is not None:
    st.image(ims[0], caption="butterfly", use_column_width=True)