import streamlit as st
from utils import cargar_modelo, genera
"Pàgina principal"
st.title('Generador de mariposas')
st.write('Modelo light GAN')
#Barra lateral
st.sidebar.subheader('Esta mariposa será generada con IA')
#st.sidebar.image('ruta', width=200)
st.sidebar.caption('Demo GAN-light')

## Modelo

repo_id = 'ceyda/butterfly_cropped_uniq1K_512'
modelo_gan = cargar_modelo(repo_id)

n_mariposas = 4

def corre():
    with st.spinner('Generando, espera 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_button = st.Button('Genera mariposas', 
                         on_click=corre,
                         help = 'Presiona para generar')
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)