Upload app.py

app.py

@@ -18,9 +18,11 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model_name = {
     "aurora": 'huggan/fastgan-few-shot-aurora',
     "painting": 'huggan/fastgan-few-shot-painting',
-    #"painting":"geninhu/fastgan-few-shot-art",
     "shell": 'huggan/fastgan-few-shot-shells',
     "fauvism": 'huggan/fastgan-few-shot-fauvism-still-life',
+    "universe": 'huggan/fastgan-few-shot-universe',
+    "grumpy cat": 'huggan/fastgan-few-shot-grumpy-cat',
+    "anime": 'huggan/fastgan-few-shot-anime-face',
 }
 
 #@st.cache(allow_output_mutation=True)
@@ -54,6 +56,34 @@ def load_lottieurl(url: str):
         return None
     return r.json()
 
+def show_model_summary(expanded):
+    with st.expander("Model gallery", expanded=expanded):
+        col1, col2, col3, col4 = st.columns(4)
+        with col1:
+            st.markdown('Fauvism GAN [model](https://huggingface.co/huggan/fastgan-few-shot-fauvism-still-life)', unsafe_allow_html=True)
+            st.image('fauvism.png', width=200)
+
+            st.markdown('Painting GAN [model](https://huggingface.co/huggan/fastgan-few-shot-painting)', unsafe_allow_html=True)
+            st.image('painting.png', width=200)
+
+        with col2:
+            st.markdown('Aurora GAN [model](https://huggingface.co/huggan/fastgan-few-shot-aurora)', unsafe_allow_html=True)
+            st.image('aurora.png', width=200)
+
+            st.markdown('Universe GAN [model](https://huggingface.co/huggan/fastgan-few-shot-universe)', unsafe_allow_html=True)
+            st.image('universe.png', width=200)
+
+        with col3:
+            st.markdown('Anime GAN [model](https://huggingface.co/huggan/fastgan-few-shot-anime-face)', unsafe_allow_html=True)
+            st.image('anime.png', width=200)
+
+            st.markdown('Shell GAN [model](https://huggingface.co/huggan/fastgan-few-shot-shells)', unsafe_allow_html=True)
+            st.image('shell.png', width=200)
+
+        with col4:
+            st.markdown('Grumpy cat GAN [model](https://huggingface.co/huggan/fastgan-few-shot-grumpy-cat)', unsafe_allow_html=True)
+            st.image('grumpy_cat.png', width=200)
+
 
 def main():
 
@@ -91,36 +121,22 @@ def main():
     )
 
     st.header("Welcome to FastGAN")
+    checked = st.checkbox('Click here if you want to create one of your own !')
 
-    col1, col2, col3, col4 = st.columns([3,3,3,3])
-    with col1:
-        st.markdown('Fauvism GAN [model](https://huggingface.co/huggan/fastgan-few-shot-fauvism-still-life)', unsafe_allow_html=True)
-        st.image('fauvism.png', width=220)
-
-    with col2:
-        st.markdown('Aurora GAN [model](https://huggingface.co/huggan/fastgan-few-shot-aurora-bs8)', unsafe_allow_html=True)
-        st.image('aurora.png', width=220)
-
-    with col3:
-        st.markdown('Painting GAN [model](https://huggingface.co/huggan/fastgan-few-shot-painting-bs8)', unsafe_allow_html=True)
-        st.image('painting.png', width=220)
-    with col4:
-        st.markdown('Shell GAN [model](https://huggingface.co/huggan/fastgan-few-shot-shells)', unsafe_allow_html=True)
-        st.image('shell.png', width=220)
-
-    st.markdown('___')
-    if st.checkbox('Click if you want to create one of your own !'):
-
-        col11, col12, col13 = st.columns([3,3.5,3.5])
+    if not checked:
+        show_model_summary(True)
+    if checked:
+        col11, col12, col13 = st.columns([2.9,3.05,3.05])
         with col11:
-            img_type = st.selectbox("Choose type of image to generate", index=0, options=["aurora", "painting", "fauvism", "shell"])
+            img_type = st.selectbox("Choose type of image to generate", index=0,
+                                    options=["aurora", "anime", "painting", "fauvism", "shell", "universe", "grumpy cat"])
         # with col12:
             number_imgs = st.slider('How many images you want to generate ?', min_value=1, max_value=5)
             if number_imgs is None:
                 st.write('Invalid number ! Please insert number of images to generate !')
                 raise ValueError('Invalid number ! Please insert number of images to generate !')
 
-            generate_button = st.button('Get Image!')
+            generate_button = st.button('Get Image')
             if generate_button:
                 st.markdown("""
                     <small><i>Predictions may take up to 1 minute under high load. Please stand by.</i></small>