import streamlit as st

from utils import load_model, genera

## Mean page
st.title("Butterflies Generator")
st.write("This is a Light GAN model trained and used by me")

## Sidebar
st.sidebar.subheader("This butterfly not is real, can you believe it?")
st.sidebar.image("assets/logo.png", width=200)
st.sidebar.caption("Demo created live")

## Load model
repo_id = "ceyda/butterfly_cropped_uniq1K_512"
model_gan = load_model(repo_id)

## We generated four butterflies
n_butterflies = 4

def run():
    with st.spinner("Generating, wait a bit..."):
        ims = genera(model_gan, 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",
    on_click=run,
    help="We are in fly"
)

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)