Duplicate from huggan/FastGan

Co-authored-by: Nhu Hoang <[email protected]>

.gitattributes

@@ -0,0 +1,47 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+aurora.png filter=lfs diff=lfs merge=lfs -text
+fauvism.png filter=lfs diff=lfs merge=lfs -text
+painting.png filter=lfs diff=lfs merge=lfs -text
+shell.png filter=lfs diff=lfs merge=lfs -text
+grumpy_cat.png filter=lfs diff=lfs merge=lfs -text
+universe.png filter=lfs diff=lfs merge=lfs -text
+moon_gate.png filter=lfs diff=lfs merge=lfs -text
+assets/video/anime.gif filter=lfs diff=lfs merge=lfs -text
+assets/video/aurora.gif filter=lfs diff=lfs merge=lfs -text
+assets/video/fauvism.gif filter=lfs diff=lfs merge=lfs -text
+assets/video/grumpy_cat.gif filter=lfs diff=lfs merge=lfs -text
+assets/video/moon_gate.gif filter=lfs diff=lfs merge=lfs -text
+assets/video/painting.gif filter=lfs diff=lfs merge=lfs -text
+assets/video/universe.gif filter=lfs diff=lfs merge=lfs -text
+assets/video/anime.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/video/fauvism.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/video/universe.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/video/moongate.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/video/painting.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: FastGan
+emoji: 😎
+colorFrom: indigo
+colorTo: pink
+sdk: streamlit
+sdk_version: 1.2.0
+app_file: app.py
+pinned: true
+duplicated_from: huggan/FastGan
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

StyleMix.py

@@ -0,0 +1,70 @@
+import torch
+from torch import nn
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.utils.data.dataloader import DataLoader
+from torchvision import transforms
+from torchvision import utils as vutils
+
+from models import Generator
+from utils import copy_G_params, load_params
+
+
+
+def get_early_features(net, noise):
+    with torch.no_grad():
+        feat_4 = net._init(noise)
+        feat_8 = net._upsample_8(feat_4)
+        feat_16 = net._upsample_16(feat_8)
+        feat_32 = net._upsample_32(feat_16)
+        feat_64 = net._upsample_64(feat_32)
+    return feat_8, feat_16, feat_32, feat_64
+
+def get_late_features(net, feat_64, feat_8, feat_16, feat_32):
+    with torch.no_grad():
+        feat_128 = net._upsample_128(feat_64)
+        feat_128 = net._sle_128(feat_8, feat_128)
+
+        feat_256 = net._upsample_256(feat_128)
+        feat_256 = net._sle_256(feat_16, feat_256)
+
+        feat_512 = net._upsample_512(feat_256)
+        feat_512 = net._sle_512(feat_32, feat_512)
+
+        feat_1024 = net._upsample_1024(feat_512)
+
+    return net._out_1024(feat_1024)
+
+def style_mix(model_name_or_path, bs, device):
+    _in_channels = 256
+    im_size = 1024
+
+    netG = Generator(in_channels=_in_channels, out_channels=3)
+    netG = netG.from_pretrained(model_name_or_path, in_channels=256, out_channels=3)
+    _ = netG.to(device)
+    _ = netG.eval()
+
+    avg_param_G = copy_G_params(netG)
+    load_params(netG, avg_param_G)
+
+    noise_a = torch.randn(bs, 256, 1, 1, device=device).to(device)
+    noise_b = torch.randn(bs, 256, 1, 1, device=device).to(device)
+
+    feat_8_a, feat_16_a, feat_32_a, feat_64_a = get_early_features(netG, noise_a)
+    feat_8_b, feat_16_b, feat_32_b, feat_64_b = get_early_features(netG, noise_b)
+
+    images_b = get_late_features(netG, feat_64_b, feat_8_b, feat_16_b, feat_32_b)
+    images_a = get_late_features(netG, feat_64_a, feat_8_a, feat_16_a, feat_32_a)
+
+    imgs = [ torch.ones(1, 3, im_size, im_size) ]
+
+    imgs.append(images_b.cpu())
+    for i in range(bs):
+        imgs.append(images_a[i].unsqueeze(0).cpu())
+        gimgs = get_late_features(netG, feat_64_a[i].unsqueeze(0).repeat(bs, 1, 1, 1), feat_8_b, feat_16_b, feat_32_b)
+        imgs.append(gimgs.cpu())
+
+    imgs = torch.cat(imgs)
+    # vutils.save_image(imgs.add(1).mul(0.5), 'style_mix/style_mix_2.jpg', nrow=bs+1)
+
+    return imgs

