add .indices_to_codes for SimVQ

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.20.2"
+version = "1.20.3"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

tests/test_readme.py

@@ -362,3 +362,17 @@ def test_latent_q():
 
     assert image_feats.shape == quantized.shape
     assert (quantized == quantizer.indices_to_codes(indices)).all()
+
+def test_sim_vq():
+    from vector_quantize_pytorch import SimVQ
+
+    sim_vq = SimVQ(
+        dim = 512,
+        codebook_size = 1024,
+    )
+
+    x = torch.randn(1, 1024, 512)
+    quantized, indices, commit_loss = sim_vq(x)
+
+    assert x.shape == quantized.shape
+    assert torch.allclose(quantized, sim_vq.indices_to_codes(indices), atol = 1e-6)

vector_quantize_pytorch/sim_vq.py

@@ -58,7 +58,7 @@ def __init__(
 
         self.codebook_to_codes = codebook_transform
 
-        self.register_buffer('codebook', codebook)
+        self.register_buffer('frozen_codebook', codebook)
 
 
         # whether to use rotation trick from Fifty et al. 
@@ -70,6 +70,23 @@ def __init__(
 
         self.input_to_quantize_commit_loss_weight = input_to_quantize_commit_loss_weight
 
+    @property
+    def codebook(self):
+        return self.codebook_to_codes(self.frozen_codebook)
+
+    def indices_to_codes(
+        self,
+        indices
+    ):
+        implicit_codebook = self.codebook
+
+        quantized = get_at('[c] d, b ... -> b ... d', implicit_codebook, indices)
+
+        if self.accept_image_fmap:
+            quantized = rearrange(quantized, 'b ... d -> b d ...')
+
+        return quantized
+
     def forward(
         self,
         x
@@ -78,7 +95,7 @@ def forward(
             x = rearrange(x, 'b d h w -> b h w d')
             x, inverse_pack = pack_one(x, 'b * d')
 
-        implicit_codebook = self.codebook_to_codes(self.codebook)
+        implicit_codebook = self.codebook
 
         with torch.no_grad():
             dist = torch.cdist(x, implicit_codebook)