Add 4x1 primal vae

README.md

@@ -9,19 +9,20 @@ tags:
 
 # 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
+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 4, diffuse.
+no latent normalization to carry around. Encode, patch, diffuse.
 
-| model | latent | patch 4 | same shape as |
+| model | latent | at 8x | 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) |
+| `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) |
 
-Both are ~80M parameters. Image sides need to be multiples of 16.
+The 2x models are ~80M parameters, the 4x1 is ~122M. Image sides need to be multiples of 16 (2x) or 32 (4x1).
 
 ## Usage
 
@@ -31,7 +32,7 @@ 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
+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]
 ```
@@ -60,24 +61,39 @@ tokens = vae.patchify(vae.encode(x), compression=8)   # [B, 32, H/8, W/8]: the F
 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 / 2`, so `compression=8` is a 4x4 patch. `vae.token_channels(8)` and
-`vae.image_multiple` give you the numbers above.
+`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
 
-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.
+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.
@@ -90,7 +106,8 @@ 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.
+  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
@@ -98,23 +115,26 @@ encoder output.
 
 ## 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.
+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 |
-|---|---|---|---|---|
-| 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 |
+| 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
@@ -123,6 +143,11 @@ is a 12:1 reduction with no slack, and it has well under half the training steps
 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
@@ -141,11 +166,11 @@ The tricks that should make the latent space easier to diffuse:
   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.
+- **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.
@@ -155,6 +180,7 @@ instead of adding a second decoder pass.
 ```
 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
 ```

config.json

@@ -6,6 +6,7 @@
   },
   "weights": [
     "primal_vae_2x1.safetensors",
-    "primal_vae_2x2.safetensors"
+    "primal_vae_2x2.safetensors",
+    "primal_vae_4x1.safetensors"
   ]
 }

modeling_primal_vae.py

@@ -1,7 +1,7 @@
 """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, compression=8)` regroups the latent into the token grid a DiT consumes at that
-total compression (1ch @ 2x -> 16ch @ 8x, 2ch @ 2x -> 32ch @ 8x).
+total compression (1ch @ 2x -> 16ch @ 8x, 2ch @ 2x -> 32ch @ 8x, 1ch @ 4x -> 4ch @ 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
@@ -113,7 +113,7 @@ class PrimalVAE(ModelMixin, ConfigMixin):
         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)
-        # the latent must tile into the U's 2^(levels-1) and a 4x4 token patch
+        # the latent must tile into the U's 2^(levels-1) and the DiT's token patch
         self.latent_multiple = max(4, 2 ** (len(widths) - 1))
         pix = 3 * self.compression ** 2
         self.encoder = UStack(pix, widths, depths, ups, self.latent_channels)
@@ -121,11 +121,11 @@ class PrimalVAE(ModelMixin, ConfigMixin):
 
     @property
     def image_multiple(self) -> int:
-        """Image sides must be a multiple of this (16 for the released models)."""
+        """Image sides must be a multiple of this (16 for 2x1 / 2x2, 32 for 4x1)."""
         return self.compression * self.latent_multiple
 
     def token_channels(self, compression: int = 8) -> int:
-        """Channels of the latent patchified to `compression`: 16 for 2x1 at 8x (Flux.1), 32 for 2x2 (Flux.2)."""
+        """Channels of the latent patchified to `compression`: 16 for 2x1 at 8x (Flux.1), 32 for 2x2 (Flux.2), 4 for 4x1 (SD 1.5 / SDXL)."""
         return self.latent_channels * self._patch(compression) ** 2
 
     def _patch(self, compression: int) -> int:
@@ -153,7 +153,8 @@ class PrimalVAE(ModelMixin, ConfigMixin):
 
     def patchify(self, z: torch.Tensor, compression: int = 8) -> torch.Tensor:
         """Latent -> token grid at `compression`x total: [B, C, h, w] -> [B, C * p^2, h/p, w/p] with
-        p = compression / model compression (space-to-depth). 8 gives the Flux 8x layout from a 2x latent."""
+        p = compression / model compression (space-to-depth). 8 gives the Flux 8x layout from a 2x latent and the
+        SD 1.5 / SDXL layout from a 4x latent."""
         return F.pixel_unshuffle(z, self._patch(compression))
 
     def unpatchify(self, tokens: torch.Tensor) -> torch.Tensor:
@@ -183,7 +184,7 @@ class PrimalVAE(ModelMixin, ConfigMixin):
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, variant: str = "2x1", weights_name: Optional[str] = None, torch_dtype=None, **kwargs):
-        """`variant` picks `primal_vae_<variant>.safetensors` in the repo ("2x1" or "2x2"); `weights_name` any file."""
+        """`variant` picks `primal_vae_<variant>.safetensors` in the repo ("2x1", "2x2" or "4x1"); `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)

primal_vae_4x1.safetensors

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