Spaces:
Runtime error
Runtime error
| 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') | |
| 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() | |