import streamlit as st

from utils import load_model, generate

## Main page
st.title("Butterfly generator")
st.write("Modelo Light GAN")

## Lateral bar
st.sidebar.subheader("No existing butterfly")
st.sidebar.image("assets/logo.png", width=200)
st.sidebar.caption("In live demo")

## Loading the model
repo_id = "ceyda/butterfly_croped_uniq1K_512"
model_gan = load_model(repo_id)

## Generate 4 butterflies
n_butterflies = 4


def execute():
    with st.spinner("Generating..."):
        ims = generate(model_gan, n_butterflies)
        st.session_state["ims"] = ims

if "ims" not in st.session_state:
    st.session_state["ims"] = None
    execute()

ims = st.session_state["ims"]

exec_button = st.button(
    "Generate!",
    on_click=execute,
    help="Help message"
)

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