import streamlit as st

from utils import carga_modelo, genera

##pagina principal
st.title("Generador de mariposas")
st.write("Esto es un modelo Light GAN entrenado y utilizado con Platzi")

## barra lateral
st.sidebar.subheader("esta mariposa no existe! Puedes creerlo??")
st.sidebar.image("asset/logo.png", width=200)
st.sidebar.caption("Demo creado en un curso de Platzi")

## Cargamos el modelo
repo_id = "https://huggingface.co/ceyda/butterfly_cropped_uniq1K_512"
modelo_gan = carga_modelo(repo_id)

## Generamos 4 mariposas
n_mariposas = 4

def corre():
    with st.spinner("Generando, espera un poco..."):
        ims = genera(modelo_gan, n_mariposas)
        st.session_state["ims"] = ims

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

ims = st.session_state["ims"]

corre_buton = st.button(
    "Genera mariposas porfa",
    on_click= corre,
    help="Estamos en vuelo, abrocha el cinturón."
)

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_sidth=True)