From cb261874504e67e39eebe83d1b0a78a6174f3aeb Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Thu, 28 Apr 2022 08:59:03 -0700 Subject: [PATCH] vqgan-vae codebook dims should be 256 or smaller --- dalle2_pytorch/vqgan_vae.py | 2 ++ setup.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/dalle2_pytorch/vqgan_vae.py b/dalle2_pytorch/vqgan_vae.py index 7e59805..76f395f 100644 --- a/dalle2_pytorch/vqgan_vae.py +++ b/dalle2_pytorch/vqgan_vae.py @@ -545,6 +545,7 @@ class VQGanVAE(nn.Module): l2_recon_loss = False, use_hinge_loss = True, vgg = None, + vq_codebook_dim = 256, vq_codebook_size = 512, vq_decay = 0.8, vq_commitment_weight = 1., @@ -579,6 +580,7 @@ class VQGanVAE(nn.Module): self.vq = VQ( dim = self.enc_dec.encoded_dim, + codebook_dim = vq_codebook_dim, codebook_size = vq_codebook_size, decay = vq_decay, commitment_weight = vq_commitment_weight, diff --git a/setup.py b/setup.py index 8a3a1e7..c171129 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ setup( 'dream = dalle2_pytorch.cli:dream' ], }, - version = '0.0.62', + version = '0.0.63', license='MIT', description = 'DALL-E 2', author = 'Phil Wang',