Upload app.py

app.py

@@ -0,0 +1,146 @@
+import pandas as pd
+import numpy as np
+import streamlit as st
+
+from models import Generator, Discriminrator
+from utils import image_to_base64
+import torch
+import torchvision.transforms as T
+from torchvision.utils import make_grid
+from PIL import Image
+
+from streamlit_lottie import st_lottie
+import requests
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+model_name = {
+    "aurora": 'huggan/fastgan-few-shot-aurora',
+    "painting": 'huggan/fastgan-few-shot-painting',
+    "shell": 'huggan/fastgan-few-shot-shells',
+    "fauvism": 'huggan/fastgan-few-shot-fauvism-still-life',
+}
+
+#@st.cache(allow_output_mutation=True)
+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.eval()
+
+    return generator
+
+def _denormalize(input: torch.Tensor) -> torch.Tensor:
+    return (input * 127.5) + 127.5
+
+
+def generate_images(generator, number_imgs):
+    noise = torch.zeros(number_imgs, 256, 1, 1, device='cuda').normal_(0.0, 1.0)
+    with torch.no_grad():
+        gan_images, _ = generator(noise)
+
+    gan_images = _denormalize(gan_images.detach()).cpu()
+    gan_images = [i for i in gan_images]
+    gan_images = [make_grid(i, nrow=1, normalize=True) for i in gan_images]
+    gan_images = [i.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() for i in gan_images]
+    gan_images = [Image.fromarray(i) for i in gan_images]
+    return gan_images
+
+def load_lottieurl(url: str):
+    r = requests.get(url)
+    if r.status_code != 200:
+        return None
+    return r.json()
+
+
+def main():
+
+    st.set_page_config(
+        page_title="FastGAN Generator",
+        page_icon="🖥️",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
+
+    lottie_penguin = load_lottieurl('https://assets7.lottiefiles.com/packages/lf20_mm4bsl3l.json')
+
+    with st.sidebar:
+        st_lottie(lottie_penguin, height=200)
+    st.sidebar.markdown(
+        """
+    ___
+    <p style='text-align: center'>
+    FastGAN is an few-shot GAN model that generates images of several types!
+    </p>
+    <p style='text-align: center'>
+    Model training and Space creation by
+    <br/>
+    <a href="https://huggingface.co/vumichien" target="_blank">Chien Vu</a> | <a href="https://huggingface.co/geninhu" target="_blank">Nhu Hoang</a>
+    <br/>
+    </p>
+
+    <p style='text-align: center'>
+    based on
+    <br/>
+    <a href="https://github.com/silentz/Towards-Faster-And-Stabilized-GAN-Training-For-High-Fidelity-Few-Shot-Image-Synthesis" target="_blank">FastGAN model</a> | <a href="https://arxiv.org/abs/2101.04775" target="_blank">Article</a>
+    </p>
+            """,
+        unsafe_allow_html=True,
+    )
+
+    st.header("Welcome to FastGAN")
+
+    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)
+
+    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)
+
+    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)
+    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.markdown('___')
+    if st.checkbox('Click if you want to create one of your own !'):
+
+        col11, col12, col13 = st.columns([3,3,3])
+        with col11:
+            img_type = st.selectbox("Choose type of image to generate", index=0, options=["aurora", "painting", "fauvism", "shell"])
+        # 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!')
+            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>
+                """,
+                unsafe_allow_html=True,)
+                
+        if generate_button:
+            generator = load_generator(model_name[img_type])
+            gan_images = generate_images(generator, number_imgs)
+            with col12:
+                st.image(gan_images[0], width=400)
+            if len(gan_images) > 1:
+                with col13:
+                    if len(gan_images) <= 2:
+                        st.image(gan_images[1], width=400)
+                    else:
+                        st.image(gan_images[1:], width=200)
+
+
+
+if __name__ == '__main__':
+    main()