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')

#Barra lateral
st.sidebar.subheader('!Esta mariposa no existe, puedes creelo')
st.sidebar.image('assets/logo.png', width=200)
st.sidebar.caption('Demo creado')

#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, 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 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)