Commit 56bcd108 authored by Phil Wang's avatar Phil Wang
Browse files

add kmeans init and expiry of stale codes as detailed in soundstream paper

parent 8b78d101

audiolm_pytorch/audiolm_pytorch.py

+4 −1
Original line number Diff line number Diff line
@@ -93,7 +93,10 @@ class SoundStream(nn.Module): 
        self.rq = ResidualVQ(

            dim = codebook_dim,

            num_quantizers = rq_num_quantizers,

            codebook_size = codebook_size

            codebook_size = codebook_size,

            kmeans_init = True,

            threshold_ema_dead_code = 2,

            sync_kmeans = False

        )



        decoder_blocks = []

setup.py

+1 −1
Original line number Diff line number Diff line
@@ -21,7 +21,7 @@ setup( 
    'einops>=0.4',

    'ema-pytorch',

    'torch>=1.6',

    'vector-quantize-pytorch'

    'vector-quantize-pytorch>=0.10.1'

  ],

  classifiers=[

    'Development Status :: 4 - Beta',