app.py

@@ -0,0 +1,217 @@
+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=True):
+        cols=st.columns(4)
+
+        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()

assets/video/anime.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:040ef94e35a65978826df636850403c72a8ad8ca97432f0e8b543db9e1474b08
+size 3398750

assets/video/fauvism.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6fe077eefbd400b8202876f8a7f1de2982a11a3e4e6e68ef2ed7f85eb398ab1
+size 1573497

assets/video/moongate.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba99ee59ea1c330dcad52c6448e726090bf741ac115d4a31765bd85d8316e85c
+size 2861613

assets/video/painting.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f172256a1373a8c513aaf097f61ec69df14dbdce08c0d765b3ecd92e132b9c56
+size 1477719

assets/video/universe.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d7ed5c0b077180d5bb9975e2a2051fca3bf67d8e1e59bfca8a8b31728c63271
+size 15562186

layers.py

@@ -0,0 +1,272 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.modules.batchnorm import BatchNorm2d
+from torch.nn.utils import spectral_norm
+
+
+class SpectralConv2d(nn.Module):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self._conv = spectral_norm(
+                nn.Conv2d(*args, **kwargs)
+            )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return self._conv(input)
+
+
+class SpectralConvTranspose2d(nn.Module):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self._conv = spectral_norm(
+                nn.ConvTranspose2d(*args, **kwargs)
+            )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return self._conv(input)
+
+
+class Noise(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self._weight = nn.Parameter(
+                torch.zeros(1),
+                requires_grad=True,
+            )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        batch_size, _, height, width = input.shape
+        noise = torch.randn(batch_size, 1, height, width, device=input.device)
+        return self._weight * noise + input
+
+
+class InitLayer(nn.Module):
+
+    def __init__(self, in_channels: int,
+                       out_channels: int):
+        super().__init__()
+
+        self._layers = nn.Sequential(
+                SpectralConvTranspose2d(
+                        in_channels=in_channels,
+                        out_channels=out_channels * 2,
+                        kernel_size=4,
+                        stride=1,
+                        padding=0,
+                        bias=False,
+                    ),
+                nn.BatchNorm2d(num_features=out_channels * 2),
+                nn.GLU(dim=1),
+            )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return self._layers(input)
+
+
+class SLEBlock(nn.Module):
+
+    def __init__(self, in_channels: int,
+                       out_channels: int):
+        super().__init__()
+
+        self._layers = nn.Sequential(
+                nn.AdaptiveAvgPool2d(output_size=4),
+                SpectralConv2d(
+                        in_channels=in_channels,
+                        out_channels=out_channels,
+                        kernel_size=4,
+                        stride=1,
+                        padding=0,
+                        bias=False,
+                    ),
+                nn.SiLU(),
+                SpectralConv2d(
+                        in_channels=out_channels,
+                        out_channels=out_channels,
+                        kernel_size=1,
+                        stride=1,
+                        padding=0,
+                        bias=False,
+                    ),
+                nn.Sigmoid(),
+            )
+
+    def forward(self, low_dim: torch.Tensor,
+                      high_dim: torch.Tensor) -> torch.Tensor:
+        return high_dim * self._layers(low_dim)
+
+
+class UpsampleBlockT1(nn.Module):
+
+    def __init__(self, in_channels: int,
+                       out_channels: int):
+        super().__init__()
+
+        self._layers = nn.Sequential(
+                nn.Upsample(scale_factor=2, mode='nearest'),
+                SpectralConv2d(
+                        in_channels=in_channels,
+                        out_channels=out_channels * 2,
+                        kernel_size=3,
+                        stride=1,
+                        padding='same',
+                        bias=False,
+                    ),
+                nn.BatchNorm2d(num_features=out_channels * 2),
+                nn.GLU(dim=1),
+            )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return self._layers(input)
+
+
+class UpsampleBlockT2(nn.Module):
+
+    def __init__(self, in_channels: int,
+                       out_channels: int):
+        super().__init__()
+
+        self._layers = nn.Sequential(
+                nn.Upsample(scale_factor=2, mode='nearest'),
+                SpectralConv2d(
+                        in_channels=in_channels,
+                        out_channels=out_channels * 2,
+                        kernel_size=3,
+                        stride=1,
+                        padding='same',
+                        bias=False,
+                    ),
+                Noise(),
+                BatchNorm2d(num_features=out_channels * 2),
+                nn.GLU(dim=1),
+                SpectralConv2d(
+                        in_channels=out_channels,
+                        out_channels=out_channels * 2,
+                        kernel_size=3,
+                        stride=1,
+                        padding='same',
+                        bias=False,
+                    ),
+                Noise(),
+                nn.BatchNorm2d(num_features=out_channels * 2),
+                nn.GLU(dim=1),
+            )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return self._layers(input)
+
+
+class DownsampleBlockT1(nn.Module):
+
+    def __init__(self, in_channels: int,
+                       out_channels: int):
+        super().__init__()
+
+        self._layers = nn.Sequential(
+                SpectralConv2d(
+                        in_channels=in_channels,
+                        out_channels=out_channels,
+                        kernel_size=4,
+                        stride=2,
+                        padding=1,
+                        bias=False,
+                    ),
+                nn.BatchNorm2d(num_features=out_channels),
+                nn.LeakyReLU(negative_slope=0.2),
+            )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return self._layers(input)
+
+
+class DownsampleBlockT2(nn.Module):
+
+    def __init__(self, in_channels: int,
+                       out_channels: int):
+        super().__init__()
+
+        self._layers_1 = nn.Sequential(
+                SpectralConv2d(
+                        in_channels=in_channels,
+                        out_channels=out_channels,
+                        kernel_size=4,
+                        stride=2,
+                        padding=1,
+                        bias=False,
+                    ),
+                nn.BatchNorm2d(num_features=out_channels),
+                nn.LeakyReLU(negative_slope=0.2),
+                SpectralConv2d(
+                        in_channels=out_channels,
+                        out_channels=out_channels,
+                        kernel_size=3,
+                        stride=1,
+                        padding='same',
+                        bias=False,
+                    ),
+                nn.BatchNorm2d(num_features=out_channels),
+                nn.LeakyReLU(negative_slope=0.2),
+            )
+
+        self._layers_2 = nn.Sequential(
+                nn.AvgPool2d(
+                        kernel_size=2,
+                        stride=2,
+                    ),
+                SpectralConv2d(
+                        in_channels=in_channels,
+                        out_channels=out_channels,
+                        kernel_size=1,
+                        stride=1,
+                        padding=0,
+                        bias=False,
+                    ),
+                nn.BatchNorm2d(num_features=out_channels),
+                nn.LeakyReLU(negative_slope=0.2),
+            )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        t1 = self._layers_1(input)
+        t2 = self._layers_2(input)
+        return (t1 + t2) / 2
+
+
+class Decoder(nn.Module):
+
+    def __init__(self, in_channels: int,
+                       out_channels: int):
+        super().__init__()
+
+        self._channels = {
+                16:   128,
+                32:   64,
+                64:   64,
+                128:  32,
+                256:  16,
+                512:  8,
+                1024: 4,
+            }
+
+        self._layers = nn.Sequential(
+                nn.AdaptiveAvgPool2d(output_size=8),
+                UpsampleBlockT1(in_channels=in_channels,        out_channels=self._channels[16]),
+                UpsampleBlockT1(in_channels=self._channels[16], out_channels=self._channels[32]),
+                UpsampleBlockT1(in_channels=self._channels[32], out_channels=self._channels[64]),
+                UpsampleBlockT1(in_channels=self._channels[64], out_channels=self._channels[128]),
+                SpectralConv2d(
+                        in_channels=self._channels[128],
+                        out_channels=out_channels,
+                        kernel_size=3,
+                        stride=1,
+                        padding='same',
+                        bias=False,
+                    ),
+                nn.Tanh(),
+            )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return self._layers(input)

models.py

@@ -0,0 +1,245 @@
+import torch
+import torch.nn as nn
+
+from typing import Any, Tuple, Union
+
+from utils import (
+    ImageType,
+    crop_image_part,
+)
+
+from layers import (
+    SpectralConv2d,
+    InitLayer,
+    SLEBlock,
+    UpsampleBlockT1,
+    UpsampleBlockT2,
+    DownsampleBlockT1,
+    DownsampleBlockT2,
+    Decoder,
+)
+
+from huggan.pytorch.huggan_mixin import HugGANModelHubMixin
+
+
+class Generator(nn.Module, HugGANModelHubMixin):
+
+    def __init__(self, in_channels: int,
+                       out_channels: int):
+        super().__init__()
+
+        self._channels = {
+                4:    1024,
+                8:    512,
+                16:   256,
+                32:   128,
+                64:   128,
+                128:  64,
+                256:  32,
+                512:  16,
+                1024: 8,
+            }
+
+        self._init = InitLayer(
+                in_channels=in_channels,
+                out_channels=self._channels[4],
+            )
+
+        self._upsample_8    = UpsampleBlockT2(in_channels=self._channels[4],   out_channels=self._channels[8]   )
+        self._upsample_16   = UpsampleBlockT1(in_channels=self._channels[8],   out_channels=self._channels[16]  )
+        self._upsample_32   = UpsampleBlockT2(in_channels=self._channels[16],  out_channels=self._channels[32]  )
+        self._upsample_64   = UpsampleBlockT1(in_channels=self._channels[32],  out_channels=self._channels[64]  )
+        self._upsample_128  = UpsampleBlockT2(in_channels=self._channels[64],  out_channels=self._channels[128] )
+        self._upsample_256  = UpsampleBlockT1(in_channels=self._channels[128], out_channels=self._channels[256] )
+        self._upsample_512  = UpsampleBlockT2(in_channels=self._channels[256], out_channels=self._channels[512] )
+        self._upsample_1024 = UpsampleBlockT1(in_channels=self._channels[512], out_channels=self._channels[1024])
+
+        self._sle_64  = SLEBlock(in_channels=self._channels[4],  out_channels=self._channels[64] )
+        self._sle_128 = SLEBlock(in_channels=self._channels[8],  out_channels=self._channels[128])
+        self._sle_256 = SLEBlock(in_channels=self._channels[16], out_channels=self._channels[256])
+        self._sle_512 = SLEBlock(in_channels=self._channels[32], out_channels=self._channels[512])
+
+        self._out_128 = nn.Sequential(
+                SpectralConv2d(
+                    in_channels=self._channels[128],
+                    out_channels=out_channels,
+                    kernel_size=1,
+                    stride=1,
+                    padding='same',
+                    bias=False,
+                ),
+                nn.Tanh(),
+            )
+
+        self._out_1024 = nn.Sequential(
+                SpectralConv2d(
+                    in_channels=self._channels[1024],
+                    out_channels=out_channels,
+                    kernel_size=3,
+                    stride=1,
+                    padding='same',
+                    bias=False,
+                ),
+                nn.Tanh(),
+            )
+
+    def forward(self, input: torch.Tensor) -> \
+            Tuple[torch.Tensor, torch.Tensor]:
+        size_4  = self._init(input)
+        size_8  = self._upsample_8(size_4)
+        size_16 = self._upsample_16(size_8)
+        size_32 = self._upsample_32(size_16)
+
+        size_64  = self._sle_64 (size_4,  self._upsample_64 (size_32) )
+        size_128 = self._sle_128(size_8,  self._upsample_128(size_64) )
+        size_256 = self._sle_256(size_16, self._upsample_256(size_128))
+        size_512 = self._sle_512(size_32, self._upsample_512(size_256))
+
+        size_1024 = self._upsample_1024(size_512)
+
+        out_128  = self._out_128 (size_128)
+        out_1024 = self._out_1024(size_1024)
+        return out_1024, out_128
+
+
+class Discriminrator(nn.Module, HugGANModelHubMixin):
+
+    def __init__(self, in_channels: int):
+        super().__init__()
+
+        self._channels = {
+                4:    1024,
+                8:    512,
+                16:   256,
+                32:   128,
+                64:   128,
+                128:  64,
+                256:  32,
+                512:  16,
+                1024: 8,
+            }
+
+        self._init = nn.Sequential(
+                SpectralConv2d(
+                        in_channels=in_channels,
+                        out_channels=self._channels[1024],
+                        kernel_size=4,
+                        stride=2,
+                        padding=1,
+                        bias=False,
+                    ),
+                nn.LeakyReLU(negative_slope=0.2),
+                SpectralConv2d(
+                        in_channels=self._channels[1024],
+                        out_channels=self._channels[512],
+                        kernel_size=4,
+                        stride=2,
+                        padding=1,
+                        bias=False,
+                    ),
+                nn.BatchNorm2d(num_features=self._channels[512]),
+                nn.LeakyReLU(negative_slope=0.2),
+            )
+
+        self._downsample_256 = DownsampleBlockT2(in_channels=self._channels[512], out_channels=self._channels[256])
+        self._downsample_128 = DownsampleBlockT2(in_channels=self._channels[256], out_channels=self._channels[128])
+        self._downsample_64  = DownsampleBlockT2(in_channels=self._channels[128], out_channels=self._channels[64] )
+        self._downsample_32  = DownsampleBlockT2(in_channels=self._channels[64],  out_channels=self._channels[32] )
+        self._downsample_16  = DownsampleBlockT2(in_channels=self._channels[32],  out_channels=self._channels[16] )
+
+        self._sle_64 = SLEBlock(in_channels=self._channels[512], out_channels=self._channels[64])
+        self._sle_32 = SLEBlock(in_channels=self._channels[256], out_channels=self._channels[32])
+        self._sle_16 = SLEBlock(in_channels=self._channels[128], out_channels=self._channels[16])
+
+        self._small_track = nn.Sequential(
+                SpectralConv2d(
+                        in_channels=in_channels,
+                        out_channels=self._channels[256],
+                        kernel_size=4,
+                        stride=2,
+                        padding=1,
+                        bias=False,
+                    ),
+                nn.LeakyReLU(negative_slope=0.2),
+                DownsampleBlockT1(in_channels=self._channels[256], out_channels=self._channels[128]),
+                DownsampleBlockT1(in_channels=self._channels[128], out_channels=self._channels[64] ),
+                DownsampleBlockT1(in_channels=self._channels[64],  out_channels=self._channels[32] ),
+            )
+
+        self._features_large = nn.Sequential(
+                SpectralConv2d(
+                        in_channels=self._channels[16] ,
+                        out_channels=self._channels[8],
+                        kernel_size=1,
+                        stride=1,
+                        padding=0,
+                        bias=False,
+                    ),
+                nn.BatchNorm2d(num_features=self._channels[8]),
+                nn.LeakyReLU(negative_slope=0.2),
+                SpectralConv2d(
+                        in_channels=self._channels[8],
+                        out_channels=1,
+                        kernel_size=4,
+                        stride=1,
+                        padding=0,
+                        bias=False,
+                    )
+            )
+
+        self._features_small = nn.Sequential(
+                SpectralConv2d(
+                        in_channels=self._channels[32],
+                        out_channels=1,
+                        kernel_size=4,
+                        stride=1,
+                        padding=0,
+                        bias=False,
+                    ),
+            )
+
+        self._decoder_large = Decoder(in_channels=self._channels[16], out_channels=3)
+        self._decoder_small = Decoder(in_channels=self._channels[32], out_channels=3)
+        self._decoder_piece = Decoder(in_channels=self._channels[32], out_channels=3)
+
+    def forward(self, images_1024: torch.Tensor,
+                      images_128: torch.Tensor,
+                      image_type: ImageType) -> \
+            Union[
+                torch.Tensor,
+                Tuple[torch.Tensor, Tuple[Any, Any, Any]]
+            ]:
+        # large track
+
+        down_512 = self._init(images_1024)
+        down_256 = self._downsample_256(down_512)
+        down_128 = self._downsample_128(down_256)
+
+        down_64 = self._downsample_64(down_128)
+        down_64 = self._sle_64(down_512, down_64)
+
+        down_32 = self._downsample_32(down_64)
+        down_32 = self._sle_32(down_256, down_32)
+
+        down_16 = self._downsample_16(down_32)
+        down_16 = self._sle_16(down_128, down_16)
+
+        # small track
+
+        down_small = self._small_track(images_128)
+
+        # features
+
+        features_large = self._features_large(down_16).view(-1)
+        features_small = self._features_small(down_small).view(-1)
+        features = torch.cat([features_large, features_small], dim=0)
+
+        # decoder
+
+        if image_type != ImageType.FAKE:
+            dec_large = self._decoder_large(down_16)
+            dec_small = self._decoder_small(down_small)
+            dec_piece = self._decoder_piece(crop_image_part(down_32, image_type))
+            return features, (dec_large, dec_small, dec_piece)
+
+        return features

requirements.txt

@@ -0,0 +1,5 @@
+git+https://github.com/huggingface/community-events@main
+streamlit==1.8.0
+torch
+streamlit-lottie
+streamlit-option-menu

utils.py

@@ -0,0 +1,87 @@
+import torch
+import torch.nn as nn
+from enum import Enum
+
+import base64
+import json
+from io import BytesIO
+from PIL import Image
+import requests
+import re
+from copy import deepcopy
+
+class ImageType(Enum):
+    REAL_UP_L = 0
+    REAL_UP_R = 1
+    REAL_DOWN_R = 2
+    REAL_DOWN_L = 3
+    FAKE = 4
+
+
+def crop_image_part(image: torch.Tensor,
+                    part: ImageType) -> torch.Tensor:
+    size = image.shape[2] // 2
+
+    if part == ImageType.REAL_UP_L:
+        return image[:, :, :size, :size]
+
+    elif part == ImageType.REAL_UP_R:
+        return image[:, :, :size, size:]
+
+    elif part == ImageType.REAL_DOWN_L:
+        return image[:, :, size:, :size]
+
+    elif part == ImageType.REAL_DOWN_R:
+        return image[:, :, size:, size:]
+
+    else:
+        raise ValueError('invalid part')
+
+
+def init_weights(module: nn.Module):
+    if isinstance(module, nn.Conv2d):
+        torch.nn.init.normal_(module.weight, 0.0, 0.02)
+
+    if isinstance(module, nn.BatchNorm2d):
+        torch.nn.init.normal_(module.weight, 1.0, 0.02)
+        module.bias.data.fill_(0)
+
+def load_image_from_local(image_path, image_resize=None):
+    image = Image.open(image_path)
+
+    if isinstance(image_resize, tuple):
+        image = image.resize(image_resize)
+    return image
+
+def load_image_from_url(image_url, rgba_mode=False, image_resize=None, default_image=None):
+    try:
+        image = Image.open(requests.get(image_url, stream=True).raw)
+
+        if rgba_mode:
+            image = image.convert("RGBA")
+
+        if isinstance(image_resize, tuple):
+            image = image.resize(image_resize)
+
+    except Exception as e:
+        image = None
+        if default_image:
+            image = load_image_from_local(default_image, image_resize=image_resize)
+
+    return image
+
+def image_to_base64(image_array):
+    buffered = BytesIO()
+    image_array.save(buffered, format="PNG")
+    image_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return f"data:image/png;base64, {image_b64}"
+
+
+def copy_G_params(model):
+    flatten = deepcopy(list(p.data for p in model.parameters()))
+    return flatten
+
+
+def load_params(model, new_param):
+    for p, new_p in zip(model.parameters(), new_param):
+        p.data.copy_(new_p)