import streamlit as st

from utils import load_model, generate


# main page
st.title("Butterflies generator")
st.write("LightGan Model to generate Butterflies")

# sidebar
st.sidebar.subheader("Butterfly generated with AI")
st.sidebar.image("assets/logo.png", width=200)
st.sidebar.caption("Demo")

# Load model
repo_id = "ceyda/butterfly_cropped_uniq1K_512"
model_version = "57d36a15546909557d9f967f47713236c8288838"
model_gan = load_model(repo_id, model_version)

# Generate 4 btfs
n_btfs = 4

def run():
    with st.spinner("Generating... "):
        ims = generate(model_gan, n_btfs)
        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 Butterfly",
    on_click = run, 
    help = "... flying ..."
)

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