Update app.py

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import streamlit as st
 
 from models import Generator, Discriminrator
-from utils import image_to_base64
+from StyleMix import style_mix
 import torch
 import torchvision.transforms as T
 from torchvision.utils import make_grid
@@ -130,22 +130,16 @@ def main():
         unsafe_allow_html=True,
     )
 
-    st.header("Welcome to FastGAN")
-
     if choose == 'Model Gallery':
+        st.header("Welcome to FastGAN")
         show_model_summary(True)
     elif choose == 'Create your own images':
-        #checked = st.checkbox('Click here if you want to create one of your own !')
-
-    # if not checked:
-    #     show_model_summary(True)
-    # if checked:
-        # show_model_summary(False)
+        st.header("Generate images")
         col11, col12, col13 = st.columns([3,3.5,3.5])
         with col11:
             img_type = st.selectbox("Choose type of image to generate", index=0,
                                     options=["aurora", "anime", "painting", "fauvism", "shell", "universe", "grumpy cat", "moon gate"])
-        # 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 !')
@@ -172,8 +166,35 @@ def main():
                         st.image(gan_images[1], width=300)
                     else:
                         st.image(gan_images[1:], width=150)
+                        
+    elif choose == 'Style mix':
+        st.header("Style mix")
+        st.markdown(
+            """
+        <p style='text-align: left'>
+        Get the style representations of 2 images generated from the model to create a new one that mixes the style of two.
+        </p>
+                """,
+            unsafe_allow_html=True,
+        )
+        st.markdown("""___""")
+        col21, col22 = st.columns([3, 6])
+        with col21:
+            img_type = st.selectbox("Choose type of image to mix", index=0,
+                                    options=["aurora", "anime", "painting", "fauvism", "shell", "universe", "grumpy cat", "moon gate"])
+            number_imgs = st.slider('How many images you want to generate ?', min_value=1, max_value=3)
+            generate_button = st.button('Mix style')
 
-
+        if generate_button:
+            with col21:
+                with st.spinner(text=f"Mixing styles..."):
+                    mix_imgs = style_mix(model_name[img_type], number_imgs, device)
+                    mix_imgs = make_grid(mix_imgs, nrow=number_imgs+1, normalize=True)
+                    mix_imgs = mix_imgs.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
+                    mix_imgs = Image.fromarray(mix_imgs)
+            with col22:
+                st.image(mix_imgs, width=600)
+                
 
 if __name__ == '__main__':
     main()