|
import torch |
|
import torch.nn as nn |
|
|
|
from typing import Any, Tuple, Union |
|
|
|
from utils import ( |
|
ImageType, |
|
crop_image_part, |
|
) |
|
|
|
from layers import ( |
|
SpectralConv2d, |
|
InitLayer, |
|
SLEBlock, |
|
UpsampleBlockT1, |
|
UpsampleBlockT2, |
|
DownsampleBlockT1, |
|
DownsampleBlockT2, |
|
Decoder, |
|
) |
|
|
|
from huggan.pytorch.huggan_mixin import HugGANModelHubMixin |
|
|
|
|
|
class Generator(nn.Module, HugGANModelHubMixin): |
|
|
|
def __init__(self, in_channels: int, |
|
out_channels: int): |
|
super().__init__() |
|
|
|
self._channels = { |
|
4: 1024, |
|
8: 512, |
|
16: 256, |
|
32: 128, |
|
64: 128, |
|
128: 64, |
|
256: 32, |
|
512: 16, |
|
1024: 8, |
|
} |
|
|
|
self._init = InitLayer( |
|
in_channels=in_channels, |
|
out_channels=self._channels[4], |
|
) |
|
|
|
self._upsample_8 = UpsampleBlockT2(in_channels=self._channels[4], out_channels=self._channels[8] ) |
|
self._upsample_16 = UpsampleBlockT1(in_channels=self._channels[8], out_channels=self._channels[16] ) |
|
self._upsample_32 = UpsampleBlockT2(in_channels=self._channels[16], out_channels=self._channels[32] ) |
|
self._upsample_64 = UpsampleBlockT1(in_channels=self._channels[32], out_channels=self._channels[64] ) |
|
self._upsample_128 = UpsampleBlockT2(in_channels=self._channels[64], out_channels=self._channels[128] ) |
|
self._upsample_256 = UpsampleBlockT1(in_channels=self._channels[128], out_channels=self._channels[256] ) |
|
self._upsample_512 = UpsampleBlockT2(in_channels=self._channels[256], out_channels=self._channels[512] ) |
|
self._upsample_1024 = UpsampleBlockT1(in_channels=self._channels[512], out_channels=self._channels[1024]) |
|
|
|
self._sle_64 = SLEBlock(in_channels=self._channels[4], out_channels=self._channels[64] ) |
|
self._sle_128 = SLEBlock(in_channels=self._channels[8], out_channels=self._channels[128]) |
|
self._sle_256 = SLEBlock(in_channels=self._channels[16], out_channels=self._channels[256]) |
|
self._sle_512 = SLEBlock(in_channels=self._channels[32], out_channels=self._channels[512]) |
|
|
|
self._out_128 = nn.Sequential( |
|
SpectralConv2d( |
|
in_channels=self._channels[128], |
|
out_channels=out_channels, |
|
kernel_size=1, |
|
stride=1, |
|
padding='same', |
|
bias=False, |
|
), |
|
nn.Tanh(), |
|
) |
|
|
|
self._out_1024 = nn.Sequential( |
|
SpectralConv2d( |
|
in_channels=self._channels[1024], |
|
out_channels=out_channels, |
|
kernel_size=3, |
|
stride=1, |
|
padding='same', |
|
bias=False, |
|
), |
|
nn.Tanh(), |
|
) |
|
|
|
def forward(self, input: torch.Tensor) -> \ |
|
Tuple[torch.Tensor, torch.Tensor]: |
|
size_4 = self._init(input) |
|
size_8 = self._upsample_8(size_4) |
|
size_16 = self._upsample_16(size_8) |
|
size_32 = self._upsample_32(size_16) |
|
|
|
size_64 = self._sle_64 (size_4, self._upsample_64 (size_32) ) |
|
size_128 = self._sle_128(size_8, self._upsample_128(size_64) ) |
|
size_256 = self._sle_256(size_16, self._upsample_256(size_128)) |
|
size_512 = self._sle_512(size_32, self._upsample_512(size_256)) |
|
|
|
size_1024 = self._upsample_1024(size_512) |
|
|
|
out_128 = self._out_128 (size_128) |
|
out_1024 = self._out_1024(size_1024) |
|
return out_1024, out_128 |
|
|
|
|
|
class Discriminrator(nn.Module, HugGANModelHubMixin): |
|
|
|
def __init__(self, in_channels: int): |
|
super().__init__() |
|
|
|
self._channels = { |
|
4: 1024, |
|
8: 512, |
|
16: 256, |
|
32: 128, |
|
64: 128, |
|
128: 64, |
|
256: 32, |
|
512: 16, |
|
1024: 8, |
|
} |
|
|
|
self._init = nn.Sequential( |
|
SpectralConv2d( |
|
in_channels=in_channels, |
|
out_channels=self._channels[1024], |
|
kernel_size=4, |
|
stride=2, |
|
padding=1, |
|
bias=False, |
|
), |
|
nn.LeakyReLU(negative_slope=0.2), |
|
SpectralConv2d( |
|
in_channels=self._channels[1024], |
|
out_channels=self._channels[512], |
|
kernel_size=4, |
|
stride=2, |
|
padding=1, |
|
bias=False, |
|
), |
|
nn.BatchNorm2d(num_features=self._channels[512]), |
|
nn.LeakyReLU(negative_slope=0.2), |
|
) |
|
|
|
self._downsample_256 = DownsampleBlockT2(in_channels=self._channels[512], out_channels=self._channels[256]) |
|
self._downsample_128 = DownsampleBlockT2(in_channels=self._channels[256], out_channels=self._channels[128]) |
|
self._downsample_64 = DownsampleBlockT2(in_channels=self._channels[128], out_channels=self._channels[64] ) |
|
self._downsample_32 = DownsampleBlockT2(in_channels=self._channels[64], out_channels=self._channels[32] ) |
|
self._downsample_16 = DownsampleBlockT2(in_channels=self._channels[32], out_channels=self._channels[16] ) |
|
|
|
self._sle_64 = SLEBlock(in_channels=self._channels[512], out_channels=self._channels[64]) |
|
self._sle_32 = SLEBlock(in_channels=self._channels[256], out_channels=self._channels[32]) |
|
self._sle_16 = SLEBlock(in_channels=self._channels[128], out_channels=self._channels[16]) |
|
|
|
self._small_track = nn.Sequential( |
|
SpectralConv2d( |
|
in_channels=in_channels, |
|
out_channels=self._channels[256], |
|
kernel_size=4, |
|
stride=2, |
|
padding=1, |
|
bias=False, |
|
), |
|
nn.LeakyReLU(negative_slope=0.2), |
|
DownsampleBlockT1(in_channels=self._channels[256], out_channels=self._channels[128]), |
|
DownsampleBlockT1(in_channels=self._channels[128], out_channels=self._channels[64] ), |
|
DownsampleBlockT1(in_channels=self._channels[64], out_channels=self._channels[32] ), |
|
) |
|
|
|
self._features_large = nn.Sequential( |
|
SpectralConv2d( |
|
in_channels=self._channels[16] , |
|
out_channels=self._channels[8], |
|
kernel_size=1, |
|
stride=1, |
|
padding=0, |
|
bias=False, |
|
), |
|
nn.BatchNorm2d(num_features=self._channels[8]), |
|
nn.LeakyReLU(negative_slope=0.2), |
|
SpectralConv2d( |
|
in_channels=self._channels[8], |
|
out_channels=1, |
|
kernel_size=4, |
|
stride=1, |
|
padding=0, |
|
bias=False, |
|
) |
|
) |
|
|
|
self._features_small = nn.Sequential( |
|
SpectralConv2d( |
|
in_channels=self._channels[32], |
|
out_channels=1, |
|
kernel_size=4, |
|
stride=1, |
|
padding=0, |
|
bias=False, |
|
), |
|
) |
|
|
|
self._decoder_large = Decoder(in_channels=self._channels[16], out_channels=3) |
|
self._decoder_small = Decoder(in_channels=self._channels[32], out_channels=3) |
|
self._decoder_piece = Decoder(in_channels=self._channels[32], out_channels=3) |
|
|
|
def forward(self, images_1024: torch.Tensor, |
|
images_128: torch.Tensor, |
|
image_type: ImageType) -> \ |
|
Union[ |
|
torch.Tensor, |
|
Tuple[torch.Tensor, Tuple[Any, Any, Any]] |
|
]: |
|
|
|
|
|
down_512 = self._init(images_1024) |
|
down_256 = self._downsample_256(down_512) |
|
down_128 = self._downsample_128(down_256) |
|
|
|
down_64 = self._downsample_64(down_128) |
|
down_64 = self._sle_64(down_512, down_64) |
|
|
|
down_32 = self._downsample_32(down_64) |
|
down_32 = self._sle_32(down_256, down_32) |
|
|
|
down_16 = self._downsample_16(down_32) |
|
down_16 = self._sle_16(down_128, down_16) |
|
|
|
|
|
|
|
down_small = self._small_track(images_128) |
|
|
|
|
|
|
|
features_large = self._features_large(down_16).view(-1) |
|
features_small = self._features_small(down_small).view(-1) |
|
features = torch.cat([features_large, features_small], dim=0) |
|
|
|
|
|
|
|
if image_type != ImageType.FAKE: |
|
dec_large = self._decoder_large(down_16) |
|
dec_small = self._decoder_small(down_small) |
|
dec_piece = self._decoder_piece(crop_image_part(down_32, image_type)) |
|
return features, (dec_large, dec_small, dec_piece) |
|
|
|
return features |
|
|
