dew.nn.autoencoders

Name	Summary
AutoEncoder	An encoder and decoder pair a latent diffusion model trains behind.
AutoencoderKL	NHWC image/latent arrays; scaling and shifts belong to AutoEncoder.
SimpleAutoEncoder	Tutorial-grade convolutional autoencoder, no pretrained weights.
StableDiffusionVAE	Frozen native AutoencoderKL variables and their latent normalization.

AutoEncoder

class source

class AutoEncoder(ABC)

An encoder and decoder pair a latent diffusion model trains behind.

A subclass encodes and decodes one batch of frames, [B, H, W, C] to [B, h, w, c] and back; encode and decode here flatten video [B, T, H, W, C] to frames around that and apply the latent normalization. Latents are normalized as (z - latent_shift) * latent_scale on the way out and inverted on the way in, the SD3 convention. The defaults are the identity; set them to the dataset’s own latent mean and 1/std so the diffusion model sees roughly unit-variance, zero-mean inputs.

The weights are an argument, as a ConditionEncoder’s are: params holds what a run loaded, and every call takes the tree to use, so the layout places the weights and the checkpoint carries them.

downscale_factor: int

H / h, the spatial factor between a frame and its latent.

latent_channels: int

c, the channels of a latent.

AutoEncoder.encode_batch

def encode_batch(params, x: jnp.ndarray, key: jax.Array | None = None) -> jnp.ndarray

Frames [B, H, W, C] to raw latents [B, h, w, c]; key draws a stochastic encoder’s sample, and None takes its mean.

AutoEncoder.decode_batch

def decode_batch(params, z: jnp.ndarray) -> jnp.ndarray

Raw latents [B, h, w, c] to frames [B, H, W, C].

AutoEncoder.encode

def encode(params, x: jnp.ndarray, key: jax.Array | None = None) -> jnp.ndarray

Images [B, H, W, C] or video [B, T, H, W, C] to normalized latents with the same leading axes.

AutoEncoder.decode

def decode(params, z: jnp.ndarray) -> jnp.ndarray

Normalized latents [B, h, w, c] or [B, T, h, w, c] back to images or video.

AutoencoderKL

Flax module source

class AutoencoderKL(
    channels: tuple[int, ...] = (128, 256, 512, 512),
    latent_channels: int = 4,
    image_channels: int = 3,
    blocks_per_level: int = 2,
    norm_groups: int = 32,
    quantize: bool = True,
    post_quantize: bool = True,
    dtype: Dtype = jnp.float32,
)

NHWC image/latent arrays; scaling and shifts belong to AutoEncoder.

encode(image, key=None) returns the posterior mean. Passing a key samples its diagonal Gaussian; decode(latents) returns unnormalized image pixels.

downscale_factor: int

Every encoder level except the last halves each spatial axis.

AutoencoderKL.setup

def setup()

AutoencoderKL.encode

def encode(image, key=None)

AutoencoderKL.decode

def decode(latents)

SimpleAutoEncoder

class source

class SimpleAutoEncoder(
    latent_channels: int = 4,
    feature_depths: Sequence[int] = (32, 64, 128),
    out_channels: int = 3,
    activation: Callable = jax.nn.silu,
    norm_groups: int = 8,
    dtype: Dtype | None = jnp.float32,
    precision: PrecisionLike = None,
    latent_shift: float = 0.0,
    latent_scale: float = 1.0,
    params=None,
    key: jax.Array | None = None,
)

Tutorial-grade convolutional autoencoder, no pretrained weights.

The encoder halves the resolution once per entry of feature_depths (stride-2 3x3 conv + GroupNorm + SiLU) and projects to latent_channels; the decoder mirrors it with nearest-neighbour upsampling. So downscale_factor == 2 ** len(feature_depths) and latent_channels is the bottleneck width, the two properties the samplers and input config read off an autoencoder (same contract as StableDiffusionVAE).

Like StableDiffusionVAE it loads a tree into params and takes the tree to use on every call; unlike it, the weights start random. Train them (or pass a trained tree as params) before the reconstructions mean anything. The latent is deterministic: there is no KL bottleneck, so the encode key is accepted and ignored. Video comes free from the AutoEncoder base class, which flattens [B, T, H, W, C] to frames.

SimpleAutoEncoder.init_params

def init_params(key: jax.Array) -> dict

Freshly initialize encoder and decoder parameters.

Convolutional, so the init resolution is irrelevant as long as it survives every downscale stage; the smallest such image is used.

SimpleAutoEncoder.encode_batch

def encode_batch(params, x: jnp.ndarray, key=None) -> jnp.ndarray

key is part of the AutoEncoder contract but unused: this encoder is deterministic.

SimpleAutoEncoder.decode_batch

def decode_batch(params, z: jnp.ndarray) -> jnp.ndarray

StableDiffusionVAE

class source

class StableDiffusionVAE(
    modelname='CompVis/stable-diffusion-v1-4',
    revision='bf16',
    dtype=jnp.bfloat16,
    latent_shift=None,
    latent_scale=None,
    params=None,
    model: AutoencoderKL | None = None,
)

Frozen native AutoencoderKL variables and their latent normalization.