import streamlit as st
import os

from utils import load_model, generate

st.title('Image Generator with GAN')
st.markdown('Butterflies generator')

st.write("Token Access:", st.secrets["access_token"])


# Sidebar
st.sidebar.subheader('This butterfly does not exist, it is completely generated!')
st.sidebar.image('assets/logo.jpeg', width=200)
st.sidebar.caption('Demo was just created.')

# Loading model
repo_id = 'ceyda/butterfly_cropped_uniq1K_512'
gan_model = load_model(repo_id)

# Generate 4 butterflies
n_butterflies = 4

def run():
    with st.spinner('Generating...'):
        ims = generate(gan_model, n_butterflies)
        st.session_state['ims'] = ims

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

ims = st.session_state['ims']

run_button = st.button(
    'Generate butterflies',
    on_click=run,
    help='We are about to start to generate'
)

if ims is not None:
    cols = st.columns(n_butterflies)
    for j, im in enumerate(ims):
        i = j % n_butterflies
        cols[i].image(im, use_column_width=True)