Upload capacitor_decoder private initial export

README.md

@@ -0,0 +1,126 @@
+---
+license: apache-2.0
+tags:
+  - diffusion
+  - autoencoder
+  - image-reconstruction
+  - decoder-only
+  - flux-compatible
+  - pytorch
+---
+
+# data-archetype/capacitor_decoder
+
+**Capacitor decoder**: a faster, lighter FLUX.2-compatible latent decoder built
+on the
+[SemDisDiffAE](https://huggingface.co/data-archetype/semdisdiffae)
+architecture.
+
+## Decode Speed
+
+| Resolution | Speedup vs FLUX.2 | Peak VRAM Reduction | capacitor_decoder (ms/image) | FLUX.2 VAE (ms/image) | capacitor_decoder peak VRAM | FLUX.2 peak VRAM |
+|---:|---:|---:|---:|---:|---:|---:|
+| `512x512` | `1.85x` | `59.3%` | `11.40` | `21.14` | `391.6 MiB` | `961.9 MiB` |
+| `1024x1024` | `3.28x` | `79.1%` | `26.31` | `86.24` | `601.4 MiB` | `2876.4 MiB` |
+| `2048x2048` | `4.70x` | `86.4%` | `86.29` | `405.84` | `1437.4 MiB` | `10531.4 MiB` |
+
+These measurements are decode-only. Each image is first encoded once with the
+same FLUX.2 encoder, latents are cached in memory, and then both decoders are
+timed over the same cached latent set.
+
+## 2k PSNR Benchmark
+
+| Model | Mean PSNR (dB) | Std (dB) | Median (dB) | Min (dB) | P5 (dB) | P95 (dB) | Max (dB) |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| FLUX.2 VAE | 36.2849 | 4.5332 | 36.0728 | 22.7343 | 28.8853 | 43.6343 | 47.3836 |
+| capacitor_decoder | 36.3399 | 4.4980 | 36.2891 | 23.2770 | 29.0597 | 43.6603 | 47.4312 |
+
+| Delta vs FLUX.2 | Mean (dB) | Std (dB) | Median (dB) | Min (dB) | P5 (dB) | P95 (dB) | Max (dB) |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| capacitor_decoder - FLUX.2 | 0.0550 | 0.5308 | 0.0618 | -1.9683 | -0.8108 | 0.8859 | 2.8071 |
+
+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.
+
+## Usage
+
+```python
+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: match the usual FLUX.2 convention.
+    # Whiten here, then let capacitor_decoder unwhiten internally before decode.
+    latents = flux2_patchify_and_whiten(latent_mean, flux2)
+    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](https://huggingface.co/data-archetype/semdisdiffae)
+- [Technical report](technical_report_capacitor_decoder.md)
+- [SemDisDiffAE technical report](https://huggingface.co/data-archetype/semdisdiffae/blob/main/technical_report_semantic.md)
+- [Results viewer](https://huggingface.co/spaces/data-archetype/capacitor_decoder-results)
+
+## Citation
+
+```bibtex
+@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},
+}
+```
+

capacitor_decoder/__init__.py

@@ -0,0 +1,30 @@
+"""capacitor_decoder: standalone decoder-only Capacitor export.
+
+The package exposes a VP-diffusion decoder that consumes FLUX.2 latents.
+
+Usage::
+
+    from capacitor_decoder import CapacitorDecoder, CapacitorDecoderInferenceConfig
+
+    model = CapacitorDecoder.from_pretrained("path/to/export", device="cuda")
+
+    # Default: expects FLUX.2 pipeline-style encoded latents
+    recon = model.decode(latents, height=H, width=W)
+
+    # Bypass unwhitening when latents are already raw decoder-space tokens
+    recon = model.decode(
+        latents,
+        height=H,
+        width=W,
+        latents_are_flux2_whitened=False,
+    )
+"""
+
+from .config import CapacitorDecoderConfig, CapacitorDecoderInferenceConfig
+from .model import CapacitorDecoder
+
+__all__ = [
+    "CapacitorDecoder",
+    "CapacitorDecoderConfig",
+    "CapacitorDecoderInferenceConfig",
+]

capacitor_decoder/adaln.py

@@ -0,0 +1,50 @@
+"""Scale+Gate AdaLN (2-way) for FCDM decoder blocks."""
+
+from __future__ import annotations
+
+from torch import Tensor, nn
+
+
+class AdaLNScaleGateZeroProjector(nn.Module):
+    """Packed 2-way AdaLN projection (SiLU -> Linear), zero-initialized.
+
+    Outputs [B, 2*d_model] packed as (scale, gate).
+    """
+
+    def __init__(self, d_model: int, d_cond: int) -> None:
+        super().__init__()
+        self.d_model: int = int(d_model)
+        self.d_cond: int = int(d_cond)
+        self.act: nn.SiLU = nn.SiLU()
+        self.proj: nn.Linear = nn.Linear(self.d_cond, 2 * self.d_model)
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward_activated(self, act_cond: Tensor) -> Tensor:
+        """Return packed modulation for a pre-activated conditioning vector."""
+        return self.proj(act_cond)
+
+    def forward(self, cond: Tensor) -> Tensor:
+        """Return packed modulation [B, 2*d_model]."""
+        return self.forward_activated(self.act(cond))
+
+
+class AdaLNScaleGateZeroLowRankDelta(nn.Module):
+    """Low-rank delta for 2-way AdaLN: down(d_cond -> rank) -> up(rank -> 2*d_model).
+
+    Zero-initialized up projection preserves zero-output semantics at init.
+    """
+
+    def __init__(self, *, d_model: int, d_cond: int, rank: int) -> None:
+        super().__init__()
+        self.d_model: int = int(d_model)
+        self.d_cond: int = int(d_cond)
+        self.rank: int = int(rank)
+        self.down: nn.Linear = nn.Linear(self.d_cond, self.rank, bias=False)
+        self.up: nn.Linear = nn.Linear(self.rank, 2 * self.d_model, bias=False)
+        nn.init.normal_(self.down.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, act_cond: Tensor) -> Tensor:
+        """Return packed delta modulation [B, 2*d_model]."""
+        return self.up(self.down(act_cond))

capacitor_decoder/config.py

@@ -0,0 +1,65 @@
+"""Frozen architecture and user-tunable inference config for capacitor_decoder."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class CapacitorDecoderConfig:
+    """Frozen architecture config stored alongside exported weights."""
+
+    in_channels: int = 3
+    patch_size: int = 16
+    model_dim: int = 896
+    decoder_depth: int = 8
+    decoder_start_blocks: int = 2
+    decoder_end_blocks: int = 2
+    bottleneck_dim: int = 128
+    mlp_ratio: float = 4.0
+    depthwise_kernel_size: int = 7
+    adaln_low_rank_rank: int = 128
+    logsnr_min: float = -10.0
+    logsnr_max: float = 10.0
+    pixel_noise_std: float = 0.558
+    flux2_batch_norm_eps: float = 1e-5
+
+    @property
+    def latent_channels(self) -> int:
+        """Channel width of the exported latent space."""
+
+        return self.bottleneck_dim
+
+    @property
+    def effective_patch_size(self) -> int:
+        """Effective spatial stride from image to latent grid."""
+
+        return self.patch_size
+
+    def save(self, path: str | Path) -> None:
+        """Save config as JSON."""
+
+        output_path = Path(path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        output_path.write_text(json.dumps(asdict(self), indent=2) + "\n")
+
+    @classmethod
+    def load(cls, path: str | Path) -> CapacitorDecoderConfig:
+        """Load config from JSON."""
+
+        data = json.loads(Path(path).read_text())
+        return cls(**data)
+
+
+@dataclass
+class CapacitorDecoderInferenceConfig:
+    """User-tunable decoder inference parameters."""
+
+    num_steps: int = 1
+    sampler: str = "ddim"
+    schedule: str = "linear"
+    pdg: bool = False
+    pdg_strength: float = 2.0
+    seed: int | None = None

capacitor_decoder/decoder.py

@@ -0,0 +1,168 @@
+"""FCDM DiffAE decoder: skip-concat topology with FCDM blocks and path-drop PDG.
+
+No outer RMSNorms (use_other_outer_rms_norms=False during training):
+norm_in, latent_norm, and norm_out are all absent.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+
+from .adaln import AdaLNScaleGateZeroLowRankDelta, AdaLNScaleGateZeroProjector
+from .fcdm_block import FCDMBlock
+from .straight_through_encoder import Patchify
+from .time_embed import SinusoidalTimeEmbeddingMLP
+
+
+class Decoder(nn.Module):
+    """VP diffusion decoder conditioned on encoder latents and timestep.
+
+    Architecture (skip-concat, 2+4+2 default):
+        Patchify x_t -> Fuse with upsampled z
+        -> Start blocks (2) -> Middle blocks (4) -> Skip fuse -> End blocks (2)
+        -> Conv1x1 -> PixelShuffle
+
+    Path-Drop Guidance (PDG) at inference:
+    - Replace middle block output with ``path_drop_mask_feature`` to create
+      an unconditional prediction, then extrapolate.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        patch_size: int,
+        model_dim: int,
+        depth: int,
+        start_block_count: int,
+        end_block_count: int,
+        bottleneck_dim: int,
+        mlp_ratio: float,
+        depthwise_kernel_size: int,
+        adaln_low_rank_rank: int,
+    ) -> None:
+        super().__init__()
+        self.patch_size = int(patch_size)
+        self.model_dim = int(model_dim)
+
+        # Input processing (no norm_in)
+        self.patchify = Patchify(in_channels, patch_size, model_dim)
+
+        # Latent conditioning path (no latent_norm)
+        self.latent_up = nn.Conv2d(bottleneck_dim, model_dim, kernel_size=1, bias=True)
+        self.fuse_in = nn.Conv2d(2 * model_dim, model_dim, kernel_size=1, bias=True)
+
+        # Time embedding
+        self.time_embed = SinusoidalTimeEmbeddingMLP(model_dim)
+
+        # 2-way AdaLN: shared base projector + per-block low-rank deltas
+        self.adaln_base = AdaLNScaleGateZeroProjector(
+            d_model=model_dim, d_cond=model_dim
+        )
+        self.adaln_deltas = nn.ModuleList(
+            [
+                AdaLNScaleGateZeroLowRankDelta(
+                    d_model=model_dim, d_cond=model_dim, rank=adaln_low_rank_rank
+                )
+                for _ in range(depth)
+            ]
+        )
+
+        # Block layout: start + middle + end
+        middle_count = depth - start_block_count - end_block_count
+        self._middle_start_idx = start_block_count
+        self._end_start_idx = start_block_count + middle_count
+
+        def _make_blocks(count: int) -> nn.ModuleList:
+            return nn.ModuleList(
+                [
+                    FCDMBlock(
+                        model_dim,
+                        mlp_ratio,
+                        depthwise_kernel_size=depthwise_kernel_size,
+                        use_external_adaln=True,
+                    )
+                    for _ in range(count)
+                ]
+            )
+
+        self.start_blocks = _make_blocks(start_block_count)
+        self.middle_blocks = _make_blocks(middle_count)
+        self.fuse_skip = nn.Conv2d(2 * model_dim, model_dim, kernel_size=1, bias=True)
+        self.end_blocks = _make_blocks(end_block_count)
+
+        # Learned mask feature for path-drop PDG
+        self.path_drop_mask_feature = nn.Parameter(torch.zeros((1, model_dim, 1, 1)))
+
+        # Output head (no norm_out)
+        self.out_proj = nn.Conv2d(
+            model_dim, in_channels * (patch_size**2), kernel_size=1, bias=True
+        )
+        self.unpatchify = nn.PixelShuffle(patch_size)
+
+    def _adaln_m_for_layer(self, cond: Tensor, layer_idx: int) -> Tensor:
+        """Compute packed AdaLN modulation = shared_base + per-layer delta."""
+        act = self.adaln_base.act(cond)
+        base_m = self.adaln_base.forward_activated(act)
+        delta_m = self.adaln_deltas[layer_idx](act)
+        return base_m + delta_m
+
+    def _run_blocks(
+        self, blocks: nn.ModuleList, x: Tensor, cond: Tensor, start_index: int
+    ) -> Tensor:
+        """Run a group of decoder blocks with per-block AdaLN modulation."""
+        for local_idx, block in enumerate(blocks):
+            adaln_m = self._adaln_m_for_layer(cond, layer_idx=start_index + local_idx)
+            x = block(x, adaln_m=adaln_m)
+        return x
+
+    def forward(
+        self,
+        x_t: Tensor,
+        t: Tensor,
+        latents: Tensor,
+        *,
+        drop_middle_blocks: bool = False,
+    ) -> Tensor:
+        """Single decoder forward pass.
+
+        Args:
+            x_t: Noised image [B, C, H, W].
+            t: Timestep [B] in [0, 1].
+            latents: Encoder latents [B, bottleneck_dim, h, w].
+            drop_middle_blocks: Replace middle block output with mask feature (PDG).
+
+        Returns:
+            x0 prediction [B, C, H, W].
+        """
+        x_feat = self.patchify(x_t)
+        z_up = self.latent_up(latents)
+
+        fused = torch.cat([x_feat, z_up], dim=1)
+        fused = self.fuse_in(fused)
+
+        cond = self.time_embed(t.to(torch.float32).to(device=x_t.device))
+
+        start_out = self._run_blocks(self.start_blocks, fused, cond, start_index=0)
+
+        if drop_middle_blocks:
+            middle_out = self.path_drop_mask_feature.to(
+                device=x_t.device, dtype=x_t.dtype
+            ).expand_as(start_out)
+        else:
+            middle_out = self._run_blocks(
+                self.middle_blocks,
+                start_out,
+                cond,
+                start_index=self._middle_start_idx,
+            )
+
+        skip_fused = torch.cat([start_out, middle_out], dim=1)
+        skip_fused = self.fuse_skip(skip_fused)
+
+        end_out = self._run_blocks(
+            self.end_blocks, skip_fused, cond, start_index=self._end_start_idx
+        )
+
+        patches = self.out_proj(end_out)
+        return self.unpatchify(patches)

capacitor_decoder/fcdm_block.py

@@ -0,0 +1,103 @@
+"""FCDM block: ConvNeXt-style conv block with GRN and scale+gate AdaLN."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from .norms import ChannelWiseRMSNorm
+
+
+class GRN(nn.Module):
+    """Global Response Normalization for NCHW tensors."""
+
+    def __init__(self, channels: int, *, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps: float = float(eps)
+        c = int(channels)
+        self.gamma = nn.Parameter(torch.zeros((1, c, 1, 1), dtype=torch.float32))
+        self.beta = nn.Parameter(torch.zeros((1, c, 1, 1), dtype=torch.float32))
+
+    def forward(self, x: Tensor) -> Tensor:
+        g = torch.linalg.vector_norm(x, ord=2, dim=(2, 3), keepdim=True)
+        g_fp32 = g.to(dtype=torch.float32)
+        n = (g_fp32 / (g_fp32.mean(dim=1, keepdim=True) + self.eps)).to(dtype=x.dtype)
+        gamma = self.gamma.to(device=x.device, dtype=x.dtype)
+        beta = self.beta.to(device=x.device, dtype=x.dtype)
+        return gamma * (x * n) + beta + x
+
+
+class FCDMBlock(nn.Module):
+    """ConvNeXt-style block with scale+gate AdaLN and GRN.
+
+    Two modes:
+    - Unconditioned (encoder): uses learned layer-scale for near-identity init.
+    - External AdaLN (decoder): receives packed [B, 2*C] modulation (scale, gate).
+      The gate is applied raw (no tanh).
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        mlp_ratio: float,
+        *,
+        depthwise_kernel_size: int = 7,
+        use_external_adaln: bool = False,
+        norm_eps: float = 1e-6,
+        layer_scale_init: float = 1e-3,
+    ) -> None:
+        super().__init__()
+        self.channels: int = int(channels)
+        self.mlp_ratio: float = float(mlp_ratio)
+
+        self.dwconv = nn.Conv2d(
+            channels,
+            channels,
+            kernel_size=depthwise_kernel_size,
+            padding=depthwise_kernel_size // 2,
+            stride=1,
+            groups=channels,
+            bias=True,
+        )
+        self.norm = ChannelWiseRMSNorm(channels, eps=float(norm_eps), affine=False)
+        hidden = max(int(float(channels) * float(mlp_ratio)), 1)
+        self.pwconv1 = nn.Conv2d(channels, hidden, kernel_size=1, bias=True)
+        self.grn = GRN(hidden, eps=1e-6)
+        self.pwconv2 = nn.Conv2d(hidden, channels, kernel_size=1, bias=True)
+
+        if not use_external_adaln:
+            self.layer_scale = nn.Parameter(
+                torch.full((channels,), float(layer_scale_init))
+            )
+        else:
+            self.register_parameter("layer_scale", None)
+
+    def forward(self, x: Tensor, *, adaln_m: Tensor | None = None) -> Tensor:
+        b, c, _, _ = x.shape
+
+        if adaln_m is not None:
+            m = adaln_m.to(device=x.device, dtype=x.dtype)
+            scale, gate = m.chunk(2, dim=-1)
+        else:
+            scale = gate = None
+
+        h = self.dwconv(x)
+        h = self.norm(h)
+
+        if scale is not None:
+            h = h * (1.0 + scale.view(b, c, 1, 1))
+
+        h = self.pwconv1(h)
+        h = F.gelu(h)
+        h = self.grn(h)
+        h = self.pwconv2(h)
+
+        if gate is not None:
+            gate_view = gate.view(b, c, 1, 1)
+        else:
+            gate_view = self.layer_scale.view(1, c, 1, 1).to(  # type: ignore[union-attr]
+                device=h.device, dtype=h.dtype
+            )
+
+        return x + gate_view * h

capacitor_decoder/model.py

@@ -0,0 +1,244 @@
+"""Standalone decoder-only Capacitor model for HuggingFace distribution."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from torch import Tensor, nn
+
+from .config import CapacitorDecoderConfig, CapacitorDecoderInferenceConfig
+from .decoder import Decoder
+from .samplers import run_ddim, run_dpmpp_2m
+from .vp_diffusion import get_schedule, make_initial_state, sample_noise
+
+
+def _resolve_model_dir(
+    path_or_repo_id: str | Path,
+    *,
+    revision: str | None,
+    cache_dir: str | Path | None,
+) -> Path:
+    """Resolve a local path or HuggingFace Hub repo ID to a local directory."""
+
+    local = Path(path_or_repo_id)
+    if local.is_dir():
+        return local
+    repo_id = str(path_or_repo_id)
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError as exc:
+        raise ImportError(
+            f"'{repo_id}' is not an existing local directory. "
+            "To download from HuggingFace Hub, install huggingface_hub: "
+            "pip install huggingface_hub"
+        ) from exc
+    cache_dir_str = str(cache_dir) if cache_dir is not None else None
+    local_dir = snapshot_download(
+        repo_id,
+        revision=revision,
+        cache_dir=cache_dir_str,
+    )
+    return Path(local_dir)
+
+
+class CapacitorDecoder(nn.Module):
+    """Decoder-only Capacitor export that reconstructs from FLUX.2 latents.
+
+    Default input convention matches the public FLUX.2 diffusers pipeline:
+    patchified latents whitened by the VAE BN running statistics. The decoder
+    therefore unwhitens by default before applying the VP-diffusion decode path.
+    """
+
+    def __init__(self, config: CapacitorDecoderConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.register_buffer(
+            "flux2_running_mean",
+            torch.zeros((config.latent_channels,), dtype=torch.float32),
+        )
+        self.register_buffer(
+            "flux2_running_var",
+            torch.ones((config.latent_channels,), dtype=torch.float32),
+        )
+        self.decoder = Decoder(
+            in_channels=config.in_channels,
+            patch_size=config.patch_size,
+            model_dim=config.model_dim,
+            depth=config.decoder_depth,
+            start_block_count=config.decoder_start_blocks,
+            end_block_count=config.decoder_end_blocks,
+            bottleneck_dim=config.bottleneck_dim,
+            mlp_ratio=config.mlp_ratio,
+            depthwise_kernel_size=config.depthwise_kernel_size,
+            adaln_low_rank_rank=config.adaln_low_rank_rank,
+        )
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        path_or_repo_id: str | Path,
+        *,
+        dtype: torch.dtype = torch.bfloat16,
+        device: str | torch.device = "cpu",
+        revision: str | None = None,
+        cache_dir: str | Path | None = None,
+    ) -> CapacitorDecoder:
+        """Load a pretrained model from a local directory or HuggingFace Hub."""
+
+        model_dir = _resolve_model_dir(
+            path_or_repo_id, revision=revision, cache_dir=cache_dir
+        )
+        config = CapacitorDecoderConfig.load(model_dir / "config.json")
+        model = cls(config)
+
+        safetensors_path = model_dir / "model.safetensors"
+        pt_path = model_dir / "model.pt"
+        if safetensors_path.exists():
+            try:
+                from safetensors.torch import load_file
+            except ImportError as exc:
+                raise ImportError(
+                    "safetensors package required to load .safetensors files. "
+                    "Install with: pip install safetensors"
+                ) from exc
+            state_dict = load_file(str(safetensors_path), device=str(device))
+        elif pt_path.exists():
+            state_dict = torch.load(str(pt_path), map_location=device, weights_only=True)
+        else:
+            raise FileNotFoundError(
+                f"No model weights found in {model_dir}. "
+                "Expected model.safetensors or model.pt."
+            )
+
+        model.load_state_dict(state_dict)
+        model = model.to(dtype=dtype, device=torch.device(device))
+        model.eval()
+        return model
+
+    def flux2_norm_stats(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[Tensor, Tensor]:
+        """Return FLUX.2 latent BN running stats as rank-4 tensors."""
+
+        mean = self.flux2_running_mean.view(1, -1, 1, 1).to(device=device, dtype=dtype)
+        var = self.flux2_running_var.view(1, -1, 1, 1).to(device=device, dtype=dtype)
+        std = torch.sqrt(var + float(self.config.flux2_batch_norm_eps))
+        return mean, std
+
+    def unwhiten_flux2_latents(self, latents: Tensor) -> Tensor:
+        """Undo FLUX.2 BN whitening on patchified latents."""
+
+        z_fp32 = latents.to(torch.float32)
+        mean, std = self.flux2_norm_stats(device=z_fp32.device, dtype=torch.float32)
+        return z_fp32 * std + mean
+
+    @torch.no_grad()
+    def decode(
+        self,
+        latents: Tensor,
+        height: int,
+        width: int,
+        *,
+        latents_are_flux2_whitened: bool = True,
+        inference_config: CapacitorDecoderInferenceConfig | None = None,
+    ) -> Tensor:
+        """Decode FLUX.2 latents to images via VP diffusion.
+
+        Args:
+            latents: [B, 128, H/16, W/16] patchified FLUX.2 latents.
+            height: Output image height.
+            width: Output image width.
+            latents_are_flux2_whitened: When True, interpret the input as
+                pipeline-style FLUX.2 encoded latents and unwhiten them using the
+                saved FLUX.2 BN running stats before decode. Set to False when the
+                input is already raw decoder-space latents.
+            inference_config: Optional inference parameters.
+        """
+
+        cfg = inference_config or CapacitorDecoderInferenceConfig()
+        config = self.config
+        if int(latents.shape[1]) != int(config.latent_channels):
+            raise ValueError(
+                "Latent channel count must match exported decoder input width: "
+                f"got C={int(latents.shape[1])}, expected={int(config.latent_channels)}"
+            )
+        if height % int(config.effective_patch_size) != 0 or width % int(
+            config.effective_patch_size
+        ) != 0:
+            raise ValueError(
+                f"height={height} and width={width} must be divisible by "
+                f"effective_patch_size={int(config.effective_patch_size)}"
+            )
+
+        try:
+            model_dtype = next(self.parameters()).dtype
+        except StopIteration:
+            model_dtype = torch.float32
+
+        latents_fp32 = (
+            self.unwhiten_flux2_latents(latents)
+            if bool(latents_are_flux2_whitened)
+            else latents.to(torch.float32)
+        )
+        latents_in = latents_fp32.to(dtype=model_dtype)
+
+        shape = (int(latents.shape[0]), config.in_channels, height, width)
+        noise = sample_noise(
+            shape,
+            noise_std=config.pixel_noise_std,
+            seed=cfg.seed,
+            device=torch.device("cpu"),
+            dtype=torch.float32,
+        )
+
+        schedule = get_schedule(cfg.schedule, cfg.num_steps).to(device=latents.device)
+        initial_state = make_initial_state(
+            noise=noise.to(device=latents.device),
+            t_start=schedule[0:1],
+            logsnr_min=config.logsnr_min,
+            logsnr_max=config.logsnr_max,
+        )
+        device_type = "cuda" if latents.device.type == "cuda" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+
+            def _forward_fn(
+                x_t: Tensor,
+                t: Tensor,
+                condition_latents: Tensor,
+                *,
+                drop_middle_blocks: bool = False,
+                mask_latent_tokens: bool = False,
+            ) -> Tensor:
+                del mask_latent_tokens
+                return self.decoder(
+                    x_t.to(dtype=model_dtype),
+                    t,
+                    condition_latents.to(dtype=model_dtype),
+                    drop_middle_blocks=drop_middle_blocks,
+                )
+
+            pdg_mode = "path_drop" if bool(cfg.pdg) else "disabled"
+            if cfg.sampler == "ddim":
+                sampler_fn = run_ddim
+            elif cfg.sampler == "dpmpp_2m":
+                sampler_fn = run_dpmpp_2m
+            else:
+                raise ValueError(
+                    f"Unsupported sampler: {cfg.sampler!r}. Use 'ddim' or 'dpmpp_2m'."
+                )
+            result = sampler_fn(
+                forward_fn=_forward_fn,
+                initial_state=initial_state,
+                schedule=schedule,
+                latents=latents_in,
+                logsnr_min=config.logsnr_min,
+                logsnr_max=config.logsnr_max,
+                pdg_mode=pdg_mode,
+                pdg_strength=float(cfg.pdg_strength),
+                device=latents.device,
+            )
+        return result.to(torch.float32).clamp(-1.0, 1.0)

capacitor_decoder/norms.py

@@ -0,0 +1,39 @@
+"""Channel-wise RMSNorm for NCHW tensors."""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+
+
+class ChannelWiseRMSNorm(nn.Module):
+    """Channel-wise RMSNorm with float32 reduction for numerical stability.
+
+    Normalizes across channels per spatial position. Supports optional
+    per-channel affine weight and bias.
+    """
+
+    def __init__(self, channels: int, eps: float = 1e-6, affine: bool = True) -> None:
+        super().__init__()
+        self.channels: int = int(channels)
+        self._eps: float = float(eps)
+        if affine:
+            self.weight = nn.Parameter(torch.ones(self.channels))
+            self.bias = nn.Parameter(torch.zeros(self.channels))
+        else:
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if x.dim() < 2:
+            return x
+        # Float32 accumulation for stability
+        ms = torch.mean(torch.square(x), dim=1, keepdim=True, dtype=torch.float32)
+        inv_rms = torch.rsqrt(ms + self._eps)
+        y = x * inv_rms.to(dtype=x.dtype)
+        if self.weight is not None:
+            shape = (1, -1) + (1,) * (x.dim() - 2)
+            y = y * self.weight.view(shape).to(dtype=x.dtype)
+            if self.bias is not None:
+                y = y + self.bias.view(shape).to(dtype=x.dtype)
+        return y

capacitor_decoder/samplers.py

@@ -0,0 +1,255 @@
+"""DDIM and DPM++2M samplers for VP diffusion with path-drop PDG support."""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+import torch
+from torch import Tensor
+
+from .vp_diffusion import (
+    alpha_sigma_from_logsnr,
+    broadcast_time_like,
+    shifted_cosine_interpolated_logsnr_from_t,
+)
+
+
+class DecoderForwardFn(Protocol):
+    """Callable that predicts x0 from (x_t, t, latents) with path-drop PDG flag."""
+
+    def __call__(
+        self,
+        x_t: Tensor,
+        t: Tensor,
+        latents: Tensor,
+        *,
+        drop_middle_blocks: bool = False,
+        mask_latent_tokens: bool = False,
+    ) -> Tensor: ...
+
+
+def _reconstruct_eps_from_x0(
+    *, x_t: Tensor, x0_hat: Tensor, alpha: Tensor, sigma: Tensor
+) -> Tensor:
+    """Reconstruct eps_hat from (x_t, x0_hat) under VP parameterization.
+
+    eps_hat = (x_t - alpha * x0_hat) / sigma. All float32.
+    """
+    alpha_view = broadcast_time_like(alpha, x_t).to(dtype=torch.float32)
+    sigma_view = broadcast_time_like(sigma, x_t).to(dtype=torch.float32)
+    x_t_f32 = x_t.to(torch.float32)
+    x0_f32 = x0_hat.to(torch.float32)
+    return (x_t_f32 - alpha_view * x0_f32) / sigma_view
+
+
+def _ddim_step(
+    *,
+    x0_hat: Tensor,
+    eps_hat: Tensor,
+    alpha_next: Tensor,
+    sigma_next: Tensor,
+    ref: Tensor,
+) -> Tensor:
+    """DDIM step: x_next = alpha_next * x0_hat + sigma_next * eps_hat."""
+    a = broadcast_time_like(alpha_next, ref).to(dtype=torch.float32)
+    s = broadcast_time_like(sigma_next, ref).to(dtype=torch.float32)
+    return a * x0_hat + s * eps_hat
+
+
+def _predict_with_pdg(
+    forward_fn: DecoderForwardFn,
+    state: Tensor,
+    t_vec: Tensor,
+    latents: Tensor,
+    *,
+    pdg_mode: str,
+    pdg_strength: float,
+) -> Tensor:
+    """Run decoder forward with optional PDG guidance.
+
+    Args:
+        forward_fn: Decoder forward function.
+        state: Current noised state [B, C, H, W].
+        t_vec: Timestep vector [B].
+        latents: Encoder latents.
+        pdg_mode: "disabled" or "path_drop".
+        pdg_strength: CFG-like strength for PDG.
+
+    Returns:
+        x0_hat prediction in float32.
+    """
+    if pdg_mode == "path_drop":
+        x0_uncond = forward_fn(state, t_vec, latents, drop_middle_blocks=True).to(
+            torch.float32
+        )
+        x0_cond = forward_fn(state, t_vec, latents, drop_middle_blocks=False).to(
+            torch.float32
+        )
+        return x0_uncond + pdg_strength * (x0_cond - x0_uncond)
+    else:
+        return forward_fn(state, t_vec, latents, drop_middle_blocks=False).to(
+            torch.float32
+        )
+
+
+def run_ddim(
+    *,
+    forward_fn: DecoderForwardFn,
+    initial_state: Tensor,
+    schedule: Tensor,
+    latents: Tensor,
+    logsnr_min: float,
+    logsnr_max: float,
+    log_change_high: float = 0.0,
+    log_change_low: float = 0.0,
+    pdg_mode: str = "disabled",
+    pdg_strength: float = 1.5,
+    device: torch.device | None = None,
+) -> Tensor:
+    """Run DDIM sampling loop with path-drop PDG support.
+
+    Args:
+        forward_fn: Decoder forward function (x_t, t, latents) -> x0_hat.
+        initial_state: Starting noised state [B, C, H, W] in float32.
+        schedule: Descending t-schedule [num_steps] in [0, 1].
+        latents: Encoder latents [B, bottleneck_dim, h, w].
+        logsnr_min, logsnr_max: VP schedule endpoints.
+        log_change_high, log_change_low: Shifted-cosine schedule parameters.
+        pdg_mode: "disabled" or "path_drop".
+        pdg_strength: CFG-like strength for PDG.
+        device: Target device.
+
+    Returns:
+        Denoised samples [B, C, H, W] in float32.
+    """
+    run_device = device or initial_state.device
+    batch_size = int(initial_state.shape[0])
+    state = initial_state.to(device=run_device, dtype=torch.float32)
+
+    # Precompute logSNR, alpha, sigma for all schedule points
+    lmb = shifted_cosine_interpolated_logsnr_from_t(
+        schedule.to(device=run_device),
+        logsnr_min=logsnr_min,
+        logsnr_max=logsnr_max,
+        log_change_high=log_change_high,
+        log_change_low=log_change_low,
+    )
+    alpha_sched, sigma_sched = alpha_sigma_from_logsnr(lmb)
+
+    for i in range(int(schedule.numel()) - 1):
+        t_i = schedule[i]
+        a_t = alpha_sched[i].expand(batch_size)
+        s_t = sigma_sched[i].expand(batch_size)
+        a_next = alpha_sched[i + 1].expand(batch_size)
+        s_next = sigma_sched[i + 1].expand(batch_size)
+
+        # Model prediction with optional PDG
+        t_vec = t_i.expand(batch_size).to(device=run_device, dtype=torch.float32)
+        x0_hat = _predict_with_pdg(
+            forward_fn,
+            state,
+            t_vec,
+            latents,
+            pdg_mode=pdg_mode,
+            pdg_strength=pdg_strength,
+        )
+
+        eps_hat = _reconstruct_eps_from_x0(
+            x_t=state, x0_hat=x0_hat, alpha=a_t, sigma=s_t
+        )
+        state = _ddim_step(
+            x0_hat=x0_hat,
+            eps_hat=eps_hat,
+            alpha_next=a_next,
+            sigma_next=s_next,
+            ref=state,
+        )
+
+    return state
+
+
+def run_dpmpp_2m(
+    *,
+    forward_fn: DecoderForwardFn,
+    initial_state: Tensor,
+    schedule: Tensor,
+    latents: Tensor,
+    logsnr_min: float,
+    logsnr_max: float,
+    log_change_high: float = 0.0,
+    log_change_low: float = 0.0,
+    pdg_mode: str = "disabled",
+    pdg_strength: float = 1.5,
+    device: torch.device | None = None,
+) -> Tensor:
+    """Run DPM++2M sampling loop with path-drop PDG support.
+
+    Multi-step solver using exponential integrator formulation in half-lambda space.
+    """
+    run_device = device or initial_state.device
+    batch_size = int(initial_state.shape[0])
+    state = initial_state.to(device=run_device, dtype=torch.float32)
+
+    # Precompute logSNR, alpha, sigma, half-lambda for all schedule points
+    lmb = shifted_cosine_interpolated_logsnr_from_t(
+        schedule.to(device=run_device),
+        logsnr_min=logsnr_min,
+        logsnr_max=logsnr_max,
+        log_change_high=log_change_high,
+        log_change_low=log_change_low,
+    )
+    alpha_sched, sigma_sched = alpha_sigma_from_logsnr(lmb)
+    half_lambda = 0.5 * lmb.to(torch.float32)
+
+    x0_prev: Tensor | None = None
+
+    for i in range(int(schedule.numel()) - 1):
+        t_i = schedule[i]
+        s_t = sigma_sched[i].expand(batch_size)
+        a_next = alpha_sched[i + 1].expand(batch_size)
+        s_next = sigma_sched[i + 1].expand(batch_size)
+
+        # Model prediction with optional PDG
+        t_vec = t_i.expand(batch_size).to(device=run_device, dtype=torch.float32)
+        x0_hat = _predict_with_pdg(
+            forward_fn,
+            state,
+            t_vec,
+            latents,
+            pdg_mode=pdg_mode,
+            pdg_strength=pdg_strength,
+        )
+
+        lam_t = half_lambda[i].expand(batch_size)
+        lam_next = half_lambda[i + 1].expand(batch_size)
+        h = (lam_next - lam_t).to(torch.float32)
+        phi_1 = torch.expm1(-h)
+
+        sigma_ratio = (s_next / s_t).to(torch.float32)
+
+        if i == 0 or x0_prev is None:
+            # First-order step
+            state = (
+                sigma_ratio.view(-1, *([1] * (state.dim() - 1))) * state
+                - broadcast_time_like(a_next, state).to(torch.float32)
+                * broadcast_time_like(phi_1, state).to(torch.float32)
+                * x0_hat
+            )
+        else:
+            # Second-order step
+            lam_prev = half_lambda[i - 1].expand(batch_size)
+            h_0 = (lam_t - lam_prev).to(torch.float32)
+            r0 = h_0 / h
+            d1_0 = (x0_hat - x0_prev) / broadcast_time_like(r0, x0_hat)
+            common = broadcast_time_like(a_next, state).to(
+                torch.float32
+            ) * broadcast_time_like(phi_1, state).to(torch.float32)
+            state = (
+                sigma_ratio.view(-1, *([1] * (state.dim() - 1))) * state
+                - common * x0_hat
+                - 0.5 * common * d1_0
+            )
+
+        x0_prev = x0_hat
+
+    return state

capacitor_decoder/straight_through_encoder.py

@@ -0,0 +1,27 @@
+"""PixelUnshuffle-based patchifier (no residual conv path)."""
+
+from __future__ import annotations
+
+from torch import Tensor, nn
+
+
+class Patchify(nn.Module):
+    """PixelUnshuffle(patch) -> Conv2d 1x1 projection.
+
+    Converts [B, C, H, W] images into [B, out_channels, H/patch, W/patch] features.
+    """
+
+    def __init__(self, in_channels: int, patch: int, out_channels: int) -> None:
+        super().__init__()
+        self.patch = int(patch)
+        self.unshuffle = nn.PixelUnshuffle(self.patch)
+        in_after = in_channels * (self.patch * self.patch)
+        self.proj = nn.Conv2d(in_after, out_channels, kernel_size=1, bias=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if x.shape[2] % self.patch != 0 or x.shape[3] % self.patch != 0:
+            raise ValueError(
+                f"Input H={x.shape[2]} and W={x.shape[3]} must be divisible by patch={self.patch}"
+            )
+        y = self.unshuffle(x)
+        return self.proj(y)

capacitor_decoder/time_embed.py

@@ -0,0 +1,83 @@
+"""Sinusoidal timestep embedding with MLP projection."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor, nn
+
+
+def _log_spaced_frequencies(
+    half: int, max_period: float, *, device: torch.device | None = None
+) -> Tensor:
+    """Log-spaced frequencies for sinusoidal embedding."""
+    return torch.exp(
+        -math.log(max_period)
+        * torch.arange(half, device=device, dtype=torch.float32)
+        / max(float(half - 1), 1.0)
+    )
+
+
+def sinusoidal_time_embedding(
+    t: Tensor,
+    dim: int,
+    *,
+    max_period: float = 10000.0,
+    scale: float | None = None,
+    freqs: Tensor | None = None,
+) -> Tensor:
+    """Sinusoidal timestep embedding (DDPM/DiT-style). Always float32."""
+    t32 = t.to(torch.float32)
+    if scale is not None:
+        t32 = t32 * float(scale)
+    half = dim // 2
+    if freqs is not None:
+        freqs = freqs.to(device=t32.device, dtype=torch.float32)
+    else:
+        freqs = _log_spaced_frequencies(half, max_period, device=t32.device)
+    angles = t32[:, None] * freqs[None, :]
+    return torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)
+
+
+class SinusoidalTimeEmbeddingMLP(nn.Module):
+    """Sinusoidal time embedding followed by Linear -> SiLU -> Linear."""
+
+    def __init__(
+        self,
+        dim: int,
+        *,
+        freq_dim: int = 256,
+        hidden_mult: float = 1.0,
+        time_scale: float = 1000.0,
+        max_period: float = 10000.0,
+    ) -> None:
+        super().__init__()
+        self.dim = int(dim)
+        self.freq_dim = int(freq_dim)
+        self.time_scale = float(time_scale)
+        self.max_period = float(max_period)
+        hidden_dim = max(int(round(int(dim) * float(hidden_mult))), 1)
+
+        freqs = _log_spaced_frequencies(self.freq_dim // 2, self.max_period)
+        self.register_buffer("freqs", freqs, persistent=True)
+
+        self.proj_in = nn.Linear(self.freq_dim, hidden_dim)
+        self.act = nn.SiLU()
+        self.proj_out = nn.Linear(hidden_dim, self.dim)
+
+    def forward(self, t: Tensor) -> Tensor:
+        freqs: Tensor = self.freqs  # type: ignore[assignment]
+        emb_freq = sinusoidal_time_embedding(
+            t.to(torch.float32),
+            self.freq_dim,
+            max_period=self.max_period,
+            scale=self.time_scale,
+            freqs=freqs,
+        )
+        dtype_in = self.proj_in.weight.dtype
+        hidden = self.proj_in(emb_freq.to(dtype_in))
+        hidden = self.act(hidden)
+        if hidden.dtype != self.proj_out.weight.dtype:
+            hidden = hidden.to(self.proj_out.weight.dtype)
+        return self.proj_out(hidden)

capacitor_decoder/vp_diffusion.py

@@ -0,0 +1,151 @@
+"""VP diffusion math: logSNR schedules, alpha/sigma computation, noise construction."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor
+
+
+def alpha_sigma_from_logsnr(lmb: Tensor) -> tuple[Tensor, Tensor]:
+    """Compute (alpha, sigma) from logSNR in float32.
+
+    VP constraint: alpha^2 + sigma^2 = 1.
+    """
+    lmb32 = lmb.to(dtype=torch.float32)
+    alpha = torch.sqrt(torch.sigmoid(lmb32))
+    sigma = torch.sqrt(torch.sigmoid(-lmb32))
+    return alpha, sigma
+
+
+def broadcast_time_like(coeff: Tensor, x: Tensor) -> Tensor:
+    """Broadcast [B] coefficient to match x for per-sample scaling."""
+    view_shape = (int(x.shape[0]),) + (1,) * (x.dim() - 1)
+    return coeff.view(view_shape)
+
+
+def _cosine_interpolated_params(
+    logsnr_min: float, logsnr_max: float
+) -> tuple[float, float]:
+    """Compute (a, b) for cosine-interpolated logSNR schedule.
+
+    logsnr(t) = -2 * log(tan(a*t + b))
+    logsnr(0) = logsnr_max, logsnr(1) = logsnr_min
+    """
+    b = math.atan(math.exp(-0.5 * logsnr_max))
+    a = math.atan(math.exp(-0.5 * logsnr_min)) - b
+    return a, b
+
+
+def cosine_interpolated_logsnr_from_t(
+    t: Tensor, *, logsnr_min: float, logsnr_max: float
+) -> Tensor:
+    """Map t in [0,1] to logSNR via cosine-interpolated schedule. Always float32."""
+    a, b = _cosine_interpolated_params(logsnr_min, logsnr_max)
+    t32 = t.to(dtype=torch.float32)
+    a_t = torch.tensor(a, device=t32.device, dtype=torch.float32)
+    b_t = torch.tensor(b, device=t32.device, dtype=torch.float32)
+    u = a_t * t32 + b_t
+    return -2.0 * torch.log(torch.tan(u))
+
+
+def shifted_cosine_interpolated_logsnr_from_t(
+    t: Tensor,
+    *,
+    logsnr_min: float,
+    logsnr_max: float,
+    log_change_high: float = 0.0,
+    log_change_low: float = 0.0,
+) -> Tensor:
+    """SiD2 "shifted cosine" schedule: logSNR with resolution-dependent shifts.
+
+    lambda(t) = (1-t) * (base(t) + log_change_high) + t * (base(t) + log_change_low)
+    """
+    base = cosine_interpolated_logsnr_from_t(
+        t, logsnr_min=logsnr_min, logsnr_max=logsnr_max
+    )
+    t32 = t.to(dtype=torch.float32)
+    high = base + float(log_change_high)
+    low = base + float(log_change_low)
+    return (1.0 - t32) * high + t32 * low
+
+
+def get_schedule(schedule_type: str, num_steps: int) -> Tensor:
+    """Generate a descending t-schedule in [0, 1] for VP diffusion sampling.
+
+    ``num_steps`` is the number of function evaluations (NFE = decoder forward
+    passes).  Internally the schedule has ``num_steps + 1`` time points
+    (including both endpoints).
+
+    Args:
+        schedule_type: "linear" or "cosine".
+        num_steps: Number of decoder forward passes (NFE), >= 1.
+
+    Returns:
+        Descending 1D tensor with ``num_steps + 1`` elements from ~1.0 to ~0.0.
+    """
+    # NOTE: the upstream training code (src/ode/time_schedules.py) uses a
+    # different convention where num_steps counts schedule *points* (so NFE =
+    # num_steps - 1).  This export package corrects the off-by-one so that
+    # num_steps means NFE directly.  TODO: align the upstream convention.
+    n = max(int(num_steps) + 1, 2)
+    if schedule_type == "linear":
+        base = torch.linspace(0.0, 1.0, n)
+    elif schedule_type == "cosine":
+        i = torch.arange(n, dtype=torch.float32)
+        base = 0.5 * (1.0 - torch.cos(math.pi * (i / (n - 1))))
+    else:
+        raise ValueError(
+            f"Unsupported schedule type: {schedule_type!r}. Use 'linear' or 'cosine'."
+        )
+    # Descending: high t (noisy) -> low t (clean)
+    return torch.flip(base, dims=[0])
+
+
+def make_initial_state(
+    *,
+    noise: Tensor,
+    t_start: Tensor,
+    logsnr_min: float,
+    logsnr_max: float,
+    log_change_high: float = 0.0,
+    log_change_low: float = 0.0,
+) -> Tensor:
+    """Construct VP initial state x_t0 = sigma_start * noise (since x0=0).
+
+    All math in float32.
+    """
+    batch = int(noise.shape[0])
+    lmb_start = shifted_cosine_interpolated_logsnr_from_t(
+        t_start.expand(batch).to(dtype=torch.float32),
+        logsnr_min=logsnr_min,
+        logsnr_max=logsnr_max,
+        log_change_high=log_change_high,
+        log_change_low=log_change_low,
+    )
+    _alpha_start, sigma_start = alpha_sigma_from_logsnr(lmb_start)
+    sigma_view = broadcast_time_like(sigma_start, noise)
+    return sigma_view * noise.to(dtype=torch.float32)
+
+
+def sample_noise(
+    shape: tuple[int, ...],
+    *,
+    noise_std: float = 1.0,
+    seed: int | None = None,
+    device: torch.device | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> Tensor:
+    """Sample Gaussian noise with optional seeding. CPU-seeded for reproducibility."""
+    if seed is None:
+        noise = torch.randn(
+            shape, device=device or torch.device("cpu"), dtype=torch.float32
+        )
+    else:
+        gen = torch.Generator(device="cpu")
+        gen.manual_seed(int(seed))
+        noise = torch.randn(shape, generator=gen, device="cpu", dtype=torch.float32)
+    noise = noise.mul(float(noise_std))
+    target_device = device if device is not None else torch.device("cpu")
+    return noise.to(device=target_device, dtype=dtype)

config.json

@@ -0,0 +1,16 @@
+{
+  "in_channels": 3,
+  "patch_size": 16,
+  "model_dim": 896,
+  "decoder_depth": 8,
+  "decoder_start_blocks": 2,
+  "decoder_end_blocks": 2,
+  "bottleneck_dim": 128,
+  "mlp_ratio": 4.0,
+  "depthwise_kernel_size": 7,
+  "adaln_low_rank_rank": 128,
+  "logsnr_min": -10.0,
+  "logsnr_max": 10.0,
+  "pixel_noise_std": 0.558,
+  "flux2_batch_norm_eps": 0.0001
+}

model.safetensors

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

technical_report_capacitor_decoder.md

@@ -0,0 +1,77 @@
+# data-archetype/capacitor_decoder — Technical Report
+
+**Capacitor decoder** is a FLUX.2-compatible latent decoder built on the
+[SemDisDiffAE](https://huggingface.co/data-archetype/semdisdiffae)
+architecture.
+
+This document only covers the export contract.
+
+- Model card: https://huggingface.co/data-archetype/capacitor_decoder
+- SemDisDiffAE technical report: https://huggingface.co/data-archetype/semdisdiffae/blob/main/technical_report_semantic.md
+
+## 1. Latent Interface
+
+The exported runtime defaults to the usual FLUX.2 pipeline convention:
+FLUX.2 patchified latents with FLUX.2 BN whitening still applied. The decoder
+first unwhitens those latents using the saved FLUX.2 BN running statistics and
+then decodes them.
+
+## 2. Two Supported Input Modes
+
+There are two valid ways to call the export:
+
+1. **Default mode**: upstream produces FLUX.2 patchified + whitened latents.
+   Use `decode(latents, ..., latents_are_flux2_whitened=True)`.
+2. **Raw-latent mode**: upstream produces raw patchified decoder-space latents
+   with no FLUX.2 BN whitening. Use
+   `decode(latents, ..., latents_are_flux2_whitened=False)`.
+
+The important invariant is that whitening must be handled consistently on both
+sides. If whitening is enabled upstream, keep the decoder default. If whitening
+is disabled upstream, disable dewhitening in the decoder too.
+
+## 3. Training
+
+This export corresponds to roughly **300k training steps**. The saved run
+configuration uses:
+
+| Parameter | Value |
+|---|---|
+| Optimizer | AdamW |
+| AdamW betas | `(0.9, 0.99)` |
+| AdamW epsilon | `1e-8` |
+| Weight decay | `0.0` |
+| Learning rate | `1e-4` |
+| LR schedule | constant after warmup |
+| Warmup steps | `2,000` |
+| Batch size | `128` |
+| Gradient accumulation | `1` |
+| Gradient clip | `1.0` max norm |
+| Precision | AMP bfloat16 |
+| FP32 matmul precision | TF32 |
+| EMA decay | `0.9995` |
+| EMA dtype | FP32 |
+| EMA update cadence | every step |
+| Compilation | `torch.compile` enabled |
+| Validation / checkpoint cadence | every `1,000` steps |
+
+## 4. Links
+
+- [SemDisDiffAE](https://huggingface.co/data-archetype/semdisdiffae)
+- [This model card](https://huggingface.co/data-archetype/capacitor_decoder)
+- [SemDisDiffAE technical report](https://huggingface.co/data-archetype/semdisdiffae/blob/main/technical_report_semantic.md)
+- [Results viewer](https://huggingface.co/spaces/data-archetype/capacitor_decoder-results)
+
+## Citation
+
+```bibtex
+@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},
+}
+```
+