import streamlit as st

from utils import carga_modelo, genera

##pagina principla
st.title('Generador de mariposas')
st.write('este es un model 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 model
repo_id = 'ceyda/butterfly_cropped_uniq1K_512'
modelo_gan = carga_modelo(repo_id)


## genera 4 mariposas

n_mariposas = 4  

def corre():
    with st.spinner('Generando, espera sentado...'):
        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_stat['ims']

corre_boton = st.button(
    'Genera mariposas',
    on_click= corre,
    help='Estamos en vuelo, abre la imaginacion'
)

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)