Adding my first config on streamlit

app.py

@@ -0,0 +1,42 @@
+import streamlit as st
+
+from utils import load_model, generate
+
+# Main page
+st.title("Butterfly Generator")
+st.write("This app generates butterfly images using a pretrained Light GAN model.")
+
+# Sibebar
+st.sidebar.subheader("This butterfly does not exist, could you believe it?")
+st.sidebar.image("assets/logo.png", width=200)
+st.sidebar.caption("Demo created on live.")
+
+# Load model
+id_rep = "ceyda/butterfly_cropped_uniq1K_512"
+gan_model = load_model(id_rep)
+
+# Generate butterflies
+# n_butterflies = st.sidebar.slider("Number of butterflies", 1, 10, 5, 1)
+n_butterflies = 4
+
+
+def run():
+    with st.spinner("Loading model, please wait..."):
+        ims = generate(model=gan_model, batch_size=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 new butterflies",
+                       on_click=run, help="Click to generate new butterflies.")
+
+if ims is not None:
+    cols = st.columns(n_butterflies)
+    for j, im in enumerate(cols):
+        i = j % n_butterflies
+        cols[i].image(im, use_column_width=True)

utils.py

@@ -0,0 +1,17 @@
+import numpy as np
+import torch
+from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
+
+
+def load_model(model_name="ceyda/butterfly_cropped_uniq1K_512", model_version=None):
+    model = LightweightGAN.from_pretrained(model_name, version=model_version)
+    model.eval()
+    return model
+
+
+def generate(model, batch_size=1):
+    with torch.no_grad():
+        ims = model.G(torch.randn(batch_size, model.latent_dim)
+                      ).clamp_(0.0, 1.0) * 255
+        ims = ims.permute(0, 2, 3, 1).detach().cpu().numpy().astype(np.uint8)
+    return ims