import streamlit as st

from utils import load_model, generate

## Main page

st.title("Butterfly generator")
st.write("This is a GAN lite model as a basic example")


## Lateral sidebar
st.sidebar.subheader("This Butterfly doesn't exists")
st.sidebar.image("assets/logo.png", width=200)
st.sidebar.caption("Demo created live.")

## Main content

repo_id = "ceyda/butterfly_cropped_uniq1K_512"
gan_model = load_model(repo_id)
n_butterfly = 4
def run():
    with st.spinner("Generating..."):
        ims = generate(gan_model, n_butterfly)
        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 Butterflies",
                        on_click=run,
                        help="Click to generate butterflies!")
if ims is not None:
    cols = st.columns(n_butterfly)
    for j, im in enumerate(ims):
        i = j % n_butterfly
        cols[i].image(im, use_column_width=True)