Primal VAE

Primal VAE от Ostris — это VAE для diffusion, который делает только 2x spatial down до 1 или 2 каналов. После patch 4 получаются 16 или 32 канала при 8x — та же форма, что у Flux.1 VAE и Flux.2 VAE. Обе версии — примерно 80M параметров.

📌 Ключевые особенности

• 📌 Два варианта: 2x1 (1 ch @ 2x → 16 ch @ 8x) и 2x2 (2 ch @ 2x → 32 ch @ 8x) после patch 4.
• ✨ Латент обучен быть похожим на гауссов шум: без scaling_factor, shift_factor и latent normalization.
• 💻 patchify/unpatchify принимают общий compression; compression=8 — это патч 4×4.
• 🖼 Conv U-Net с RMSNorm по каналам, без attention; тайлы кодируются и декодируются так же, как полное изображение, при наличии overlap.
• 🔥 Через decode можно пускать backprop с gradient checkpointing; инференс в bf16 стоит около 0.05 dB.
• 📚 Лицензия Apache 2.0, веса fp32, архитектура читается из формы тензоров.

💻 Пример

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")
z = vae.encode(x)       # [B, 1, H/2, W/2], ~N(0, 1)
x_hat = vae.decode(z)   # [B, 3, H, W]

Подробнее

PSNR на трёх 512 px изображениях: Primal 2x2 — 37.5 / 27.1 / 27.4 dB, Flux.2 VAE — 37.2 / 26.5 / 25.9 dB; Primal 2x1 уступает Flux.1 на текстурах: 22.0 против 24.2 dB на img02.
2x2 обучен 147k шагов, 2x1 — 62k; у 2x1 std латента ещё дрейфует (1.07–1.21), у 2x2 — в пределах нескольких процентов от 1.
Стороны изображения должны быть кратны 16. Для инференса bf16 стоит ~0.05 dB. Лицензия Apache 2.0.
Три изображения и PSNR — не полноценный бенчмарк.