import streamlit as st
from utils import cargar_mdoel, general

## Pagina principal

st.title("Generaedor de mariposas")
st.write("Este es un modelo Linght GAN entrenado y utilizado con platzi!!!")

##barra lateral
st.sidebar.subheader("!Esta mariposa no existe, ¿Puedes creerlo?")
st.sidebar.image("assets/logo.png", width=200)
st.sidebar.caption("Demo creado en vivo.")

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

#Generamos 4 mariposas
n_mariposas = 4

def corre():
    with st.spinner("Generando, espera un poco..."):
        ims = general( 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_boton = st.button(
    "Genera mariposas porfa",
    on_click=corre,
    help="Estamos en vuelo, abrocha tu 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(ims, use_column_width=True)