File size: 9,756 Bytes
5e5faf7
 
 
 
 
98c3ab4
5e5faf7
 
 
 
 
 
e196b57
5e5faf7
 
 
 
 
 
 
4e206a5
5e5faf7
 
8d454d8
 
 
5c9af2e
5e5faf7
 
 
 
 
 
e728a84
5e5faf7
 
 
 
 
 
 
 
 
0262342
5e5faf7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8d454d8
c3b9254
9ffffed
8d454d8
 
 
9ffffed
8d454d8
9ffffed
8d454d8
 
 
9ffffed
8d454d8
9ffffed
8d454d8
 
 
9ffffed
8d454d8
9ffffed
8d454d8
 
 
9ffffed
5c9af2e
9ffffed
 
 
 
8f812f4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9ffffed
8f812f4
9ffffed
8f812f4
9ffffed
 
8f812f4
9ffffed
8f812f4
9ffffed
 
8f812f4
9ffffed
8f812f4
9ffffed
 
8f812f4
8d454d8
5e5faf7
 
 
 
 
 
 
 
 
 
 
 
 
 
661c6e3
 
 
5c9af2e
661c6e3
5c9af2e
 
e196b57
5e5faf7
 
 
 
661c6e3
 
 
5e5faf7
661c6e3
 
5e5faf7
9ffffed
5e5faf7
 
 
 
 
 
e196b57
98c3ab4
8d454d8
661c6e3
98c3ab4
15822d1
5e5faf7
8d454d8
c3b9254
9ffffed
5e5faf7
 
 
 
 
8d454d8
5e5faf7
 
e728a84
5e5faf7
 
0262342
5e5faf7
5c9af2e
 
 
 
 
5e5faf7
e728a84
5e5faf7
 
 
e728a84
5e5faf7
e728a84
9ffffed
661c6e3
 
98c3ab4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5e5faf7
98c3ab4
 
 
 
 
 
 
 
 
9ffffed
5e5faf7
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
import pandas as pd
import numpy as np
import streamlit as st

from models import Generator, Discriminrator
from StyleMix import style_mix
import torch
import torchvision.transforms as T
from torchvision.utils import make_grid
from PIL import Image

from streamlit_lottie import st_lottie
from streamlit_option_menu import option_menu
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',
    "universe": 'huggan/fastgan-few-shot-universe',
    "grumpy cat": 'huggan/fastgan-few-shot-grumpy-cat',
    "anime": 'huggan/fastgan-few-shot-anime-face',
    "moon gate": 'huggan/fastgan-few-shot-moongate',
}

#@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(device)
    _ = 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=device).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 show_model_summary(expanded):
    st.subheader("Model gallery")
    with st.expander('Image 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('assets/image/fauvism.png', width=200)
            st.markdown('Painting GAN [model](https://huggingface.co/huggan/fastgan-few-shot-painting)', unsafe_allow_html=True)
            st.image('assets/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('assets/image/aurora.png', width=200)
            st.markdown('Universe GAN [model](https://huggingface.co/huggan/fastgan-few-shot-universe)', unsafe_allow_html=True)
            st.image('assets/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('assets/image/anime.png', width=200)
            st.markdown('Shell GAN [model](https://huggingface.co/huggan/fastgan-few-shot-shells)', unsafe_allow_html=True)
            st.image('assets/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('assets/image/grumpy_cat.png', width=200)
            st.markdown('Moon gate GAN [model](https://huggingface.co/huggan/fastgan-few-shot-moongate)', unsafe_allow_html=True)
            st.image('assets/image/moon_gate.png', width=200)

    with st.expander('Video gallery', expanded=False):
        cols=st.columns(4)
        # cols[0].write("Universe GAN")
        # cols[0].image('assets/video/universe.gif')
        # cols[0].write("Fauvism still life GAN")
        # cols[0].image('assets/video/fauvism.gif')
        #
        # cols[1].write("Aurora GAN")
        # cols[1].image('assets/video/aurora.gif')
        # cols[1].write("Moon gate GAN")
        # cols[1].image('assets/video/moon_gate.gif')
        #
        # cols[2].write("Anime GAN")
        # cols[2].image('assets/video/anime.gif')
        # cols[2].write("Painting GAN")
        # cols[2].image('assets/video/painting.gif')
        #
        # cols[3].write("Grumpy cat GAN")
        # cols[3].image('assets/video/grumpy_cat.gif')

        cols[0].write("Universe GAN")
        cols[0].video('assets/video/universe.mp4')
        cols[0].write("Fauvism still life GAN")
        cols[0].video('assets/video/fauvism.mp4')

        cols[1].write("Aurora GAN")
        cols[1].video('assets/video/aurora.mp4')
        cols[1].write("Moon gate GAN")
        cols[1].video('assets/video/moongate.mp4')

        cols[2].write("Anime GAN")
        cols[2].video('assets/video/anime.mp4')
        cols[2].write("Painting GAN")
        cols[2].video('assets/video/painting.mp4')

        cols[3].write("Grumpy cat GAN")
        cols[3].video('assets/video/grumpy.mp4')


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)
        choose = option_menu("FastGAN", ["Model Gallery", "Generate images", "Mix style"],
                             icons=['collection', 'file-plus', 'intersect'],
                             menu_icon="infinity", default_index=0,
                             styles={
                                "container": {"padding": ".0rem", "font-size": "14px"},
                                "nav-link-selected": {"color": "#000000", "font-size": "16px"},
                            }
                             )
    st.sidebar.markdown(
        """
    ___
    <p style='text-align: center'>
    FastGAN is a few-shot GAN model trained on high-fidelity images which requires less computation resource and samples for training.
    <br/>
    <a href="https://arxiv.org/abs/2101.04775" target="_blank">Article</a>
    </p>
    <p style='text-align: center; font-size: 14px;'>
    Model training and Spaces creating by
    <br/>
    <a href="https://www.linkedin.com/in/vumichien/" target="_blank">Chien Vu</a> | <a href="https://www.linkedin.com/in/nhu-hoang/" target="_blank">Nhu Hoang</a>
    <br/>
    </p>
            """,
        unsafe_allow_html=True,
    )

    if choose == 'Model Gallery':
        st.header("Welcome to FastGAN")
        show_model_summary(True)
    elif choose == 'Generate images':
        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"])

            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. Please stand by.</i></small>
                """,
                unsafe_allow_html=True,)

        if generate_button:
            with col11:
                with st.spinner(text=f"Loading selected model..."):
                    generator = load_generator(model_name[img_type])
                with st.spinner(text=f"Generating images..."):
                    gan_images = generate_images(generator, number_imgs)
            with col12:
                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=300)
                    else:
                        st.image(gan_images[1:], width=150)

    elif choose == 'Mix style':
        st.header("Mix style")
        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()