import streamlit as st
from utils import carga_modelo, genera

## Página principal
st.title("Generador de mariposas")
st.write(
    "Este es un generador de mariposas que utiliza una red neuronal generativa (GAN)"
)

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

# Numero de mariposas a generar
n_mariposas = 4


def corre():
    with st.spinner("Generando mariposa..."):
        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(
    "Generar mariposas", on_click=corre, help="Genera mariposas aleatorias."
)

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)