Upload 13 files

models/__init__.py

@@ -0,0 +1,3 @@
+from .anime_gan import GeneratorV1
+from .anime_gan_v2 import GeneratorV2
+from .anime_gan_v3 import GeneratorV3

models/anime_gan.py

@@ -0,0 +1,112 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import spectral_norm
+from .conv_blocks import DownConv
+from .conv_blocks import UpConv
+from .conv_blocks import SeparableConv2D
+from .conv_blocks import InvertedResBlock
+from .conv_blocks import ConvBlock
+from .layers import get_norm
+from utils.common import initialize_weights
+
+
+class GeneratorV1(nn.Module):
+    def __init__(self, dataset=''):
+        super(GeneratorV1, self).__init__()
+        self.name = f'{self.__class__.__name__}_{dataset}'
+        bias = False
+
+        self.encode_blocks = nn.Sequential(
+            ConvBlock(3, 64, bias=bias),
+            ConvBlock(64, 128, bias=bias),
+            DownConv(128, bias=bias),
+            ConvBlock(128, 128, bias=bias),
+            SeparableConv2D(128, 256, bias=bias),
+            DownConv(256, bias=bias),
+            ConvBlock(256, 256, bias=bias),
+        )
+
+        self.res_blocks = nn.Sequential(
+            InvertedResBlock(256, 256, bias=bias),
+            InvertedResBlock(256, 256, bias=bias),
+            InvertedResBlock(256, 256, bias=bias),
+            InvertedResBlock(256, 256, bias=bias),
+            InvertedResBlock(256, 256, bias=bias),
+            InvertedResBlock(256, 256, bias=bias),
+            InvertedResBlock(256, 256, bias=bias),
+            InvertedResBlock(256, 256, bias=bias),
+        )
+
+        self.decode_blocks = nn.Sequential(
+            ConvBlock(256, 128, bias=bias),
+            UpConv(128, bias=bias),
+            SeparableConv2D(128, 128, bias=bias),
+            ConvBlock(128, 128, bias=bias),
+            UpConv(128, bias=bias),
+            ConvBlock(128, 64, bias=bias),
+            ConvBlock(64, 64, bias=bias),
+            nn.Conv2d(64, 3, kernel_size=1, stride=1, padding=0, bias=bias),
+            nn.Tanh(),
+        )
+
+        initialize_weights(self)
+
+    def forward(self, x):
+        out = self.encode_blocks(x)
+        out = self.res_blocks(out)
+        img = self.decode_blocks(out)
+
+        return img
+
+
+class Discriminator(nn.Module):
+    def __init__(
+        self,
+        dataset=None,
+        num_layers=1,
+        use_sn=False,
+        norm_type="instance",
+    ):
+        super(Discriminator, self).__init__()
+        self.name = f'discriminator_{dataset}'
+        self.bias = False
+        channels = 32
+
+        layers = [
+            nn.Conv2d(3, channels, kernel_size=3, stride=1, padding=1, bias=self.bias),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        in_channels = channels
+        for i in range(num_layers):
+            layers += [
+                nn.Conv2d(in_channels, channels * 2, kernel_size=3, stride=2, padding=1, bias=self.bias),
+                nn.LeakyReLU(0.2, True),
+                nn.Conv2d(channels * 2, channels * 4, kernel_size=3, stride=1, padding=1, bias=self.bias),
+                get_norm(norm_type)(channels * 4),
+                nn.LeakyReLU(0.2, True),
+            ]
+            in_channels = channels * 4
+            channels *= 2
+
+        channels *= 2
+        layers += [
+            nn.Conv2d(channels, channels, kernel_size=3, stride=1, padding=1, bias=self.bias),
+            get_norm(norm_type)(channels),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(channels, 1, kernel_size=3, stride=1, padding=1, bias=self.bias),
+        ]
+
+        if use_sn:
+            for i in range(len(layers)):
+                if isinstance(layers[i], nn.Conv2d):
+                    layers[i] = spectral_norm(layers[i])
+
+        self.discriminate = nn.Sequential(*layers)
+
+        initialize_weights(self)
+
+    def forward(self, img):
+        logits = self.discriminate(img)
+        return logits

models/anime_gan_v2.py

@@ -0,0 +1,61 @@
+
+import torch.nn as nn
+from models.conv_blocks import InvertedResBlock
+from models.conv_blocks import ConvBlock
+from models.conv_blocks import UpConvLNormLReLU
+from utils.common import initialize_weights
+
+
+class GeneratorV2(nn.Module):
+    def __init__(self, dataset=''):
+        super(GeneratorV2, self).__init__()
+        self.name = f'{self.__class__.__name__}_{dataset}'
+        bias = False
+
+        self.conv_block1 = nn.Sequential(
+            ConvBlock(3, 32, kernel_size=7, stride=1, norm_type="layer", bias=bias),
+            ConvBlock(32, 64, kernel_size=3, stride=2, norm_type="layer", bias=bias),
+            ConvBlock(64, 64, kernel_size=3, stride=1, norm_type="layer", bias=bias),
+        )
+
+        self.conv_block2 = nn.Sequential(
+            ConvBlock(64, 128, kernel_size=3, stride=2, norm_type="layer", bias=bias),
+            ConvBlock(128, 128, kernel_size=3, stride=1, norm_type="layer", bias=bias),
+        )
+
+        self.res_blocks = nn.Sequential(
+            ConvBlock(128, 128, kernel_size=3, stride=1, norm_type="layer", bias=bias),
+            InvertedResBlock(128, 256, expand_ratio=2, norm_type="layer", bias=bias),
+            InvertedResBlock(256, 256, expand_ratio=2, norm_type="layer", bias=bias),
+            InvertedResBlock(256, 256, expand_ratio=2, norm_type="layer", bias=bias),
+            InvertedResBlock(256, 256, expand_ratio=2, norm_type="layer", bias=bias),
+            ConvBlock(256, 128, kernel_size=3, stride=1, norm_type="layer", bias=bias),
+        )
+
+        self.upsample1 = nn.Sequential(
+            UpConvLNormLReLU(128, 128),
+            ConvBlock(128, 128, kernel_size=3, stride=1, norm_type="layer", bias=bias),
+        )
+
+        self.upsample2 = nn.Sequential(
+            UpConvLNormLReLU(128, 64),
+            ConvBlock(64, 64, kernel_size=3, stride=1, norm_type="layer", bias=bias),
+            ConvBlock(64, 32, kernel_size=7, stride=1, norm_type="layer", bias=bias),
+        )
+
+        self.decode_blocks = nn.Sequential(
+            nn.Conv2d(32, 3, kernel_size=1, stride=1, padding=0, bias=bias),
+            nn.Tanh(),
+        )
+
+        initialize_weights(self)
+
+    def forward(self, x):
+        out = self.conv_block1(x)
+        out = self.conv_block2(out)
+        out = self.res_blocks(out)
+        out = self.upsample1(out)
+        out = self.upsample2(out)
+        img = self.decode_blocks(out)
+
+        return img

models/anime_gan_v3.py

@@ -0,0 +1,14 @@
+
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import spectral_norm
+from models.conv_blocks import DownConv
+from models.conv_blocks import UpConv
+from models.conv_blocks import SeparableConv2D
+from models.conv_blocks import InvertedResBlock
+from models.conv_blocks import ConvBlock
+from utils.common import initialize_weights
+
+
+class GeneratorV3(nn.Module):
+    pass

models/conv_blocks.py

@@ -0,0 +1,185 @@
+import torch.nn as nn
+import torch.nn.functional as F
+from utils.common import initialize_weights
+from .layers import LayerNorm2d
+
+
+class DownConv(nn.Module):
+
+    def __init__(self, channels, bias=False):
+        super(DownConv, self).__init__()
+
+        self.conv1 = SeparableConv2D(channels, channels, stride=2, bias=bias)
+        self.conv2 = SeparableConv2D(channels, channels, stride=1, bias=bias)
+
+    def forward(self, x):
+        out1 = self.conv1(x)
+        out2 = F.interpolate(x, scale_factor=0.5, mode='bilinear')
+        out2 = self.conv2(out2)
+
+        return out1 + out2
+
+
+class UpConv(nn.Module):
+    def __init__(self, channels, bias=False):
+        super(UpConv, self).__init__()
+
+        self.conv = SeparableConv2D(channels, channels, stride=1, bias=bias)
+
+    def forward(self, x):
+        out = F.interpolate(x, scale_factor=2.0, mode='bilinear')
+        out = self.conv(out)
+        return out
+
+
+class UpConvLNormLReLU(nn.Module):
+    """Upsample Conv block with Layer Norm and Leaky ReLU"""
+    def __init__(self, in_channels, out_channels, bias=False):
+        super(UpConvLNormLReLU, self).__init__()
+
+        self.conv_block = ConvBlock(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            bias=bias,
+        )
+
+    def forward(self, x):
+        out = F.interpolate(x, scale_factor=2.0, mode='bilinear')
+        out = self.conv_block(out)
+        return out
+
+class SeparableConv2D(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1, bias=False):
+        super(SeparableConv2D, self).__init__()
+        self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size=3,
+            stride=stride, padding=1, groups=in_channels, bias=bias)
+        self.pointwise = nn.Conv2d(in_channels, out_channels,
+            kernel_size=1, stride=1, bias=bias)
+        # self.pad = 
+        self.ins_norm1 = nn.InstanceNorm2d(in_channels)
+        self.activation1 = nn.LeakyReLU(0.2, True)
+        self.ins_norm2 = nn.InstanceNorm2d(out_channels)
+        self.activation2 = nn.LeakyReLU(0.2, True)
+
+        initialize_weights(self)
+
+    def forward(self, x):
+        out = self.depthwise(x)
+        out = self.ins_norm1(out)
+        out = self.activation1(out)
+
+        out = self.pointwise(out)
+        out = self.ins_norm2(out)
+
+        return self.activation2(out)
+
+
+class ConvBlock(nn.Module):
+    """Stack of Conv2D + Norm + LeakyReLU"""
+    def __init__(
+        self,
+        channels,
+        out_channels,
+        kernel_size=3,
+        stride=1,
+        padding="valid",
+        bias=False,
+        norm_type="instance"
+    ):
+        super(ConvBlock, self).__init__()
+
+        if kernel_size == 3 and stride == 1:
+            self.pad = nn.ReflectionPad2d((1, 1, 1, 1))
+        elif kernel_size == 7 and stride == 1:
+            self.pad = nn.ReflectionPad2d((3, 3, 3, 3))
+        elif stride == 2:
+            self.pad = nn.ReflectionPad2d((0, 1, 1, 0))
+        else:
+            self.pad = None
+
+        self.conv = nn.Conv2d(
+            channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=bias
+        )
+        if norm_type == "instance":
+            self.ins_norm = nn.InstanceNorm2d(out_channels)
+        elif norm_type == "layer":
+            self.ins_norm = LayerNorm2d(out_channels)
+        self.activation = nn.LeakyReLU(0.2, True)
+
+        initialize_weights(self)
+
+    def forward(self, x):
+        if self.pad is not None:
+            x = self.pad(x)
+        out = self.conv(x)
+        out = self.ins_norm(out)
+        out = self.activation(out)
+        return out
+
+
+
+class InvertedResBlock(nn.Module):
+    def __init__(
+        self,
+        channels=256,
+        out_channels=256,
+        expand_ratio=2,
+        bias=False,
+        norm_type="instance",
+    ):
+        super(InvertedResBlock, self).__init__()
+        bottleneck_dim = round(expand_ratio * channels)
+        self.conv_block = ConvBlock(
+            channels,
+            bottleneck_dim,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=bias
+        )
+        self.depthwise_conv = nn.Conv2d(
+            bottleneck_dim,
+            bottleneck_dim,
+            kernel_size=3,
+            groups=bottleneck_dim,
+            stride=1,
+            padding=1,
+            bias=bias
+        )
+        self.conv = nn.Conv2d(
+            bottleneck_dim,
+            out_channels,
+            kernel_size=1,
+            stride=1,
+            bias=bias
+        )
+
+        if norm_type == "instance":
+            self.ins_norm1 = nn.InstanceNorm2d(out_channels)
+            self.ins_norm2 = nn.InstanceNorm2d(out_channels)
+        elif norm_type == "layer":
+            # Keep var name as is for v1 compatibility.
+            self.ins_norm1 = LayerNorm2d(bottleneck_dim)
+            self.ins_norm2 = LayerNorm2d(out_channels)
+        self.activation = nn.LeakyReLU(0.2, True)
+
+        initialize_weights(self)
+
+    def forward(self, x):
+        out = self.conv_block(x)
+        out = self.depthwise_conv(out)
+        out = self.ins_norm1(out)
+        out = self.activation(out)
+        out = self.conv(out)
+        out = self.ins_norm2(out)
+
+        if out.shape[1] != x.shape[1]:
+            # Only concate if same shape
+            return out
+        return out + x

