import streamlit as st
from utils import carga_modelo, genera

## pagina principal

st.title("Generador de Mariposas")
st.write("Este es un modelo Light GAN pre entrenado")

##Barra lateral

st.sidebar.subheader("Esta Mariposa no existe! haha")
st.sidebar.image("assets/mariposa.png", width = 200)
st.sidebar.caption("Dema creado usando 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("Generando Mariposa... :"):
        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(
"Genera Mariposa",
on_click = run_model(),
help="Estamos generando la Mariposa"
)

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)