app.py

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()