app.py

import streamlit as st
import torch

from utils import load_model, generate

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

# Sidebar
st.sidebar.subheader('This butterfly does not exist, it is completely generated!')
st.sidebar.image('assets/logo.png', 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)