| """Discriminator implementation.""" |
| import functools |
| import math |
| from typing import Tuple |
|
|
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .maskgit_vqgan import Conv2dSame |
|
|
|
|
| class BlurBlock(torch.nn.Module): |
| def __init__(self, |
| kernel: Tuple[int] = (1, 3, 3, 1) |
| ): |
| super().__init__() |
|
|
| kernel = torch.tensor(kernel, dtype=torch.float32, requires_grad=False) |
| kernel = kernel[None, :] * kernel[:, None] |
| kernel /= kernel.sum() |
| kernel = kernel.unsqueeze(0).unsqueeze(0) |
| self.register_buffer("kernel", kernel) |
|
|
| def calc_same_pad(self, i: int, k: int, s: int) -> int: |
| return max((math.ceil(i / s) - 1) * s + (k - 1) + 1 - i, 0) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| ic, ih, iw = x.size()[-3:] |
| pad_h = self.calc_same_pad(i=ih, k=4, s=2) |
| pad_w = self.calc_same_pad(i=iw, k=4, s=2) |
| if pad_h > 0 or pad_w > 0: |
| x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]) |
|
|
| weight = self.kernel.expand(ic, -1, -1, -1) |
|
|
| out = F.conv2d(input=x, weight=weight, stride=2, groups=x.shape[1]) |
| return out |
|
|
|
|
| class NLayerDiscriminator(torch.nn.Module): |
| def __init__( |
| self, |
| num_channels: int = 3, |
| hidden_channels: int = 128, |
| num_stages: int = 3, |
| blur_resample: bool = True, |
| blur_kernel_size: int = 4 |
| ): |
| """ Initializes the NLayerDiscriminator. |
| |
| Args: |
| num_channels -> int: The number of input channels. |
| hidden_channels -> int: The number of hidden channels. |
| num_stages -> int: The number of stages. |
| blur_resample -> bool: Whether to use blur resampling. |
| blur_kernel_size -> int: The blur kernel size. |
| """ |
| super().__init__() |
| assert num_stages > 0, "Discriminator cannot have 0 stages" |
| assert (not blur_resample) or (blur_kernel_size >= 3 and blur_kernel_size <= 5), "Blur kernel size must be in [3,5] when sampling]" |
|
|
| in_channel_mult = (1,) + tuple(map(lambda t: 2**t, range(num_stages))) |
| init_kernel_size = 5 |
| activation = functools.partial(torch.nn.LeakyReLU, negative_slope=0.1) |
|
|
| self.block_in = torch.nn.Sequential( |
| Conv2dSame( |
| num_channels, |
| hidden_channels, |
| kernel_size=init_kernel_size |
| ), |
| activation(), |
| ) |
|
|
| BLUR_KERNEL_MAP = { |
| 3: (1,2,1), |
| 4: (1,3,3,1), |
| 5: (1,4,6,4,1), |
| } |
|
|
| discriminator_blocks = [] |
| for i_level in range(num_stages): |
| in_channels = hidden_channels * in_channel_mult[i_level] |
| out_channels = hidden_channels * in_channel_mult[i_level + 1] |
| block = torch.nn.Sequential( |
| Conv2dSame( |
| in_channels, |
| out_channels, |
| kernel_size=3, |
| ), |
| torch.nn.AvgPool2d(kernel_size=2, stride=2) if not blur_resample else BlurBlock(BLUR_KERNEL_MAP[blur_kernel_size]), |
| torch.nn.GroupNorm(32, out_channels), |
| activation(), |
| ) |
| discriminator_blocks.append(block) |
|
|
| self.blocks = torch.nn.ModuleList(discriminator_blocks) |
|
|
| self.pool = torch.nn.AdaptiveMaxPool2d((16, 16)) |
|
|
| self.to_logits = torch.nn.Sequential( |
| Conv2dSame(out_channels, out_channels, 1), |
| activation(), |
| Conv2dSame(out_channels, 1, kernel_size=5) |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """ Forward pass. |
| |
| Args: |
| x -> torch.Tensor: The input tensor. |
| |
| Returns: |
| output -> torch.Tensor: The output tensor. |
| """ |
| hidden_states = self.block_in(x) |
| for block in self.blocks: |
| hidden_states = block(hidden_states) |
|
|
| hidden_states = self.pool(hidden_states) |
|
|
| return self.to_logits(hidden_states) |
|
|
