Primal VAE 2x1 and 2x2

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/img01.png filter=lfs diff=lfs merge=lfs -text
+assets/img02.png filter=lfs diff=lfs merge=lfs -text
+assets/img03.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,157 @@
 ---
-license: apache-2.0
+library_name: diffusers
+tags:
+  - vae
+  - autoencoder
+  - latent-diffusion
 ---
+
+# Primal VAE
+
+Primal VAE is a new VAE concept I am working on. It 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 and nearly lossless.
+
+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))      # [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))      # [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` are just `pixel_unshuffle` / `pixel_shuffle` with `token_patch = 4`.
+`vae.token_channels`, `vae.token_compression` and `vae.image_multiple` give you the numbers above.
+
+## 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** (both): `mean^2 + (std - 1)^2`, weight 0.1. On `2x1` the mean and std are pooled over the
+  batch. On `2x2` they are per image and per channel, so every single image comes out mean 0 / std 1 in each
+  channel, not just the dataset on average.
+- **Gaussian order-statistic loss** (`2x2`): 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.2 | 24.2 | 27.1 | 26.5 |
+| img03 (UI screenshot) | 21.0 | 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 about a third of the training steps of the 2x2. Its per-image
+latent std also drifts (0.9 to 1.4 here) because its moment loss was batch pooled. The per-image normalization
+on the 2x2 holds std within a few percent of 1.
+
+## Training
+
+I am training these 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** (`2x2`): 5% of images are swapped for a random flat color. Flat regions are the hard case
+  for a per-image normalized latent.
+- **Critic** (`2x1`): 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.

config.json

@@ -0,0 +1,12 @@
+{
+  "_class_name": "PrimalVAE",
+  "_diffusers_version": "0.36.0",
+  "auto_map": {
+    "AutoModel": "modeling_primal_vae.PrimalVAE"
+  },
+  "weights": [
+    "primal_vae_2x1.safetensors",
+    "primal_vae_2x2.safetensors"
+  ],
+  "token_patch": 4
+}

modeling_primal_vae.py

