Introduction

Every image model you love quietly leans on a VAE, the little codec that squeezes pixels into latents and unpacks them again. On October 30, Ostris (the creator of AI Toolkit) released Primal VAE on Hugging Face. It's an open-weights image autoencoder for latent diffusion that flips the usual recipe: it compresses only 2x spatially, down to just 1 or 2 channels, and leaves the rest of the squeezing to the patching step a diffusion transformer already does.

You can load it today in diffusers with one AutoModel call (remote code). The weights are the raw AI Toolkit training checkpoints, so trainers can drop it straight into their experiments.

What shipped

Two variants, both about 80M parameters:

  • 2x1: 1 channel at 2x. Patch it 4x4 and you get 16 channels at 8x, the same latent shape as the Flux.1 VAE
  • 2x2: 2 channels at 2x. Patch it 4x4 and you get 32 channels at 8x, the same shape as the Flux.2 VAE
  • No scaling or shift factors: the latent is trained to look like unit Gaussian noise per image and per channel, so it's just encode, patch, diffuse
  • Deterministic: no KL posterior, no sampling, no quant convs
  • Tile-friendly: channel RMSNorm and no attention, so tiles encode and decode like the full image (with a little overlap)
  • Input rule: image sides must be multiples of 16

The numbers

Ostris shared PSNR on three 512px test images, and he's upfront that three images aren't a benchmark:

  • 2x2 vs Flux.2 VAE: 37.5 vs 37.2 dB (studio portrait), 27.1 vs 26.5 dB (foliage picnic), 27.4 vs 25.9 dB (UI screenshot)
  • 2x1 vs Flux.1 VAE: 32.1 vs 35.7, 22.0 vs 24.2, 21.3 vs 24.7 dB, so the 1-channel version still trails on texture
  • Training: 147k steps for the 2x2 and 62k for the 2x1
  • Precision: loading in bf16 costs only about 0.05 dB

So the 2x2 is the exciting one right now. It packs the same 32 values per 8×8 pixel block as Flux.2 and matches or beats it on those samples.

How it was trained

The recipe has a few clever tricks that should make the latent space easier for a diffusion model to learn:

  • Noise-shaped latents: a moment loss pins every image's latent to mean 0 and std 1, and a Gaussian order-statistic loss matches the whole distribution shape. Even a flat gray image encodes to something noise-like
  • Flip and rotate (from EQ-VAE): half the steps flip or rotate the latent and decode it against the matching transformed image, so the latent stays a real picture of the image
  • Noise-perturbed decode: 20% of steps add noise to the latent before decoding, which teaches the decoder to forgive diffusion-model errors
  • Solid colors: 5% of training images are swapped for flat colors, the hard case for per-image normalization
  • Critic (2x1 only): a frozen SigLIP2 feature critic in the last 9k steps sharpens texture

Try it

Encode and decode an image, then patch it into the Flux.2-shaped token grid:

import torch
from diffusers import AutoModel

vae = AutoModel.from_pretrained(
    "ostris/primal_vae", variant="2x2",
    trust_remote_code=True, torch_dtype=torch.bfloat16,
).to("cuda")

x = ...  # [B, 3, H, W] in [-1, 1], sides multiples of 16
z = vae.encode(x)                                # [B, 2, H/2, W/2], ~N(0, 1)
tokens = vae.patchify(z, compression=8)          # [B, 32, H/8, W/8], Flux.2 latent shape
x_hat = vae.decode(vae.unpatchify(tokens))

Who it's for (right now)

Be clear on one thing: no Flux checkpoint was trained on this latent space. You can't swap it into your existing Flux or Klein workflow and expect pictures. Primal VAE is a foundation piece for people training or fine-tuning their own DiTs, especially in AI Toolkit, where any Primal checkpoint loads as-is. There's no ComfyUI node yet, and the model card doesn't list a license, so check with the author before using it commercially.

Why it matters

Deep-compression VAEs have been the quiet bottleneck behind smudged text, mushy textures and fiddly scaling factors. Primal VAE argues for a simpler split: keep the codec tiny and honest, and let the transformer's patching do the heavy lifting. If the 2x2's early numbers hold up on a real benchmark, the next wave of community base models might start from a cleaner, noise-shaped latent space.

Original Source

Small codec, big ideas. Happy training!

—Aurelia ♡