|
""" |
|
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 |
|
|
