vaezip 0.1.0__py3-none-any.whl

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/_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/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/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/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)

vaezip/codecs/base.py

@@ -0,0 +1,119 @@
+"""Abstract codec API."""
+
+from __future__ import annotations
+
+import json
+import os
+from abc import ABC, abstractmethod
+from typing import Any, TypeVar
+
+import torch
+
+from vaezip._version import __version__
+from vaezip.codecs.config import ConfigSource, split_config
+from vaezip.codecs.shortcuts import CODEC_CLASS_KEY, VERSION_KEY
+from vaezip.encoded_tensor import CODEC_CONFIG_KEY, EncodedTensor
+
+CodecType = TypeVar("CodecType", bound="Codec")
+
+
+class Codec(ABC):
+    """
+    Compression algorithm for tensors.
+
+    The configuration is fixed at construction. Encoding takes a tensor and returns an
+    `EncodedTensor` on the same device, decoding restores a float tensor of the
+    original shape.
+    """
+
+    def encode(self, tensor: torch.Tensor) -> EncodedTensor:
+        """
+        args:
+            tensor: [B, ...], floating point, finite values
+        returns:
+            encoded: tensors on the device of the input, metadata with the codec config
+        """
+        if not isinstance(tensor, torch.Tensor):
+            raise TypeError(f"Expected a torch.Tensor, got {type(tensor).__name__}")
+        if not tensor.is_floating_point():
+            raise TypeError(f"Expected a floating point tensor, got {tensor.dtype}")
+        if tensor.numel() == 0:
+            raise ValueError("Can not encode an empty tensor")
+        tensors = self._encode_tensors(tensor.detach())
+        metadata = {CODEC_CONFIG_KEY: json.dumps(self.get_config())}
+        return EncodedTensor(tensors=tensors, metadata=metadata)
+
+    def decode(
+        self, encoded: EncodedTensor, dtype: torch.dtype = torch.float32
+    ) -> torch.Tensor:
+        """
+        args:
+            encoded: output of `encode()` of the same codec configuration
+            dtype: floating point dtype of the result
+        returns:
+            tensor: [B, ...], original shape, dtype: dtype, device of encoded tensors
+        """
+        if not dtype.is_floating_point:
+            raise TypeError(f"dtype must be a floating point dtype, got {dtype}")
+        return self._decode_tensors(encoded.tensors, dtype)
+
+    @abstractmethod
+    def _encode_tensors(self, tensor: torch.Tensor) -> dict[str, torch.Tensor]:
+        """
+        args:
+            tensor: [B, ...], floating point
+        returns:
+            tensors: compressed representation, per-sample tensors have the batch first
+        """
+
+    @abstractmethod
+    def _decode_tensors(
+        self, tensors: dict[str, torch.Tensor], dtype: torch.dtype
+    ) -> torch.Tensor:
+        """
+        args:
+            tensors: output of `_encode_tensors()`
+            dtype: floating point dtype of the result
+        returns:
+            tensor: [B, ...], original shape, dtype: dtype
+        """
+
+    @abstractmethod
+    def get_config_parameters(self) -> dict[str, Any]:
+        """JSON-serializable constructor parameters, `from_config` accepts them back."""
+
+    def get_config(self) -> dict[str, Any]:
+        """Diffusers-style config: class name, package version and the parameters."""
+        return {
+            CODEC_CLASS_KEY: type(self).__name__,
+            VERSION_KEY: __version__,
+            **self.get_config_parameters(),
+        }
+
+    def save_config(self, path: str | os.PathLike[str]) -> None:
+        with open(path, "w", encoding="utf-8") as config_file:
+            json.dump(self.get_config(), config_file, indent=2)
+            config_file.write("\n")
+
+    @classmethod
+    def from_config(
+        cls: type[CodecType], config: ConfigSource, **overrides: Any
+    ) -> CodecType:
+        """
+        args:
+            config: path to a JSON config, a shortcut name 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 not None and class_name != cls.__name__:
+            raise ValueError(
+                f"Config describes {class_name}, "
+                f"can not construct {cls.__name__} from it"
+            )
+        return cls(**parameters)
+
+    def __repr__(self) -> str:
+        parameters = ", ".join(
+            f"{name}={value!r}" for name, value in self.get_config_parameters().items()
+        )
+        return f"{type(self).__name__}({parameters})"

vaezip/codecs/block_range_quant.py

@@ -0,0 +1,232 @@
+"""Blockwise range quantization to low-bit integers."""
+
+import math
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+
+from vaezip.codecs.base import Codec
+from vaezip.codecs.range_quant import (
+    CHANNEL_MAX_KEY,
+    CHANNEL_MIN_KEY,
+    CODES_KEY,
+    MIN_RANGE,
+    SHAPE_KEY,
+)
+from vaezip.utils.bit_packing import check_bits, max_code, pack_bits, unpack_bits
+from vaezip.utils.blocks import merge_blocks, padded_size, split_into_blocks
+from vaezip.utils.layout import (
+    channel_from_second,
+    channel_to_second,
+    check_channel_dim,
+    shape_with_channel_second,
+)
+
+RANGE_LOWER_KEY = "range_lower"
+RANGE_UPPER_KEY = "range_upper"
+
+
+def parse_block_size(block_size: Sequence[int]) -> tuple[int, ...]:
+    if isinstance(block_size, (str, bytes)) or not isinstance(block_size, Sequence):
+        raise TypeError("block_size must be a sequence of ints")
+    parsed = tuple(block_size)
+    if not parsed:
+        raise ValueError("block_size must have at least one dimension")
+    for extent in parsed:
+        if isinstance(extent, bool) or not isinstance(extent, int) or extent < 1:
+            raise ValueError(f"block_size must contain positive ints, got {block_size}")
+    return parsed
+
+
+class BlockRangeQuantCodec(Codec):
+    """
+    Quantizes each block of each channel between its own min and max.
+    The block ranges are quantized to `range_bits` relative to the min and max of the
+    whole channel, the values to `value_bits` inside the block range. Everything is
+    packed into bytes.
+
+    The first dimension is the batch, one other dimension is the channel, the remaining
+    dimensions are spatial and `block_size` has one entry for each of them. Spatial
+    sizes that are not divisible by the block size are padded internally, the decoded
+    tensor has the original shape.
+    """
+
+    def __init__(
+        self,
+        value_bits: int = 4,
+        range_bits: int = 4,
+        channel_dim: int = 1,
+        block_size: Sequence[int] = (4, 4),
+    ) -> None:
+        super().__init__()
+        check_bits(value_bits, "value_bits")
+        check_bits(range_bits, "range_bits")
+        check_channel_dim(channel_dim)
+        self.value_bits = value_bits
+        self.range_bits = range_bits
+        self.channel_dim = channel_dim
+        self.block_size = parse_block_size(block_size)
+
+    def _encode_tensors(self, tensor: torch.Tensor) -> dict[str, torch.Tensor]:
+        """
+        args:
+            tensor: [B, ...] with ndim == len(block_size) + 2, floating point
+        returns:
+            tensors: codes: [B, ceil(C * G * V * value_bits / 8)],
+            range_lower, range_upper: [B, ceil(C * G * range_bits / 8)],
+            channel_min, channel_max: [B, C], shape: [ndim], dtype: torch.int64,
+            G: number of blocks per channel, V: number of values per block
+        """
+        if tensor.ndim != len(self.block_size) + 2:
+            raise ValueError(
+                f"Expected a tensor with {len(self.block_size) + 2} dimensions "
+                f"(batch, channel and {len(self.block_size)} spatial), got shape "
+                f"{tuple(tensor.shape)}"
+            )
+        channel_first = channel_to_second(tensor.float(), self.channel_dim)
+        blocks = split_into_blocks(channel_first, self.block_size)  # [B, C, G, V]
+        batch_size = blocks.shape[0]
+        range_levels = max_code(self.range_bits)
+
+        channel_min = blocks.amin(dim=(2, 3), keepdim=True)  # [B, C, 1, 1]
+        channel_max = blocks.amax(dim=(2, 3), keepdim=True)
+        range_step = _range_step(channel_min, channel_max, range_levels)
+        block_min = blocks.amin(dim=3, keepdim=True)  # [B, C, G, 1]
+        block_max = blocks.amax(dim=3, keepdim=True)
+        lower = torch.floor((block_min - channel_min) / range_step)
+        lower = lower.clamp(0, range_levels - 1)
+        upper = torch.ceil((block_max - channel_min) / range_step).clamp(
+            max=range_levels
+        )
+        upper = torch.maximum(upper, lower + 1)
+        start, span = _block_grid(lower, upper, channel_min, range_step)
+
+        normalized = ((blocks - start) / span).clamp(0, 1)
+        codes, quantizer_tensors = self._quantize_normalized(normalized, span)
+        return {
+            CODES_KEY: pack_bits(codes.reshape(batch_size, -1), self.value_bits),
+            RANGE_LOWER_KEY: pack_bits(
+                lower.reshape(batch_size, -1).to(torch.uint8), self.range_bits
+            ),
+            RANGE_UPPER_KEY: pack_bits(
+                upper.reshape(batch_size, -1).to(torch.uint8), self.range_bits
+            ),
+            CHANNEL_MIN_KEY: channel_min.reshape(batch_size, -1),
+            CHANNEL_MAX_KEY: channel_max.reshape(batch_size, -1),
+            SHAPE_KEY: torch.tensor(
+                tensor.shape, dtype=torch.int64, device=tensor.device
+            ),
+            **quantizer_tensors,
+        }
+
+    def _decode_tensors(
+        self, tensors: dict[str, torch.Tensor], dtype: torch.dtype
+    ) -> torch.Tensor:
+        """
+        args:
+            tensors: output of `_encode_tensors()`
+        returns:
+            tensor: [B, ...], dtype: dtype
+        """
+        shape = [int(size) for size in tensors[SHAPE_KEY].tolist()]
+        if len(shape) != len(self.block_size) + 2:
+            raise ValueError(
+                f"Encoded tensor has {len(shape)} dimensions, "
+                f"this codec expects {len(self.block_size) + 2}"
+            )
+        batch_size, channels, *spatial_shape = shape_with_channel_second(
+            shape, self.channel_dim
+        )
+        blocks_per_channel = math.prod(
+            padded_size(size, extent) // extent
+            for size, extent in zip(spatial_shape, self.block_size, strict=True)
+        )
+        values_per_block = math.prod(self.block_size)
+        range_levels = max_code(self.range_bits)
+
+        ranges_count = channels * blocks_per_channel
+        lower = unpack_bits(tensors[RANGE_LOWER_KEY], ranges_count, self.range_bits)
+        upper = unpack_bits(tensors[RANGE_UPPER_KEY], ranges_count, self.range_bits)
+        lower = lower.reshape(batch_size, channels, blocks_per_channel, 1).float()
+        upper = upper.reshape(batch_size, channels, blocks_per_channel, 1).float()
+        channel_min = tensors[CHANNEL_MIN_KEY].reshape(batch_size, channels, 1, 1)
+        channel_max = tensors[CHANNEL_MAX_KEY].reshape(batch_size, channels, 1, 1)
+        range_step = _range_step(channel_min, channel_max, range_levels)
+        start, span = _block_grid(lower, upper, channel_min, range_step)
+
+        codes = unpack_bits(
+            tensors[CODES_KEY], ranges_count * values_per_block, self.value_bits
+        )
+        codes = codes.reshape(
+            batch_size, channels, blocks_per_channel, values_per_block
+        )
+        normalized = self._dequantize_normalized(codes, tensors)
+        blocks = normalized * span + start  # [B, C, G, V]
+        channel_first = merge_blocks(blocks, spatial_shape, self.block_size)
+        return (
+            channel_from_second(channel_first, self.channel_dim).to(dtype).contiguous()
+        )
+
+    def _quantize_normalized(
+        self, normalized: torch.Tensor, span: torch.Tensor
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        """
+        args:
+            normalized: [B, C, G, V], values: [0...1], position inside the block range
+            span: [B, C, G, 1], width of the block range
+        returns:
+            codes: [B, C, G, V], dtype: torch.uint8, values: [0...2^value_bits - 1]
+            tensors: additional tensors needed for dequantization
+        """
+        codes = torch.round(normalized * max_code(self.value_bits))
+        return codes.to(torch.uint8), {}
+
+    def _dequantize_normalized(
+        self, codes: torch.Tensor, tensors: dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        """
+        args:
+            codes: [B, C, G, V], dtype: torch.uint8
+            tensors: output of `_encode_tensors()`
+        returns:
+            normalized: [B, C, G, V], dtype: torch.float32, values: [0...1]
+        """
+        return codes.float() / max_code(self.value_bits)
+
+    def get_config_parameters(self) -> dict[str, Any]:
+        return {
+            "value_bits": self.value_bits,
+            "range_bits": self.range_bits,
+            "channel_dim": self.channel_dim,
+            "block_size": list(self.block_size),
+        }
+
+
+def _range_step(
+    channel_min: torch.Tensor, channel_max: torch.Tensor, range_levels: int
+) -> torch.Tensor:
+    """
+    args:
+        channel_min, channel_max: [B, C, 1, 1]
+    returns:
+        range_step: [B, C, 1, 1], distance between neighbouring block range positions
+    """
+    return ((channel_max - channel_min) / range_levels).clamp(min=MIN_RANGE)
+
+
+def _block_grid(
+    lower: torch.Tensor,
+    upper: torch.Tensor,
+    channel_min: torch.Tensor,
+    range_step: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    args:
+        lower, upper: [B, C, G, 1], dtype: torch.float32, block range positions
+        channel_min, range_step: [B, C, 1, 1]
+    returns:
+        start: [B, C, G, 1], lower end of the block range
+        span: [B, C, G, 1], width of the block range
+    """
+    return channel_min + lower * range_step, (upper - lower) * range_step

