vaezip 0.1.0__py3-none-any.whl

vaezip/utils/lloyd.py

@@ -0,0 +1,97 @@
+"""Lloyd-Max quantization tables for values normalized to [0, 1]."""
+
+import torch
+
+MAX_FIT_POINTS = 2**17
+TINY = 1e-30
+
+
+def nearest_level(values: torch.Tensor, tables: torch.Tensor) -> torch.Tensor:
+    """
+    args:
+        values: [B, N], values: [0...1]
+        tables: [B, K], sorted ascending
+    returns:
+        codes: [B, N], dtype: torch.uint8, index of the nearest table entry
+    """
+    midpoints = (tables[:, 1:] + tables[:, :-1]) / 2
+    return torch.searchsorted(midpoints.contiguous(), values.contiguous()).to(
+        torch.uint8
+    )
+
+
+def fit_lloyd_table(
+    values: torch.Tensor,
+    weights: torch.Tensor,
+    levels: int,
+    iterations: int = 20,
+) -> torch.Tensor:
+    """
+    args:
+        values: [N], values: [0...1]
+        weights: [N], positive importance of each value
+        levels: number of table entries
+    returns:
+        table: [levels], dtype: torch.float32, sorted, fp16-representable
+    """
+    stride = max(1, values.numel() // MAX_FIT_POINTS)
+    values, weights = values[::stride], weights[::stride]
+    weights = weights / weights.mean().clamp(min=TINY)
+    table = torch.linspace(0, 1, levels, device=values.device)
+    for _ in range(iterations):
+        assignment = nearest_level(values.unsqueeze(0), table.unsqueeze(0))[0].long()
+        weighted_sum = torch.zeros_like(table).index_add_(
+            0, assignment, values * weights
+        )
+        weight_total = torch.zeros_like(table).index_add_(0, assignment, weights)
+        table = torch.where(
+            weight_total > 0, weighted_sum / weight_total.clamp(min=TINY), table
+        )
+    return table.to(torch.float16).float().sort().values
+
+
+def fit_lloyd_tables_fast(
+    values: torch.Tensor,
+    weights: torch.Tensor,
+    levels: int,
+    iterations: int = 2,
+    bins: int = 4096,
+) -> torch.Tensor:
+    """
+    Batched approximation of `fit_lloyd_table` on a weighted histogram of the values.
+    The start is companded quantiles, followed by a few Lloyd steps on bin centroids.
+
+    args:
+        values: [R, N], values: [0...1]
+        weights: [R, N], positive importance of each value
+        levels: number of table entries
+    returns:
+        tables: [R, levels], dtype: torch.float32, sorted, fp16-representable
+    """
+    rows = values.shape[0]
+    weights = weights / weights.mean(dim=1, keepdim=True).clamp(min=TINY)
+    bin_index = (values * bins).long().clamp_(0, bins - 1)
+    mass = torch.zeros(rows, bins, device=values.device).scatter_add_(
+        1, bin_index, weights
+    )
+    moment = torch.zeros_like(mass).scatter_add_(1, bin_index, values * weights)
+    bin_centers = (torch.arange(bins, device=values.device) + 0.5) / bins
+    centroid = torch.where(mass > 0, moment / mass.clamp(min=TINY), bin_centers)
+
+    cumulative_mass = mass.pow(1 / 3).cumsum(dim=1)
+    targets = (torch.arange(levels, device=values.device) + 0.5) / levels
+    positions = torch.searchsorted(
+        cumulative_mass / cumulative_mass[:, -1:],
+        targets.expand(rows, -1).contiguous(),
+    )
+    tables = centroid.gather(1, positions.clamp(max=bins - 1))
+    for _ in range(iterations):
+        assignment = torch.searchsorted(
+            ((tables[:, 1:] + tables[:, :-1]) / 2).contiguous(), centroid
+        )
+        weighted_sum = torch.zeros_like(tables).scatter_add_(1, assignment, moment)
+        weight_total = torch.zeros_like(tables).scatter_add_(1, assignment, mass)
+        tables = torch.where(
+            weight_total > 0, weighted_sum / weight_total.clamp(min=TINY), tables
+        )
+    return tables.to(torch.float16).float().sort(dim=1).values

vaezip-0.1.0.dist-info/METADATA

@@ -0,0 +1,160 @@
+Metadata-Version: 2.4
+Name: vaezip
+Version: 0.1.0
+Summary: Compression of VAE latents: low-bit quantization codecs on PyTorch and safetensors
+License-Expression: Unlicense
+Project-URL: Homepage, https://github.com/tensorforger/VaeZip
+Project-URL: Issues, https://github.com/tensorforger/VaeZip/issues
+Project-URL: Changelog, https://github.com/tensorforger/VaeZip/blob/main/CHANGELOG.md
+Project-URL: Benchmark, https://github.com/tensorforger/LatentStorageBenchmark
+Keywords: vae,latents,compression,quantization,pytorch,safetensors
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.2
+Requires-Dist: safetensors>=0.4.3
+Requires-Dist: numpy>=1.24
+Dynamic: license-file
+
+![main_logo](https://raw.githubusercontent.com/tensorforger/tensorforger/main/assets/VaeZip_main_logo.png)
+
+
+**VaeZip** is the simple and efficient way to store latent tensors. It brings in low-bit quantization codecs, built around **PyTorch** and **safetensors** and exploits how [VAEs store there latents](https://github.com/tensorforger/LatentStorageBenchmark).
+
+
+The codecs were selected from the results of the
+[LatentStorageBenchmark](https://github.com/tensorforger/LatentStorageBenchmark),
+which compares many latent storage schemes by reconstruction quality, size and I/O time.
+
+![demo](https://raw.githubusercontent.com/tensorforger/tensorforger/main/assets/VaeZip_demo.png)
+
+*10x file size reduction with almost no quality loss. [reproduce these results](https://github.com/tensorforger/VaeZip/blob/main/examples/demo_reproduction.ipynb)*.
+
+**Use cases**:
+* Storage of large-scale encoded datasets for pretraining.
+* Network transfer of latents in efficient inference pipelines, where VAE and generative transformer/VLM are on different nodes.
+* Pipeline parallelism: Transfer activations/latents between pipeline stages over slow interconnects (PCIe, Ethernet, cross-datacenter).
+
+## Installation
+
+```bash
+pip install vaezip
+```
+
+Requires Python 3.10+ and PyTorch 2.2+. Install PyTorch for your platform first if you need a specific build.
+
+## Quick start
+
+```python
+import torch
+import vaezip
+
+latents = torch.randn([4, 32, 128, 128], device="cuda")
+
+vaezip.save(latents, "latents.vz.safetensors", codec="lbrq_4bit")
+restored = vaezip.load("latents.vz.safetensors", device="cuda")
+```
+
+The codec config is stored in the file metadata, so `load()` needs no codec.
+Pass `save_metadata=False` to drop it, then `load()` requires an explicit `codec=`.
+
+## Key concepts
+
+**Codec** is a compression algorithm configured at construction.
+`encode()` returns an `EncodedTensor`, `decode()` returns a `float32` tensor of the original shape
+(or another floating dtype via `decode(encoded, dtype=...)`).
+
+```python
+from vaezip.codecs import BlockRangeQuantCodec
+
+codec = BlockRangeQuantCodec(
+    value_bits=4, range_bits=4, channel_dim=1, block_size=(8, 8)
+)
+encoded = codec.encode(latents)
+restored = codec.decode(encoded)
+```
+
+**EncodedTensor** holds named tensors plus string metadata with the codec config.
+It serializes to `.safetensors` (default) or `.pth`:
+
+```python
+data = encoded.serialize()  # or vaezip.serialize(encoded)
+encoded = vaezip.deserialize(data)
+data = encoded.serialize(save_metadata=False, format="pth")
+```
+
+**Configs** are diffusers-style JSON files and can be created from files, shortcut names or dicts.
+Keyword arguments override config values:
+
+```python
+from vaezip.codecs import AutoCodec
+
+codec.save_config("codec.json")
+codec = AutoCodec.from_config("codec.json")
+codec = AutoCodec.from_config("brq_8bit", block_size=(4, 4, 4))
+```
+
+## Codecs
+
+| Codec | Description |
+|---|---|
+| `CastCodec(cast_dtype)` | Cast to a floating dtype, e.g. `torch.float16`, `torch.float8_e4m3fn`. |
+| `RangeQuantCodec(bits, channel_dim)` | Per-channel min/max normalization, `bits`-bit integers. |
+| `BlockRangeQuantCodec(value_bits, range_bits, channel_dim, block_size)` | Min/max per block of a channel, `value_bits` inside the block range, block ranges stored with `range_bits`. |
+| `LloydBlockRangeQuantCodec(..., lloyd_fast, cache_table)` | Same with Lloyd-Max level tables instead of uniform levels. |
+
+The first tensor dimension is always the batch. Block codecs take one more dimension as the channel
+(`channel_dim`, negative values count from the end) and treat the rest as spatial, `block_size` has one entry per spatial dimension:
+
+```python
+# [B, C, H, W]
+BlockRangeQuantCodec(value_bits=4, range_bits=4, channel_dim=1, block_size=(8, 8))
+# [B, C, T, H, W], no temporal blocks
+BlockRangeQuantCodec(value_bits=4, range_bits=4, channel_dim=1, block_size=(1, 8, 8))
+# [B, T, C, H, W]
+BlockRangeQuantCodec(value_bits=4, range_bits=4, channel_dim=2, block_size=(4, 8, 8))
+```
+
+Sizes that are not divisible by the block size are padded internally, decoded tensors keep the original shape.
+Low-bit values are bit-packed, so the size follows the number of bits.
+Quantization ranges are computed per sample, samples in a batch are encoded independently.
+
+`LloydBlockRangeQuantCodec` always stores its level table in the encoded tensor.
+`cache_table=True` fits the table once on the first batch and reuses it to speed up encoding,
+`codec.flush_cache()` resets it.
+
+## Shortcuts
+
+| Name | Codec | Approx. bits per value |
+|---|---|---|
+| `cast_fp32` (default), `cast_fp16`, `cast_bf16`, `cast_fp8_e4m3` | `CastCodec` | 32, 16, 16, 8 |
+| `rq_8bit`, `rq_6bit`, `rq_4bit` | `RangeQuantCodec` | 8, 6, 4 |
+| `brq_8bit`, `brq_6bit`, `brq_5bit`, `brq_4bit`, `brq_3bit`, `brq_2bit` | `BlockRangeQuantCodec` | 8.1, 6.1, 5.1, 4.1, 3.1, 2.1 |
+| `lbrq_8bit`, `lbrq_6bit`, `lbrq_5bit`, `lbrq_4bit`, `lbrq_3bit`, `lbrq_2bit` | `LloydBlockRangeQuantCodec` | 8.1, 6.1, 5.1, 4.1, 3.1, 2.1 |
+
+The block shortcuts use `value_bits` from the name, `range_bits=4`, `block_size=(8, 8)` and `channel_dim=1`.
+These settings are the points of the MSE Pareto front on FLUX.2 latents (benchmark run 1.2) closest to the target bit count.
+`vaezip.codecs.list_shortcuts()` lists all names.
+
+## Development
+
+```bash
+uv sync
+uv run pytest
+uv run ruff check . && uv run ruff format --check . && uv run mypy
+```
+
+## License
+
+[Unlicense](LICENSE)

vaezip-0.1.0.dist-info/RECORD

@@ -0,0 +1,24 @@
+vaezip/__init__.py,sha256=WUX4-5GhUmr8696GkeR5f4kerXiphGiRna9ThWQTNho,344
+vaezip/_version.py,sha256=ZXJqehvwFXAGJKgzeyjtkk4Q4Zm39gPOgD2cEGPL0wM,190
+vaezip/api.py,sha256=4NtwuyRFRox7UVDDpbgV1YHWKC-9Rn5O2ctdL8luPCY,4422
+vaezip/encoded_tensor.py,sha256=Q_7JnVk1Htq0FyxijD_wWq3TzKXd7maf3sXSPVD-6UI,4343
+vaezip/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vaezip/codecs/__init__.py,sha256=9__t5zGCkUTb2f1myzz7VcBZrB1egKw64kbHSpVu8j4,621
+vaezip/codecs/auto.py,sha256=8-Ei9-HO9Evu8oFiwiPzNiw8ApnLDlPtxybZe_CzrYY,1561
+vaezip/codecs/base.py,sha256=TvwcirA89mWhbB6Ul2BUBVWuyajBCFpXhnbQXtc7CJ4,4276
+vaezip/codecs/block_range_quant.py,sha256=2dmzXj9UV29tyIKdM8buFIxn3rfnRf9PCWgiw5cFLFg,9259
+vaezip/codecs/cast.py,sha256=cGPJgYTyOFk9WmvFMlcVN9FllPxSbBZzhjnFee872U0,2060
+vaezip/codecs/config.py,sha256=eBreKKHrVpAgsX75we8Ub0vmPJGyP_-FlQdc9q0WiuI,1744
+vaezip/codecs/lloyd_block_range_quant.py,sha256=QGs25sP6yBg3s_dXoXkuAbvEWZhAQCEU9TLQvnHbXko,4181
+vaezip/codecs/range_quant.py,sha256=RRb8MLOtC1mTFJQEJl_Q1JJeAzZp_C0hujUSv2E4Yb8,3506
+vaezip/codecs/shortcuts.py,sha256=nrtgJ95oJlRNJoFpbYELaPfvxDfF1HWQMk_-8EILOXk,1624
+vaezip/utils/__init__.py,sha256=tntvtMeE1lu8n7XMAtZAGXy5MmHaecc3hsVKWBhDS4k,43
+vaezip/utils/bit_packing.py,sha256=aJJj9uQJtI9UVUlLHsyajmFmX9TRX2UD_uolHxSkJ5o,2267
+vaezip/utils/blocks.py,sha256=o_BKl8Vcw1t0VMBSo7orksKdV0X6bbvG06mnNBFZNBk,3147
+vaezip/utils/layout.py,sha256=WIySLbmcdVa_iRYFTcs7w8Mz_xRbiukMBv2Bkb4kkQY,1817
+vaezip/utils/lloyd.py,sha256=6AV1my6B69lOPbFRewnbmltTLxPEZnbCdha09wBb8tI,3594
+vaezip-0.1.0.dist-info/licenses/LICENSE,sha256=b5ovsAkiXeXMocKiTf7XpovCJlfkVbSj1j8GUli0VdE,1211
+vaezip-0.1.0.dist-info/METADATA,sha256=80BOroDlG8di_rM3aB5rCuK4ago3GvmwaQyH-CiVQGU,6699
+vaezip-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+vaezip-0.1.0.dist-info/top_level.txt,sha256=QZ0uPyZhNWehXV3pZqVa3ZjEFQ232ONVszOYzgva4r8,7
+vaezip-0.1.0.dist-info/RECORD,,

vaezip-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

vaezip-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,24 @@
+This is free and unencumbered software released into the public domain.
+
+Anyone is free to copy, modify, publish, use, compile, sell, or
+distribute this software, either in source code form or as a compiled
+binary, for any purpose, commercial or non-commercial, and by any
+means.
+
+In jurisdictions that recognize copyright laws, the author or authors
+of this software dedicate any and all copyright interest in the
+software to the public domain. We make this dedication for the benefit
+of the public at large and to the detriment of our heirs and
+successors. We intend this dedication to be an overt act of
+relinquishment in perpetuity of all present and future rights to this
+software under copyright law.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
+OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+OTHER DEALINGS IN THE SOFTWARE.
+
+For more information, please refer to <https://unlicense.org/>

vaezip-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+vaezip