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=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 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( """ ___

FastGAN is an few-shot GAN model that generates images of several types!

Model training and Space creation by
Chien Vu | Nhu Hoang

based on
FastGAN model | Article

""", 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(""" Predictions may take up to 1 minute under high load.
Please stand by. """, 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()