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 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 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..."):
    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)