ostris/primal_vae

Primal VAE

Primal VAE is a VAE that only does 2x spatial down to 1 channel (or 2). If you patch 4 the latent, you end up with 8x at 16 channels, the same 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 4, diffuse.

model	latent	patch 4	same shape as
2x1	1 ch @ 2x	16 ch @ 8x	Flux.1 VAE (16 ch @ 8x)
2x2	2 ch @ 2x	32 ch @ 8x	Flux.2 VAE (32 ch @ 8x)

Both are ~80M parameters. Image sides need to be multiples of 16.

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 16
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))

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 / 2, so compression=8 is a 4x4 patch. 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

Both models are the same network, they only differ in output channels.

• Input is pixel_unshuffle(2): RGB becomes 12 channels at 1/2 resolution. 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, 4 levels at 1/2, 1/4, 1/8 and 1/16 of the image, widths 96 / 192 / 384 / 384, residual block depths 1 / 1 / 2 / 4 down and 1 / 1 / 2 up. Downs and ups are pixel_unshuffle / pixel_shuffle plus a 3x3 conv. Residual blocks cycle dilations 1, 2, 4.
• 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.
• 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 is the Flux VAE with the same latent layout, on the same image. 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 Flux latents show their first 3 of 16 / 32 channels after the scaling the DiT actually sees (Flux.1: shift and scale factor; Flux.2: batch-norm statistics). Mean and std are over the full latent tensor.

PSNR on these three 512 px images (dB):

image	Primal 2x1	Flux.1 VAE	Primal 2x2	Flux.2 VAE
img01 (people, studio)	32.1	35.7	37.5	37.2
img02 (picnic, foliage)	22.0	24.2	27.1	26.5
img03 (UI screenshot)	21.3	24.7	27.4	25.9

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.

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 only, last 9k steps): 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)
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.