Primal VAE: 2x Downsampling That Matches Flux Latent Shapes

Most modern VAEs compress the image 8x and then hand the diffusion model a latent that has to be rescaled with a measured scaling factor, a shift factor, or batch-norm statistics. Primal VAE takes the other route: the VAE only performs a 2x spatial reduction, and the DiT's own patchify step supplies the remaining 4x. A 1-channel latent at 2x patched 4 ways becomes 16 channels at 8x, the same shape a Flux.1 VAE produces; the 2-channel variant lands on 32 channels at 8x, the Flux.2 VAE shape.

Both checkpoints are about 80M parameters, and image sides need to be multiples of 16.

Each row is input, reconstruction, latent and |error| x 4. Under every Primal VAE row sits the Flux VAE with the same latent layout, on the same image. Latents are drawn at 3 sigma; the two-channel latent puts channel 0 on red and channel 1 on blue.

Reconstruction quality on the three 512 px images in the repository, in dB:

The author states plainly that three images and PSNR are not a benchmark, and that a KL VAE pays a reconstruction tax for its posterior noise that a deterministic one does not. The 2x2 variant matches or beats the Flux.2 VAE at the same information budget of 32 values per 8x8 pixel block. The 2x1 variant trails the Flux.1 VAE on textured content: one channel at 2x is a 12:1 reduction with no slack, and it received well under half the training steps of the 2x2 run (62k against 147k).

The 1-channel comparison. The single-channel latent is shown in gray and its per-image standard deviation still drifts, from 1.07 to 1.21, because it only received the per-image normalization for its last 9k steps.

Architecture

Both variants are the same network and differ only in output channels:

• Input goes through pixel_unshuffle(2), so RGB becomes 12 channels at half resolution. Everything runs at latent resolution or below, which keeps the latent a spatially honest map of the image.
• Encoder and decoder are mirrored convolutional U-Nets with four levels at 1/2, 1/4, 1/8 and 1/16 of the image size, widths 96 / 192 / 384 / 384, and residual blocks with dilations cycling 1, 2, 4.
• Normalization is RMSNorm over channels at each position. There is no spatial statistic and no attention, so a tile encodes and decodes the same as the full image given some overlap.
• The encoder ends in a plain 3x3 convolution to 1 or 2 channels. No posterior, no KL term, no quantization convolutions.

Keeping the latent in line with noise

A diffusion model is trained against Gaussian noise, so the latent should already look like Gaussian noise. Instead of adding a KL term and measuring a scaling factor afterwards, the statistics go directly into the loss on the encoder output:

• Moment loss: mean^2 + (std - 1)^2 at weight 0.1, computed per image and per channel, so every image comes out at mean 0 and standard deviation 1 rather than only the dataset on average.
• Gaussian order-statistic loss: the latent values of each image and channel are sorted and matched against the expected order statistics of a standard normal of the same size, using Blom's approximation. This pins the whole distribution shape instead of the first two moments, so a flat gray image still encodes to something that resembles noise rather than a constant.

Because of that, there is no scaling factor, no shift factor and no latent normalization to carry around, and encoding returns a plain tensor instead of a distribution: encode gives [B, 1, H/2, W/2] with no .latent_dist and no sampling step.

Training

Reconstruction is MAE at 0.5 and LPIPS at 1.0, plus an L1 term on the log-magnitude 2D spectrum at 0.1 that pushes high-frequency energy without demanding exact phase. On top of that, several tricks aim at making the latent easier to diffuse:

• Latent flips and rotations taken from EQ-VAE: on 50% of steps the latent is flipped or rotated by 90 degrees and decoded against the equally transformed image, so the latent has to be a picture of the image rather than an arbitrary code.
• Noise-perturbed decode: on 20% of steps the decoder receives z + sigma * noise with sigma drawn per image from U(0, 0.5) and still has to reconstruct the clean image, which teaches the decoder to tolerate diffusion error.
• 5% of images are replaced by a random flat color, because flat regions are the hard case for a per-image normalized latent.
• The 2x1 variant spent its last 9k steps under a critic: a frozen SigLIP2 NaFlex ViT supplies patch features from layers 4, 8 and 12 of the real and reconstructed image at native resolution, small spectral-normalized convolutional heads train with a hinge loss plus R1, and the decoder receives the generator term at weight 0.1. The 2x2 was trained without a critic.

The published weight files are the training checkpoints, unchanged, in fp32. Loading them with torch_dtype=torch.bfloat16 for inference costs about 0.05 dB.

Availability

Primal VAE is not usable inside ComfyUI yet. There is no custom VAE loader node, no Comfy-Org repack and no core support: the repository ships a diffusers AutoModel implementation with trust_remote_code=True (modeling_primal_vae.py plus config.json), and the documented path is AutoModel.from_pretrained("ostris/primal_vae", variant="2x1") followed by patchify / unpatchify to reach the Flux.1 or Flux.2 token layout.

A repository discussion opened on October 31 asks directly whether it can run in ComfyUI and has no working answer yet; a separate community discussion offers a single-file Rust build for local inference on GPU or CPU.

Anything else stays in the Python lane: vae.encode(x) returns the latent, vae.patchify(latent, compression=8) turns it into Flux.1-shaped or Flux.2-shaped tokens, and vae.decode(vae.unpatchify(tokens)) brings it back, with no scaling or shift applied on either side. The architecture is read from tensor shapes, so a fine-tune exported from ai-toolkit can be loaded with weights_name="my_finetune.safetensors".