data-archetype/dinac_ae

DINAC-AE is a DINO-Aligned Class-token AutoEncoder. It follows the SemDisDiffAE family: patch-16 spatial latents, a VP diffusion decoder, and DINO-aligned representations.

Relative to SemDisDiffAE, DINAC-AE changes the encoder from FCDM blocks to a 6-block ViT/DiT-style transformer encoder and uses DINOv3 ViT-B/16 alignment. The latent-to-DINO alignment head is extended to predict the DINO class token as well as patch tokens. predict_class(latents) exposes that class-token feature directly from latents.

2k PSNR Benchmark

Evaluated on 2000 validation images.

DINAC-AE targets a compromise between high reconstruction quality, a learnable latent space with KL-like variance expansion, DINOv3 alignment, and robustness to local token errors.

Results viewer shows the 39-image reconstruction set with DINAC-AE and FLUX.2 VAE reconstructions, RGB differences, and latent PCA. The released export recheck on that 39-image set gives 35.15 dB mean PSNR (25.73 min, 45.99 max).

Full technical report

Encode Throughput

Measured on an NVIDIA GeForce RTX 5090 in bfloat16, averaging repeated batches per resolution.

The transformer encoder is slightly slower and larger than the full_capacitor FCDM encoder, but remains much faster and much smaller than the FLUX.2 VAE encoder.

Latent Interface

• encode() returns DINAC-AE's own whitened latent space.
• decode() expects that same whitened latent space and dewhitens internally.
• predict_class() expects the same whitened latent space, dewhitens internally, and predicts a DINOv3 ViT-B/16 class-token feature.
• whiten() and dewhiten() are exposed for explicit control.
• encode_posterior() returns the raw exported posterior before whitening.
• DinacAEInferenceConfig.num_steps counts decoder evaluations directly: num_steps=1 means one NFE.

The export ships weights in float32. The recommended and default runtime path is bfloat16 AMP for the main encoder, decoder, and class-token path. The loader retains normalization affine parameters, GRN/residual gates, the final pixel projection, latent statistics, RoPE/time frequencies, sampler state, and whitening/dewhitening in float32. These tensors are loaded from the original FP32 weights before ordinary parameters are converted to BF16.

Usage

import torch

from dinac_ae import DinacAE, DinacAEInferenceConfig


device = "cuda"
model = DinacAE.from_pretrained(
    "data-archetype/dinac_ae",
    device=device,
    dtype=torch.bfloat16,
)

image = ...  # [1, 3, H, W] in [-1, 1], H and W divisible by 16

with torch.inference_mode():
    latents = model.encode(image.to(device=device, dtype=torch.bfloat16))
    class_token = model.predict_class(latents)
    recon = model.decode(
        latents,
        height=int(image.shape[-2]),
        width=int(image.shape[-1]),
        inference_config=DinacAEInferenceConfig(num_steps=1),
    )

Details

• DINAC-AE uses a 6-block ViT/DiT-style transformer encoder and an 8-block FCDM decoder.
• Patch size is 16, model width is 896, and latent width is 128.
• The DINO alignment head predicts spatial patch tokens and is extended with a class-token output in DINOv3 ViT-B/16 feature space.
• The class-token output is used to improve semantic organization of the latent space and to support FD-loss / Representation Frechet Distance objectives directly in latent space.
• predict_class(latents) reaches mean cosine similarity 0.757458 against the frozen DINOv3 ViT-B/16 teacher class token on the same 2000 images.
• DINO alignment is applied directly to clean latent tokens. Robustness to local token errors is handled by random-token logSNR offset regularization.
• Results viewer: https://huggingface.co/spaces/data-archetype/dinac_ae-results
• Related: SemDisDiffAE, full_capacitor, capacitor_decoder

Citation

@misc{dinac_ae,
  title   = {DINAC-AE: a DINO-aligned class-token diffusion autoencoder},
  author  = {data-archetype},
  email   = {data-archetype@proton.me},
  year    = {2026},
  month   = may,
  url     = {https://huggingface.co/data-archetype/dinac_ae},
}