Upload models.py

models.py

@@ -0,0 +1,245 @@
+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]]
+            ]:
+        # large track
+
+        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)
+
+        # small track
+
+        down_small = self._small_track(images_128)
+
+        # features
+
+        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)
+
+        # decoder
+
+        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