import streamlit as st
from utils import carga_modelo, genera
import os


## Página principal
st.title('Generador de Mariposas')
st.write('Este es un modelo light GAN entrenado para generación de mariposas')

## Barra Lateral
st.sidebar.subheader('Esta mariposa no existe, ¿Puedes creerlo?')
logo_path = 'assets/logo.png'

if os.path.exists(logo_path):
    st.sidebar.image(logo_path, width=200)
else:
    st.sidebar.write("⚠️ Logo not found.")

st.sidebar.caption('Demo creado en vivo')

## Cargamos el modelo
repo_id = 'ceyda/butterfly_cropped_uniq1K_512'
modelo_gan = carga_modelo(repo_id)

## No de Mariposas
n_mariposas = 4

## Core de la app
def corre():
    with st.spinner('Generando espera un poco....'):
        ims = genera(modelo_gan, n_mariposas)
        st.session_state['ims'] = ims  # Corrected line

if 'ims' not in st.session_state:
    st.session_state['ims'] = None
    corre()

ims = st.session_state['ims']

corre_boton = st.button(
    'Genera mariposa por favor',
    on_click=corre, 
    help='Estamos en vuelo, abrocha tu cinturón'
)

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)