import streamlit as st # HF spaces at v1.2.0
from demo import load_model,generate

# TODOs
# Add markdown short readme project intro
# project setup:

# git clone https://github.com/huggingface/community-events.git
# cd community-events
# pip install .
 
st.title("ButterflyGAN")
st.write("## This butterfly does not exist! ")
st.write("Demo prep still in progress!!")

@st.experimental_singleton
def load_model_intocache(model_name):
    
    # model_name='ceyda/butterfly_512_base'
    gan = load_model(model_name)

    return gan

model_name='ceyda/butterfly_cropped_uniq1K_512'
model=load_model_intocache(model_name)

st.write(f"Model {model_name} is loaded")
st.write(f"Latent dimension: {model.latent_dim}, Image size:{model.image_size}")

run=st.button("Generate")
if run:
    with st.spinner("Generating..."):
    
        batch_size=4 #generate 4 butterflies 
        ims=generate(model,batch_size)

        cols=st.columns(batch_size)
        for i,im in enumerate(ims):
            cols[i].image(im)