zea.models.hvae

Hierarchical Variational Auto-Encoder for image generation, posterior sampling and inference tasks. To try this model, simply load one of the available presets:

>>> from zea.models.hvae import HierarchicalVAE

>>> model = HierarchicalVAE.from_preset("hvae")

Important

This is a zea implementation of the model. For the original code, see here.

Reference

S. W. Penninga, H. van Gorp, and R. J. G. van Sloun, “Deep Sylvester Posterior Inference for Adaptive Compressed Sensing in Ultrasound Imaging,” in 2025 IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP 2025, Hyderabad, India, April 6-11, 2025, pp. 1–5, 2025. doi.org/10.1109/ICASSP49660.2025.10888253

See also

A tutorial notebook where this model is used: Hierarchical VAEs for ultrasound image generation and inpainting.

Classes

HierarchicalVAE(*args, **kwargs)

Hierarchical Variational Autoencoder (HVAE) model.
class zea.models.hvae.HierarchicalVAE(*args, **kwargs)[source]

Bases: DeepGenerativeModel

Hierarchical Variational Autoencoder (HVAE) model. The network as defined here is a snippet of the complete model at: https://github.com/swpenninga/hvae

The lvh versions are trained on EchoNetLVH at 256x256 resolution with 3 channels. (video-frames as channel dimension) The ur(.) versions denote retraining with a UniformRandom agent with (.)/256 lines.

Unlike the other models, this network is built when the weights are loaded.

Parameters:

  • name (str) – Name of the model.

  • version (str) – Version of the HVAE model to use. Supported versions are: “lvh”, “lvh_ur24”, “lvh_ur16”, “lvh_ur8”, “lvh_ur4”, “lvh_ge24”, “lvh_ge16”, “lvh_ge8”, “lvh_ge4”.

call(measurements, seed=None)[source]

Returns a reconstruction of the input, together with the latent samples and KL divergences.

Parameters:

  • measurements (tensor) – Input measurements of shape [B, 256, 256, 3].

  • seed – Random seed generator. When None, the model falls back to the internal seed generators of its blocks.

Returns:

Reconstructed output of shape [B, 256, 256, 3], List of latent samples from the decoder, and list of KL divergences at each latent layer.

Return type:

recon (tensor)

custom_load_weights(preset, load_weights=True, **kwargs)[source]

Load the pretrained weights of the HVAE model from a preset. First builds the model architecture from args.pkl, then loads the weights into the model.

Parameters:

  • preset (str) – Preset identifier or path.

  • load_weights (bool) – If False, only the model architecture is built from args.pkl without downloading or loading the (large) weights file. Useful for testing.

log_density(measurements, **kwargs)[source]

Calculates the log density (ELBO) of the data under the model.

Parameters:

measurements (tensor) – Input measurements of shape [B, 256, 256, 3].

Returns:

negative ELBO of the input measurements, averaged over the batch.

Return type:

-elbo (tensor)

partial_inference(measurements, num_layers=0.5, n_samples=1, seed=None, **kwargs)[source]

Performs TopDown inference with the HVAE up until a certain layer, after which it continues in the decoder with multiple prior streams.

Parameters:

  • measurements (tensor) – Input measurements of shape [256, 256, 3], or [n_samples, 256, 256, 3] to condition each sample on a different measurement.

  • num_layers (float or int) – If float, fraction of total layers to use from the top. If int, number of layers to use from the top.

  • n_samples (int) – Number of posterior samples to generate.

  • seed – Random seed generator. When None, the model falls back to the internal seed generators of its blocks.

Returns:

Posterior samples of shape [n_samples, 256, 256, 3].

Return type:

output (tensor)

posterior_sample(measurements, n_samples=1, seed=None)[source]

Performs posterior sampling for a single measurement.

The encoder is deterministic, so one measurement needs only a single encoder pass; the n_samples decoder passes are what produce distinct posterior samples.

Parameters:

  • measurements (tensor) – Input measurements of shape [256, 256, 3] in [-1, 1], or [n_samples, 256, 256, 3] to condition each sample on a different measurement (which costs one encoder pass each).

  • n_samples (int, optional) – Number of posterior samples to generate. Defaults to 1.

  • seed – Random seed generator. When None, the model falls back to the internal seed generators of its blocks, which are shared across a zea.func.vmap() batch; pass one seed per mapped measurement to keep their samples independent.

Returns:

Posterior samples of shape [n_samples, 256, 256, 3].

Return type:

output (tensor)

sample(n_samples=1, seed=None, **kwargs)[source]

Samples from the prior distribution.

Parameters:

  • n_samples (int) – Number of samples to generate.

  • seed – Random seed generator. When None, the model falls back to the internal seed generators of its blocks.

Returns:

Generated samples of shape (n_samples, 256, 256, 3) in [-1, 1].

Return type:

tensor