add blurs

Author: lucidrains

setup.py

@@ -8,7 +8,7 @@
           'stylegan2_pytorch = stylegan2_pytorch.cli:main',
       ],
   },
-  version = '1.0.2',
+  version = '1.0.3',
   license='GPLv3+',
   description = 'StyleGan2 in Pytorch',
   author = 'Phil Wang',
@@ -17,15 +17,16 @@
   download_url = 'https://github.com/lucidrains/stylegan2-pytorch/archive/v_036.tar.gz',
   keywords = ['generative adversarial networks', 'artificial intelligence'],
   install_requires=[
+      'contrastive_learner>=0.1.0',
       'fire',
+      'kornia',
+      'linear_attention_transformer',
       'numpy',
       'retry',
       'tqdm',
       'torch',
       'torchvision',
       'pillow',
-      'contrastive_learner>=0.1.0',
-      'linear_attention_transformer',
       'vector-quantize-pytorch'
   ],
   classifiers=[

stylegan2_pytorch/stylegan2_pytorch.py

@@ -23,6 +23,8 @@
 from torch.utils.data.distributed import DistributedSampler
 from torch.nn.parallel import DistributedDataParallel as DDP
 
+from kornia.filters import filter2D
+
 import torchvision
 from torchvision import transforms
 from stylegan2_pytorch.diff_augment import DiffAugment
@@ -101,6 +103,16 @@ def forward(self, x):
         out = out.permute(0, 3, 1, 2)
         return out, loss
 
+class Blur(nn.Module):
+    def __init__(self):
+        super().__init__()
+        f = torch.Tensor([1, 2, 1])
+        self.register_buffer('f', f)
+    def forward(self, x):
+        f = self.f
+        f = f[None, None, :] * f [None, :, None]
+        return filter2D(x, f, normalized=True)
+
 # one layer of self-attention and feedforward, for images
 
 attn_and_ff = lambda chan: nn.Sequential(*[
@@ -352,7 +364,10 @@ def __init__(self, latent_dim, input_channel, upsample, rgba = False):
         out_filters = 3 if not rgba else 4
         self.conv = Conv2DMod(input_channel, out_filters, 1, demod=False)
 
-        self.upsample = nn.Upsample(scale_factor = 2, mode='bilinear', align_corners=False) if upsample else None
+        self.upsample = nn.Sequential(
+            nn.Upsample(scale_factor = 2, mode='bilinear', align_corners=False),
+            Blur()
+        ) if upsample else None
 
     def forward(self, x, prev_rgb, istyle):
         b, c, h, w = x.shape
@@ -450,7 +465,10 @@ def __init__(self, input_channels, filters, downsample=True):
             leaky_relu()
         )
 
-        self.downsample = nn.Conv2d(filters, filters, 3, padding = 1, stride = 2) if downsample else None
+        self.downsample = nn.Sequential(
+            Blur(),
+            nn.Conv2d(filters, filters, 3, padding = 1, stride = 2)
+        ) if downsample else None
 
     def forward(self, x):
         res = self.conv_res(x)