Primal VAE

Primal VAE is a VAE that only does 2x (or 4x) spatial down to 1 channel (or 2). If you patch the latent up to 8x, you end up with the same latent shape as most modern VAEs, without the need for deep compression. The DiT does the patching, like it already does, and the VAE stays small.

The latent is trained to keep its norm in line with noise, so there is no scaling factor, no shift factor, and no latent normalization to carry around. Encode, patch, diffuse.

The 2x models are ~80M parameters, the 4x1 is ~122M. Image sides need to be multiples of 16 (2x) or 32 (4x1).

Usage

import torch
from diffusers import AutoModel

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

x = ...                      # [B, 3, H, W] in [-1, 1], H and W multiples of vae.image_multiple (16, or 32 for 4x1)
z = vae.encode(x)            # [B, 1, H/2, W/2], ~N(0, 1). No .latent_dist, no sampling.
x_hat = vae.decode(z)        # [B, 3, H, W]

Matching the Flux.1 latent (2x1)

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

tokens = vae.patchify(vae.encode(x), compression=8)   # [B, 16, H/8, W/8]: the Flux.1 VAE latent shape
# drop `tokens` in wherever the Flux.1 VAE latent went. No scaling_factor / shift_factor, it is already unit normal.
# e.g. FluxPipeline._pack_latents(tokens, B, 16, H // 8, W // 8) -> [B, (H/16) * (W/16), 64]

x_hat = vae.decode(vae.unpatchify(tokens))

Matching the Flux.2 latent (2x2)

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

tokens = vae.patchify(vae.encode(x), compression=8)   # [B, 32, H/8, W/8]: the Flux.2 VAE latent shape
# drop `tokens` in where the Flux.2 VAE latent went, before its own 2x2 patchify to 128-channel tokens.
# Skip the Flux.2 batch-norm latent normalization, it is already unit normal per image and channel.

x_hat = vae.decode(vae.unpatchify(tokens))

Matching the SD 1.5 / SDXL latent (4x1)

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

tokens = vae.patchify(vae.encode(x), compression=8)   # [B, 4, H/8, W/8]: the SD 1.5 / SDXL VAE latent shape
# drop `tokens` in wherever the SD VAE latent went. No scaling_factor (0.18215 / 0.13025), it is already unit normal.

x_hat = vae.decode(vae.unpatchify(tokens))

variant picks the weight file, primal_vae_<variant>.safetensors. The architecture is read from the tensor shapes, so weights_name="my_finetune.safetensors" loads any Primal VAE checkpoint out of ai-toolkit as-is. patchify / unpatchify take the total compression you want and are just pixel_unshuffle / pixel_shuffle by compression / <model compression>, so compression=8 is a 4x4 patch on the 2x models and a 2x2 patch on the 4x1. vae.token_channels(8) and vae.image_multiple give you the numbers above.

