forward lfq kwargs

lucidrains · Oct 11, 2023 · 3c721ac · 3c721ac
1 parent 6e97bef
commit 3c721ac
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/muse_maskgit_pytorch/vqgan_vae.py b/muse_maskgit_pytorch/vqgan_vae.py
@@ -330,7 +330,8 @@ def __init__(
         if lookup_free_quantization:
             self.quantizer = LFQ(
                 dim = self.enc_dec.encoded_dim,
-                codebook_size = codebook_size
+                codebook_size = codebook_size,
+                **lfq_kwargs
             )
         else:
             self.quantizer = VQ(

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'muse-maskgit-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.3.0',
+  version = '0.3.1',
   license='MIT',
   description = 'MUSE - Text-to-Image Generation via Masked Generative Transformers, in Pytorch',
   author = 'Phil Wang',