Upload folder using huggingface_hub

README.md

@@ -0,0 +1,109 @@
+---
+license: apache-2.0
+tags:
+  - diffusion
+  - autoencoder
+  - image-reconstruction
+  - image-tokenizer
+  - pytorch
+  - fcdm
+  - semantic-alignment
+library_name: capacitor_diffae
+---
+
+# data-archetype/semdisdiffae
+
+**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 v2, 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.
+
+## Key Features
+
+- **Fast**: ~3 ms/img encode, ~6 ms/img decode (1 step) on H200 — 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 v2
+- **Pure convolutional**: no attention in encoder/decoder, O(n) in spatial resolution
+- **VP diffusion decoder**: single-step DDIM for PSNR-optimal, 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 capacitor_diffae import CapacitorDiffAE, CapacitorDiffAEInferenceConfig
+
+model = CapacitorDiffAE.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 = CapacitorDiffAEInferenceConfig(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)
+
+## License
+
+Apache 2.0

_results_appendix_semantic.md

@@ -0,0 +1,72 @@
+## 7. Results
+
+Reconstruction quality evaluated on a curated set of test images covering photographs, book covers, and documents. Flux.1 VAE (patch 8, 16 channels) is included as a reference at the same 12x compression ratio as the c64 variant.
+
+### 7.1 Interactive Viewer
+
+**[Open full-resolution comparison viewer](https://huggingface.co/spaces/data-archetype/irdiffae-results)** — side-by-side reconstructions, RGB deltas, and latent PCA with adjustable image size.
+
+### 7.2 Inference Settings
+
+| Setting | Value |
+|---------|-------|
+| Sampler | ddim |
+| Steps | 1 |
+| Schedule | linear |
+| Seed | 42 |
+| PDG | no_path_dropg |
+| Batch size (timing) | 4 |
+
+> All models run in bfloat16. Timings measured on an NVIDIA RTX Pro 6000 (Blackwell).
+
+### 7.3 Global Metrics
+
+| Metric | semdisdiffae (1 step) | Flux.2 VAE |
+|--------|--------|--------|
+| Avg PSNR (dB) | 35.78 | 34.16 |
+| Avg encode (ms/image) | 2.5 | 46.1 |
+| Avg decode (ms/image) | 5.5 | 91.8 |
+
+### 7.4 Per-Image PSNR (dB)
+
+| Image | semdisdiffae (1 step) | Flux.2 VAE |
+|-------|--------|--------|
+| p640x1536:94623 | 35.44 | 33.50 |
+| p640x1536:94624 | 31.33 | 30.03 |
+| p640x1536:94625 | 35.05 | 33.98 |
+| p640x1536:94626 | 33.21 | 31.53 |
+| p640x1536:94627 | 32.54 | 30.53 |
+| p640x1536:94628 | 29.80 | 28.88 |
+| p960x1024:216264 | 46.37 | 45.39 |
+| p960x1024:216265 | 29.70 | 27.80 |
+| p960x1024:216266 | 47.15 | 46.20 |
+| p960x1024:216267 | 40.99 | 39.23 |
+| p960x1024:216268 | 38.47 | 36.13 |
+| p960x1024:216269 | 32.74 | 30.24 |
+| p960x1024:216270 | 36.23 | 34.18 |
+| p960x1024:216271 | 44.41 | 42.18 |
+| p704x1472:94699 | 43.80 | 41.79 |
+| p704x1472:94700 | 32.83 | 32.08 |
+| p704x1472:94701 | 39.00 | 37.90 |
+| p704x1472:94702 | 34.52 | 32.50 |
+| p704x1472:94703 | 32.81 | 31.35 |
+| p704x1472:94704 | 33.38 | 31.84 |
+| p704x1472:94705 | 39.70 | 37.44 |
+| p704x1472:94706 | 35.12 | 33.66 |
+| r256_p1344x704:15577 | 31.02 | 29.98 |
+| r256_p1344x704:15578 | 32.38 | 30.79 |
+| r256_p1344x704:15579 | 33.27 | 31.83 |
+| r256_p1344x704:15580 | 37.84 | 36.03 |
+| r256_p1344x704:15581 | 38.57 | 36.94 |
+| r256_p1344x704:15582 | 33.41 | 32.10 |
+| r256_p1344x704:15583 | 36.67 | 34.54 |
+| r256_p1344x704:15584 | 33.23 | 31.76 |
+| r256_p896x1152:144131 | 35.30 | 33.60 |
+| r256_p896x1152:144132 | 36.99 | 35.32 |
+| r256_p896x1152:144133 | 39.69 | 37.33 |
+| r256_p896x1152:144134 | 36.01 | 34.47 |
+| r256_p896x1152:144135 | 31.20 | 29.87 |
+| r256_p896x1152:144136 | 37.51 | 35.68 |
+| r256_p896x1152:144137 | 33.83 | 32.86 |
+| r256_p896x1152:144138 | 27.39 | 25.63 |
+| VAE_accuracy_test_image | 36.64 | 35.25 |

capacitor_diffae/__init__.py

@@ -0,0 +1,33 @@
+"""CapacitorDiffAE: Standalone diffusion autoencoder with FCDM blocks.
+
+Capacitor DiffAE — a fast diffusion autoencoder with a 128-channel spatial
+bottleneck and a VP-parameterized diagonal Gaussian posterior. Built on FCDM
+(Fully Convolutional Diffusion Model) blocks with GRN and scale+gate AdaLN.
+
+Usage::
+
+    from capacitor_diffae import CapacitorDiffAE, CapacitorDiffAEInferenceConfig
+
+    model = CapacitorDiffAE.from_pretrained("path/to/weights", device="cuda")
+
+    # Encode (returns posterior mode by default)
+    latents = model.encode(images)  # images: [B,3,H,W] in [-1,1]
+
+    # Decode — PSNR-optimal (1 step, default)
+    recon = model.decode(latents, height=H, width=W)
+
+    # Decode — perceptual sharpness (10 steps + path-drop PDG)
+    cfg = CapacitorDiffAEInferenceConfig(num_steps=10, pdg=True, pdg_strength=2.0)
+    recon = model.decode(latents, height=H, width=W, inference_config=cfg)
+"""
+
+from .config import CapacitorDiffAEConfig, CapacitorDiffAEInferenceConfig
+from .encoder import EncoderPosterior
+from .model import CapacitorDiffAE
+
+__all__ = [
+    "CapacitorDiffAE",
+    "CapacitorDiffAEConfig",
+    "CapacitorDiffAEInferenceConfig",
+    "EncoderPosterior",
+]

capacitor_diffae/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_diffae/config.py

@@ -0,0 +1,62 @@
+"""Frozen model architecture and user-tunable inference configuration."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class CapacitorDiffAEConfig:
+    """Frozen model architecture config. Stored alongside weights as config.json."""
+
+    in_channels: int = 3
+    patch_size: int = 16
+    model_dim: int = 896
+    encoder_depth: int = 4
+    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
+    # Encoder posterior kind: "diagonal_gaussian" or "deterministic"
+    bottleneck_posterior_kind: str = "diagonal_gaussian"
+    # Post-bottleneck normalization: "channel_wise" or "disabled"
+    bottleneck_norm_mode: str = "disabled"
+    # VP diffusion schedule endpoints
+    logsnr_min: float = -10.0
+    logsnr_max: float = 10.0
+    # Pixel-space noise std for VP diffusion initialization
+    pixel_noise_std: float = 0.558
+
+    def save(self, path: str | Path) -> None:
+        """Save config as JSON."""
+        p = Path(path)
+        p.parent.mkdir(parents=True, exist_ok=True)
+        p.write_text(json.dumps(asdict(self), indent=2) + "\n")
+
+    @classmethod
+    def load(cls, path: str | Path) -> CapacitorDiffAEConfig:
+        """Load config from JSON."""
+        data = json.loads(Path(path).read_text())
+        return cls(**data)
+
+
+@dataclass
+class CapacitorDiffAEInferenceConfig:
+    """User-tunable inference parameters with sensible defaults.
+
+    PDG (Path-Drop Guidance) sharpens reconstructions by degrading conditioning
+    in one pass and amplifying the difference. When enabled, uses 2 NFE per step.
+    Recommended: ``pdg=True, pdg_strength=2.0, num_steps=10``.
+    """
+
+    num_steps: int = 1  # number of denoising steps (NFE)
+    sampler: str = "ddim"  # "ddim" or "dpmpp_2m"
+    schedule: str = "linear"  # "linear" or "cosine"
+    pdg: bool = False  # enable PDG for perceptual sharpening
+    pdg_strength: float = 2.0  # CFG-like strength when pdg=True
+    seed: int | None = None

capacitor_diffae/decoder.py

@@ -0,0 +1,169 @@
+"""Capacitor decoder: skip-concat topology with FCDM blocks and dual 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
+
+    Dual PDG at inference:
+    - Path drop: replace middle block output with ``path_drop_mask_feature``.
+    - Token mask: replace a fraction of upsampled latent tokens with
+      ``latent_mask_feature`` before fusion.
+    """
+
+    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_diffae/encoder.py

@@ -0,0 +1,129 @@
+"""Capacitor encoder: patchify -> FCDMBlocks -> diagonal Gaussian posterior.
+
+No input RMSNorm (use_other_outer_rms_norms=False during training).
+Post-bottleneck RMSNorm (affine=False) on the mean branch.
+Encoder outputs posterior mode by default: alpha * RMSNorm(mean).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+
+from .fcdm_block import FCDMBlock
+from .norms import ChannelWiseRMSNorm
+from .straight_through_encoder import Patchify
+
+
+@dataclass(frozen=True)
+class EncoderPosterior:
+    """VP-parameterized diagonal Gaussian posterior.
+
+    mean: Clean signal branch mu [B, bottleneck_dim, h, w]
+    logsnr: Per-element log signal-to-noise ratio [B, bottleneck_dim, h, w]
+    """
+
+    mean: Tensor
+    logsnr: Tensor
+
+    @property
+    def alpha(self) -> Tensor:
+        """VP signal coefficient: sqrt(sigmoid(logsnr))."""
+        return torch.sigmoid(self.logsnr).sqrt()
+
+    @property
+    def sigma(self) -> Tensor:
+        """VP noise coefficient: sqrt(sigmoid(-logsnr))."""
+        return torch.sigmoid(-self.logsnr).sqrt()
+
+    def mode(self) -> Tensor:
+        """Posterior mode in token space: alpha * mean."""
+        return self.alpha.to(dtype=self.mean.dtype) * self.mean
+
+    def sample(self, *, generator: torch.Generator | None = None) -> Tensor:
+        """Sample from posterior: alpha * mean + sigma * eps."""
+        eps = torch.randn_like(self.mean, generator=generator)  # type: ignore[call-overload]
+        alpha = self.alpha.to(dtype=self.mean.dtype)
+        sigma = self.sigma.to(dtype=self.mean.dtype)
+        return alpha * self.mean + sigma * eps
+
+
+class Encoder(nn.Module):
+    """Encoder: Image [B,3,H,W] -> latents [B,bottleneck_dim,h,w].
+
+    With diagonal_gaussian posterior, the to_bottleneck projection outputs
+    2 * bottleneck_dim channels, split into mean and logsnr. The default
+    encode() returns the posterior mode: alpha * RMSNorm(mean).
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        patch_size: int,
+        model_dim: int,
+        depth: int,
+        bottleneck_dim: int,
+        mlp_ratio: float,
+        depthwise_kernel_size: int,
+        bottleneck_posterior_kind: str = "diagonal_gaussian",
+        bottleneck_norm_mode: str = "disabled",
+    ) -> None:
+        super().__init__()
+        self.bottleneck_dim = int(bottleneck_dim)
+        self.bottleneck_posterior_kind = bottleneck_posterior_kind
+        self.bottleneck_norm_mode = bottleneck_norm_mode
+        self.patchify = Patchify(in_channels, patch_size, model_dim)
+        self.blocks = nn.ModuleList(
+            [
+                FCDMBlock(
+                    model_dim,
+                    mlp_ratio,
+                    depthwise_kernel_size=depthwise_kernel_size,
+                    use_external_adaln=False,
+                )
+                for _ in range(depth)
+            ]
+        )
+        out_dim = (
+            2 * bottleneck_dim
+            if bottleneck_posterior_kind == "diagonal_gaussian"
+            else bottleneck_dim
+        )
+        self.to_bottleneck = nn.Conv2d(model_dim, out_dim, kernel_size=1, bias=True)
+        if bottleneck_norm_mode == "channel_wise":
+            self.norm_out = ChannelWiseRMSNorm(bottleneck_dim, eps=1e-6, affine=False)
+        else:
+            self.norm_out = nn.Identity()
+
+    def encode_posterior(self, images: Tensor) -> EncoderPosterior:
+        """Encode images and return the full posterior (mean + logsnr).
+
+        Only valid when bottleneck_posterior_kind == "diagonal_gaussian".
+        """
+        z = self.patchify(images)
+        for block in self.blocks:
+            z = block(z)
+        projection = self.to_bottleneck(z)
+        mean, logsnr = projection.chunk(2, dim=1)
+        mean = self.norm_out(mean)
+        return EncoderPosterior(mean=mean, logsnr=logsnr)
+
+    def forward(self, images: Tensor) -> Tensor:
+        """Encode images [B,3,H,W] in [-1,1] to latents [B,bottleneck_dim,h,w].
+
+        Returns posterior mode (alpha * mean) for diagonal_gaussian,
+        or deterministic latents otherwise.
+        """
+        z = self.patchify(images)
+        for block in self.blocks:
+            z = block(z)
+        projection = self.to_bottleneck(z)
+        if self.bottleneck_posterior_kind == "diagonal_gaussian":
+            mean, logsnr = projection.chunk(2, dim=1)
+            mean = self.norm_out(mean)
+            alpha = torch.sigmoid(logsnr).sqrt().to(dtype=mean.dtype)
+            return alpha * mean
+        z = self.norm_out(projection)
+        return z

capacitor_diffae/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_diffae/model.py

@@ -0,0 +1,304 @@
+"""CapacitorDiffAE: standalone HuggingFace-compatible diffusion autoencoder."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from torch import Tensor, nn
+
+from .config import CapacitorDiffAEConfig, CapacitorDiffAEInferenceConfig
+from .decoder import Decoder
+from .encoder import Encoder, EncoderPosterior
+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:
+        raise ImportError(
+            f"'{repo_id}' is not an existing local directory. "
+            "To download from HuggingFace Hub, install huggingface_hub: "
+            "pip install huggingface_hub"
+        )
+    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 CapacitorDiffAE(nn.Module):
+    """Standalone Capacitor DiffAE model for HuggingFace distribution.
+
+    A diffusion autoencoder built on FCDM (Fully Convolutional Diffusion Model)
+    blocks. Encodes images to compact 128-channel spatial latents via a
+    VP-parameterized diagonal Gaussian posterior, and decodes them back via
+    iterative VP diffusion with a skip-concat decoder.
+
+    Usage::
+
+        model = CapacitorDiffAE.from_pretrained("path/to/weights")
+        model = model.to("cuda", dtype=torch.bfloat16)
+
+        # Encode (returns posterior mode by default)
+        latents = model.encode(images)  # images: [B,3,H,W] in [-1,1]
+
+        # Decode (1 step by default — PSNR-optimal)
+        recon = model.decode(latents, height=H, width=W)
+
+        # Reconstruct (encode + 1-step decode)
+        recon = model.reconstruct(images)
+    """
+
+    def __init__(self, config: CapacitorDiffAEConfig) -> None:
+        super().__init__()
+        self.config = config
+
+        self.encoder = Encoder(
+            in_channels=config.in_channels,
+            patch_size=config.patch_size,
+            model_dim=config.model_dim,
+            depth=config.encoder_depth,
+            bottleneck_dim=config.bottleneck_dim,
+            mlp_ratio=config.mlp_ratio,
+            depthwise_kernel_size=config.depthwise_kernel_size,
+            bottleneck_posterior_kind=config.bottleneck_posterior_kind,
+            bottleneck_norm_mode=config.bottleneck_norm_mode,
+        )
+
+        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,
+    ) -> CapacitorDiffAE:
+        """Load a pretrained model from a local directory or HuggingFace Hub.
+
+        The directory (or repo) should contain:
+        - config.json: Model architecture config.
+        - model.safetensors (preferred) or model.pt: Model weights.
+
+        Args:
+            path_or_repo_id: Local directory path or HuggingFace Hub repo ID.
+            dtype: Load weights in this dtype (float32 or bfloat16).
+            device: Target device.
+            revision: Git revision for Hub downloads.
+            cache_dir: Where to cache Hub downloads.
+
+        Returns:
+            Loaded model in eval mode.
+        """
+        model_dir = _resolve_model_dir(
+            path_or_repo_id, revision=revision, cache_dir=cache_dir
+        )
+        config = CapacitorDiffAEConfig.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
+
+                state_dict = load_file(str(safetensors_path), device=str(device))
+            except ImportError:
+                raise ImportError(
+                    "safetensors package required to load .safetensors files. "
+                    "Install with: pip install safetensors"
+                )
+        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 encode(self, images: Tensor) -> Tensor:
+        """Encode images to latents (posterior mode).
+
+        Args:
+            images: [B, 3, H, W] in [-1, 1], H and W divisible by patch_size.
+
+        Returns:
+            Latents [B, bottleneck_dim, H/patch, W/patch].
+        """
+        try:
+            model_dtype = next(self.parameters()).dtype
+        except StopIteration:
+            model_dtype = torch.float32
+        return self.encoder(images.to(dtype=model_dtype))
+
+    def encode_posterior(self, images: Tensor) -> EncoderPosterior:
+        """Encode images and return the full posterior (mean + logsnr).
+
+        Args:
+            images: [B, 3, H, W] in [-1, 1], H and W divisible by patch_size.
+
+        Returns:
+            EncoderPosterior with mean and logsnr tensors.
+        """
+        try:
+            model_dtype = next(self.parameters()).dtype
+        except StopIteration:
+            model_dtype = torch.float32
+        return self.encoder.encode_posterior(images.to(dtype=model_dtype))
+
+    @torch.no_grad()
+    def decode(
+        self,
+        latents: Tensor,
+        height: int,
+        width: int,
+        *,
+        inference_config: CapacitorDiffAEInferenceConfig | None = None,
+    ) -> Tensor:
+        """Decode latents to images via VP diffusion.
+
+        Args:
+            latents: [B, bottleneck_dim, h, w] encoder latents.
+            height: Output image height (divisible by patch_size).
+            width: Output image width (divisible by patch_size).
+            inference_config: Optional inference parameters.
+
+        Returns:
+            Reconstructed images [B, 3, H, W] in float32.
+        """
+        cfg = inference_config or CapacitorDiffAEInferenceConfig()
+        config = self.config
+        batch = int(latents.shape[0])
+        device = latents.device
+
+        try:
+            model_dtype = next(self.parameters()).dtype
+        except StopIteration:
+            model_dtype = torch.float32
+
+        if height % config.patch_size != 0 or width % config.patch_size != 0:
+            raise ValueError(
+                f"height={height} and width={width} must be divisible by "
+                f"patch_size={config.patch_size}"
+            )
+
+        shape = (batch, 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=device)
+        initial_state = make_initial_state(
+            noise=noise.to(device=device),
+            t_start=schedule[0:1],
+            logsnr_min=config.logsnr_min,
+            logsnr_max=config.logsnr_max,
+        )
+
+        device_type = "cuda" if device.type == "cuda" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            latents_in = latents.to(device=device)
+
+            def _forward_fn(
+                x_t: Tensor,
+                t: Tensor,
+                latents: Tensor,
+                *,
+                drop_middle_blocks: bool = False,
+                mask_latent_tokens: bool = False,
+            ) -> Tensor:
+                return self.decoder(
+                    x_t.to(dtype=model_dtype),
+                    t,
+                    latents.to(dtype=model_dtype),
+                    drop_middle_blocks=drop_middle_blocks,
+                )
+
+            pdg_mode = "path_drop" if 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=cfg.pdg_strength,
+                device=device,
+            )
+
+        return result
+
+    @torch.no_grad()
+    def reconstruct(
+        self,
+        images: Tensor,
+        *,
+        inference_config: CapacitorDiffAEInferenceConfig | None = None,
+    ) -> Tensor:
+        """Encode then decode. Convenience wrapper.
+
+        Args:
+            images: [B, 3, H, W] in [-1, 1].
+            inference_config: Optional inference parameters.
+
+        Returns:
+            Reconstructed images [B, 3, H, W] in float32.
+        """
+        latents = self.encode(images)
+        _, _, h, w = images.shape
+        return self.decode(
+            latents, height=h, width=w, inference_config=inference_config
+        )

