import streamlit as st

from utils import load_model,generate

st.title("butterfly_generator")
st.write("this is a ligth GAN generator model")

st.sidebar.subheader("this butterfly is created by IA")
st.sidebar.image("assets/logo.png",width=200)

model_id="ceyda/butterfly_cropped_uniq1K_512"
model_gan=load_model(model_id)
n_fly=4
def run():
    with st.spinner("generating"):
        ims=generate(model_gan,n_fly)
        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(
    "click me",
    on_click=run
    )

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