| """ |
| Ported from Paella |
| """ |
|
|
| import torch |
| from torch import nn |
|
|
| from diffusers.configuration_utils import ConfigMixin, register_to_config |
| from diffusers.models.modeling_utils import ModelMixin |
|
|
|
|
| |
| class Discriminator(ModelMixin, ConfigMixin): |
| @register_to_config |
| def __init__(self, in_channels=3, cond_channels=0, hidden_channels=512, depth=6): |
| super().__init__() |
| d = max(depth - 3, 3) |
| layers = [ |
| nn.utils.spectral_norm( |
| nn.Conv2d(in_channels, hidden_channels // (2**d), kernel_size=3, stride=2, padding=1) |
| ), |
| nn.LeakyReLU(0.2), |
| ] |
| for i in range(depth - 1): |
| c_in = hidden_channels // (2 ** max((d - i), 0)) |
| c_out = hidden_channels // (2 ** max((d - 1 - i), 0)) |
| layers.append(nn.utils.spectral_norm(nn.Conv2d(c_in, c_out, kernel_size=3, stride=2, padding=1))) |
| layers.append(nn.InstanceNorm2d(c_out)) |
| layers.append(nn.LeakyReLU(0.2)) |
| self.encoder = nn.Sequential(*layers) |
| self.shuffle = nn.Conv2d( |
| (hidden_channels + cond_channels) if cond_channels > 0 else hidden_channels, 1, kernel_size=1 |
| ) |
| self.logits = nn.Sigmoid() |
|
|
| def forward(self, x, cond=None): |
| x = self.encoder(x) |
| if cond is not None: |
| cond = cond.view( |
| cond.size(0), |
| cond.size(1), |
| 1, |
| 1, |
| ).expand(-1, -1, x.size(-2), x.size(-1)) |
| x = torch.cat([x, cond], dim=1) |
| x = self.shuffle(x) |
| x = self.logits(x) |
| return x |
|
|
