Add 4x1 primal vae

b0b828d verified 3 days ago

10.2 kB

---
license: apache-2.0
library_name: diffusers
tags:
  - vae
  - autoencoder
  - latent-diffusion
---
# 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.

| model | latent | at 8x | same shape as |
|---|---|---|---|
| `2x1` | 1 ch @ 2x | 16 ch @ 8x (4×4 patch) | Flux.1 VAE (16 ch @ 8x) |
| `2x2` | 2 ch @ 2x | 32 ch @ 8x (4×4 patch) | Flux.2 VAE (32 ch @ 8x) |
| `4x1` | 1 ch @ 4x | 4 ch @ 8x (2×2 patch) | SD 1.5 / SDXL VAE (4 ch @ 8x) |

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

## Usage

```python
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)

```python
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)

```python
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)

```python
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.

![2x1 vs Flux.1](assets/compare_2x1.svg)

![2x2 vs Flux.2](assets/compare_2x2.svg)

![4x1 vs SD 1.5 / SDXL](assets/compare_4x1.svg)

PSNR on these three 512 px images (dB):

| image | Primal 2x1 | Flux.1 VAE | Primal 2x2 | Flux.2 VAE | Primal 4x1 | SD 1.5 VAE | SDXL VAE |
|---|---|---|---|---|---|---|---|
| img01 (people, studio) | 32.1 | 35.7 | 37.5 | 37.2 | 25.8 | 27.8 | 29.4 |
| img02 (picnic, foliage) | 22.0 | 24.2 | 27.1 | 26.5 | 17.3 | 17.5 | 18.6 |
| img03 (UI screenshot) | 21.3 | 24.7 | 27.4 | 25.9 | 17.2 | 17.3 | 19.6 |

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](https://github.com/ostris/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`.

	---
	license: apache-2.0
	library_name: diffusers
	tags:
	- vae
	- autoencoder
	- latent-diffusion
	---

	# 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.

	\| model \| latent \| at 8x \| same shape as \|
	\|---\|---\|---\|---\|
	\| `2x1` \| 1 ch @ 2x \| 16 ch @ 8x (4×4 patch) \| Flux.1 VAE (16 ch @ 8x) \|
	\| `2x2` \| 2 ch @ 2x \| 32 ch @ 8x (4×4 patch) \| Flux.2 VAE (32 ch @ 8x) \|
	\| `4x1` \| 1 ch @ 4x \| 4 ch @ 8x (2×2 patch) \| SD 1.5 / SDXL VAE (4 ch @ 8x) \|

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

	## Usage

	```python
	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)

	```python
	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)

	```python
	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)

	```python
	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.

	![2x1 vs Flux.1](assets/compare_2x1.svg)

	![2x2 vs Flux.2](assets/compare_2x2.svg)

	![4x1 vs SD 1.5 / SDXL](assets/compare_4x1.svg)

	PSNR on these three 512 px images (dB):

	\| image \| Primal 2x1 \| Flux.1 VAE \| Primal 2x2 \| Flux.2 VAE \| Primal 4x1 \| SD 1.5 VAE \| SDXL VAE \|
	\|---\|---\|---\|---\|---\|---\|---\|---\|
	\| img01 (people, studio) \| 32.1 \| 35.7 \| 37.5 \| 37.2 \| 25.8 \| 27.8 \| 29.4 \|
	\| img02 (picnic, foliage) \| 22.0 \| 24.2 \| 27.1 \| 26.5 \| 17.3 \| 17.5 \| 18.6 \|
	\| img03 (UI screenshot) \| 21.3 \| 24.7 \| 27.4 \| 25.9 \| 17.2 \| 17.3 \| 19.6 \|

	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](https://github.com/ostris/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`.