vaezip/codecs/cast.py

@@ -0,0 +1,62 @@
+"""Codec that casts values to a lower precision floating point dtype."""
+
+from typing import Any
+
+import torch
+
+from vaezip.codecs.base import Codec
+
+VALUES_KEY = "values"
+
+
+def parse_dtype(dtype: torch.dtype | str) -> torch.dtype:
+    """Accepts a torch dtype or its name such as "float16"."""
+    if isinstance(dtype, str):
+        name = dtype.removeprefix("torch.")
+        resolved = getattr(torch, name, None)
+        if not isinstance(resolved, torch.dtype):
+            raise ValueError(f"Unknown torch dtype {dtype!r}")
+        return resolved
+    if isinstance(dtype, torch.dtype):
+        return dtype
+    raise TypeError(
+        f"Expected a torch.dtype or a dtype name, got {type(dtype).__name__}"
+    )
+
+
+class CastCodec(Codec):
+    """Stores the tensor in a floating point dtype, out-of-range values are clamped."""
+
+    def __init__(self, cast_dtype: torch.dtype | str = torch.float16) -> None:
+        super().__init__()
+        self.cast_dtype = parse_dtype(cast_dtype)
+        if not self.cast_dtype.is_floating_point:
+            raise ValueError(
+                f"cast_dtype must be floating point, got {self.cast_dtype}"
+            )
+
+    def _encode_tensors(self, tensor: torch.Tensor) -> dict[str, torch.Tensor]:
+        """
+        args:
+            tensor: [B, ...], floating point
+        returns:
+            tensors: values: [B, ...], dtype: cast_dtype
+        """
+        cast_max = torch.finfo(self.cast_dtype).max
+        if cast_max < torch.finfo(tensor.dtype).max:
+            tensor = tensor.clamp(-cast_max, cast_max)
+        return {VALUES_KEY: tensor.to(self.cast_dtype)}
+
+    def _decode_tensors(
+        self, tensors: dict[str, torch.Tensor], dtype: torch.dtype
+    ) -> torch.Tensor:
+        """
+        args:
+            tensors: values: [B, ...], dtype: cast_dtype
+        returns:
+            tensor: [B, ...], dtype: dtype
+        """
+        return tensors[VALUES_KEY].to(dtype)
+
+    def get_config_parameters(self) -> dict[str, Any]:
+        return {"cast_dtype": str(self.cast_dtype).removeprefix("torch.")}