# Streamlit Library
import streamlit as st 

from utils import load_model, generate

# Main Page
# Title
st.title("Butterflies Generator")
# Sumamry Description Model
st.write("This is Light GAN trained model")

# Sidebar
st.sidebar.subheader("This butterfly does not exist!, could you believe it?")
# Logo folder
st.sidebar.image("assets/logo.png", width=200)
#
st.sidebar.caption("Demo live created")

# Loading & Model's Name
repo_id = "ceyda/butterfly_cropped_uniq1K_512"
gan_model = load_model(repo_id)

# Butterflies Generator (4)
n_butterflies = 4


def run():
    with st.spinner("Loading, Please! be patient..."):
        # Inner Processing 
        ims = generate(gan_model, n_butterflies)
        # To save on
        st.session_state["ims"]
        
     
if "ims" not in st.session_state:
    st.session_state["ims"] = None
    run()
    
# Extracting info from ims    
ims = st.session_state["ims"]

#
run_button = st.button(
    "Please! Generate buttlerflies",
    on_click= run(),
    help="we are in flight, fasten your seatbelt."
)

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)