Add gradient checkpointing to the vae

5dcf26e verified 1 day ago

8.49 kB

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

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

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

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

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

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

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](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` 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`.

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

	# 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

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

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

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

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

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

	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](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` 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`.