butterfly-gan / app.py
Ceyda Cinarel
add nearest neighbor
9fbe234
raw
history blame
1.91 kB
import re
import streamlit as st # HF spaces at v1.2.0
from demo import load_model,generate,get_dataset,embed
# 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
@st.experimental_singleton
def load_dataset():
dataset=get_dataset()
return dataset
model_name='ceyda/butterfly_cropped_uniq1K_512'
model=load_model_intocache(model_name)
dataset=load_dataset()
st.write(f"Model {model_name} is loaded")
st.write(f"Latent dimension: {model.latent_dim}, Image size:{model.image_size}")
if 'ims' not in st.session_state:
st.session_state['ims'] = None
ims=st.session_state["ims"]
batch_size=4 #generate 4 butterflies
def run():
with st.spinner("Generating..."):
ims=generate(model,batch_size)
st.session_state['ims'] = ims
runb=st.button("Generate", on_click=run)
if ims is not None:
cols=st.columns(batch_size)
picks=[False]*batch_size
for i,im in enumerate(ims):
cols[i].image(im)
picks[i]=cols[i].button("Find Nearest",key="pick_"+str(i))
# if picks[i]:
# scores, retrieved_examples=dataset.get_nearest_examples('beit_embeddings', embed(im), k=5)
# for r in retrieved_examples["image"]:
# st.image(r)
if any(picks):
# st.write("Nearest butterflies:")
for i,pick in enumerate(picks):
if pick:
scores, retrieved_examples=dataset.get_nearest_examples('beit_embeddings', embed(ims[i]), k=5)
for r in retrieved_examples["image"]:
cols[i].image(r)