complete vit-vqgan from https://arxiv.org/abs/2110.04627

dalle2_pytorch/vqgan_vae.py

@@ -495,10 +495,12 @@ class ViTEncDec(nn.Module):
                 layers = layers
             ),
             nn.Sequential(
-                nn.Linear(dim, dim * 4, bias = False),
+                nn.Linear(dim, dim * 2, bias = False),
                 nn.Tanh(),
-                nn.Linear(dim * 4, input_dim, bias = False),
+                nn.Linear(dim * 2, dim, bias = False),
             ),
+            nn.LayerNorm(dim),
+            nn.Linear(dim, input_dim),
             RearrangeImage(),
             Rearrange('b h w (p1 p2 c) -> b c (h p1) (w p2)', p1 = patch_size, p2 = patch_size)
         )

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.0.54',
+  version = '0.0.53',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',