vector-quantize-pytorch 1.14.2__tar.gz → 1.14.5__tar.gz

{vector_quantize_pytorch-1.14.2/vector_quantize_pytorch.egg-info → vector_quantize_pytorch-1.14.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vector_quantize_pytorch
-Version: 1.14.2
+Version: 1.14.5
 Summary: Vector Quantization - Pytorch
 Home-page: https://github.com/lucidrains/vector-quantizer-pytorch
 Author: Phil Wang

{vector_quantize_pytorch-1.14.2 → vector_quantize_pytorch-1.14.5}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '1.14.2',
+  version = '1.14.5',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   long_description_content_type = 'text/markdown',

{vector_quantize_pytorch-1.14.2 → vector_quantize_pytorch-1.14.5}/vector_quantize_pytorch/finite_scalar_quantization.py

@@ -3,7 +3,7 @@ Finite Scalar Quantization: VQ-VAE Made Simple - https://arxiv.org/abs/2309.1550
 Code adapted from Jax version in Appendix A.1
 """
 
-from typing import List, Optional
+from typing import List, Tuple, Optional
 
 import torch
 import torch.nn as nn
@@ -46,7 +46,8 @@ class FSQ(Module):
         dim: Optional[int] = None,
         num_codebooks = 1,
         keep_num_codebooks_dim: Optional[bool] = None,
-        scale: Optional[float] = None
+        scale: Optional[float] = None,
+        allowed_dtypes: Tuple[torch.dtype, ...] = (torch.float32, torch.float64)
     ):
         super().__init__()
         _levels = torch.tensor(levels, dtype=int32)
@@ -80,6 +81,8 @@ class FSQ(Module):
         implicit_codebook = self.indices_to_codes(torch.arange(self.codebook_size), project_out = False)
         self.register_buffer("implicit_codebook", implicit_codebook, persistent = False)
 
+        self.allowed_dtypes = allowed_dtypes
+
     def bound(self, z: Tensor, eps: float = 1e-3) -> Tensor:
         """Bound `z`, an array of shape (..., d)."""
         half_l = (self._levels - 1) * (1 + eps) / 2
@@ -141,8 +144,14 @@ class FSQ(Module):
         c - number of codebook dim
         """
 
+        orig_dtype = z.dtype
         is_img_or_video = z.ndim >= 4
 
+        # make sure allowed dtype
+
+        if z.dtype not in self.allowed_dtypes:
+            z = z.float()
+
         # standardize image or video into (batch, seq, dimension)
 
         if is_img_or_video:
@@ -173,4 +182,11 @@ class FSQ(Module):
         if not self.keep_num_codebooks_dim:
             indices = rearrange(indices, '... 1 -> ...')
 
+        # cast back to original dtype
+
+        if out.dtype != orig_dtype:
+            out = out.type(orig_dtype)
+
+        # return quantized output and indices
+
         return out, indices

{vector_quantize_pytorch-1.14.2 → vector_quantize_pytorch-1.14.5}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -826,17 +826,21 @@ class VectorQuantize(nn.Module):
         is_multiheaded = codebook.ndim > 2
 
         if not is_multiheaded:
-            return codebook[indices]
+            codes = codebook[indices]
+        else:
+            indices, ps = pack_one(indices, 'b * h')
+            indices = rearrange(indices, 'b n h -> b h n')
+
+            indices = repeat(indices, 'b h n -> b h n d', d = codebook.shape[-1])
+            codebook = repeat(codebook, 'h n d -> b h n d', b = indices.shape[0])
 
-        indices, ps = pack_one(indices, 'b * h')
-        indices = rearrange(indices, 'b n h -> b h n')
+            codes = codebook.gather(2, indices)
+            codes = rearrange(codes, 'b h n d -> b n (h d)')
+            codes = unpack_one(codes, ps, 'b * d')
 
-        indices = repeat(indices, 'b h n -> b h n d', d = codebook.shape[-1])
-        codebook = repeat(codebook, 'h n d -> b h n d', b = indices.shape[0])
+        if not self.channel_last:
+            codes = rearrange(codes, 'b ... d -> b d ...')
 
-        codes = codebook.gather(2, indices)
-        codes = rearrange(codes, 'b h n d -> b n (h d)')
-        codes = unpack_one(codes, ps, 'b * d')
         return codes
 
     def get_output_from_indices(self, indices):

{vector_quantize_pytorch-1.14.2 → vector_quantize_pytorch-1.14.5/vector_quantize_pytorch.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vector_quantize_pytorch
-Version: 1.14.2
+Version: 1.14.5
 Summary: Vector Quantization - Pytorch
 Home-page: https://github.com/lucidrains/vector-quantizer-pytorch
 Author: Phil Wang