raiiaf.chunks.latent

Latent chunk utilities for RAIIAF.

This module provides functionality to validate, pack, compress, and parse latent arrays as RAIIAF LATN chunks. It supports compressed (zstd + zfpy) and uncompressed representations and exposes helpers for lazy loading.

Classes

raiiafLatent()

Operations for RAIIAF latent (LATN) chunks.
class raiiaf.chunks.latent.raiiafLatent[source]

Bases: object

Operations for RAIIAF latent (LATN) chunks.

dtype_from_flags(flags: bytes) → dtype[source]

Map 4-byte chunk flags to a NumPy dtype.

Handles null-padding and validates known types.

Parameters:

flags (bytes) – 4-byte flag field from the chunk header (e.g., b”F16").

Returns:

Corresponding dtype (float16 or float32).

Return type:

numpy.dtype

Raises:

ValueError – If the flag does not correspond to a known dtype.

iter_lazy_latents(filename: str, chunk_records: list)[source]

Yield callables for each latent chunk to enable lazy loading.

Parameters:

  • filename (str) – Path to the RAIIAF file.

  • chunk_records (list) – List of manifest chunk records.

Yields:

Callable[[], np.ndarray] – A function that loads and returns each latent array.

latent_compressor(chunk_type: bytes, chunk_flags: bytes, data: ndarray, tolerance: float = 0.0) → dict[source]

Compress a latent array into a LATN chunk.

The header is compressed with Zstd and payload with ZFPY (optionally lossy).

Parameters:

  • chunk_type (bytes) – 4-byte chunk identifier (b”LATN”).

  • chunk_flags (bytes) – 4-byte dtype flags (b”F16" or b”F32").

  • data (np.ndarray) – NumPy array to compress (float16/float32).

  • tolerance (float) – zfpy lossy tolerance; 0.0 for lossless. Defaults to 0.0.

Returns:

Dictionary with keys ‘chunk’, ‘len_header’, ‘len_data’.

Return type:

dict

latent_dtype_validator(latent_array: ndarray) → bool[source]

Validate latent data type and values.

Parameters:

latent_array (np.ndarray) – Latent array to validate.

Returns:

True if valid.

Return type:

bool

Raises:

raiiafLatentError – If dtype or values are invalid (NaN/Inf or unsupported dtype).

latent_packer(latent: Dict[str, ndarray], file_offset: int = 0, chunk_records=None, should_compress: bool = True, convert_float16: bool = True) → list[source]

Pack latent arrays into chunk bytes and append manifest records.

Parameters:

  • latent (Dict[str, np.ndarray]) – Mapping of latent key to array.

  • file_offset (int) – Current file offset where the first latent chunk will be placed.

  • chunk_records (Optional[list]) – Mutable list to append chunk records to.

  • should_compress (bool) – Whether to compress using zstd+zfpy. Defaults to True.

  • convert_float16 (bool) – Convert arrays to float16 before storage. Defaults to True.

Returns:

List of chunk byte strings (to be concatenated externally).

Return type:

list[bytes]

Raises:

raiiafLatentError – If inputs are invalid or unsupported dtypes are used.

latent_parser(chunk: bytes | dict, shape: tuple, compressed: bool = True)[source]

Parse a latent chunk and return the latent array.

Parameters:

  • chunk (Union[bytes, dict]) – Raw bytes for uncompressed chunks or dict with keys ‘chunk’, ‘len_header’, ‘len_data’ for compressed chunks.

  • shape (tuple) – Expected array shape.

  • compressed (bool) – Whether the chunk is compressed. Defaults to True.

Returns:

Parsed latent array with the expected dtype and shape.

Return type:

np.ndarray

Raises:

raiiafLatentError – If the chunk is malformed or decompression fails.

latent_shape_validator(latent_array: ndarray, expected_dims: int = 4, max_dimension_size: int = 8192) → bool[source]

Validate latent array shape and dimensions.

Parameters:

  • latent_array (np.ndarray) – Latent array to validate.

  • expected_dims (int) – Expected number of dimensions (default: 4 for NCHW).

  • max_dimension_size (int) – Maximum allowed size for any single dimension.

Returns:

True if valid.

Return type:

bool

Raises:

raiiafLatentError – If validation fails.

make_lazy_latent_loader(filename: str, chunk_record: dict)[source]

Create a callable that lazily loads the latent array on demand.

Parameters:

  • filename (str) – Path to the RAIIAF file.

  • chunk_record (dict) – Manifest record describing the latent chunk.

Returns:

A function that loads and returns the latent array.

Return type:

Callable[[], np.ndarray]

Raises:

raiiafLatentError – If the chunk cannot be loaded from the file.