models/layers.py

@@ -0,0 +1,24 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LayerNorm2d(nn.LayerNorm):
+    """ LayerNorm for channels of '2D' spatial NCHW tensors """
+    def __init__(self, num_channels, eps=1e-6, affine=True):
+        super().__init__(num_channels, eps=eps, elementwise_affine=affine)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.permute(0, 2, 3, 1)
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        x = x.permute(0, 3, 1, 2)
+        return x
+
+
+def get_norm(norm_type):
+    if norm_type == "instance":
+        return nn.InstanceNorm2d
+    elif norm_type == "layer":
+        return LayerNorm2d
+    else:
+        raise ValueError(norm_type)

models/vgg.py

@@ -0,0 +1,80 @@
+from numpy.lib.arraysetops import isin
+import torchvision.models as models
+import torch.nn as nn
+import torch
+
+
+
+class Vgg19(nn.Module):
+    def __init__(self):
+        super(Vgg19, self).__init__()
+        self.vgg19 = self.get_vgg19().eval()
+        vgg_mean = torch.tensor([0.485, 0.456, 0.406]).float()
+        vgg_std = torch.tensor([0.229, 0.224, 0.225]).float()
+        self.mean = vgg_mean.view(-1, 1 ,1)
+        self.std = vgg_std.view(-1, 1, 1)
+
+    def to(self, device):
+        new_self = super(Vgg19, self).to(device)
+        new_self.mean = new_self.mean.to(device)
+        new_self.std = new_self.std.to(device)
+        return new_self
+
+    def forward(self, x):
+        return self.vgg19(self.normalize_vgg(x))
+
+    @staticmethod
+    def get_vgg19(last_layer='conv4_4'):
+        vgg = models.vgg19(weights=models.VGG19_Weights.IMAGENET1K_V1).features
+        model_list = []
+
+        i = 0
+        j = 1
+        for layer in vgg.children():
+            if isinstance(layer, nn.MaxPool2d):
+                i = 0
+                j += 1
+
+            elif isinstance(layer, nn.Conv2d):
+                i += 1
+
+            name = f'conv{j}_{i}'
+
+            if name == last_layer:
+                model_list.append(layer)
+                break
+
+            model_list.append(layer)
+
+
+        model = nn.Sequential(*model_list)
+        return model
+
+
+    def normalize_vgg(self, image):
+        '''
+        Expect input in range -1 1
+        '''
+        image = (image + 1.0) / 2.0
+        return (image - self.mean) / self.std
+
+
+if __name__ == '__main__':
+    from PIL import Image
+    import numpy as np
+    from utils.image_processing import normalize_input
+
+    image = Image.open("example/10.jpg")
+    image = image.resize((224, 224))
+    np_img = np.array(image).astype('float32')
+    np_img = normalize_input(np_img)
+
+    img = torch.from_numpy(np_img)
+    img = img.permute(2, 0, 1)
+    img = img.unsqueeze(0)
+
+    vgg = Vgg19()
+
+    feat = vgg(img)
+
+    print(feat.shape)