Field noteOct 31, 20265 min read1 source

Ostris releases Primal VAE, an ~80M-parameter 2x autoencoder with 1- or 2-channel latents that patch into Flux.1/Flux.2 latent shapes

In one minute

Ostris has released Primal VAE, an approximately 80-million-parameter autoencoder intended for diffusion pipelines using Flux.1 or Flux.2-shaped latents. Its main proposition is architectural: retain a shallow, low-channel latent representation, let the diffusion transformer perform patching, and eliminate latent scaling and shift constants.

What changed: Ostris released Primal VAE, a small 2× autoencoder whose 1- or 2-channel latents can be patched into Flux.1/Flux.2 latent shapes; the archive covers other Flux model releases.

What happened

Ostris published Primal VAE on Hugging Face under the Apache-2.0 licence. It is available in two variants: 2x1, which downsamples an image spatially by 2x and produces one latent channel, and 2x2, which uses the same 2x spatial reduction but produces two channels. Both models have roughly 80 million parameters. Images must have side lengths that are multiples of 16.

The release is designed to fit into existing Flux-shaped latent pipelines after patching. A 2x1 latent, patched with a 4x4 operation to reach total 8x compression, becomes 16 channels: the shape used by the Flux.1 VAE. The 2x2 variant becomes 32 channels, matching Flux.2. In the supplied example, a developer encodes an image, calls patchify with compression 8, then places the resulting tokens where the corresponding Flux latent would otherwise enter. The model card says the Diffusion Transformer, rather than the VAE, performs the effective deep patching.

That split is the release’s central technical claim. Conventional modern VAEs commonly carry more deeply compressed latent representations. Primal VAE keeps a spatially less-compressed map at the VAE boundary, then relies on pixel unshuffle and pixel shuffle operations to convert it to the shape expected downstream. The author describes patchify and unpatchify as these operations with a compression factor divided by two. The package also exposes helpers for token-channel counts and required image multiples.

Primal VAE is trained so its encoder outputs are approximately standard normal on every image and channel. The card describes a moment loss for mean and standard deviation plus a Gaussian order-statistic loss that targets the distribution’s full shape. The intended consequence is that users do not need a scaling_factor, shift_factor, or latent normalization step. For Flux.2 integration, the instructions specifically say to skip Flux.2’s batch-normalization latent normalization because Primal VAE is already unit-normal per image and channel.

The architecture is a mirrored convolutional U-Net with four resolution levels, RMSNorm, residual blocks, and no attention or spatial statistics. There is no posterior, KL term, or quantization convolution. The release can be loaded through Diffusers’ AutoModel, using trust_remote_code=True; the variant selects the weight file. It also supports gradient checkpointing for pixel-space losses propagated through decoding, while inference under torch.no_grad() is said to be unaffected.

The evidence on image quality is explicitly preliminary. The card reports PSNR results on three 512-pixel images. On those examples, 2x2 matched or exceeded Flux.2 VAE PSNR at the same stated information budget; 2x1 trailed Flux.1 VAE on textured content. The author cautions that three images and PSNR alone are not a benchmark, and notes that the deterministic design does not pay the reconstruction cost associated with posterior noise in a KL VAE.

What the archive adds

This is a new technical event within an existing Flux implementation thread. A prior Field Note covered Comfy-Org’s single-file Flux.1-dev packaging for ComfyUI, including bundled text encoders and a smaller checkpoint for machines with less than 24 GB of VRAM. Primal VAE does not repackage that model; it proposes a substitute encoder-decoder component that preserves the latent shape expected by Flux-oriented pipelines.

Trend signal

Generative media becomes application infrastructure is modestly supported here. Primal VAE treats an image-generation component as a swappable, explicitly shaped building block rather than a fixed part of one end-user model. But the claim remains architectural: the source does not demonstrate a broadly deployed application integration or production reliability.

Why it matters

For teams experimenting with custom diffusion training, a compatible latent boundary can reduce the amount of surrounding pipeline code that must change. The 2x1 and 2x2 variants target the shapes Flux.1 and Flux.2 expect after patching, while their unit-normal output is intended to remove scaling, shift, and normalization bookkeeping. That can make experiments easier to reason about and make a custom checkpoint simpler to load into a Flux-shaped harness.

The trade-off is not established quality. A VAE is not merely a file-format adapter: reconstruction behavior and latent distribution affect downstream diffusion training and generated output. The release offers a plausible integration path, but teams should regard it as an experimental component until it has been evaluated on their images, objectives, and training runs.

Limits and uncertainty

The model card supplies only three 512-pixel reconstruction examples and PSNR figures, not a broad benchmark, downstream generation comparison, throughput measurement, memory profile, or independent evaluation. It says 2x1 is weaker on textured content and attributes some variation to its shorter training schedule and late move to per-image normalization. The source does not establish that either variant improves a trained Flux model, works equivalently across all Flux implementations, or is suitable for production workloads. Loading requires remote custom code, which teams should review under their own dependency and security practices.

Skill to build now

Create a small VAE-adapter test before replacing anything in a training stack. Encode a representative image set with 2x1 or 2x2, patch to the target 8x latent layout, unpatch and decode, then record reconstruction artifacts, latent mean and standard deviation, VRAM use, and throughput. Run the same test with the incumbent Flux VAE. This makes the advertised compatibility claim testable before committing compute to diffusion training.

Reader takeaway

Primal VAE is an open, small-latent autoencoder designed to slot into Flux.1- and Flux.2-shaped pipelines after patching. Its normalization-free interface may simplify custom experimentation, but its current quality evidence is a model-card demonstration rather than a sufficient benchmark for deployment.

Open question

Will the unit-normal latent design remain stable and useful after full diffusion-model training across diverse datasets, rather than only in reconstruction tests?

Evidence / sources

  1. 1huggingface.co/ostris/primal_vae
All Field Notes

Next reads

Read next

Open the graph
  1. 01topicGenerative MediaExplore the generative-media topic hub.
  2. 02topicMultimodalExplore the multimodal topic hub.
  3. 03archiveField NotesOpen the latest editorial analysis.
  4. 04source ledgerSource LedgerInspect the public source evidence graph.
Who read this page?Machine requests, hidden until opened

Loading the privacy-safe route aggregate...

Open the JSON contract