| --- |
| license: apache-2.0 |
| tags: |
| - diffusion |
| - autoencoder |
| - image-reconstruction |
| - image-tokenizer |
| - pytorch |
| - fcdm |
| - semantic-alignment |
| library_name: fcdm_diffae |
| --- |
| |
| # data-archetype/semdisdiffae |
|
|
| > **Retired:** new projects should use |
| > [data-archetype/dinac_ae](https://huggingface.co/data-archetype/dinac_ae). |
|
|
| ### Version History |
|
|
| | Date | Change | |
| |------|--------| |
| | 2026-04-18 | Stabilize VP posterior and schedule alpha/sigma computation for extreme log-SNR values while keeping posterior interpolation in float32 | |
| | 2026-04-10 | Refresh standalone package: fix bf16 RMSNorm precision path in both encoder and decoder to match training code; local export tooling now preserves fp32 EMA weights for future re-exports | |
| | 2026-04-08 | Fix posterior VP interpolation to use float32 precision (was using model dtype) | |
| | 2026-04-07 | Rename package `capacitor_diffae` → `fcdm_diffae`, class `FCDMDiffAE`; encode() now returns whitened latents, decode() dewhitens internally | |
| | 2026-04-06 | Initial release | |
|
|
| **SemDisDiffAE** (**Sem**antically **Dis**entangled **Diff**usion **A**uto**E**ncoder) |
| — a fast image tokenizer with semantically structured 128-channel latents, built |
| on FCDM (Fully Convolutional Diffusion Model) blocks with a VP-parameterized |
| diagonal Gaussian posterior. |
|
|
| Trained with DINOv2 semantic alignment, this VAE was empirically found to |
| offer comparable downstream diffusion convergence speed to other semantically |
| aligned VAEs such as Flux.2 and PS-VAE, while being much faster to encode |
| and decode and achieving very high reconstruction quality (38.6 dB mean PSNR |
| on 2k images). |
|
|
| Built on a pure convolutional architecture with no attention layers in the |
| encoder or decoder, enabling efficient inference at any resolution. |
|
|
| **[Technical Report](technical_report_semantic.md)** · |
| **[Interactive Results Viewer](https://huggingface.co/spaces/data-archetype/semdisdiffae-results)** |
|
|
| ## Key Features |
|
|
| - **Fast**: ~3 ms/img encode, ~6 ms/img decode (1 step) on Blackwell RTX Pro 6000 — significantly |
| faster than Flux.2 VAE |
| - **High fidelity**: 38.6 dB mean PSNR (2k images), exceeding Flux.2 VAE (37.0 dB) |
| - **Semantically structured latents**: DINOv2-aligned, producing latents with |
| clear semantic segmentation visible in PCA projections |
| - **Comparable downstream convergence**: empirically matches the downstream |
| diffusion training convergence speed of Flux.2 and PS-VAE |
| - **Pure convolutional**: no attention in encoder/decoder, O(n) in spatial resolution |
| - **VP diffusion decoder**: single-step DDIM for PSNR-optimal, optional multi-step |
| with PDG for perceptual sharpening |
|
|
| ## Architecture |
|
|
| | Property | Value | |
| |----------|-------| |
| | Parameters | 88.8M | |
| | Patch size | 16 | |
| | Model dim | 896 | |
| | Encoder depth | 4 blocks | |
| | Decoder depth | 8 blocks (2+4+2 skip-concat) | |
| | Bottleneck | 128 channels | |
| | Compression | 16x spatial, 6.0x total | |
| | Posterior | Diagonal Gaussian (VP log-SNR) | |
| | Block type | FCDM (ConvNeXt + GRN + scale/gate AdaLN) | |
|
|
| ## Quick Start |
|
|
| ```python |
| from fcdm_diffae import FCDMDiffAE, FCDMDiffAEInferenceConfig |
| |
| model = FCDMDiffAE.from_pretrained("data-archetype/semdisdiffae", device="cuda") |
| |
| # Encode (returns posterior mode by default) |
| latents = model.encode(images) # [B,3,H,W] in [-1,1] -> [B,128,H/16,W/16] |
| |
| # Decode — PSNR-optimal (1 step, default) |
| recon = model.decode(latents, height=H, width=W) |
| |
| # Decode — perceptual sharpening (10 steps + PDG) |
| cfg = FCDMDiffAEInferenceConfig(num_steps=10, pdg=True, pdg_strength=2.0) |
| recon = model.decode(latents, height=H, width=W, inference_config=cfg) |
| |
| # Full posterior access |
| posterior = model.encode_posterior(images) |
| z_sampled = posterior.sample() |
| ``` |
|
|
| ## Recommended Settings |
|
|
| | Use case | Steps | PDG | Notes | |
| |----------|-------|-----|-------| |
| | PSNR-optimal | 1 | off | Default, fastest | |
| | Perceptual | 10 | on (2.0) | Sharper, ~15x slower | |
|
|
| PDG is primarily useful for more compressed bottlenecks (32 or 64 channels) |
| and is rarely necessary for 128-channel models where reconstruction quality |
| is already high. |
|
|
| ## Training |
|
|
| Trained with: |
| - Pixel-space VP diffusion reconstruction loss (x-prediction, SiD2 weighting) |
| - DINOv2-S semantic alignment (negative cosine, weight 0.01) |
| - VP posterior variance expansion (weight 1e-5) |
| - Latent scale regularization (weight 0.0001) |
| - AdamW optimizer, bf16 mixed precision, EMA decay 0.9995 |
| - 251k steps on a single GPU |
|
|
| See the [technical report](technical_report_semantic.md) for full details. |
|
|
| ## Dependencies |
|
|
| - PyTorch >= 2.0 |
| - safetensors (for loading weights) |
|
|
| ## Citation |
|
|
| ```bibtex |
| @misc{semdisdiffae, |
| title = {SemDisDiffAE: A Semantically Disentangled Diffusion Autoencoder}, |
| author = {data-archetype}, |
| email = {data-archetype@proton.me}, |
| year = {2026}, |
| month = apr, |
| url = {https://huggingface.co/data-archetype/semdisdiffae}, |
| } |
| ``` |
|
|
| ## License |
|
|
| Apache 2.0 |
|
|
