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 entrenando y utilizado con Platzi")(

## Barra lateral
st.sidebar.subheader("Esta mariposa no existe, ¿Puedes creeerlo?")
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=carga_modelo(repo_id)

##Generamos 4 mariposas
n_mariposas=4

def corre():
    with st.spinner("Generando,espera un poco"):
        ims=genera(model_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 cinturon"
)

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)