@@ -0,0 +1,192 @@
+"""Primal VAE for diffusers (AutoModel remote code). Deterministic autoencoder: RGB in [-1, 1] ->
+latent_channels at `compression`x spatial compression, latent ~ N(0, 1) by training, no scaling or
+shift factor. `patchify(z)` regroups the latent into the token grid a DiT consumes (4x4 by default:
+1ch @ 2x -> 16ch @ 8x, 2ch @ 2x -> 32ch @ 8x).
+
+Loads the ai-toolkit training checkpoints directly: `from_pretrained(repo, variant="2x1")` reads
+`primal_vae_2x1.safetensors` and derives the architecture from the tensor shapes; `weights_name=` picks
+any other file. Standalone: torch + diffusers + safetensors.
+"""
+
+import math
+import os
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from huggingface_hub import hf_hub_download
+from safetensors import safe_open
+
+# dilation of block i's first conv; widens the receptive field without downsampling
+DILATION_CYCLE = (1, 2, 4)
+
+
+class ChannelRMSNorm(nn.Module):
+    """RMSNorm over channels at each position: no spatial statistics, so tiles encode/decode like the whole image."""
+
+    def __init__(self, ch: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(ch))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.rms_norm(x.movedim(1, -1), (x.shape[1],), self.weight.to(x.dtype), self.eps).movedim(-1, 1)
+
+
+class ResBlock(nn.Module):
+    def __init__(self, width: int, dilation: int):
+        super().__init__()
+        self.norm1 = ChannelRMSNorm(width)
+        self.conv1 = nn.Conv2d(width, width, 3, padding=dilation, dilation=dilation)
+        self.norm2 = ChannelRMSNorm(width)
+        self.conv2 = nn.Conv2d(width, width, 3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.conv1(F.silu(self.norm1(x)))
+        h = self.conv2(F.silu(self.norm2(h)))
+        return x + h
+
+
+class UStack(nn.Module):
+    """conv_in -> [blocks, pixel-unshuffle down] per level -> bottom blocks -> [up, + skip, blocks] per level
+    -> norm/act -> conv_out."""
+
+    def __init__(self, in_ch: int, widths: List[int], depths: List[int], up_depths: List[int], out_ch: int):
+        super().__init__()
+        L = len(widths)
+        if len(depths) != L or len(up_depths) != L - 1:
+            raise ValueError("level_depths must match level_widths; up_depths has one entry less")
+        self.conv_in = nn.Conv2d(in_ch, widths[0], 3, padding=1)
+        self.down_blocks = nn.ModuleList([nn.ModuleList([ResBlock(widths[i], DILATION_CYCLE[j % len(DILATION_CYCLE)]) for j in range(depths[i])]) for i in range(L)])
+        self.downs = nn.ModuleList([nn.Conv2d(4 * widths[i], widths[i + 1], 3, padding=1) for i in range(L - 1)])
+        self.ups = nn.ModuleList([nn.Conv2d(widths[i + 1], 4 * widths[i], 3, padding=1) for i in range(L - 1)])
+        self.up_blocks = nn.ModuleList([nn.ModuleList([ResBlock(widths[i], DILATION_CYCLE[j % len(DILATION_CYCLE)]) for j in range(up_depths[i])]) for i in range(L - 1)])
+        self.norm_out = ChannelRMSNorm(widths[0])
+        self.conv_out = nn.Conv2d(widths[0], out_ch, 3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv_in(x)
+        skips = []
+        L = len(self.down_blocks)
+        for i in range(L):
+            for blk in self.down_blocks[i]:
+                x = blk(x)
+            if i < L - 1:
+                skips.append(x)
+                x = self.downs[i](F.pixel_unshuffle(x, 2))
+        for i in reversed(range(L - 1)):
+            x = F.pixel_shuffle(self.ups[i](x), 2) + skips[i]
+            for blk in self.up_blocks[i]:
+                x = blk(x)
+        return self.conv_out(F.silu(self.norm_out(x)))
+
+
+class PrimalVAE(ModelMixin, ConfigMixin):
+    """Primal VAE: `encode` -> [B, latent_channels, H/compression, W/compression] (~N(0,1)), `decode` back to
+    [B, 3, H, W] in [-1, 1]. No posterior, no scaling/shift factor."""
+
+    @register_to_config
+    def __init__(
+        self,
+        compression: int = 2,
+        latent_channels: int = 1,
+        level_widths: List[int] = (96, 192, 384, 384),
+        level_depths: List[int] = (1, 1, 2, 4),
+        up_depths: Optional[List[int]] = (1, 1, 2),
+        token_patch: int = 4,
+    ):
+        super().__init__()
+        widths, depths = [int(v) for v in level_widths], [int(v) for v in level_depths]
+        ups = [int(v) for v in up_depths] if up_depths is not None else depths[:-1]
+        self.compression = int(compression)
+        self.latent_channels = int(latent_channels)
+        self.token_patch = int(token_patch)
+        # the latent must tile into the U's 2^(levels-1) and the DiT's token patch
+        self.latent_multiple = max(self.token_patch, 2 ** (len(widths) - 1))
+        pix = 3 * self.compression ** 2
+        self.encoder = UStack(pix, widths, depths, ups, self.latent_channels)
+        self.decoder = UStack(self.latent_channels, widths, depths, ups, pix)
+
+    @property
+    def image_multiple(self) -> int:
+        """Image sides must be a multiple of this (16 for the released models)."""
+        return self.compression * self.latent_multiple
+
+    @property
+    def token_channels(self) -> int:
+        """Channels of the patchified latent: 16 for 2x1 (Flux.1 layout), 32 for 2x2 (Flux.2 layout)."""
+        return self.latent_channels * self.token_patch ** 2
+
+    @property
+    def token_compression(self) -> int:
+        """Spatial compression of the patchified latent (8 for the released models)."""
+        return self.compression * self.token_patch
+
+    def _check_image(self, x: torch.Tensor):
+        H, W = x.shape[-2:]
+        m = self.image_multiple
+        if H % m or W % m:
+            raise ValueError(f"image sides must be multiples of {m}, got {H}x{W}")
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """x: [B, 3, H, W] in [-1, 1] -> z: [B, latent_channels, H/compression, W/compression]."""
+        self._check_image(x)
+        return self.encoder(F.pixel_unshuffle(x.to(self.dtype), self.compression))
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """z -> x_hat: [B, 3, H, W], nominally in [-1, 1] (not clamped)."""
+        return F.pixel_shuffle(self.decoder(z.to(self.dtype)), self.compression)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.decode(self.encode(x))
+
+    def patchify(self, z: torch.Tensor) -> torch.Tensor:
+        """Latent -> DiT token grid: [B, C, h, w] -> [B, C * p^2, h/p, w/p] (space-to-depth, p = token_patch)."""
+        return F.pixel_unshuffle(z, self.token_patch)
+
+    def unpatchify(self, tokens: torch.Tensor) -> torch.Tensor:
+        """Inverse of `patchify`."""
+        return F.pixel_shuffle(tokens, self.token_patch)
+
+    # ---- loading -----------------------------------------------------------
+    HUB_KWARGS = ("cache_dir", "force_download", "proxies", "local_files_only", "token", "revision", "subfolder")
+
+    @staticmethod
+    def config_from_shapes(shapes: dict) -> dict:
+        """Architecture from the tensor shapes alone; the file's header metadata is not consulted."""
+        widths = [shapes["encoder.conv_in.weight"][0]]
+        while f"encoder.downs.{len(widths) - 1}.weight" in shapes:
+            widths.append(shapes[f"encoder.downs.{len(widths) - 1}.weight"][0])
+        count = lambda pre, i: len({k.split(".")[3] for k in shapes if k.startswith(f"{pre}.{i}.")})
+        return {
+            "compression": math.isqrt(shapes["encoder.conv_in.weight"][1] // 3),
+            "latent_channels": shapes["encoder.conv_out.weight"][0],
+            "level_widths": widths,
+            "level_depths": [count("encoder.down_blocks", i) for i in range(len(widths))],
+            "up_depths": [count("encoder.up_blocks", i) for i in range(len(widths) - 1)],
+        }
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, variant: str = "2x1", weights_name: Optional[str] = None, torch_dtype=None, token_patch: int = 4, **kwargs):
+        """`variant` picks `primal_vae_<variant>.safetensors` in the repo ("2x1" or "2x2"); `weights_name` any file."""
+        name = weights_name or f"primal_vae_{variant}.safetensors"
+        hub = {k: v for k, v in kwargs.items() if k in cls.HUB_KWARGS and v is not None}
+        path = str(pretrained_model_name_or_path)
+        if os.path.isdir(path):
+            file = os.path.join(path, hub.get("subfolder") or "", name)
+        else:
+            file = hf_hub_download(path, name, **hub)
+        with safe_open(file, framework="pt", device="cpu") as f:
+            keys = [k for k in f.keys() if k != "arch_spec"]  # arch_spec: ai-toolkit's own derived buffer
+            cfg = cls.config_from_shapes({k: f.get_slice(k).get_shape() for k in keys})
+            with torch.device("meta"):
+                model = cls(**cfg, token_patch=token_patch)
+            state = {k: f.get_tensor(k) for k in keys}
+        if torch_dtype is not None:
+            state = {k: v.to(torch_dtype) for k, v in state.items()}
+        model.load_state_dict(state, assign=True)
+        model.register_to_config(_name_or_path=path)
+        return model.eval()

primal_vae_2x1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9cf37dd0ff80edde481ee2cae858107a71855e6ed5fb495aa095a09b33b9189a
+size 321415508

primal_vae_2x2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f90813d9a7b2aa4f9334a93dc22c28867f5204b850357fe224029431727be3e8
+size 321422432