import streamlit as st
from utils import carga_modelo, genera

## pagina principal

st.title("Butterfly Generator")
st.write("Generate your butterflies with one click!")

##Barra lateral

st.sidebar.subheader("ButterflyGAN")
st.sidebar.image("assets/mariposa.png", width = 200)
st.sidebar.caption("My first demo using Streamlit")


## Cargamos el modelo
repo_id = "ceyda/butterfly_cropped_uniq1K_512"
modelo_gan = carga_modelo(repo_id)

##Generamos 4 mariposas
n_mariposas = 4


def run_model():
    with st.spinner("Crafting your new butterfly :"):
        ims = genera(modelo_gan, n_mariposas)
    st.session_state["ims"] = ims
if "ims" not in st.session_state:
    st.session_state["ims"] = None
    run_model()
ims = st.session_state["ims"]

run_model_button = st.button(
"Generate",
on_click = run_model(),
help="Still in process"
)

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