From 3c721ac874cd3baabcbf66ed99cf77b3822bab84 Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Wed, 11 Oct 2023 10:07:40 -0700
Subject: [PATCH] forward lfq kwargs

---
 muse_maskgit_pytorch/vqgan_vae.py | 3 ++-
 setup.py                          | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/muse_maskgit_pytorch/vqgan_vae.py b/muse_maskgit_pytorch/vqgan_vae.py
index 11acb12..6b0e255 100644
--- 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
index 58bd9f0..3efbeb8 100644
--- 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',