FoundationVision/LlamaGen

1# Modified from:2#   stylegan2-pytorch: https://github.com/lucidrains/stylegan2-pytorch/blob/master/stylegan2_pytorch/stylegan2_pytorch.py3#   stylegan2-pytorch: https://github.com/rosinality/stylegan2-pytorch/blob/master/model.py4#   maskgit: https://github.com/google-research/maskgit/blob/main/maskgit/nets/discriminator.py5import math6import torch7import torch.nn as nn8try:9    from kornia.filters import filter2d10except:11    pass12 13class Discriminator(nn.Module):14    def __init__(self, input_nc=3, ndf=64, n_layers=3, channel_multiplier=1, image_size=256):15        super().__init__()16        channels = {17            4: 512,18            8: 512,19            16: 512,20            32: 512,21            64: 256 * channel_multiplier,22            128: 128 * channel_multiplier,23            256: 64 * channel_multiplier,24            512: 32 * channel_multiplier,25            1024: 16 * channel_multiplier,26        }27        28        log_size = int(math.log(image_size, 2))29        in_channel = channels[image_size]30 31        blocks = [nn.Conv2d(input_nc, in_channel, 3, padding=1), leaky_relu()]32        for i in range(log_size, 2, -1):33            out_channel = channels[2 ** (i - 1)]34            blocks.append(DiscriminatorBlock(in_channel, out_channel))35            in_channel = out_channel36        self.blocks = nn.ModuleList(blocks)37 38        self.final_conv = nn.Sequential(39            nn.Conv2d(in_channel, channels[4], 3, padding=1),40            leaky_relu(),41        )42        self.final_linear = nn.Sequential(43            nn.Linear(channels[4] * 4 * 4, channels[4]),44            leaky_relu(),45            nn.Linear(channels[4], 1)46        )47    48    def forward(self, x):49        for block in self.blocks:50            x = block(x)51        x = self.final_conv(x)52        x = x.view(x.shape[0], -1)53        x = self.final_linear(x)54        return x55 56 57class DiscriminatorBlock(nn.Module):58    def __init__(self, input_channels, filters, downsample=True):59        super().__init__()60        self.conv_res = nn.Conv2d(input_channels, filters, 1, stride = (2 if downsample else 1))61 62        self.net = nn.Sequential(63            nn.Conv2d(input_channels, filters, 3, padding=1),64            leaky_relu(),65            nn.Conv2d(filters, filters, 3, padding=1),66            leaky_relu()67        )68 69        self.downsample = nn.Sequential(70            Blur(),71            nn.Conv2d(filters, filters, 3, padding = 1, stride = 2)72        ) if downsample else None73 74    def forward(self, x):75        res = self.conv_res(x)76        x = self.net(x)77        if exists(self.downsample):78            x = self.downsample(x)79        x = (x + res) * (1 / math.sqrt(2))80        return x81 82 83 84class Blur(nn.Module):85    def __init__(self):86        super().__init__()87        f = torch.Tensor([1, 2, 1])88        self.register_buffer('f', f)89    90    def forward(self, x):91        f = self.f92        f = f[None, None, :] * f [None, :, None]93        return filter2d(x, f, normalized=True)94 95 96def leaky_relu(p=0.2):97    return nn.LeakyReLU(p, inplace=True)98 99 100def exists(val):101    return val is not None102