vaezip 0.2.0

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.

The codecs were selected from the results of the LatentStorageBenchmark, which compares many latent storage schemes by reconstruction quality, size and I/O time.

10x file size reduction with almost no quality loss. reproduce these results.

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

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

import torch
import vaezip

latents = torch.randn([4, 32, 128, 128], device="cuda")

vaezip.save_tensor(latents, "latents.vz.safetensors", codec="lbrq_4bit")
restored = vaezip.load_tensor("latents.vz.safetensors", device="cuda")

The codec config is stored in the file, so load_tensor() needs no codec.

Several tensors, each with its own codec, go to one file as a dict of EncodedTensor:

from vaezip.codecs import AutoCodec

tensor_a = torch.randn([4, 32, 128, 128])
tensor_b = torch.randn([4, 64, 256])
codec_b = AutoCodec.from_config("brq_8bit", block_size=(4,))

batch_dict = {
    "tensor_a": vaezip.encode(tensor_a, codec="lbrq_3bit"),
    "tensor_b": vaezip.encode(tensor_b, codec=codec_b),
}
vaezip.save(batch_dict, "batch_dict.vz.safetensors")

batch_dict = vaezip.load("batch_dict.vz.safetensors")
tensor_a = batch_dict["tensor_a"].decode(device="cuda")

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=...)).

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. encoded.decode() restores the tensor with the codec from the metadata, or with an explicit codec=. vaezip.encode() and vaezip.decode() are the same with codec shortcuts and configs.

Files hold a dict of EncodedTensor, vaezip.save() and vaezip.load() work with paths or bytes. The format is .safetensors by default, .pth only with format="pth". vaezip.save_tensor() and vaezip.load_tensor() are shortcuts for a file with a single tensor stored as "tensor".

Configs are diffusers-style JSON files and can be created from files, shortcut names or dicts. Keyword arguments override config values:

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

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:

# [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

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

uv sync
uv run pytest
uv run ruff check . && uv run ruff format --check . && uv run mypy

License

Unlicense