data-archetype/capacitor_decoder

Capacitor decoder: a faster, lighter FLUX.2-compatible latent decoder built on the SemDisDiffAE architecture.

Decode Speed

RTX 5090

These measurements are decode-only, were run on an NVIDIA GeForce RTX 5090 in bfloat16, and time sequential batch-1 decode over the same cached latent set for both decoders.

2k PSNR Benchmark

Evaluated on 2000 validation images: roughly 2/3 photographs and 1/3 book covers. Each image is encoded once with FLUX.2 and reused for both decoders.

Results viewer

Usage

import torch
from diffusers.models import AutoencoderKLFlux2

from capacitor_decoder import CapacitorDecoder, CapacitorDecoderInferenceConfig


def flux2_patchify_and_whiten(
    latents: torch.Tensor,
    vae: AutoencoderKLFlux2,
) -> torch.Tensor:
    b, c, h, w = latents.shape
    if h % 2 != 0 or w % 2 != 0:
        raise ValueError(f"Expected even FLUX.2 latent grid, got H={h}, W={w}")
    z = latents.reshape(b, c, h // 2, 2, w // 2, 2)
    z = z.permute(0, 1, 3, 5, 2, 4).reshape(b, c * 4, h // 2, w // 2)
    mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device=z.device, dtype=torch.float32)
    var = vae.bn.running_var.view(1, -1, 1, 1).to(device=z.device, dtype=torch.float32)
    std = torch.sqrt(var + float(vae.config.batch_norm_eps))
    return (z.to(torch.float32) - mean) / std


device = "cuda"
flux2 = AutoencoderKLFlux2.from_pretrained(
    "BiliSakura/VAEs",
    subfolder="FLUX2-VAE",
    torch_dtype=torch.bfloat16,
).to(device)
decoder = CapacitorDecoder.from_pretrained(
    "data-archetype/capacitor_decoder",
    device=device,
    dtype=torch.bfloat16,
)

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

with torch.inference_mode():
    posterior = flux2.encode(image.to(device=device, dtype=torch.bfloat16))
    latent_mean = posterior.latent_dist.mean

    # Default path: whiten in float32, then cast back to model dtype before decode.
    latents = flux2_patchify_and_whiten(latent_mean, flux2).to(dtype=torch.bfloat16)
    recon = decoder.decode(
        latents,
        height=int(image.shape[-2]),
        width=int(image.shape[-1]),
        inference_config=CapacitorDecoderInferenceConfig(num_steps=1),
    )

Whitening and dewhitening are optional, but they must stay consistent. The default above matches the usual FLUX.2 pipeline behavior. If your upstream path already gives you raw patchified decoder-space latents instead, skip whitening upstream and call decode(..., latents_are_flux2_whitened=False).

Details

• Default input contract: FLUX.2 patchified latents with FLUX.2 BN whitening still applied.
• Default decoder behavior: unwhiten with saved FLUX.2 BN running stats, then decode.
• Optional raw-latent mode: disable whitening upstream and call decode(..., latents_are_flux2_whitened=False).
• Reused decoder architecture: SemDisDiffAE
• Technical report
• SemDisDiffAE technical report

Citation

@misc{capacitor_decoder,
  title   = {Capacitor Decoder: A Faster, Lighter FLUX.2-Compatible Latent Decoder},
  author  = {data-archetype},
  email   = {data-archetype@proton.me},
  year    = {2026},
  month   = apr,
  url     = {https://huggingface.co/data-archetype/capacitor_decoder},
}