zea.models.taesd

Tiny Autoencoder (TAESD) model.

>>> from zea.models.taesd import TinyAutoencoder

>>> model = TinyAutoencoder.from_preset("taesdxl")

Important

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

See also

A tutorial notebook where this model is used: Autoencoder for ultrasound images.

Functions

convert_original_weights([model_name, ...])

Convert the original PyTorch TAESD weights to TensorFlow / Keras v3 models.

Classes

TinyAutoencoder(*args, **kwargs)	Tiny Autoencoder model.
TinyBase(*args, **kwargs)	Base class for TAESD encoder and decoder.
TinyDecoder(*args, **kwargs)	Decoder from TAESD model.
TinyEncoder(*args, **kwargs)	Encoder from TAESD model.

class zea.models.taesd.TinyAutoencoder(*args, **kwargs)

Bases: BaseModel

Tiny Autoencoder model.

Note

This model currently only supports TensorFlow and Jax backends.

Initializes the TAESD model with the given parameters.

Parameters:

**kwargs – Additional keyword arguments to pass to the superclass initializer.

call(inputs)

Applies the full autoencoder to the input.

custom_load_weights(preset, **kwargs)

Load the weights for the encoder and decoder.

decode(inputs)

Decode the encoded images.

Parameters:

inputs (tensor) – Input images of shape (batch_size, height, width, 4).

encode(inputs)

Encode the input images.

Parameters:

inputs (tensor) – Input images of shape (batch_size, height, width, channels).

class zea.models.taesd.TinyBase(*args, **kwargs)

Bases: BaseModel

Base class for TAESD encoder and decoder.

build(input_shape)

Builds the network.

call(inputs)

Applies the network to the input.

custom_load_weights(preset, **kwargs)

Load the weights for the encoder or decoder.

maybe_convert_to_jax(input_shape)

Converts the network to Jax if backend is Jax.

class zea.models.taesd.TinyDecoder(*args, **kwargs)

Bases: TinyBase

Decoder from TAESD model.

Initializes the TAESD decoder.

Parameters:

**kwargs – Additional keyword arguments passed to the superclass initializer.

class zea.models.taesd.TinyEncoder(*args, **kwargs)

Bases: TinyBase

Encoder from TAESD model.

Initializes the TAESD encoder.

Parameters:

**kwargs – Additional keyword arguments passed to the superclass initializer.

zea.models.taesd.convert_original_weights(model_name='madebyollin/taesdxl', output_dir=None, revision=None)

Convert the original PyTorch TAESD weights to TensorFlow / Keras v3 models.

This is how the taesdxl presets on the Hugging Face Hub were created; it is kept here for reproducibility and is not needed to use the model.

The conversion goes PyTorch -> ONNX -> TensorFlow and therefore needs a few extra packages that are not part of the zea dependencies:

pip install torch diffusers[torch] onnx==1.16.1 onnxruntime==1.18.1 onnx2tf \
    onnx-graphsurgeon onnxsim==0.4.33 sne4onnx sng4onnx tf-keras

Note

torch and TensorFlow have to coexist in one process here, which not every combination of wheels survives. Prefer the zeahub/all container, and import torch before this module if you hit a crash while building the torch model.

Parameters:

• model_name (str, optional) – Hugging Face model id of the original PyTorch autoencoder. Defaults to "madebyollin/taesdxl".

• revision (str, optional) – Git revision of model_name to convert. Defaults to None, i.e. whatever upstream currently has on its default branch; pass a commit or tag to reproduce one specific set of weights.

• output_dir (str | Path, optional) – Folder to write the converted models to. Defaults to a timestamped folder under ./temp/zea.

Returns:

Folder containing the converted encoder and decoder models.

Return type:

Path