vector-quantize-pytorch 1.28.3__tar.gz → 1.29.0__tar.gz

{vector_quantize_pytorch-1.28.3 → vector_quantize_pytorch-1.29.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vector-quantize-pytorch
-Version: 1.28.3
+Version: 1.29.0
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch
@@ -884,3 +884,15 @@ assert loss.item() >= 0
     url     = {https://arxiv.org/abs/2307.13304},
 }
 ```
+
+```bibtex
+@misc{tian2024visualautoregressivemodelingscalable,
+    title   = {Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction},
+    author  = {Keyu Tian and Yi Jiang and Zehuan Yuan and Bingyue Peng and Liwei Wang},
+    year    = {2024},
+    eprint  = {2404.02905},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV},
+    url     = {https://arxiv.org/abs/2404.02905},
+}
+```

{vector_quantize_pytorch-1.28.3 → vector_quantize_pytorch-1.29.0}/README.md

@@ -840,3 +840,15 @@ assert loss.item() >= 0
     url     = {https://arxiv.org/abs/2307.13304},
 }
 ```
+
+```bibtex
+@misc{tian2024visualautoregressivemodelingscalable,
+    title   = {Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction},
+    author  = {Keyu Tian and Yi Jiang and Zehuan Yuan and Bingyue Peng and Liwei Wang},
+    year    = {2024},
+    eprint  = {2404.02905},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV},
+    url     = {https://arxiv.org/abs/2404.02905},
+}
+```

{vector_quantize_pytorch-1.28.3 → vector_quantize_pytorch-1.29.0}/pyproject.toml

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

{vector_quantize_pytorch-1.28.3 → vector_quantize_pytorch-1.29.0}/vector_quantize_pytorch/__init__.py

@@ -11,5 +11,6 @@ from vector_quantize_pytorch.sim_vq import SimVQ
 from vector_quantize_pytorch.residual_sim_vq import ResidualSimVQ
 
 from vector_quantize_pytorch.binary_mapper import BinaryMapper
+from vector_quantize_pytorch.hierarchical_vq import HierarchicalVQ
 
 from vector_quantize_pytorch.utils import Sequential

vector_quantize_pytorch-1.29.0/vector_quantize_pytorch/hierarchical_vq.py

@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+from typing import List, Optional, Sequence
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from vector_quantize_pytorch.vector_quantize_pytorch import VectorQuantize
+
+
+def exists(v):
+    return v is not None
+
+
+class _Phi2D(nn.Module):
+    def __init__(self, dim: int, resi_ratio: float):
+        super().__init__()
+        self.resi_ratio = float(abs(resi_ratio))
+        self.conv = nn.Conv2d(dim, dim, 3, padding = 1)
+
+    def forward(self, x):
+        if self.resi_ratio <= 1e-8:
+            return x
+        return (1. - self.resi_ratio) * x + self.resi_ratio * self.conv(x)
+
+
+class HierarchicalVQ(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim: int,
+        codebook_size: int,
+        scales: Sequence[int],
+        decay: float = 0.99,
+        commitment_weight: float = 1.,
+        rotation_trick: bool = False,
+        kmeans_init: bool = True,
+        kmeans_iters: int = 10,
+        threshold_ema_dead_code: int = 2,
+        stochastic_sample_codes: bool = False,
+        sample_codebook_temp: float = 0.1,
+        orthogonal_reg_weight: float = 0.,
+        orthogonal_reg_max_codes: int = 128,
+        orthogonal_reg_active_codes_only: bool = False,
+        quant_resi: float = 0.5,
+        share_quant_resi: int = 1,
+        accept_image_fmap: bool = False
+    ):
+        super().__init__()
+        assert accept_image_fmap, 'HierarchicalVQ currently expects accept_image_fmap = True'
+
+        scales = [int(scale) for scale in scales]
+        assert len(scales) > 0
+        assert scales == sorted(scales)
+        assert all(scale > 0 for scale in scales)
+
+        self.dim = dim
+        self.scales = tuple(scales)
+        self.accept_image_fmap = True
+
+        self.vq = VectorQuantize(
+            dim = dim,
+            codebook_size = codebook_size,
+            decay = decay,
+            commitment_weight = commitment_weight,
+            rotation_trick = rotation_trick,
+            kmeans_init = kmeans_init,
+            kmeans_iters = kmeans_iters,
+            threshold_ema_dead_code = threshold_ema_dead_code,
+            stochastic_sample_codes = stochastic_sample_codes,
+            sample_codebook_temp = sample_codebook_temp,
+            orthogonal_reg_weight = orthogonal_reg_weight,
+            orthogonal_reg_max_codes = orthogonal_reg_max_codes,
+            orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only,
+            accept_image_fmap = True
+        )
+
+        if share_quant_resi == 1:
+            self.phi_shared = _Phi2D(dim, quant_resi)
+            self.phi_levels = None
+        else:
+            num_phi_levels = len(self.scales) if share_quant_resi <= 0 else min(len(self.scales), int(share_quant_resi))
+            self.phi_shared = None
+            self.phi_levels = nn.ModuleList([_Phi2D(dim, quant_resi) for _ in range(num_phi_levels)])
+
+    def _choose_phi(self, scale_index: int):
+        if exists(self.phi_shared):
+            return self.phi_shared
+
+        assert exists(self.phi_levels)
+
+        if len(self.phi_levels) == len(self.scales):
+            return self.phi_levels[scale_index]
+
+        if len(self.scales) == 1:
+            return self.phi_levels[0]
+
+        position = scale_index / float(len(self.scales) - 1)
+        phi_index = round(position * (len(self.phi_levels) - 1))
+        phi_index = max(0, min(len(self.phi_levels) - 1, phi_index))
+        return self.phi_levels[phi_index]
+
+    def _upsample_to_full(self, q, full_hw, scale_index: int):
+        if q.shape[-2:] != full_hw:
+            q = F.interpolate(q, size = full_hw, mode = 'bilinear', align_corners = False)
+
+        phi = self._choose_phi(scale_index)
+        if exists(phi):
+            q = phi(q)
+
+        return q
+
+    def forward(
+        self,
+        x,
+        indices = None,
+        sample_codebook_temp = None,
+        **kwargs
+    ):
+        assert indices is None, 'reconstruction-from-indices path not implemented in forward'
+        del kwargs
+
+        assert x.ndim == 4, 'expected image fmap of shape (batch, channels, height, width)'
+        batch, dim, height, width = x.shape
+        assert dim == self.dim
+
+        residual = x
+        reconstruction = torch.zeros_like(x)
+        all_indices = []
+        all_commit_losses = []
+
+        for scale_index, scale in enumerate(self.scales):
+            residual_down = F.adaptive_avg_pool2d(residual, output_size = (scale, scale))
+
+            vq_kwargs = {}
+            if exists(sample_codebook_temp):
+                vq_kwargs.update(sample_codebook_temp = sample_codebook_temp)
+
+            quantized, scale_indices, commit_loss = self.vq(residual_down, **vq_kwargs)
+            quantized = self._upsample_to_full(quantized, (height, width), scale_index)
+
+            reconstruction = reconstruction + quantized
+            residual = residual - quantized
+
+            all_indices.append(scale_indices)
+            all_commit_losses.append(commit_loss)
+
+        mean_commit_loss = torch.stack(all_commit_losses).mean()
+        return reconstruction, tuple(all_indices), mean_commit_loss
+
+    def get_output_from_indices(self, indices):
+        assert isinstance(indices, (tuple, list))
+        assert len(indices) == len(self.scales)
+
+        first = indices[0]
+        assert first.ndim == 3
+        batch, height, width = first.shape
+        full_hw = (self.scales[-1], self.scales[-1])
+        if (height, width) != full_hw:
+            full_hw = (self.scales[-1], self.scales[-1])
+
+        reconstructed = None
+
+        for scale_index, scale_indices in enumerate(indices):
+            q = self.vq.get_output_from_indices(scale_indices)
+            q = self._upsample_to_full(q, full_hw, scale_index)
+            reconstructed = q if reconstructed is None else reconstructed + q
+
+        return reconstructed

{vector_quantize_pytorch-1.28.3 → vector_quantize_pytorch-1.29.0}/vector_quantize_pytorch/utils.py

@@ -14,6 +14,7 @@ from vector_quantize_pytorch.residual_fsq import ResidualFSQ, GroupedResidualFSQ
 from vector_quantize_pytorch.latent_quantization import LatentQuantize
 from vector_quantize_pytorch.sim_vq import SimVQ
 from vector_quantize_pytorch.residual_sim_vq import ResidualSimVQ
+from vector_quantize_pytorch.hierarchical_vq import HierarchicalVQ
 
 QUANTIZE_KLASSES = (
     VectorQuantize,
@@ -28,7 +29,8 @@ QUANTIZE_KLASSES = (
     GroupedResidualLFQ,
     ResidualFSQ,
     GroupedResidualFSQ,
-    LatentQuantize
+    LatentQuantize,
+    HierarchicalVQ
 )
 
 # classes