Miguel Angel commited on
Commit
e63b271
1 Parent(s): cb57041

Adding my first config on streamlit

Browse files
Files changed (2) hide show
  1. app.py +42 -0
  2. utils.py +17 -0
app.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+
3
+ from utils import load_model, generate
4
+
5
+ # Main page
6
+ st.title("Butterfly Generator")
7
+ st.write("This app generates butterfly images using a pretrained Light GAN model.")
8
+
9
+ # Sibebar
10
+ st.sidebar.subheader("This butterfly does not exist, could you believe it?")
11
+ st.sidebar.image("assets/logo.png", width=200)
12
+ st.sidebar.caption("Demo created on live.")
13
+
14
+ # Load model
15
+ id_rep = "ceyda/butterfly_cropped_uniq1K_512"
16
+ gan_model = load_model(id_rep)
17
+
18
+ # Generate butterflies
19
+ # n_butterflies = st.sidebar.slider("Number of butterflies", 1, 10, 5, 1)
20
+ n_butterflies = 4
21
+
22
+
23
+ def run():
24
+ with st.spinner("Loading model, please wait..."):
25
+ ims = generate(model=gan_model, batch_size=n_butterflies)
26
+ st.session_state["ims"] = ims
27
+
28
+
29
+ if "ims" not in st.session_state:
30
+ st.session_state["ims"] = None
31
+ run()
32
+
33
+ ims = st.session_state["ims"]
34
+
35
+ run_button = st.button("Generate new butterflies",
36
+ on_click=run, help="Click to generate new butterflies.")
37
+
38
+ if ims is not None:
39
+ cols = st.columns(n_butterflies)
40
+ for j, im in enumerate(cols):
41
+ i = j % n_butterflies
42
+ cols[i].image(im, use_column_width=True)
utils.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import torch
3
+ from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
4
+
5
+
6
+ def load_model(model_name="ceyda/butterfly_cropped_uniq1K_512", model_version=None):
7
+ model = LightweightGAN.from_pretrained(model_name, version=model_version)
8
+ model.eval()
9
+ return model
10
+
11
+
12
+ def generate(model, batch_size=1):
13
+ with torch.no_grad():
14
+ ims = model.G(torch.randn(batch_size, model.latent_dim)
15
+ ).clamp_(0.0, 1.0) * 255
16
+ ims = ims.permute(0, 2, 3, 1).detach().cpu().numpy().astype(np.uint8)
17
+ return ims