To backprop through decode (a pixel-space loss on a diffusion model's prediction), vae.enable_gradient_checkpointing() checkpoints every residual block of both U-Nets; inference under torch.no_grad() is unaffected.

Architecture

The 2x1 and 2x2 are the same network, they only differ in output channels. The 4x1 is the same design with a 4x input patch and a wider, deeper U.

• Input is pixel_unshuffle(compression): RGB becomes 12 channels at 1/2 resolution (2x) or 48 channels at 1/4 (4x1). Everything runs at latent resolution or below, so the latent is a spatially honest map of the image.
• Encoder and decoder are mirrored conv U-Nets with 4 levels, each level halving again from the latent (2x: 1/2 to 1/16 of the image, 4x1: 1/4 to 1/32). Downs and ups are pixel_unshuffle / pixel_shuffle plus a 3x3 conv. Residual blocks cycle dilations 1, 2, 4.
  • 2x: widths 96 / 192 / 384 / 384, residual block depths 1 / 1 / 2 / 4 down and 1 / 1 / 2 up.
  • 4x1: widths 192 / 384 / 384 / 384, depths 2 / 2 / 2 / 4 down and 2 / 2 / 2 up.
• Normalization is RMSNorm over channels at each position. No spatial statistics, no attention, so tiles encode and decode the same as the full image given some overlap.
• The encoder ends in a plain 3x3 conv to 1 or 2 channels. No posterior, no KL, no quant convs.

Keeping the latent in line with noise

A diffusion model is trained against Gaussian noise, so the latent should look like Gaussian noise. Instead of a KL term and a scaling factor someone measures after the fact, I put the statistics directly in the loss on the encoder output.

• Moment loss: mean^2 + (std - 1)^2, weight 0.1, taken per image and per channel, so every single image comes out mean 0 / std 1 in each channel, not just the dataset on average. The 2x1 spent its first 53k steps with this pooled over the batch and switched to per image for the last 9k. The 2x2 and 4x1 had it per image throughout.
• Gaussian order-statistic loss: sort the latent values of each image and channel and MSE them against the expected order statistics of a standard normal of the same size (Blom's approximation). This pins the whole distribution shape, not just the first two moments, so a flat gray image still encodes to something that looks like noise instead of a constant.

Examples

Each row is input, reconstruction, latent, |error| × 4. Under each Primal VAE row are the reference VAEs with the same latent layout, on the same image: Flux.1 for the 2x1, Flux.2 for the 2x2, SD 1.5 and SDXL for the 4x1. Latents are shown at ±3σ. The 1-channel latent is gray; the 2-channel one has channel 0 on red and channel 1 on blue. The reference latents show their first 3 channels after the scaling the DiT actually sees (Flux.1: shift and scale factor; Flux.2: batch-norm statistics; SD: the scaling factor). Mean and std are over the full latent tensor.

PSNR on these three 512 px images (dB):

The 2x2 matches or beats the Flux.2 VAE here at the same information budget (both are 32 values per 8×8 pixel block). Three images and PSNR alone are not a benchmark, and a KL VAE pays a reconstruction tax for its posterior noise that a deterministic one does not. The 2x1 trails the Flux.1 VAE on textured content. One channel at 2x is a 12:1 reduction with no slack, and it has well under half the training steps of the 2x2 (62k vs 147k). Its per-image latent std still drifts (1.07 to 1.21 here) because it only got the per-image normalization for its last 9k steps. The 2x2, trained with it throughout, holds std within a few percent of 1.

The 4x1 is the same 4 values per 8×8 block as the SD VAEs, a 48:1 reduction, and lands within a fraction of a dB of the SD 1.5 VAE on the textured images and 2 dB behind on the clean studio shot. SDXL's VAE, a retrain of the same architecture, stays 1 to 2.5 dB ahead. Its latent holds per-image std within about 5% of 1 on these (1.03 to 1.11).

Training

Trained in ai-toolkit with a custom extension on aspect-bucketed images. The weight files here are the training checkpoints, unchanged.

Reconstruction is MAE (0.5) and LPIPS (1.0), plus an L1 on the log-magnitude 2D spectrum (0.1) that pushes high-frequency energy without demanding exact phase.

The tricks that should make the latent space easier to diffuse:

• Latent flips and rotations (from EQ-VAE). On 50% of steps I flip or rotate the latent by 90 degrees and decode it against the same flipped or rotated image. The latent has to be a picture of the image, not an arbitrary code. I tried downscale equivariance too and dropped it, it slowed down color learning.
• Noise-perturbed decode. On 20% of steps the decoder gets z + sigma * noise, sigma drawn per image from U(0, 0.5), and still has to reconstruct the clean image. The decoder learns to tolerate diffusion model error and the latent gets smooth: nearby codes decode to nearby images.
• Solid colors: 5% of images are swapped for a random flat color. Flat regions are the hard case for a per-image normalized latent.
• Critic (2x1 for its last 9k steps, 4x1 from step 0 for all 155k): a frozen SigLIP2 NaFlex ViT gives patch features from layers 4, 8 and 12 of the real and reconstructed image at native resolution. Small spectral-normalized conv heads on those features train with a hinge loss plus R1, and the decoder gets the generator term at weight 0.1, ramped in over the first 1000 steps. It sharpens the texture that MAE and LPIPS leave soft. The 2x2 was trained without a critic, it seemed to converge fine without one.

Every step is one encode and one decode. On a flip/rotate step the transformed pair replaces the clean pair instead of adding a second decoder pass.

Files

primal_vae_2x1.safetensors   1 ch @ 2x (fp32)
primal_vae_2x2.safetensors   2 ch @ 2x (fp32)
primal_vae_4x1.safetensors   1 ch @ 4x (fp32)
modeling_primal_vae.py       the diffusers model (remote code)
config.json                  points AutoModel at it

Load with torch_dtype=torch.bfloat16 for inference, it costs ~0.05 dB.

License

Apache 2.0, Copyright 2026 Ostris, LLC. See LICENSE.