import streamlit as st

from utils import load_model, generate

# Main page
st.title("Butterfly Generator")
st.write("This app generates butterfly images using a pretrained Light GAN model.")

# Sibebar
st.sidebar.subheader("This butterfly does not exist, could you believe it?")
st.sidebar.image("assets/logo.png", width=200)
st.sidebar.caption("Demo created on live.")

# Load model
id_rep = "ceyda/butterfly_cropped_uniq1K_512"
gan_model = load_model(id_rep)

# Generate butterflies
# n_butterflies = st.sidebar.slider("Number of butterflies", 1, 10, 5, 1)
n_butterflies = 4


def run():
    with st.spinner("Loading model, please wait..."):
        ims = generate(model=gan_model, batch_size=n_butterflies)
        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 butterflies",
                       on_click=run, help="Click to generate new butterflies.")

if ims is not None:
    cols = st.columns(n_butterflies)
    for j, im in enumerate(cols):
        i = j % n_butterflies
        cols[i].image(im, use_column_width=True)