Upload 2 files

animegan2/model.py

@@ -0,0 +1,52 @@
+
+import torch.nn as nn
+
+class ConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride):
+        super(ConvLayer, self).__init__()
+        reflection_padding = kernel_size // 2
+        self.layer = nn.Sequential(
+            nn.ReflectionPad2d(reflection_padding),
+            nn.Conv2d(in_channels, out_channels, kernel_size, stride),
+            nn.InstanceNorm2d(out_channels, affine=True),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        return self.layer(x)
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super(ResidualBlock, self).__init__()
+        self.block = nn.Sequential(
+            ConvLayer(channels, channels, 3, 1),
+            ConvLayer(channels, channels, 3, 1)
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        self.encoder = nn.Sequential(
+            ConvLayer(3, 32, 7, 1),
+            ConvLayer(32, 64, 3, 2),
+            ConvLayer(64, 128, 3, 2),
+        )
+        self.res_blocks = nn.Sequential(*[ResidualBlock(128) for _ in range(5)])
+        self.decoder = nn.Sequential(
+            nn.Upsample(scale_factor=2),
+            ConvLayer(128, 64, 3, 1),
+            nn.Upsample(scale_factor=2),
+            ConvLayer(64, 32, 3, 1),
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(32, 3, 7, 1),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.res_blocks(x)
+        x = self.decoder(x)
+        return x