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 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 = carga_modelo(repo_id)

#Generamos 4 mariposas
n_mariposas = 4

def corre():
    with st.spinner("Generando, espero 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_boton =st.button(
    "Genera mariposas",
    on_click=corre,
    help = "Estamos en vuelo, abrocha tu cinturon."
)           

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