Upload app.py

app.py

@@ -26,7 +26,7 @@ model_name = {
 def load_generator(model_name_or_path):
     generator = Generator(in_channels=256, out_channels=3)
     generator = generator.from_pretrained(model_name_or_path, in_channels=256, out_channels=3)
-    _ = generator.to('cuda')
+    _ = generator.to(device)
     _ = generator.eval()
 
     return generator
@@ -94,18 +94,18 @@ def main():
     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=250)
+        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=250)
+        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=250)
+        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=250)
+        st.image('shell.png', width=220)
 
     st.markdown('___')
     if st.checkbox('Click if you want to create one of your own !'):
@@ -122,9 +122,7 @@ def main():
             generate_button = st.button('Get Image!')
             if generate_button:
                 st.markdown("""
-                    <small><i>Predictions may take up to 1 minute under high load.</i></small>
-                    </br>
-                    <small><i>Please stand by.</i></small>
+                    <small><i>Predictions may take up to 1 minute under high load. Please stand by.</i></small>
                 """,
                 unsafe_allow_html=True,)
 
@@ -132,13 +130,13 @@ def main():
             generator = load_generator(model_name[img_type])
             gan_images = generate_images(generator, number_imgs)
             with col12:
-                st.image(gan_images[0], width=400)
+                st.image(gan_images[0], width=300)
             if len(gan_images) > 1:
                 with col13:
                     if len(gan_images) <= 2:
-                        st.image(gan_images[1], width=400)
+                        st.image(gan_images[1], width=300)
                     else:
-                        st.image(gan_images[1:], width=200)
+                        st.image(gan_images[1:], width=150)