vaezip 0.1.0__tar.gz

vaezip-0.1.0/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/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,131 @@
+![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/pyproject.toml

@@ -0,0 +1,78 @@
+[build-system]
+requires = ["setuptools>=77.0.3"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "vaezip"
+version = "0.1.0"
+description = "Compression of VAE latents: low-bit quantization codecs on PyTorch and safetensors"
+readme = "README.md"
+license = "Unlicense"
+license-files = ["LICENSE"]
+requires-python = ">=3.10"
+keywords = [
+    "vae",
+    "latents",
+    "compression",
+    "quantization",
+    "pytorch",
+    "safetensors",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Typing :: Typed",
+]
+dependencies = [
+    "torch>=2.2",
+    "safetensors>=0.4.3",
+    "numpy>=1.24",
+]
+
+[project.urls]
+Homepage = "https://github.com/tensorforger/VaeZip"
+Issues = "https://github.com/tensorforger/VaeZip/issues"
+Changelog = "https://github.com/tensorforger/VaeZip/blob/main/CHANGELOG.md"
+Benchmark = "https://github.com/tensorforger/LatentStorageBenchmark"
+
+[dependency-groups]
+dev = [
+    "build>=1.2",
+    "mypy>=1.10",
+    "pre-commit>=3.7",
+    "pytest>=8.0",
+    "ruff>=0.5",
+    "twine>=5.1",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+vaezip = ["py.typed"]
+
+[tool.ruff]
+line-length = 88
+target-version = "py310"
+src = ["src", "tests"]
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I", "UP", "B", "SIM", "RUF"]
+
+[tool.mypy]
+strict = true
+mypy_path = "src"
+explicit_package_bases = true
+files = ["src/vaezip"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = ["-ra", "--strict-markers", "--strict-config"]

vaezip-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

vaezip-0.1.0/src/vaezip/__init__.py

@@ -0,0 +1,16 @@
+"""VaeZip: compression of VAE latents."""
+
+from vaezip import codecs
+from vaezip._version import __version__
+from vaezip.api import load, save
+from vaezip.encoded_tensor import EncodedTensor, deserialize, serialize
+
+__all__ = [
+    "EncodedTensor",
+    "__version__",
+    "codecs",
+    "deserialize",
+    "load",
+    "save",
+    "serialize",
+]

vaezip-0.1.0/src/vaezip/_version.py

@@ -0,0 +1,8 @@
+"""Package version."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("vaezip")
+except PackageNotFoundError:
+    __version__ = "0.0.0+unknown"

vaezip-0.1.0/src/vaezip/api.py

@@ -0,0 +1,139 @@
+"""High level save and load of tensors."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import overload
+
+import torch
+
+from vaezip.codecs.auto import AutoCodec
+from vaezip.codecs.base import Codec
+from vaezip.codecs.config import without_version
+from vaezip.codecs.shortcuts import DEFAULT_SHORTCUT
+from vaezip.encoded_tensor import (
+    PTH_FORMAT,
+    SAFETENSORS_FORMAT,
+    EncodedTensor,
+)
+
+CodecSpecification = Codec | str | os.PathLike[str]
+PTH_SUFFIXES = (".pth", ".pt")
+
+
+def _resolve_codec(codec: CodecSpecification | None) -> Codec:
+    if codec is None:
+        return AutoCodec.from_config(DEFAULT_SHORTCUT)
+    if isinstance(codec, Codec):
+        return codec
+    return AutoCodec.from_config(codec)
+
+
+def _infer_format(path: str | os.PathLike[str], format: str | None) -> str:
+    if format is not None:
+        return format
+    return PTH_FORMAT if Path(path).suffix in PTH_SUFFIXES else SAFETENSORS_FORMAT
+
+
+@overload
+def save(
+    tensor: torch.Tensor,
+    path: None = None,
+    codec: CodecSpecification | None = None,
+    save_metadata: bool = True,
+    format: str | None = None,
+) -> bytes: ...
+
+
+@overload
+def save(
+    tensor: torch.Tensor,
+    path: str | os.PathLike[str],
+    codec: CodecSpecification | None = None,
+    save_metadata: bool = True,
+    format: str | None = None,
+) -> None: ...
+
+
+def save(
+    tensor: torch.Tensor,
+    path: str | os.PathLike[str] | None = None,
+    codec: CodecSpecification | None = None,
+    save_metadata: bool = True,
+    format: str | None = None,
+) -> bytes | None:
+    """
+    Encodes the tensor and writes it to a file, or returns bytes if there is no path.
+
+    args:
+        tensor: [B, ...], floating point
+        path: output file, `None` to return bytes
+        codec: codec object, shortcut name or config path, default is "cast_fp32"
+        save_metadata: store the codec config, required by `load()` without a codec
+        format: "safetensors" or "pth", by default inferred from the path suffix
+        (safetensors if unknown)
+    """
+    resolved_codec = _resolve_codec(codec)
+    if path is None:
+        serialization_format = format or SAFETENSORS_FORMAT
+    else:
+        serialization_format = _infer_format(path, format)
+    data = resolved_codec.encode(tensor).serialize(
+        save_metadata=save_metadata, format=serialization_format
+    )
+    if path is None:
+        return data
+    Path(path).write_bytes(data)
+    return None
+
+
+def load(
+    source: str | os.PathLike[str] | bytes | bytearray,
+    codec: CodecSpecification | None = None,
+    dtype: torch.dtype = torch.float32,
+    device: torch.device | str | None = None,
+    format: str | None = None,
+) -> torch.Tensor:
+    """
+    Reads a file or bytes written by `save()` and decodes the tensor.
+
+    args:
+        source: file path or bytes
+        codec: codec object, shortcut name or config path, by default the codec
+        from the metadata. If both are available they must have the same config.
+        dtype: floating point dtype of the result
+        device: device to decode on, default is CPU
+        format: "safetensors" or "pth", by default inferred from the path suffix
+        (safetensors for bytes)
+    returns:
+        tensor: [B, ...], original shape, dtype: dtype
+    """
+    if isinstance(source, (bytes, bytearray)):
+        data = bytes(source)
+        serialization_format = format or SAFETENSORS_FORMAT
+    else:
+        data = Path(source).read_bytes()
+        serialization_format = _infer_format(source, format)
+    encoded = EncodedTensor.deserialize(data, format=serialization_format)
+    stored_config = encoded.codec_config
+
+    if codec is None:
+        if stored_config is None:
+            raise ValueError(
+                "The data has no codec config in its metadata, "
+                "pass the codec explicitly"
+            )
+        resolved_codec = AutoCodec.from_config(stored_config)
+    else:
+        resolved_codec = _resolve_codec(codec)
+        if stored_config is not None and without_version(
+            stored_config
+        ) != without_version(resolved_codec.get_config()):
+            raise ValueError(
+                f"The codec {resolved_codec!r} does not match the config stored in the "
+                f"data: {stored_config}"
+            )
+    if device is not None:
+        encoded = encoded.to(device)
+    return resolved_codec.decode(encoded, dtype=dtype)

vaezip-0.1.0/src/vaezip/codecs/__init__.py

@@ -0,0 +1,20 @@
+"""Codecs for compression of VAE latents."""
+
+from vaezip.codecs.auto import AutoCodec
+from vaezip.codecs.base import Codec
+from vaezip.codecs.block_range_quant import BlockRangeQuantCodec
+from vaezip.codecs.cast import CastCodec
+from vaezip.codecs.lloyd_block_range_quant import LloydBlockRangeQuantCodec
+from vaezip.codecs.range_quant import RangeQuantCodec
+from vaezip.codecs.shortcuts import DEFAULT_SHORTCUT, list_shortcuts
+
+__all__ = [
+    "DEFAULT_SHORTCUT",
+    "AutoCodec",
+    "BlockRangeQuantCodec",
+    "CastCodec",
+    "Codec",
+    "LloydBlockRangeQuantCodec",
+    "RangeQuantCodec",
+    "list_shortcuts",
+]

vaezip-0.1.0/src/vaezip/codecs/auto.py

@@ -0,0 +1,44 @@
+"""Construction of codecs from configs without knowing the codec class."""
+
+from typing import Any
+
+from vaezip.codecs.base import Codec
+from vaezip.codecs.block_range_quant import BlockRangeQuantCodec
+from vaezip.codecs.cast import CastCodec
+from vaezip.codecs.config import ConfigSource, split_config
+from vaezip.codecs.lloyd_block_range_quant import LloydBlockRangeQuantCodec
+from vaezip.codecs.range_quant import RangeQuantCodec
+
+CODEC_CLASSES: dict[str, type[Codec]] = {
+    codec_class.__name__: codec_class
+    for codec_class in (
+        CastCodec,
+        RangeQuantCodec,
+        BlockRangeQuantCodec,
+        LloydBlockRangeQuantCodec,
+    )
+}
+
+
+class AutoCodec:
+    """Builds the right codec class from a config."""
+
+    def __init__(self) -> None:
+        raise TypeError("AutoCodec is not instantiable, use AutoCodec.from_config()")
+
+    @staticmethod
+    def from_config(config: ConfigSource, **overrides: Any) -> Codec:
+        """
+        args:
+            config: path to a JSON config, a shortcut name such as "cast_fp16",
+            or a config dict
+            overrides: constructor parameters that replace the ones from the config
+        """
+        class_name, parameters = split_config(config, overrides)
+        if class_name is None:
+            raise ValueError("Config has no '_class_name' entry")
+        if class_name not in CODEC_CLASSES:
+            raise ValueError(
+                f"Unknown codec class {class_name!r}, available: {list(CODEC_CLASSES)}"
+            )
+        return CODEC_CLASSES[class_name](**parameters)