vaezip 0.1.0__tar.gz → 0.2.0__tar.gz

{vaezip-0.1.0/src/vaezip.egg-info → vaezip-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vaezip
-Version: 0.1.0
+Version: 0.2.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
@@ -62,12 +62,30 @@ 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")
+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 metadata, so `load()` needs no codec.
-Pass `save_metadata=False` to drop it, then `load()` requires an explicit `codec=`.
+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`:
+
+```python
+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
 
@@ -86,13 +104,12 @@ restored = codec.decode(encoded)
 ```
 
 **EncodedTensor** holds named tensors plus string metadata with the codec config.
-It serializes to `.safetensors` (default) or `.pth`:
+`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.
 
-```python
-data = encoded.serialize()  # or vaezip.serialize(encoded)
-encoded = vaezip.deserialize(data)
-data = encoded.serialize(save_metadata=False, format="pth")
-```
+**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:

{vaezip-0.1.0 → vaezip-0.2.0}/README.md

@@ -33,12 +33,30 @@ 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")
+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 metadata, so `load()` needs no codec.
-Pass `save_metadata=False` to drop it, then `load()` requires an explicit `codec=`.
+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`:
+
+```python
+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
 
@@ -57,13 +75,12 @@ restored = codec.decode(encoded)
 ```
 
 **EncodedTensor** holds named tensors plus string metadata with the codec config.
-It serializes to `.safetensors` (default) or `.pth`:
+`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.
 
-```python
-data = encoded.serialize()  # or vaezip.serialize(encoded)
-encoded = vaezip.deserialize(data)
-data = encoded.serialize(save_metadata=False, format="pth")
-```
+**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:

{vaezip-0.1.0 → vaezip-0.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vaezip"
-version = "0.1.0"
+version = "0.2.0"
 description = "Compression of VAE latents: low-bit quantization codecs on PyTorch and safetensors"
 readme = "README.md"
 license = "Unlicense"

{vaezip-0.1.0 → vaezip-0.2.0}/src/vaezip/__init__.py

@@ -2,15 +2,17 @@
 
 from vaezip import codecs
 from vaezip._version import __version__
-from vaezip.api import load, save
-from vaezip.encoded_tensor import EncodedTensor, deserialize, serialize
+from vaezip.api import decode, encode, load, load_tensor, save, save_tensor
+from vaezip.encoded_tensor import EncodedTensor
 
 __all__ = [
     "EncodedTensor",
     "__version__",
     "codecs",
-    "deserialize",
+    "decode",
+    "encode",
     "load",
+    "load_tensor",
     "save",
-    "serialize",
+    "save_tensor",
 ]

vaezip-0.2.0/src/vaezip/api.py

@@ -0,0 +1,188 @@
+"""High level encoding, decoding, saving and loading 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.shortcuts import DEFAULT_SHORTCUT
+from vaezip.encoded_tensor import EncodedTensor
+from vaezip.serialization import (
+    SAFETENSORS_FORMAT,
+    SerializationFormat,
+    deserialize,
+    serialize,
+)
+
+CodecSpecification = Codec | str | os.PathLike[str]
+SINGLE_TENSOR_NAME = "tensor"
+
+
+def encode(
+    tensor: torch.Tensor, codec: CodecSpecification | None = None
+) -> EncodedTensor:
+    """
+    args:
+        tensor: [B, ...], floating point
+        codec: codec object, shortcut name or config path, default is "cast_fp32"
+    returns:
+        encoded: tensors on the device of the input, metadata with the codec config
+    """
+    if codec is None:
+        codec = DEFAULT_SHORTCUT
+    resolved_codec = codec if isinstance(codec, Codec) else AutoCodec.from_config(codec)
+    return resolved_codec.encode(tensor)
+
+
+def decode(
+    encoded: EncodedTensor,
+    codec: CodecSpecification | None = None,
+    dtype: torch.dtype = torch.float32,
+    device: torch.device | str | None = None,
+) -> torch.Tensor:
+    """
+    Same as `EncodedTensor.decode()`.
+
+    returns:
+        tensor: [B, ...], original shape, dtype: dtype
+    """
+    return encoded.decode(codec=codec, dtype=dtype, device=device)
+
+
+@overload
+def save(
+    encoded_tensors: dict[str, EncodedTensor],
+    path: None = None,
+    format: SerializationFormat = SAFETENSORS_FORMAT,
+) -> bytes: ...
+
+
+@overload
+def save(
+    encoded_tensors: dict[str, EncodedTensor],
+    path: str | os.PathLike[str],
+    format: SerializationFormat = SAFETENSORS_FORMAT,
+) -> None: ...
+
+
+def save(
+    encoded_tensors: dict[str, EncodedTensor],
+    path: str | os.PathLike[str] | None = None,
+    format: SerializationFormat = SAFETENSORS_FORMAT,
+) -> bytes | None:
+    """
+    Writes named encoded tensors to a file, or returns bytes if there is no path.
+
+    args:
+        encoded_tensors: name -> encoded tensor, tensors may use different codecs
+        path: output file, `None` to return bytes
+        format: "safetensors" or "pth", never inferred from the path
+    """
+    return _write_or_return(serialize(encoded_tensors, format=format), path)
+
+
+def load(
+    source: str | os.PathLike[str] | bytes | bytearray,
+    device: torch.device | str | None = None,
+    format: SerializationFormat = SAFETENSORS_FORMAT,
+) -> dict[str, EncodedTensor]:
+    """
+    Reads a file or bytes written by `save()`.
+
+    args:
+        source: file path or bytes
+        device: device of the loaded tensors, default is CPU
+        format: "safetensors" or "pth", never inferred from the path
+    returns:
+        encoded_tensors: name -> encoded tensor, sorted by name
+    """
+    if isinstance(source, (bytes, bytearray)):
+        data = bytes(source)
+    else:
+        data = Path(source).read_bytes()
+    encoded_tensors = deserialize(data, format=format)
+    if device is not None:
+        encoded_tensors = {
+            name: encoded.to(device) for name, encoded in encoded_tensors.items()
+        }
+    return encoded_tensors
+
+
+@overload
+def save_tensor(
+    tensor: torch.Tensor,
+    path: None = None,
+    codec: CodecSpecification | None = None,
+    format: SerializationFormat = SAFETENSORS_FORMAT,
+) -> bytes: ...
+
+
+@overload
+def save_tensor(
+    tensor: torch.Tensor,
+    path: str | os.PathLike[str],
+    codec: CodecSpecification | None = None,
+    format: SerializationFormat = SAFETENSORS_FORMAT,
+) -> None: ...
+
+
+def save_tensor(
+    tensor: torch.Tensor,
+    path: str | os.PathLike[str] | None = None,
+    codec: CodecSpecification | None = None,
+    format: SerializationFormat = SAFETENSORS_FORMAT,
+) -> bytes | None:
+    """
+    Encodes a single tensor and saves it as `{"tensor": encoded}`.
+
+    args:
+        tensor: [B, ...], floating point
+        path: output file, `None` to return bytes
+        codec: codec object, shortcut name or config path, default is "cast_fp32"
+        format: "safetensors" or "pth", never inferred from the path
+    """
+    encoded_tensors = {SINGLE_TENSOR_NAME: encode(tensor, codec=codec)}
+    return _write_or_return(serialize(encoded_tensors, format=format), path)
+
+
+def load_tensor(
+    source: str | os.PathLike[str] | bytes | bytearray,
+    codec: CodecSpecification | None = None,
+    dtype: torch.dtype = torch.float32,
+    device: torch.device | str | None = None,
+    format: SerializationFormat = SAFETENSORS_FORMAT,
+) -> torch.Tensor:
+    """
+    Loads and decodes a file that holds exactly one entry named "tensor".
+
+    args:
+        source: file path or bytes
+        codec: codec object, shortcut name or config path, by default the codec
+        from the metadata
+        dtype: floating point dtype of the result
+        device: device to decode on, default is CPU
+        format: "safetensors" or "pth", never inferred from the path
+    returns:
+        tensor: [B, ...], original shape, dtype: dtype
+    """
+    encoded_tensors = load(source, format=format)
+    if list(encoded_tensors) != [SINGLE_TENSOR_NAME]:
+        raise ValueError(
+            f"Expected a single entry named {SINGLE_TENSOR_NAME!r}, "
+            f"found: {list(encoded_tensors)}"
+        )
+    return encoded_tensors[SINGLE_TENSOR_NAME].decode(
+        codec=codec, dtype=dtype, device=device
+    )
+
+
+def _write_or_return(data: bytes, path: str | os.PathLike[str] | None) -> bytes | None:
+    if path is None:
+        return data
+    Path(path).write_bytes(data)
+    return None

vaezip-0.2.0/src/vaezip/encoded_tensor.py

@@ -0,0 +1,93 @@
+"""Container for compressed latents."""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+if TYPE_CHECKING:
+    from vaezip.codecs.base import Codec
+
+CODEC_CONFIG_KEY = "codec_config"
+
+
+@dataclass
+class EncodedTensor:
+    """
+    Output of `Codec.encode()` and input of `Codec.decode()`.
+
+    tensors: named tensors with the compressed representation.
+    metadata: optional string metadata, normally holds the JSON config of the codec.
+    """
+
+    tensors: dict[str, torch.Tensor]
+    metadata: dict[str, str] = field(default_factory=dict)
+
+    @property
+    def codec_config(self) -> dict[str, Any] | None:
+        """Parsed codec config from the metadata, `None` if there is none."""
+        raw_config = self.metadata.get(CODEC_CONFIG_KEY)
+        if raw_config is None:
+            return None
+        config = json.loads(raw_config)
+        if not isinstance(config, dict):
+            raise ValueError("Codec config in the metadata must be a JSON object")
+        return config
+
+    @property
+    def nbytes(self) -> int:
+        """Total size of the stored tensors in bytes."""
+        return sum(
+            tensor.numel() * tensor.element_size() for tensor in self.tensors.values()
+        )
+
+    def to(self, device: torch.device | str) -> EncodedTensor:
+        return EncodedTensor(
+            tensors={name: tensor.to(device) for name, tensor in self.tensors.items()},
+            metadata=dict(self.metadata),
+        )
+
+    def decode(
+        self,
+        codec: Codec | str | os.PathLike[str] | None = None,
+        dtype: torch.dtype = torch.float32,
+        device: torch.device | str | None = None,
+    ) -> torch.Tensor:
+        """
+        args:
+            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, by default the device of the stored tensors
+        returns:
+            tensor: [B, ...], original shape, dtype: dtype
+        """
+        from vaezip.codecs.auto import AutoCodec
+        from vaezip.codecs.base import Codec
+        from vaezip.codecs.config import without_version
+
+        stored_config = self.codec_config
+        if codec is None:
+            if stored_config is None:
+                raise ValueError(
+                    "The encoded tensor has no codec config in its metadata, "
+                    "pass the codec explicitly"
+                )
+            resolved_codec = AutoCodec.from_config(stored_config)
+        else:
+            resolved_codec = (
+                codec if isinstance(codec, Codec) else AutoCodec.from_config(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 "
+                    f"in the encoded tensor: {stored_config}"
+                )
+        encoded = self if device is None else self.to(device)
+        return resolved_codec.decode(encoded, dtype=dtype)

vaezip-0.2.0/src/vaezip/serialization.py

@@ -0,0 +1,189 @@
+"""Storage of named encoded tensors in safetensors and pth containers."""
+
+from __future__ import annotations
+
+import io
+import json
+import pickle
+import struct
+from typing import Any, Literal, get_args
+
+import torch
+from safetensors import SafetensorError
+from safetensors.torch import load as load_safetensors
+from safetensors.torch import save as save_safetensors
+
+from vaezip._version import __version__
+from vaezip.encoded_tensor import EncodedTensor
+
+SerializationFormat = Literal["safetensors", "pth"]
+SAFETENSORS_FORMAT: SerializationFormat = "safetensors"
+PTH_FORMAT: SerializationFormat = "pth"
+SUPPORTED_FORMATS: tuple[str, ...] = get_args(SerializationFormat)
+
+FORMAT_VERSION = "1"
+FORMAT_VERSION_KEY = "vaezip_format_version"
+PACKAGE_VERSION_KEY = "vaezip_version"
+KEY_DELIMITER = "."
+
+_PTH_TENSORS_KEY = "tensors"
+_PTH_METADATA_KEY = "metadata"
+_SAFETENSORS_HEADER_SIZE_BYTES = 8
+_SAFETENSORS_METADATA_KEY = "__metadata__"
+
+
+def check_format(format: str) -> None:
+    if format not in SUPPORTED_FORMATS:
+        raise ValueError(
+            f"Unsupported format {format!r}, expected one of {SUPPORTED_FORMATS}"
+        )
+
+
+def flatten(
+    encoded_tensors: dict[str, EncodedTensor],
+) -> tuple[dict[str, torch.Tensor], dict[str, str]]:
+    """
+    Joins entry names with inner names: tensor "codes" of entry "latents" becomes
+    "latents.codes", metadata is flattened the same way.
+
+    args:
+        encoded_tensors: entry name -> encoded tensor, names are non-empty
+    returns:
+        tensors: flat name -> tensor, on CPU, contiguous
+        metadata: flat key -> string, includes the format and package versions
+    """
+    if not isinstance(encoded_tensors, dict):
+        raise TypeError(
+            f"Expected a dict of EncodedTensor, got {type(encoded_tensors).__name__}"
+        )
+    if not encoded_tensors:
+        raise ValueError("Can not save an empty dict")
+    tensors: dict[str, torch.Tensor] = {}
+    metadata: dict[str, str] = {
+        FORMAT_VERSION_KEY: FORMAT_VERSION,
+        PACKAGE_VERSION_KEY: __version__,
+    }
+    for name, encoded in encoded_tensors.items():
+        if not isinstance(name, str) or not name:
+            raise ValueError(f"Names must be non-empty strings, got {name!r}")
+        if not isinstance(encoded, EncodedTensor):
+            raise TypeError(
+                f"Expected an EncodedTensor for {name!r}, got {type(encoded).__name__}"
+            )
+        if not encoded.tensors:
+            raise ValueError(f"EncodedTensor {name!r} has no tensors")
+        for inner_name, tensor in encoded.tensors.items():
+            _check_inner_key(name, inner_name)
+            tensors[f"{name}{KEY_DELIMITER}{inner_name}"] = (
+                tensor.detach().to("cpu", copy=True).contiguous()
+            )
+        for key, value in encoded.metadata.items():
+            _check_inner_key(name, key)
+            if not isinstance(value, str):
+                raise ValueError(
+                    f"Metadata {key!r} of {name!r} must be a string, "
+                    f"got {type(value).__name__}"
+                )
+            metadata[f"{name}{KEY_DELIMITER}{key}"] = value
+    return tensors, metadata
+
+
+def unflatten(
+    tensors: dict[str, torch.Tensor], metadata: dict[str, str]
+) -> dict[str, EncodedTensor]:
+    """
+    Inverse of `flatten()`, splits keys on the last delimiter.
+
+    returns:
+        entry name -> encoded tensor, sorted by name
+    """
+    format_version = metadata.get(FORMAT_VERSION_KEY)
+    if format_version != FORMAT_VERSION:
+        raise ValueError(
+            f"Not a valid VaeZip file, expected {FORMAT_VERSION_KEY}="
+            f"{FORMAT_VERSION!r} in the metadata, got {format_version!r}"
+        )
+    grouped_tensors: dict[str, dict[str, torch.Tensor]] = {}
+    for flat_key, tensor in tensors.items():
+        name, inner_name = _split_key(flat_key)
+        grouped_tensors.setdefault(name, {})[inner_name] = tensor
+    grouped_metadata: dict[str, dict[str, str]] = {}
+    for flat_key, value in metadata.items():
+        if flat_key in (FORMAT_VERSION_KEY, PACKAGE_VERSION_KEY):
+            continue
+        name, key = _split_key(flat_key)
+        if name not in grouped_tensors:
+            raise ValueError(
+                f"Not a valid VaeZip file, metadata {flat_key!r} has no tensors"
+            )
+        grouped_metadata.setdefault(name, {})[key] = value
+    return {
+        name: EncodedTensor(
+            tensors=grouped_tensors[name], metadata=grouped_metadata.get(name, {})
+        )
+        for name in sorted(grouped_tensors)
+    }
+
+
+def serialize(
+    encoded_tensors: dict[str, EncodedTensor], format: str = SAFETENSORS_FORMAT
+) -> bytes:
+    check_format(format)
+    tensors, metadata = flatten(encoded_tensors)
+    if format == SAFETENSORS_FORMAT:
+        return save_safetensors(tensors, metadata=metadata)
+    buffer = io.BytesIO()
+    torch.save({_PTH_TENSORS_KEY: tensors, _PTH_METADATA_KEY: metadata}, buffer)
+    return buffer.getvalue()
+
+
+def deserialize(
+    data: bytes, format: str = SAFETENSORS_FORMAT
+) -> dict[str, EncodedTensor]:
+    """returns: entry name -> encoded tensor with tensors on CPU, sorted by name"""
+    check_format(format)
+    try:
+        if format == SAFETENSORS_FORMAT:
+            tensors = load_safetensors(data)
+            metadata = _read_safetensors_metadata(data)
+        else:
+            content = torch.load(
+                io.BytesIO(data), map_location="cpu", weights_only=True
+            )
+            tensors = dict(content[_PTH_TENSORS_KEY])
+            metadata = dict(content[_PTH_METADATA_KEY])
+    except (
+        SafetensorError,
+        struct.error,
+        ValueError,
+        KeyError,
+        TypeError,
+        EOFError,
+        RuntimeError,
+        pickle.UnpicklingError,
+    ) as error:
+        raise ValueError(f"Not a valid VaeZip {format} file: {error}") from error
+    return unflatten(tensors, metadata)
+
+
+def _check_inner_key(name: str, key: Any) -> None:
+    if not isinstance(key, str) or not key or KEY_DELIMITER in key:
+        raise ValueError(
+            f"Tensor and metadata names inside {name!r} must be non-empty strings "
+            f"without {KEY_DELIMITER!r}, got {key!r}"
+        )
+
+
+def _split_key(flat_key: str) -> tuple[str, str]:
+    name, delimiter, inner_name = flat_key.rpartition(KEY_DELIMITER)
+    if not delimiter or not name or not inner_name:
+        raise ValueError(f"Not a valid VaeZip file, unexpected key {flat_key!r}")
+    return name, inner_name
+
+
+def _read_safetensors_metadata(data: bytes) -> dict[str, str]:
+    (header_size,) = struct.unpack_from("<Q", data, 0)
+    header_end = _SAFETENSORS_HEADER_SIZE_BYTES + header_size
+    header = json.loads(bytes(data[_SAFETENSORS_HEADER_SIZE_BYTES:header_end]))
+    metadata: dict[str, str] = header.get(_SAFETENSORS_METADATA_KEY, {})
+    return metadata

{vaezip-0.1.0 → vaezip-0.2.0/src/vaezip.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vaezip
-Version: 0.1.0
+Version: 0.2.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
@@ -62,12 +62,30 @@ 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")
+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 metadata, so `load()` needs no codec.
-Pass `save_metadata=False` to drop it, then `load()` requires an explicit `codec=`.
+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`:
+
+```python
+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
 
@@ -86,13 +104,12 @@ restored = codec.decode(encoded)
 ```
 
 **EncodedTensor** holds named tensors plus string metadata with the codec config.
-It serializes to `.safetensors` (default) or `.pth`:
+`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.
 
-```python
-data = encoded.serialize()  # or vaezip.serialize(encoded)
-encoded = vaezip.deserialize(data)
-data = encoded.serialize(save_metadata=False, format="pth")
-```
+**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: