increase default codebook size

lucidrains · Oct 24, 2023 · 8fd178d · 8fd178d
1 parent 29a57ef
commit 8fd178d
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -65,7 +65,7 @@ Otherwise, to stay more true to the original paper, you can use `SoundStream`. F
 from audiolm_pytorch import SoundStream, SoundStreamTrainer
 
 soundstream = SoundStream(
-    codebook_size = 1024,
+    codebook_size = 4096,
     rq_num_quantizers = 8,
     rq_groups = 2,                      # this paper proposes using multi-headed residual vector quantization - https://arxiv.org/abs/2305.02765
     use_lookup_free_quantizer = True,   # whether to use residual lookup free quantization

diff --git a/audiolm_pytorch/soundstream.py b/audiolm_pytorch/soundstream.py
@@ -428,7 +428,7 @@ def __init__(
         strides = (2, 4, 5, 8),
         channel_mults = (2, 4, 8, 16),
         codebook_dim = 512,
-        codebook_size = 1024,
+        codebook_size = 4096,
         rq_num_quantizers = 8,
         rq_commitment_weight = 1.,
         rq_ema_decay = 0.95,

diff --git a/audiolm_pytorch/version.py b/audiolm_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.6.1'
+__version__ = '1.6.2'