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

# Load environment variables from .env
load_dotenv()

# Access the variables
token = os.getenv("HF_AUTH_TOKEN")
api_key = os.getenv("API_KEY")
repo_id = 'ceyda/butterfly_cropped_uniq1K_512'
modelo_gan = carga_modelo(repo_id, token)


# 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?')
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'
try:
    modelo_gan = carga_modelo(repo_id)
except Exception as e:
    st.error(f"Error loading model: {str(e)}")
    st.stop()

# 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

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

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

ims = st.session_state['ims']

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)