capacitor_diffae/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
+        if self.weight is not None:
+            shape = (1, -1) + (1,) * (x.dim() - 2)
+            y = y * self.weight.view(shape).to(dtype=y.dtype)
+            if self.bias is not None:
+                y = y + self.bias.view(shape).to(dtype=y.dtype)
+        return y.to(dtype=x.dtype)

capacitor_diffae/samplers.py

@@ -0,0 +1,263 @@
+"""DDIM and DPM++2M samplers for VP diffusion with dual 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 dual PDG flags."""
+
+    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", "path_drop", or "token_mask".
+        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)
+    elif pdg_mode == "token_mask":
+        x0_uncond = forward_fn(state, t_vec, latents, mask_latent_tokens=True).to(
+            torch.float32
+        )
+        x0_cond = forward_fn(state, t_vec, latents, mask_latent_tokens=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 dual 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", "path_drop", or "token_mask".
+        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 dual 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_diffae/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_diffae/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_diffae/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,18 @@
+{
+  "in_channels": 3,
+  "patch_size": 16,
+  "model_dim": 896,
+  "encoder_depth": 4,
+  "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,
+  "bottleneck_posterior_kind": "diagonal_gaussian",
+  "bottleneck_norm_mode": "disabled",
+  "logsnr_min": -10.0,
+  "logsnr_max": 10.0,
+  "pixel_noise_std": 0.558
+}

model.safetensors

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

technical_report_semantic.md

@@ -0,0 +1,669 @@
+# SemDisDiffAE — Technical Report
+
+**SemDisDiffAE** (**Sem**antically **Dis**entangled **Diff**usion **A**uto**E**ncoder)
+— a fast diffusion autoencoder with a 128-channel spatial bottleneck built on
+FCDM (Fully Convolutional Diffusion Model) blocks. The encoder uses a
+VP-parameterized diagonal Gaussian posterior (learned log-SNR output head),
+and the decoder reconstructs via single-step VP diffusion.
+
+This checkpoint is trained with DINOv2 semantic alignment and variance
+expansion regularization. The name is a nod to DRA (Page et al., 2026) whose
+disentangled representation alignment approach inspired the semantic alignment
+method used here.
+
+## Contents
+
+1. [Architecture](#1-architecture)
+   - [FCDM Block](#11-fcdm-block) · [Encoder](#12-encoder) · [Decoder](#13-decoder) · [AdaLN](#14-adaln-shared-base--low-rank-deltas) · [PDG](#15-path-drop-guidance-pdg)
+2. [Decoder VP Diffusion Parameterization](#2-decoder-vp-diffusion-parameterization)
+   - [Forward Process](#21-forward-process) · [Log SNR](#22-log-signal-to-noise-ratio) · [Schedule](#23-cosine-interpolated-schedule) · [X-Prediction](#24-x-prediction-objective) · [Sampling](#25-sampling)
+3. [Stochastic Posterior](#3-stochastic-posterior)
+   - [VP Log-SNR Parameterization](#31-vp-log-snr-parameterization) · [Variance Expansion Loss](#32-variance-expansion-loss) · [Posterior Mode](#33-posterior-mode-for-inference)
+4. [Semantic Alignment](#4-semantic-alignment)
+5. [Design Choices](#5-design-choices)
+   - [Convolutional Architecture](#51-convolutional-architecture) · [Single-Stride Encoder](#52-single-stride-encoder) · [Diffusion Decoding](#53-diffusion-decoding) · [Skip Connection and PDG](#54-skip-connection-and-path-drop-guidance)
+6. [Training](#6-training)
+   - [Loss Functions](#61-loss-functions) · [Optimizer](#62-optimizer-and-hyperparameters) · [Data](#63-data)
+7. [Model Configuration](#7-model-configuration)
+8. [Inference](#8-inference)
+9. [Results](#9-results)
+
+**References:**
+
+- **FCDM** — Kwon et al., *Reviving ConvNeXt for Efficient Convolutional Diffusion Models*, [arXiv:2603.09408](https://arxiv.org/abs/2603.09408), 2026.
+- **SiD2** — Hoogeboom et al., *Simpler Diffusion (SiD2): 1.5 FID on ImageNet512 with pixel-space diffusion*, [arXiv:2410.19324](https://arxiv.org/abs/2410.19324), ICLR 2025.
+- **DiTo** — Yin et al., *Diffusion Autoencoders are Scalable Image Tokenizers*, [arXiv:2501.18593](https://arxiv.org/abs/2501.18593), 2025.
+- **DiCo** — Ai et al., *DiCo: Revitalizing ConvNets for Scalable and Efficient Diffusion Modeling*, [arXiv:2505.11196](https://arxiv.org/abs/2505.11196), 2025.
+- **ConvNeXt V2** — Woo et al., *ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoders*, [arXiv:2301.00808](https://arxiv.org/abs/2301.00808), CVPR 2023.
+- **Z-image** — Cai et al., *Z-Image: An Efficient Image Generation Foundation Model with Single-Stream Diffusion Transformer*, [arXiv:2511.22699](https://arxiv.org/abs/2511.22699), 2025.
+- **SPRINT** — Park et al., *Sprint: Sparse-Dense Residual Fusion for Efficient Diffusion Transformers*, [arXiv:2510.21986](https://arxiv.org/abs/2510.21986), 2025.
+- **DINOv2** — Oquab et al., *DINOv2: Learning Robust Visual Features without Supervision*, [arXiv:2304.07193](https://arxiv.org/abs/2304.07193), 2023. Register variant: Darcet et al., *Vision Transformers Need Registers*, [arXiv:2309.16588](https://arxiv.org/abs/2309.16588), ICLR 2024.
+- **iREPA** — Singh et al., *What matters for Representation Alignment: Global Information or Spatial Structure?*, [arXiv:2512.10794](https://arxiv.org/abs/2512.10794), 2025.
+- **DRA** — Page et al., *Boosting Latent Diffusion Models via Disentangled Representation Alignment*, [arXiv:2601.05823](https://arxiv.org/abs/2601.05823), 2026.
+- **VEL** — Li et al., *Taming Sampling Perturbations with Variance Expansion Loss for Latent Diffusion Models*, [arXiv:2603.21085](https://arxiv.org/abs/2603.21085), 2026.
+- **iRDiffAE** — [data-archetype/irdiffae-v1](https://huggingface.co/data-archetype/irdiffae-v1) — predecessor model using DiCo blocks.
+
+---
+
+## 1. Architecture
+
+### 1.1 FCDM Block
+
+SemDisDiffAE uses **FCDM blocks** — ConvNeXt-style convolutional blocks
+adapted for diffusion models (Li et al., 2026). Each block follows a single
+unified residual path:
+
+```
+x ──► DWConv 7×7 ──► RMSNorm ──► [Scale] ──► Conv 1×1 ──► GELU ──► GRN ──► Conv 1×1 ──► [Gate] ──► + ──► out
+│                                                                                                    ▲
+└────────────────────────────────────────────────────────────────────────────────────────────────────────┘
+```
+
+This differs from DiCo blocks (used in the predecessor
+[iRDiffAE](https://huggingface.co/data-archetype/irdiffae-v1)) which use two
+separate residual paths (conv + MLP) with Compact Channel Attention (CCA).
+FCDM consolidates into one path, replacing CCA with Global Response
+Normalization (GRN).
+
+Key components:
+
+- **Depthwise convolution** (7×7, groups=channels): spatial mixing without
+  cross-channel interaction. The depthwise conv output feeds directly into
+  RMSNorm (no intermediate activation).
+
+- **RMSNorm** (non-affine, per-channel): normalizes activations before the
+  pointwise MLP, replacing LayerNorm used in standard ConvNeXt.
+
+- **Global Response Normalization (GRN)** (from ConvNeXt V2, Woo et al. 2023):
+  applied between the two pointwise convolutions. GRN computes per-channel L2
+  norms across the spatial dimensions and normalizes by the cross-channel mean:
+  ```
+  g = ||x||_2 over (H, W)
+  n = g / mean(g over channels)
+  GRN(x) = gamma * (x * n) + beta + x
+  ```
+  This encourages feature diversity across channels and prevents channel
+  collapse during training.
+
+- **Scale+Gate modulation** (decoder only): FCDM blocks use a 2-way modulation
+  `(scale, gate)` from the timestep embedding, in contrast to DiCo's 4-way
+  `(shift_conv, gate_conv, shift_mlp, gate_mlp)`. Scale is applied after
+  RMSNorm: `h = h * (1 + scale)`. Gate is applied to the residual:
+  `out = x + gate * h`. The gate is used **raw** (no tanh activation), giving
+  unbounded gating — this differs from DiCo which applies tanh to constrain
+  the gate to [-1, 1].
+
+- **Layer Scale** (encoder only): for unconditioned encoder blocks, a learnable
+  per-channel scale (initialized to 1e-3) gates the residual for near-identity
+  initialization, following ConvNeXt.
+
+### 1.2 Encoder
+
+The encoder uses a single spatial stride (via PixelUnshuffle at the input)
+followed by FCDM blocks at constant spatial resolution, then a bottleneck
+projection that outputs both the posterior mean and per-element log-SNR:
+
+```
+Image [B, 3, H, W]
+  ──► PixelUnshuffle(p=16) + Conv 1×1 (3·16² → 896)     [Patchify]
+  ──► 4 × FCDMBlock (unconditioned, layer-scale gated)
+  ──► Conv 1×1 (896 → 256)                               [Bottleneck projection]
+  ──► Split → mean [B, 128, h, w] + logsnr [B, 128, h, w]
+  ──► α(logsnr) · mean                                    [Posterior mode]
+```
+
+The single-stride design ensures all encoder blocks see the full spatial
+resolution and full channel width simultaneously. The information bottleneck
+is imposed only at the very end, where a single linear projection selects
+which channels to retain. See Section 4.2 for the rationale.
+
+**Note:** This checkpoint uses `bottleneck_norm_mode=disabled`, so no
+post-bottleneck RMSNorm is applied to the mean branch. The posterior mode
+output is simply `α · μ` where `α = √σ(λ)`.
+
+### 1.3 Decoder
+
+The decoder predicts x̂₀ from noisy input x_t, conditioned on encoder
+latents z and timestep t:
+
+```
+Noised image x_t [B, 3, H, W]
+  ──► PixelUnshuffle(p=16) + Conv 1×1 (3·16² → 896)     [Patchify]
+  ──► Concatenate with Conv 1×1(latents, 128 → 896)      [Latent fusion]
+  ──► Conv 1×1 (2·896 → 896)
+  ──► 2 × FCDMBlock (AdaLN conditioned)                   [Start blocks]
+  ──► 4 × FCDMBlock (AdaLN conditioned)                   [Middle blocks]
+  ──► Concat(start_out, middle_out) + Conv 1×1            [Skip fusion]
+  ──► 2 × FCDMBlock (AdaLN conditioned)                   [End blocks]
+  ──► Conv 1×1 (896 → 3·16²) + PixelShuffle(16)          [Unpatchify]
+  ──► x̂₀ prediction [B, 3, H, W]
+```
+
+The skip-concat topology with 2+4+2 blocks is inspired by SPRINT's
+sparse-dense residual fusion (Park et al., 2025). See Section 4.4 for the
+design rationale.
+
+### 1.4 AdaLN: Shared Base + Low-Rank Deltas
+
+Timestep conditioning follows the Z-image style AdaLN
+([Cai et al., 2025](https://arxiv.org/abs/2511.22699)): a shared base
+projection plus a low-rank delta per layer.
+
+A single base projector is shared across all 8 decoder layers, and each
+layer adds a low-rank correction:
+
+```
+m_i = Base(SiLU(cond)) + Δ_i(SiLU(cond))
+```
+
+where `Base: ℝ^D → ℝ^{2D}` is a linear projection (zero-initialized) and
+`Δ_i: ℝ^D → ℝ^r → ℝ^{2D}` is a low-rank factorization with rank r = 128
+(zero-initialized up-projection).
+
+The packed modulation `m_i ∈ ℝ^{B × 2D}` is split into `(scale, gate)` which
+modulate the FCDM block (no shift term):
+
+```
+ĥ = RMSNorm(x) · (1 + scale)
+x ← x + gate · f(ĥ)
+```
+
+### 1.5 Path-Drop Guidance (PDG)
+
+At inference, optional PDG sharpens reconstructions by exploiting the
+skip-concat structure — a classifier-free guidance analogue that does not
+require training with conditioning dropout:
+
+1. **Conditional pass:** run all blocks normally → x̂₀^cond
+2. **Unconditional pass:** replace the middle block output with a learned
+   mask feature m ∈ ℝ^{1×D×1×1} (initialized to zero), effectively dropping
+   the deep processing path → x̂₀^uncond
+3. **Guided prediction:** x̂₀ = x̂₀^uncond + s · (x̂₀^cond - x̂₀^uncond)
+
+where s is the guidance strength.
+
+For PSNR-optimal reconstruction, PDG is disabled (1 NFE). For perceptual
+sharpening, use 10 steps with PDG strength 2.0. Note that PDG is primarily
+useful for more compressed bottlenecks (e.g. 32 or 64 channels) and is
+rarely necessary for 128-channel models where reconstruction quality is
+already high.
+
+---
+
+## 2. Decoder VP Diffusion Parameterization
+
+The decoder uses the variance-preserving (VP) diffusion framework from
+SiD2 with an x-prediction objective.
+
+### 2.1 Forward Process
+
+Given a clean image x₀, the forward process constructs a noisy sample at
+continuous time t ∈ [0, 1]:
+
+```
+x_t = α_t · x₀ + σ_t · ε,    ε ~ N(0, s²I)
+```
+
+where s = 0.558 is the pixel-space noise standard deviation (estimated from
+the dataset image distribution) and the VP constraint holds: α²_t + σ²_t = 1.
+
+### 2.2 Log Signal-to-Noise Ratio
+
+The schedule is parameterized through the log signal-to-noise ratio:
+
+```
+λ_t = log(α²_t / σ²_t)
+```
+
+which monotonically decreases as t → 1 (pure noise). From λ_t we recover
+α_t and σ_t via the sigmoid function:
+
+```
+α_t = √σ(λ_t),    σ_t = √σ(-λ_t)
+```
+
+### 2.3 Cosine-Interpolated Schedule
+
+Following SiD2, the logSNR schedule uses cosine interpolation:
+
+```
+λ(t) = -2 log tan(a·t + b)
+```
+
+where a and b are computed to satisfy the boundary conditions
+λ(0) = λ_max = 10 and λ(1) = λ_min = -10.
+
+### 2.4 X-Prediction Objective
+
+The model predicts the clean image x̂₀ = f_θ(x_t, t, z) conditioned on
+encoder latents z.
+
+**Schedule-invariant loss.** Following SiD2, the training loss is defined as
+an integral over logSNR λ, making it invariant to the choice of noise schedule.
+Since timesteps are sampled uniformly t ~ U(0,1), the change of variable
+introduces a Jacobian factor:
+
+```
+L = E_{t ~ U(0,1)} [ (-dλ/dt) · w(λ(t)) · ||x₀ - x̂₀||² ]
+```
+
+**Sigmoid weighting.** The weighting function uses a sigmoid centered at bias
+b = -2.0, converting from ε-prediction to x-prediction form:
+
+```
+weight(t) = -(1/2) · (dλ/dt) · e^b · σ(λ(t) - b)
+```
+
+### 2.5 Sampling
+
+Decoding uses DDIM by default. With 1 NFE (default), the model runs a single
+evaluation at t_start ≈ 1 (near pure noise) and directly outputs the x₀
+prediction. This is equivalent to a denoising autoencoder that maps
+`σ_start · noise → x̂₀` conditioned on encoder latents.
+
+DPM++2M is also supported as an alternative sampler, using a half-lambda
+exponential integrator for faster convergence with more steps.
+
+---
+
+## 3. Stochastic Posterior
+
+### 3.1 VP Log-SNR Parameterization
+
+Instead of a KL-divergence penalty on a Gaussian encoder, SemDisDiffAE
+parameterizes the bottleneck posterior using the VP interpolation convention.
+This approach uses a VP-style noise interpolation in the encoder bottleneck
+as an alternative to the traditional VAE KL penalty.
+
+The encoder outputs two sets of 128 channels:
+
+- **μ** — the clean signal (posterior mean)
+- **λ** — per-element log signal-to-noise ratio
+
+The posterior distribution is:
+
+```
+z = α(λ) · μ + σ(λ) · ε,    ε ~ N(0, I)
+```
+
+where α = √σ(λ) and σ = √σ(-λ) (sigmoid parameterization). This is
+equivalent to a Gaussian with mean α·μ and variance σ².
+
+Using a VP interpolation rather than simple additive noise decouples token
+scale from stochasticity. With additive noise (`z = μ + σε`), the encoder
+faces gradient pressure to scale latents up to counter the noise — the SNR
+depends on the magnitude of μ. The VP formulation (`z = α·μ + σ·ε` with
+`α² + σ² = 1`) removes this coupling: the noise level is controlled
+entirely by the predicted log-SNR, independent of the latent magnitude.
+
+### 3.2 Variance Expansion Loss
+
+To prevent posterior collapse (where the encoder learns to set σ → 0 and
+ignore the stochastic component entirely), we adopt a **variance expansion
+loss** inspired by VEL (Li et al., 2026,
+[arXiv:2603.21085](https://arxiv.org/abs/2603.21085)):
+
+```
+L_var = -mean(log(σ² + δ))
+```
+
+where σ² is the posterior variance derived from the predicted log-SNR and
+δ is a small epsilon (1e-6) for numerical stability. This loss encourages
+non-zero posterior variance by penalizing small σ².
+
+VEL proposes the form `1/(σ² + δ)` for variance expansion. We found this to
+be too aggressive — the `1/σ²` gradient pushes variance up very rapidly,
+leading to excessive high-frequency noise in the latent space. We use the
+`-log(σ² + δ)` form instead, which provides a gentler, logarithmic penalty
+that stabilizes training.
+
+**For this checkpoint:** the variance expansion loss is active with weight
+**1e-5**.
+
+> **Key finding: latent spectral structure matters for downstream diffusion.**
+>
+> Reconstruction quality is not very sensitive to the posterior noise level —
+> good PSNR is achievable even with log-SNR as low as -2. However, the
+> posterior noise level has a strong effect on the **spatial frequency
+> content** of the latent space. When variance expansion is too aggressive,
+> the latent space develops excessive high-frequency content; when it is
+> too weak or absent, latents become overly smooth.
+>
+> We found empirically that downstream diffusion models converge best when
+> the latent space has a **radial power spectral density (PSD) decay
+> exponent of approximately 1.5** — deviating significantly in either
+> direction (too smooth or too high-frequency) consistently yields worse
+> downstream training convergence. We monitor this metric during training
+> validation to guide the variance expansion weight.
+>
+> The weight of 1e-5 for this checkpoint was chosen to target this spectral
+> sweet spot.
+
+### 3.3 Posterior Mode for Inference
+
+At inference, the encoder returns the **posterior mode**: `z = α(λ) · μ`. For
+this checkpoint, the posterior log-SNR is very high (posterior variance is
+negligible), so sampling and mode are nearly identical.
+
+The `encode_posterior()` method is available for users who need the full
+posterior distribution.
+
+---
+
+## 4. Semantic Alignment
+
+This checkpoint uses **semantic alignment** to encourage semantically
+structured latent representations. The approach is inspired by DRA
+(Page et al., 2026, [arXiv:2601.05823](https://arxiv.org/abs/2601.05823))
+which aligns autoencoder latents with frozen vision encoder features. Our
+implementation differs in the projection architecture and noise schedule.
+
+### 4.1 Teacher
+
+A frozen DINOv2-S with registers
+(timm: `vit_small_patch16_dinov3.lvd_1689m`, 384-dim patch tokens) provides
+the target spatial semantic features.
+
+### 4.2 Projection Head
+
+The student projection head maps noisy encoder latents to the teacher's
+token space. It consists of:
+
+```
+Noisy latents z_noisy ∈ ℝ^{B×128×h×w}
+  ──► Conv 1×1 (128 → 384)                     [Channel projection]
+  ──► Flatten to tokens [B, T, 384]
+  ──► DiT transformer block                     [Single block, 6 heads × 64 dim]
+      (self-attention with axial RoPE 2D + AdaLN conditioned on τ)
+  ──► RMSNorm
+  ──► student tokens ∈ ℝ^{B×T×384}
+```
+
+The DiT block uses standard multi-head self-attention with 2D axial
+rotary position embeddings (RoPE) and AdaLN-Zero timestep conditioning.
+This gives the projection head global spatial reasoning — important for
+matching the teacher's self-attention-based representations — while the
+main encoder/decoder remain purely convolutional.
+
+### 4.3 Noisy Alignment
+
+Unlike standard representation alignment which operates on clean latents,
+we align **noisy** latent versions. The noise level τ is sampled from a
+Beta(2,2) distribution (concentrated around τ=0.5) using flow matching
+linear interpolation:
+
+```
+z_noisy = (1 - τ) · z + τ · ε,    ε ~ N(0, I),    τ ~ Beta(2, 2)
+```
+
+The projection head receives both the noisy latents and the noise level τ
+(via its AdaLN conditioning). This trains the head to extract semantic
+information even from partially corrupted latents, improving robustness
+for downstream diffusion models which operate on noised latent inputs.
+
+### 4.4 Training Details
+
+The alignment loss is the mean negative cosine similarity between student
+and teacher tokens, weighted at **0.01** throughout training. The student
+projection head operates on all 128 bottleneck channels, unlike the
+predecessor iRDiffAE which aligned only the first 64 of 128 channels.
+
+Note that the projection head is a training-only component — it is not
+included in the exported model weights.
+
+---
+
+## 5. Design Choices
+
+### 5.1 Convolutional Architecture
+
+SemDisDiffAE uses a fully convolutional architecture rather than a vision
+transformer. For an autoencoder whose goal is faithful pixel-level
+reconstruction (not global semantic understanding), convolutions offer:
+
+- **Resolution generalization.** Convolutions operate on local patches and
+  generalize naturally to arbitrary image dimensions without interpolating
+  position embeddings or suffering attention distribution shift.
+- **Translation invariance.** Weight sharing across spatial positions is well
+  matched to reconstruction, where the same local patterns (edges, textures)
+  conditioned on the low-frequency latent recur throughout the image.
+- **Locality.** Reconstruction quality depends on preserving fine spatial
+  detail. Convolutions are inherently local operators, avoiding the quadratic
+  cost of global attention while focusing computation where it matters most.
+
+### 5.2 Single-Stride Encoder with Final Bottleneck
+
+The encoder uses a single spatial stride (PixelUnshuffle at the input)
+followed by blocks at constant spatial resolution, then a final 1×1 convolution
+to project to the bottleneck. This differs from classical VAE encoders that use
+progressive downsampling with channel expansion at each stage.
+
+The single-stride design ensures that all encoder blocks see the full spatial
+resolution and full channel width simultaneously. The information bottleneck is
+imposed only at the very end, where a single linear projection selects which
+channels to retain.
+
+### 5.3 Diffusion Decoding
+
+The main advantage of diffusion decoding over the standard GAN + LPIPS
+approach is **simplicity and speed of experimentation**. The training
+objective is a straightforward weighted MSE — no discriminator, no LPIPS
+perceptual loss, no delicate adversarial balancing. This makes it very fast to train and easy to iterate on — typically a few
+hours on a single GPU is sufficient. This checkpoint was trained for 251k
+steps. By contrast, GAN + LPIPS-based VAEs require many days of large-GPU
+time and are notoriously difficult to stabilize from scratch.
+
+This simplicity enables rapid experimentation with latent space shaping to
+get it as diffusion-friendly as possible, while still achieving excellent
+reconstruction quality.
+
+### 5.4 Skip Connection and Path-Drop Guidance
+
+The decoder's start → middle → skip-fuse → end architecture is inspired by
+SPRINT's sparse-dense residual fusion (Park et al., 2025). The design serves
+three purposes:
+
+1. **Regularization.** The skip path ensures that even if the middle blocks
+   are dropped or poorly conditioned, the end blocks still receive meaningful
+   features from the start blocks.
+2. **High-frequency preservation.** The start blocks (which see the input most
+   directly) pass fine detail through the skip to the end blocks.
+3. **Path-Drop Guidance.** At inference, replacing the middle block output
+   with a learned mask feature creates an "unconditional" prediction that
+   preserves the skip path but drops the deep processing. Interpolating
+   between conditional and unconditional predictions (as in classifier-free
+   guidance) sharpens the output without requiring training-time dropout.
+
+---
+
+## 6. Training
+
+### 6.1 Loss Functions
+
+The total training loss is:
+
+```
+L_total = L_recon + 0.01 · L_semantic + 0.0001 · L_scale + 1e-5 · L_var
+```
+
+| Loss | Weight | Description |
+|------|--------|-------------|
+| **Reconstruction** (L_recon) | 1.0 | SiD2 sigmoid-weighted x-prediction MSE (bias b = -2.0). Per-pixel `(x̂₀ - x₀)²` averaged over (C, H, W) per sample, multiplied by the SiD2 per-sample weight `w(t) = -½ · dλ/dt · e^b · σ(λ-b)`, then averaged over the batch |
+| **Semantic alignment** (L_semantic) | 0.01 | Per-token `(1 - cosine(student, teacher))` averaged over all tokens and batch (see §4) |
+| **Latent scale penalty** (L_scale) | 0.0001 | Per-channel variance `var_c` estimated over the batch and spatial dims (B, H, W), then `(log(var_c + ε) - log(target))²` averaged over channels. Target variance = 1.0 |
+| **Posterior variance expansion** (L_var) | 1e-5 | Per-element `-log(σ² + δ)` where σ² is the posterior variance derived from the predicted log-SNR, averaged over all dims (B, C, H, W). See §3.2 |
+
+**Note on loss scales:** The decoder reconstruction loss has a small
+effective magnitude due to the SiD2 VP x-prediction weighting (the Jacobian
+dλ/dt and sigmoid weighting compress the per-sample loss scale). As a
+result, all auxiliary loss weights must be kept correspondingly small to
+avoid dominating the reconstruction objective.
+
+### 6.2 Optimizer and Hyperparameters
+
+| Parameter | Value |
+|-----------|-------|
+| Optimizer | AdamW (β₁=0.9, β₂=0.99) |
+| Learning rate | 1e-4 (constant after warmup) |
+| Weight decay | 0.0 |
+| Warmup steps | 2,000 |
+| Gradient clip | 1.0 (max norm) |
+| Precision | AMP bfloat16 (FP32 master weights, TF32 matmul) |
+| EMA decay | 0.9995 (updated every step) |
+| Batch size | 128 |
+| Timestep sampling | Uniform with SiD2 logSNR shift -1.0 |
+| Compilation | `torch.compile` enabled |
+| Training steps | 251k |
+| Hardware | Single GPU |
+
+Convergence is fast — training is stopped when the training loss starts
+plateauing, which typically occurs within a few hours on a single GPU.
+
+### 6.3 Data
+
+Training uses ~5M images at various resolutions: mostly photographs, with
+a significant proportion of illustrations and text-heavy images (documents,
+screenshots, book covers, diagrams) to encourage crisp line and edge
+reconstruction. Images are loaded via two strategies in a 50/50 mix:
+
+- **Full-image downsampling:** images are bucketed by aspect ratio and
+  downsampled to ~256² resolution (preserving aspect ratio).
+- **Random 256×256 crops:** deterministic patches extracted from images
+  stored at ≥512px resolution.
+
+This mixed strategy exposes the model to both global scene composition (via
+downsampled full images) and fine local detail (via crops from higher-resolution
+sources).
+
+---
+
+## 7. Model Configuration
+
+| Parameter | Value |
+|-----------|-------|
+| Patch size | 16 |
+| Model dimension | 896 |
+| Encoder depth | 4 blocks |
+| Decoder depth | 8 blocks (2 start + 4 middle + 2 end) |
+| Bottleneck dimension | 128 channels |
+| Spatial compression | 16× (H/16 × W/16) |
+| Total compression | 6.0× (3·256 / 128) |
+| MLP ratio | 4.0 |
+| Depthwise kernel | 7×7 |
+| AdaLN per-block delta rank | 128 |
+| Block type | FCDM (ConvNeXt + GRN + scale/gate AdaLN) |
+| Posterior | Diagonal Gaussian (VP log-SNR), variance expansion weight 1e-5 |
+| Bottleneck norm | Disabled |
+| λ_min, λ_max | -10, +10 |
+| Sigmoid bias b | -2.0 |
+| Pixel noise std s | 0.558 |
+| Parameters | 88.8M |
+
+---
+
+## 8. Inference
+
+### Recommended Settings
+
+| Use case | Steps (NFE) | PDG | Sampler | Notes |
+|----------|-------------|-----|---------|-------|
+| **PSNR-optimal** | 1 | off | DDIM | Default. Fastest. |
+| **Perceptual** | 10 | on (2.0) | DDIM | Sharper details, ~15× slower (PDG skips middle blocks) |
+
+### Usage
+
+```python
+from capacitor_diffae import CapacitorDiffAE, CapacitorDiffAEInferenceConfig
+
+# Load model
+model = CapacitorDiffAE.from_pretrained("data-archetype/semdisdiffae", device="cuda")
+
+# Encode (returns posterior mode by default)
+latents = model.encode(images)  # [B,3,H,W] → [B,128,H/16,W/16]
+
+# Decode (1 step)
+recon = model.decode(latents, height=H, width=W)
+
+# Full posterior access
+posterior = model.encode_posterior(images)
+print(posterior.mean.shape, posterior.logsnr.shape)
+z_sampled = posterior.sample()
+```
+
+---
+
+## 9. Results
+
+## 7. Results
+
+Reconstruction quality evaluated on a curated set of test images covering photographs, book covers, and documents. Flux.1 VAE (patch 8, 16 channels) is included as a reference at the same 12x compression ratio as the c64 variant.
+
+### 7.1 Interactive Viewer
+
+**[Open full-resolution comparison viewer](https://huggingface.co/spaces/data-archetype/irdiffae-results)** — side-by-side reconstructions, RGB deltas, and latent PCA with adjustable image size.
+
+### 7.2 Inference Settings
+
+| Setting | Value |
+|---------|-------|
+| Sampler | ddim |
+| Steps | 1 |
+| Schedule | linear |
+| Seed | 42 |
+| PDG | no_path_dropg |
+| Batch size (timing) | 4 |
+
+> All models run in bfloat16. Timings measured on an NVIDIA RTX Pro 6000 (Blackwell).
+
+### 7.3 Global Metrics
+
+| Metric | semdisdiffae (1 step) | Flux.2 VAE |
+|--------|--------|--------|
+| Avg PSNR (dB) | 35.78 | 34.16 |
+| Avg encode (ms/image) | 2.5 | 46.1 |
+| Avg decode (ms/image) | 5.5 | 91.8 |
+
+### 7.4 Per-Image PSNR (dB)
+
+| Image | semdisdiffae (1 step) | Flux.2 VAE |
+|-------|--------|--------|
+| p640x1536:94623 | 35.44 | 33.50 |
+| p640x1536:94624 | 31.33 | 30.03 |
+| p640x1536:94625 | 35.05 | 33.98 |
+| p640x1536:94626 | 33.21 | 31.53 |
+| p640x1536:94627 | 32.54 | 30.53 |
+| p640x1536:94628 | 29.80 | 28.88 |
+| p960x1024:216264 | 46.37 | 45.39 |
+| p960x1024:216265 | 29.70 | 27.80 |
+| p960x1024:216266 | 47.15 | 46.20 |
+| p960x1024:216267 | 40.99 | 39.23 |
+| p960x1024:216268 | 38.47 | 36.13 |
+| p960x1024:216269 | 32.74 | 30.24 |
+| p960x1024:216270 | 36.23 | 34.18 |
+| p960x1024:216271 | 44.41 | 42.18 |
+| p704x1472:94699 | 43.80 | 41.79 |
+| p704x1472:94700 | 32.83 | 32.08 |
+| p704x1472:94701 | 39.00 | 37.90 |
+| p704x1472:94702 | 34.52 | 32.50 |
+| p704x1472:94703 | 32.81 | 31.35 |
+| p704x1472:94704 | 33.38 | 31.84 |
+| p704x1472:94705 | 39.70 | 37.44 |
+| p704x1472:94706 | 35.12 | 33.66 |
+| r256_p1344x704:15577 | 31.02 | 29.98 |
+| r256_p1344x704:15578 | 32.38 | 30.79 |
+| r256_p1344x704:15579 | 33.27 | 31.83 |
+| r256_p1344x704:15580 | 37.84 | 36.03 |
+| r256_p1344x704:15581 | 38.57 | 36.94 |
+| r256_p1344x704:15582 | 33.41 | 32.10 |
+| r256_p1344x704:15583 | 36.67 | 34.54 |
+| r256_p1344x704:15584 | 33.23 | 31.76 |
+| r256_p896x1152:144131 | 35.30 | 33.60 |
+| r256_p896x1152:144132 | 36.99 | 35.32 |
+| r256_p896x1152:144133 | 39.69 | 37.33 |
+| r256_p896x1152:144134 | 36.01 | 34.47 |
+| r256_p896x1152:144135 | 31.20 | 29.87 |
+| r256_p896x1152:144136 | 37.51 | 35.68 |
+| r256_p896x1152:144137 | 33.83 | 32.86 |
+| r256_p896x1152:144138 | 27.39 | 25.63 |
+| VAE_accuracy_test_image | 36.64 | 35.25 |
+