Initial upload: iRDiffAE v1.0 (p16_c128, EMA weights)

README.md

@@ -0,0 +1,118 @@
+---
+license: apache-2.0
+tags:
+  - diffusion
+  - autoencoder
+  - image-reconstruction
+  - pytorch
+library_name: irdiffae
+---
+
+# irdiffae_v1
+
+**iRDiffAE** — **i**REPA **Diff**usion **A**uto**E**ncoder.
+A fast, single-GPU-trainable diffusion autoencoder with spatially structured
+latents for rapid downstream model convergence. Encoding runs ~5× faster than
+Flux VAE; single-step decoding runs ~3× faster.
+
+## Model Variants
+
+| Variant | Patch | Channels | Compression | |
+|---------|-------|----------|-------------|---|
+| **irdiffae_v1** | 16x16 | 128 | 6x | recommended |
+
+This variant (irdiffae_v1): 121.0M parameters, 461.4 MB.
+
+## Documentation
+
+- [Technical Report](technical_report.md) — diffusion math, architecture, training, and results
+- [Results — interactive viewer](https://huggingface.co/spaces/data-archetype/irdiffae-results) — full-resolution side-by-side comparison
+- [Results — summary stats](technical_report.md#7-results) — metrics and per-image PSNR
+
+## Quick Start
+
+```python
+import torch
+from ir_diffae import IRDiffAE
+
+# Load from HuggingFace Hub (or a local path)
+model = IRDiffAE.from_pretrained("irdiffae_v1", device="cuda")
+
+# Encode
+images = ...  # [B, 3, H, W] in [-1, 1], H and W divisible by 16
+latents = model.encode(images)
+
+# Decode
+recon = model.decode(latents, height=H, width=W)
+
+# Reconstruct (encode + decode)
+recon = model.reconstruct(images)
+```
+
+> **Note:** Requires `pip install huggingface_hub safetensors` for Hub downloads.
+> You can also pass a local directory path to `from_pretrained()`.
+
+## Architecture
+
+| Property | Value |
+|---|---|
+| Parameters | 120,957,440 |
+| File size | 461.4 MB |
+| Patch size | 16 |
+| Model dim | 896 |
+| Encoder depth | 4 |
+| Decoder depth | 8 |
+| Bottleneck dim | 128 |
+| MLP ratio | 4.0 |
+| Depthwise kernel | 7 |
+| AdaLN rank | 128 |
+
+**Encoder**: Deterministic. Patchify (PixelUnshuffle + 1x1 conv) followed by
+DiCo blocks (depthwise conv + compact channel attention + GELU MLP) with
+learned residual gates.
+
+**Decoder**: VP diffusion conditioned on encoder latents and timestep via
+shared-base + per-layer low-rank AdaLN-Zero. Start blocks (2) -> middle
+blocks (4) -> skip fusion -> end blocks (2). Supports
+Path-Drop Guidance (PDG) at inference for quality/speed tradeoff.
+
+## Recommended Settings
+
+Best quality is achieved with just **1 DDIM step** and PDG disabled,
+making inference extremely fast. PDG (strength 2-4) can optionally
+increase perceptual sharpness but is easy to overdo.
+
+| Setting | Default |
+|---|---|
+| Sampler | DDIM |
+| Steps | 1 |
+| PDG | Disabled |
+
+```python
+from ir_diffae import IRDiffAEInferenceConfig
+
+# PSNR-optimal (fast, 1 step)
+cfg = IRDiffAEInferenceConfig(num_steps=1, sampler="ddim")
+recon = model.decode(latents, height=H, width=W, inference_config=cfg)
+```
+
+## Citation
+
+```bibtex
+@misc{irdiffae_v1,
+  title   = {iRDiffAE: A Fast, Representation Aligned Diffusion Autoencoder with DiCo Blocks},
+  author  = {data-archetype},
+  year    = {2026},
+  month   = feb,
+  url     = {https://huggingface.co/irdiffae_v1},
+}
+```
+
+## Dependencies
+
+- PyTorch >= 2.0
+- safetensors (for loading weights)
+
+## License
+
+Apache 2.0

config.json

@@ -0,0 +1,14 @@
+{
+  "in_channels": 3,
+  "patch_size": 16,
+  "model_dim": 896,
+  "encoder_depth": 4,
+  "decoder_depth": 8,
+  "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
+}

ir_diffae/__init__.py

@@ -0,0 +1,26 @@
+"""iRDiffAE: Standalone diffusion autoencoder for HuggingFace distribution.
+
+iREPA Diffusion AutoEncoder — a compact diffusion autoencoder
+that encodes images to spatial latents and decodes via iterative VP diffusion.
+
+Usage::
+
+    from ir_diffae import IRDiffAE, IRDiffAEInferenceConfig
+
+    model = IRDiffAE.from_pretrained("path/to/weights", device="cuda")
+
+    # Encode
+    latents = model.encode(images)  # images: [B,3,H,W] in [-1,1]
+
+    # Decode with custom settings
+    cfg = IRDiffAEInferenceConfig(num_steps=50, sampler="dpmpp_2m")
+    recon = model.decode(latents, height=512, width=512, inference_config=cfg)
+
+    # Reconstruct (encode + decode)
+    recon = model.reconstruct(images)
+"""
+
+from .config import IRDiffAEConfig, IRDiffAEInferenceConfig
+from .model import IRDiffAE
+
+__all__ = ["IRDiffAE", "IRDiffAEConfig", "IRDiffAEInferenceConfig"]

ir_diffae/adaln.py

@@ -0,0 +1,51 @@
+"""AdaLN-Zero modules for shared-base + low-rank-delta conditioning."""
+
+from __future__ import annotations
+
+from torch import Tensor, nn
+
+
+class AdaLNZeroProjector(nn.Module):
+    """Shared base AdaLN projection: SiLU -> Linear(d_cond -> 4*d_model).
+
+    Returns packed modulation tensor [B, 4*d_model]. Zero-initialized.
+    """
+
+    def __init__(self, d_model: int, d_cond: int) -> None:
+        super().__init__()
+        self.d_model = int(d_model)
+        self.d_cond = int(d_cond)
+        self.act = nn.SiLU()
+        self.proj = nn.Linear(self.d_cond, 4 * self.d_model)
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward(self, cond: Tensor) -> Tensor:
+        """Return packed modulation [B, 4*d_model] from conditioning [B, d_cond]."""
+        act = self.act(cond)
+        return self.proj(act)
+
+    def forward_activated(self, act_cond: Tensor) -> Tensor:
+        """Return packed modulation from pre-activated conditioning."""
+        return self.proj(act_cond)
+
+
+class AdaLNZeroLowRankDelta(nn.Module):
+    """Per-layer low-rank delta: down(d_cond -> rank) -> up(rank -> 4*d_model).
+
+    Zero-initialized up-projection preserves AdaLN "zero output" at init.
+    """
+
+    def __init__(self, *, d_model: int, d_cond: int, rank: int) -> None:
+        super().__init__()
+        self.d_model = int(d_model)
+        self.d_cond = int(d_cond)
+        self.rank = int(rank)
+        self.down = nn.Linear(self.d_cond, self.rank, bias=False)
+        self.up = nn.Linear(self.rank, 4 * 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, 4*d_model] from activated cond."""
+        return self.up(self.down(act_cond))

ir_diffae/compact_channel_attention.py

@@ -0,0 +1,21 @@
+"""Compact Channel Attention (CCA) module."""
+
+from __future__ import annotations
+
+from torch import Tensor, nn
+
+
+class CompactChannelAttention(nn.Module):
+    """Global average pool -> 1x1 Conv2d -> Sigmoid channel gate."""
+
+    def __init__(self, channels: int) -> None:
+        super().__init__()
+        c = int(channels)
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.proj = nn.Conv2d(c, c, kernel_size=1, padding=0, stride=1, bias=True)
+        self.act = nn.Sigmoid()
+
+    def forward(self, x: Tensor) -> Tensor:
+        w = self.pool(x)
+        w = self.proj(w)
+        return self.act(w)

ir_diffae/config.py

@@ -0,0 +1,51 @@
+"""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 IRDiffAEConfig:
+    """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
+    bottleneck_dim: int = 128
+    mlp_ratio: float = 4.0
+    depthwise_kernel_size: int = 7
+    adaln_low_rank_rank: int = 128
+    # 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) -> IRDiffAEConfig:
+        """Load config from JSON."""
+        data = json.loads(Path(path).read_text())
+        return cls(**data)
+
+
+@dataclass
+class IRDiffAEInferenceConfig:
+    """User-tunable inference parameters with sensible defaults."""
+
+    num_steps: int = 1  # decoder forward passes (NFE)
+    sampler: str = "ddim"  # "ddim" or "dpmpp_2m"
+    schedule: str = "linear"  # "linear" or "cosine"
+    pdg_enabled: bool = False
+    pdg_strength: float = 2.0
+    seed: int | None = None

ir_diffae/conv_mlp.py

@@ -0,0 +1,26 @@
+"""Conv-based MLP with GELU activation for DiCo blocks."""
+
+from __future__ import annotations
+
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from .norms import ChannelWiseRMSNorm
+
+
+class ConvMLP(nn.Module):
+    """1x1 Conv-based MLP: RMSNorm -> Conv1x1 -> GELU -> Conv1x1."""
+
+    def __init__(
+        self, channels: int, hidden_channels: int, norm_eps: float = 1e-6
+    ) -> None:
+        super().__init__()
+        self.norm = ChannelWiseRMSNorm(channels, eps=norm_eps, affine=False)
+        self.conv_in = nn.Conv2d(channels, hidden_channels, kernel_size=1, bias=True)
+        self.conv_out = nn.Conv2d(hidden_channels, channels, kernel_size=1, bias=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        y = self.norm(x)
+        y = self.conv_in(y)
+        y = F.gelu(y)
+        return self.conv_out(y)

ir_diffae/decoder.py

@@ -0,0 +1,175 @@
+"""iRDiffAE decoder: conditioned DiCoBlocks with AdaLN + skip connection."""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+
+from .adaln import AdaLNZeroLowRankDelta, AdaLNZeroProjector
+from .dico_block import DiCoBlock
+from .norms import ChannelWiseRMSNorm
+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:
+        Patchify x_t -> Norm -> Fuse with upsampled z
+        -> Start blocks (2) -> Middle blocks (depth-4) -> Skip fuse -> End blocks (2)
+        -> Norm -> Conv1x1 -> PixelShuffle
+
+    Middle blocks support path-drop for PDG (inference-time guidance).
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        patch_size: int,
+        model_dim: int,
+        depth: 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
+        self.patchify = Patchify(in_channels, patch_size, model_dim)
+        self.norm_in = ChannelWiseRMSNorm(model_dim, eps=1e-6, affine=True)
+
+        # Latent conditioning path
+        self.latent_up = nn.Conv2d(bottleneck_dim, model_dim, kernel_size=1, bias=True)
+        self.latent_norm = ChannelWiseRMSNorm(model_dim, eps=1e-6, affine=True)
+        self.fuse_in = nn.Conv2d(2 * model_dim, model_dim, kernel_size=1, bias=True)
+
+        # Time embedding
+        self.time_embed = SinusoidalTimeEmbeddingMLP(model_dim)
+
+        # AdaLN: shared base projector + per-block low-rank deltas
+        self.adaln_base = AdaLNZeroProjector(d_model=model_dim, d_cond=model_dim)
+        self.adaln_deltas = nn.ModuleList(
+            [
+                AdaLNZeroLowRankDelta(
+                    d_model=model_dim, d_cond=model_dim, rank=adaln_low_rank_rank
+                )
+                for _ in range(depth)
+            ]
+        )
+
+        # Block layout: start(2) + middle(depth-4) + end(2)
+        start_count = 2
+        end_count = 2
+        middle_count = depth - start_count - end_count
+        self._middle_start_idx = start_count
+        self._end_start_idx = start_count + middle_count
+
+        def _make_blocks(count: int) -> nn.ModuleList:
+            return nn.ModuleList(
+                [
+                    DiCoBlock(
+                        model_dim,
+                        mlp_ratio,
+                        depthwise_kernel_size=depthwise_kernel_size,
+                        use_external_adaln=True,
+                    )
+                    for _ in range(count)
+                ]
+            )
+
+        self.start_blocks = _make_blocks(start_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_count)
+
+        # Learned mask feature for path-drop guidance
+        self.mask_feature = nn.Parameter(torch.zeros((1, model_dim, 1, 1)))
+
+        # Output head
+        self.norm_out = ChannelWiseRMSNorm(model_dim, eps=1e-6, affine=True)
+        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: If True, replace middle block output with mask_feature (for PDG).
+
+        Returns:
+            x0 prediction [B, C, H, W].
+        """
+        # Patchify and normalize x_t
+        x_feat = self.patchify(x_t)
+        x_feat = self.norm_in(x_feat)
+
+        # Upsample and normalize latents, fuse with x_feat
+        z_up = self.latent_up(latents)
+        z_up = self.latent_norm(z_up)
+        fused = torch.cat([x_feat, z_up], dim=1)
+        fused = self.fuse_in(fused)
+
+        # Time conditioning
+        cond = self.time_embed(t.to(torch.float32).to(device=x_t.device))
+
+        # Start blocks
+        start_out = self._run_blocks(self.start_blocks, fused, cond, start_index=0)
+
+        # Middle blocks (or mask feature for PDG)
+        if drop_middle_blocks:
+            middle_out = self.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 fusion
+        skip_fused = torch.cat([start_out, middle_out], dim=1)
+        skip_fused = self.fuse_skip(skip_fused)
+
+        # End blocks
+        end_out = self._run_blocks(
+            self.end_blocks, skip_fused, cond, start_index=self._end_start_idx
+        )
+
+        # Output head
+        end_out = self.norm_out(end_out)
+        patches = self.out_proj(end_out)
+        return self.unpatchify(patches)

ir_diffae/dico_block.py

@@ -0,0 +1,111 @@
+"""DiCo block: conv path (1x1 -> depthwise -> SiLU -> CCA -> 1x1) + GELU MLP."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from .compact_channel_attention import CompactChannelAttention
+from .conv_mlp import ConvMLP
+from .norms import ChannelWiseRMSNorm
+
+
+class DiCoBlock(nn.Module):
+    """DiCo-style conv block with optional external AdaLN conditioning.
+
+    Two modes:
+    - Unconditioned (encoder): uses learned per-channel residual gates.
+    - External AdaLN (decoder): receives packed modulation [B, 4*C] via adaln_m.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        mlp_ratio: float,
+        *,
+        depthwise_kernel_size: int = 7,
+        use_external_adaln: bool = False,
+        norm_eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        self.channels = int(channels)
+        self.use_external_adaln = bool(use_external_adaln)
+
+        # Pre-norm for conv and MLP paths (no affine)
+        self.norm1 = ChannelWiseRMSNorm(self.channels, eps=norm_eps, affine=False)
+        self.norm2 = ChannelWiseRMSNorm(self.channels, eps=norm_eps, affine=False)
+
+        # Conv path: 1x1 -> depthwise kxk -> SiLU -> CCA -> 1x1
+        self.conv1 = nn.Conv2d(self.channels, self.channels, kernel_size=1, bias=True)
+        self.conv2 = nn.Conv2d(
+            self.channels,
+            self.channels,
+            kernel_size=depthwise_kernel_size,
+            padding=depthwise_kernel_size // 2,
+            groups=self.channels,
+            bias=True,
+        )
+        self.conv3 = nn.Conv2d(self.channels, self.channels, kernel_size=1, bias=True)
+        self.cca = CompactChannelAttention(self.channels)
+
+        # MLP path: GELU activation
+        hidden_channels = max(int(round(float(self.channels) * mlp_ratio)), 1)
+        self.mlp = ConvMLP(self.channels, hidden_channels, norm_eps=norm_eps)
+
+        # Conditioning: learned gates (encoder) or external adaln_m (decoder)
+        if not self.use_external_adaln:
+            self.gate_attn = nn.Parameter(torch.zeros(self.channels))
+            self.gate_mlp = nn.Parameter(torch.zeros(self.channels))
+
+    def forward(self, x: Tensor, *, adaln_m: Tensor | None = None) -> Tensor:
+        b, c = x.shape[:2]
+
+        if self.use_external_adaln:
+            if adaln_m is None:
+                raise ValueError(
+                    "adaln_m required for externally-conditioned DiCoBlock"
+                )
+            adaln_m_cast = adaln_m.to(device=x.device, dtype=x.dtype)
+            scale_a, gate_a, scale_m, gate_m = adaln_m_cast.chunk(4, dim=-1)
+        elif adaln_m is not None:
+            raise ValueError("adaln_m must be None for unconditioned DiCoBlock")
+
+        residual = x
+
+        # Conv path
+        x_att = self.norm1(x)
+        if self.use_external_adaln:
+            x_att = x_att * (1.0 + scale_a.view(b, c, 1, 1))  # type: ignore[possibly-undefined]
+        y = self.conv1(x_att)
+        y = self.conv2(y)
+        y = F.silu(y)
+        y = y * self.cca(y)
+        y = self.conv3(y)
+
+        if self.use_external_adaln:
+            gate_a_view = torch.tanh(gate_a).view(b, c, 1, 1)  # type: ignore[possibly-undefined]
+            x = residual + gate_a_view * y
+        else:
+            gate = self.gate_attn.view(1, self.channels, 1, 1).to(
+                dtype=y.dtype, device=y.device
+            )
+            x = residual + gate * y
+
+        # MLP path
+        residual_mlp = x
+        x_mlp = self.norm2(x)
+        if self.use_external_adaln:
+            x_mlp = x_mlp * (1.0 + scale_m.view(b, c, 1, 1))  # type: ignore[possibly-undefined]
+        y_mlp = self.mlp(x_mlp)
+
+        if self.use_external_adaln:
+            gate_m_view = torch.tanh(gate_m).view(b, c, 1, 1)  # type: ignore[possibly-undefined]
+            x = residual_mlp + gate_m_view * y_mlp
+        else:
+            gate = self.gate_mlp.view(1, self.channels, 1, 1).to(
+                dtype=y_mlp.dtype, device=y_mlp.device
+            )
+            x = residual_mlp + gate * y_mlp
+
+        return x

ir_diffae/encoder.py

@@ -0,0 +1,55 @@
+"""iRDiffAE encoder: patchify -> DiCoBlocks -> bottleneck projection."""
+
+from __future__ import annotations
+
+from torch import Tensor, nn
+
+from .dico_block import DiCoBlock
+from .norms import ChannelWiseRMSNorm
+from .straight_through_encoder import Patchify
+
+
+class Encoder(nn.Module):
+    """Deterministic encoder: Image [B,3,H,W] -> latents [B,bottleneck_dim,h,w].
+
+    Pipeline: Patchify -> RMSNorm -> DiCoBlocks (unconditioned) -> Conv1x1 -> RMSNorm(no affine)
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        patch_size: int,
+        model_dim: int,
+        depth: int,
+        bottleneck_dim: int,
+        mlp_ratio: float,
+        depthwise_kernel_size: int,
+    ) -> None:
+        super().__init__()
+        self.patchify = Patchify(in_channels, patch_size, model_dim)
+        self.norm_in = ChannelWiseRMSNorm(model_dim, eps=1e-6, affine=True)
+        self.blocks = nn.ModuleList(
+            [
+                DiCoBlock(
+                    model_dim,
+                    mlp_ratio,
+                    depthwise_kernel_size=depthwise_kernel_size,
+                    use_external_adaln=False,
+                )
+                for _ in range(depth)
+            ]
+        )
+        self.to_bottleneck = nn.Conv2d(
+            model_dim, bottleneck_dim, kernel_size=1, bias=True
+        )
+        self.norm_out = ChannelWiseRMSNorm(bottleneck_dim, eps=1e-6, affine=False)
+
+    def forward(self, images: Tensor) -> Tensor:
+        """Encode images [B,3,H,W] in [-1,1] to latents [B,bottleneck_dim,h,w]."""
+        z = self.patchify(images)
+        z = self.norm_in(z)
+        for block in self.blocks:
+            z = block(z)
+        z = self.to_bottleneck(z)
+        z = self.norm_out(z)
+        return z

ir_diffae/model.py

@@ -0,0 +1,291 @@
+"""IRDiffAE: standalone HuggingFace-compatible iRDiffAE model."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from torch import Tensor, nn
+
+from .config import IRDiffAEConfig, IRDiffAEInferenceConfig
+from .decoder import Decoder
+from .encoder import Encoder
+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
+    # Not a local directory — try HuggingFace Hub
+    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 IRDiffAE(nn.Module):
+    """Standalone iRDiffAE model for HuggingFace distribution.
+
+    A diffusion autoencoder that encodes images to compact latents and
+    decodes them back via iterative VP diffusion.
+
+    Usage::
+
+        model = IRDiffAE.from_pretrained("path/to/weights")
+        model = model.to("cuda", dtype=torch.bfloat16)
+
+        # Encode
+        latents = model.encode(images)  # images: [B,3,H,W] in [-1,1]
+
+        # Decode
+        recon = model.decode(latents, height=H, width=W)
+
+        # Reconstruct (encode + decode)
+        recon = model.reconstruct(images)
+    """
+
+    def __init__(self, config: IRDiffAEConfig) -> 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,
+        )
+
+        self.decoder = Decoder(
+            in_channels=config.in_channels,
+            patch_size=config.patch_size,
+            model_dim=config.model_dim,
+            depth=config.decoder_depth,
+            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,
+    ) -> IRDiffAE:
+        """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
+                (e.g. ``"data-archetype/irdiffae-v1"``).
+            dtype: Load weights in this dtype (float32 or bfloat16).
+            device: Target device.
+            revision: Git revision (branch, tag, or commit) for Hub downloads.
+            cache_dir: Where to cache Hub downloads. Uses HF default if None.
+
+        Returns:
+            Loaded model in eval mode.
+        """
+        model_dir = _resolve_model_dir(
+            path_or_repo_id, revision=revision, cache_dir=cache_dir
+        )
+        config = IRDiffAEConfig.load(model_dir / "config.json")
+        model = cls(config)
+
+        # Try safetensors first, fall back to .pt
+        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.
+
+        Args:
+            images: [B, 3, H, W] in [-1, 1], H and W must be 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))
+
+    @torch.no_grad()
+    def decode(
+        self,
+        latents: Tensor,
+        height: int,
+        width: int,
+        *,
+        inference_config: IRDiffAEInferenceConfig | None = None,
+    ) -> Tensor:
+        """Decode latents to images via VP diffusion.
+
+        Args:
+            latents: [B, bottleneck_dim, h, w] encoder latents.
+            height: Output image height (must be divisible by patch_size).
+            width: Output image width (must be divisible by patch_size).
+            inference_config: Optional inference parameters. Uses defaults if None.
+
+        Returns:
+            Reconstructed images [B, 3, H, W] in float32.
+        """
+        cfg = inference_config or IRDiffAEInferenceConfig()
+        config = self.config
+        batch = int(latents.shape[0])
+        device = latents.device
+
+        # Determine model dtype from parameters
+        try:
+            model_dtype = next(self.parameters()).dtype
+        except StopIteration:
+            model_dtype = torch.float32
+
+        # Validate dimensions
+        if height % config.patch_size != 0 or width % config.patch_size != 0:
+            raise ValueError(
+                f"height={height} and width={width} must be divisible by patch_size={config.patch_size}"
+            )
+
+        # Generate initial noise
+        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,
+        )
+
+        # Build schedule
+        schedule = get_schedule(cfg.schedule, cfg.num_steps).to(device=device)
+
+        # Construct initial state: sigma_start * noise
+        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,
+        )
+
+        # Disable autocast for numerical precision
+        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,
+            ) -> Tensor:
+                return self.decoder(
+                    x_t.to(dtype=model_dtype),
+                    t,
+                    latents.to(dtype=model_dtype),
+                    drop_middle_blocks=drop_middle_blocks,
+                )
+
+            # Select sampler
+            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_enabled=cfg.pdg_enabled,
+                pdg_strength=cfg.pdg_strength,
+                device=device,
+            )
+
+        return result
+
+    @torch.no_grad()
+    def reconstruct(
+        self,
+        images: Tensor,
+        *,
+        inference_config: IRDiffAEInferenceConfig | 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
+        )

ir_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)

ir_diffae/samplers.py

@@ -0,0 +1,226 @@
+"""DDIM and DPM++2M samplers for VP diffusion with x-prediction objective."""
+
+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)."""
+
+    def __call__(
+        self,
+        x_t: Tensor,
+        t: Tensor,
+        latents: Tensor,
+        *,
+        drop_middle_blocks: 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 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_enabled: bool = False,
+    pdg_strength: float = 1.5,
+    device: torch.device | None = None,
+) -> Tensor:
+    """Run DDIM sampling loop.
+
+    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_enabled: Whether to use Path-Drop Guidance.
+        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
+        t_vec = t_i.expand(batch_size).to(device=run_device, dtype=torch.float32)
+        if pdg_enabled:
+            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
+            )
+            x0_hat = x0_uncond + pdg_strength * (x0_cond - x0_uncond)
+        else:
+            x0_hat = forward_fn(state, t_vec, latents, drop_middle_blocks=False).to(
+                torch.float32
+            )
+
+        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_enabled: bool = False,
+    pdg_strength: float = 1.5,
+    device: torch.device | None = None,
+) -> Tensor:
+    """Run DPM++2M sampling loop.
+
+    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
+        t_vec = t_i.expand(batch_size).to(device=run_device, dtype=torch.float32)
+        if pdg_enabled:
+            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
+            )
+            x0_hat = x0_uncond + pdg_strength * (x0_cond - x0_uncond)
+        else:
+            x0_hat = forward_fn(state, t_vec, latents, drop_middle_blocks=False).to(
+                torch.float32
+            )
+
+        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

ir_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)

ir_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)

ir_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)

model.safetensors

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

technical_report.md

@@ -0,0 +1,709 @@
+# iRDiffAE v1.0 — Technical Report
+
+**i**REPA **Diff**usion **A**uto**E**ncoder = **iRDiffAE**
+
+A fast, single-GPU-trainable diffusion autoencoder with spatially structured
+latents for rapid downstream model convergence. Encoding runs ~5× faster than
+Flux VAE; single-step decoding runs ~3× faster.
+
+## Contents
+
+1. [VP Diffusion Parameterization](#1-vp-diffusion-parameterization)
+   - [Forward Process](#11-forward-process) · [Log SNR](#12-log-signal-to-noise-ratio) · [Cosine Schedule](#13-cosine-interpolated-schedule) · [X-Prediction](#14-x-prediction-objective) · [Sampling](#15-sampling)
+2. [Architecture](#2-architecture)
+   - [Overview](#21-overview) · [DiCo Block](#22-dico-block) · [Encoder](#23-encoder) · [Decoder](#24-decoder) · [AdaLN](#25-adaln-shared-base--low-rank-deltas) · [PDG](#26-path-drop-guidance-pdg)
+3. [Design Choices](#3-design-choices)
+   - [Convolutional Architecture](#31-convolutional-architecture) · [Single-Stride Encoder](#32-single-stride-encoder-with-final-bottleneck) · [Diffusion vs GAN Decoding](#33-diffusion-decoding-vs-gan-based-decoding) · [Skip Connection & PDG](#34-skip-connection-and-path-drop-guidance) · [iREPA](#35-half-channel-representation-alignment-irepa)
+4. [Model Configuration](#4-model-configuration)
+5. [Training](#5-training)
+   - [Data](#51-data) · [Timestep Sampling](#52-timestep-sampling) · [Latent Noise Sync](#53-latent-noise-synchronization-dito-regularization) · [Noise Standards](#54-pixel-vs-latent-noise-standards) · [Optimizer](#55-optimizer-and-hyperparameters) · [Loss](#56-loss)
+6. [Inference](#6-inference)
+   - [Sampling Pipeline](#61-sampling-pipeline) · [Recommended Settings](#62-recommended-settings) · [Usage](#63-usage)
+7. [Results](#7-results)
+   - [Interactive Viewer](#71-interactive-viewer) · [Inference Settings](#72-inference-settings) · [Global Metrics](#73-global-metrics) · [Per-Image PSNR](#74-per-image-psnr-db) · [Latent Smoothness](#75-latent-space-smoothness)
+
+**References:**
+
+- **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.
+- **SPRINT** — Park et al., *Sprint: Sparse-Dense Residual Fusion for Efficient Diffusion Transformers*, [arXiv:2510.21986](https://arxiv.org/abs/2510.21986), 2025.
+- **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.
+- **iREPA** — Singh et al., *What matters for Representation Alignment: Global Information or Spatial Structure?*, [arXiv:2512.10794](https://arxiv.org/abs/2512.10794), 2025.
+
+---
+
+## 1. VP Diffusion Parameterization
+
+iRDiffAE uses the variance-preserving (VP) diffusion framework from SiD2
+with an x-prediction objective.
+
+### 1.1 Forward Process
+
+Given a clean image \\(x_0\\), the forward process constructs a noisy sample at
+continuous time \\(t \in [0, 1]\\):
+
+$$x_t = \alpha_t \, x_0 + \sigma_t \, \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, s^2 I)$$
+
+where \\(s = 0.558\\) is the pixel-space noise standard deviation (estimated from
+the dataset image distribution) and the VP constraint holds:
+
+$$\alpha_t^2 + \sigma_t^2 = 1$$
+
+### 1.2 Log Signal-to-Noise Ratio
+
+The schedule is parameterized through the log signal-to-noise ratio:
+
+$$\lambda_t = \log \frac{\alpha_t^2}{\sigma_t^2}$$
+
+which monotonically decreases as \\(t \to 1\\) (pure noise). From \\(\lambda_t\\) we
+recover \\(\alpha_t\\) and \\(\sigma_t\\) via the sigmoid function:
+
+$$\alpha_t = \sqrt{\sigma(\lambda_t)}, \qquad \sigma_t = \sqrt{\sigma(-\lambda_t)}$$
+
+where \\(\sigma(\cdot)\\) is the logistic sigmoid.
+
+### 1.3 Cosine-Interpolated Schedule
+
+Following SiD2, the logSNR schedule uses cosine interpolation:
+
+$$\lambda(t) = -2 \log \tan(a \cdot t + b)$$
+
+where \\(a\\) and \\(b\\) are computed to satisfy the boundary conditions
+\\(\lambda(0) = \lambda_\text{max}\\) and \\(\lambda(1) = \lambda_\text{min}\\):
+
+$$b = \arctan\!\bigl(e^{-\lambda_\text{max}/2}\bigr), \qquad a = \arctan\!\bigl(e^{-\lambda_\text{min}/2}\bigr) - b$$
+
+SiD2 also defines a "shifted cosine" variant with resolution-dependent additive
+shifts \\(\Delta_\text{high}\\) and \\(\Delta_\text{low}\\):
+
+$$\lambda_\text{shifted}(t) = (1 - t) \cdot [\lambda(t) + \Delta_\text{high}] + t \cdot [\lambda(t) + \Delta_\text{low}]$$
+
+iRDiffAE uses \\(\lambda_\text{min} = -10\\), \\(\lambda_\text{max} = 10\\),
+\\(\Delta_\text{high} = 0\\), and \\(\Delta_\text{low} = 0\\) (no resolution-dependent
+shift), so the schedule reduces to the unshifted cosine interpolation.
+
+### 1.4 X-Prediction Objective
+
+The model predicts the clean image \\(\hat{x}_0 = f_\theta(x_t, t, z)\\)
+conditioned on the encoder latents \\(z\\).
+
+**Schedule-invariant loss.** Following SiD2, the training loss is defined as an
+integral over logSNR \\(\lambda\\), making it invariant to the choice of noise
+schedule:
+
+$$\mathcal{L}(x) = \int w(\lambda) \, \| x_0 - \hat{x}_0 \|^2 \, d\lambda$$
+
+Since timesteps are sampled uniformly \\(t \sim \mathcal{U}(0,1)\\) rather than
+integrated over \\(\lambda\\) directly, the change of variable
+\\(d\lambda = \frac{d\lambda}{dt} \, dt\\) introduces a Jacobian factor:
+
+$$\mathcal{L} = \mathbb{E}_{t \sim \mathcal{U}(0,1)} \left[ \left(-\frac{d\lambda}{dt}\right) \cdot w(\lambda(t)) \cdot \| x_0 - \hat{x}_0 \|^2 \right]$$
+
+**Sigmoid weighting.** SiD2 defines the weighting function in \\(\varepsilon\\)-prediction
+form as \\(\sigma(b - \lambda)\\) — a sigmoid centered at bias \\(b\\). Converting from
+\\(\varepsilon\\)-prediction to \\(x\\)-prediction MSE via
+\\(\|\varepsilon - \hat{\varepsilon}\|^2 = e^{\lambda} \|x_0 - \hat{x}_0\|^2\\)
+gives:
+
+$$\sigma(b - \lambda) \cdot e^{\lambda} = e^b \cdot \sigma(\lambda - b)$$
+
+Combining the Jacobian with the weighting, the per-sample weight used in
+training is:
+
+$$\text{weight}(t) = -\frac{1}{2} \frac{d\lambda}{dt} \cdot e^b \cdot \sigma(\lambda(t) - b)$$
+
+The bias \\(b = -2.0\\) controls the relative emphasis on high-SNR (low-noise) vs
+low-SNR (high-noise) timesteps. A more negative \\(b\\) shifts emphasis toward
+noisier timesteps.
+
+### 1.5 Sampling
+
+At inference, each timestep \\(t\\) in the schedule is first mapped to logSNR via
+the cosine-interpolated schedule (Section 1.3), then to diffusion coefficients:
+
+$$t \;\xrightarrow{\text{schedule}}\; \lambda(t) \;\xrightarrow{\text{sigmoid}}\; \alpha_t = \sqrt{\sigma(\lambda)}, \quad \sigma_t = \sqrt{\sigma(-\lambda)}$$
+
+**DDIM.** The default sampler uses a descending time schedule
+\\(t_0 > t_1 > \cdots > t_N\\) with \\(N\\) denoising steps. At each step:
+
+1. Predict \\(\hat{x}_0 = f_\theta(x_{t_i}, t_i, z)\\)
+2. Reconstruct \\(\hat{\varepsilon} = \frac{x_{t_i} - \alpha_{t_i} \hat{x}_0}{\sigma_{t_i}}\\)
+3. Step: \\(x_{t_{i+1}} = \alpha_{t_{i+1}} \hat{x}_0 + \sigma_{t_{i+1}} \hat{\varepsilon}\\)
+
+**DPM++2M.** Also supported as an alternative sampler, using a half-lambda
+(\\(\lambda/2\\)) exponential integrator for faster convergence with fewer steps.
+
+---
+
+## 2. Architecture
+
+### 2.1 Overview
+
+iRDiffAE consists of a deterministic encoder and an iterative VP diffusion
+decoder. The encoder maps an image to a compact spatial latent, and the decoder
+reconstructs the image by iteratively denoising from Gaussian noise,
+conditioned on both the latents and the diffusion timestep.
+
+```
+Encoder:  x ∈ ℝ^{B×3×H×W}  →  z ∈ ℝ^{B×C×h×w}     (deterministic, single pass)
+Decoder:  (z, t, x_t)       →  x̂₀ ∈ ℝ^{B×3×H×W}    (iterative, N diffusion steps)
+```
+
+where \\(h = H / p\\), \\(w = W / p\\), \\(p\\) is the patch size, and \\(C\\) is the
+bottleneck dimension.
+
+### 2.2 DiCo Block
+
+Both encoder and decoder use DiCo blocks (from the [DiCo paper](https://arxiv.org/abs/2505.11196)),
+a convolution-based alternative to transformer blocks. Each block consists of
+two residual paths:
+
+**Conv path:**
+
+$$y = \text{Conv}_{1 \times 1} \to \text{DWConv}_{k \times k} \to \text{SiLU} \to \text{CCA} \to \text{Conv}_{1 \times 1}$$
+
+**MLP path:**
+
+$$y = \text{Conv}_{1 \times 1} \to \text{GELU} \to \text{Conv}_{1 \times 1}$$
+
+where \\(\text{DWConv}_{k \times k}\\) is a depthwise convolution (default \\(k = 7\\))
+and \\(\text{CCA}\\) is Compact Channel Attention:
+
+$$\text{CCA}(y) = y \odot \sigma\bigl(\text{Conv}_{1 \times 1}(\text{AvgPool}(y))\bigr)$$
+
+Both paths use channel-wise RMSNorm (without affine parameters) as pre-norm.
+Residual connections use gating:
+
+- **Encoder (unconditioned):** learned per-channel gate parameters
+  \\(x \leftarrow x + g \cdot y\\), where \\(g\\) is a learnable vector initialized to zero.
+- **Decoder (conditioned):** AdaLN-Zero gating via
+  \\(x \leftarrow x + \tanh(g_\text{adaln}) \cdot y\\), where \\(g_\text{adaln}\\) comes
+  from the timestep conditioning.
+
+### 2.3 Encoder
+
+The encoder is deterministic — no variational posterior, no KL loss. Latent
+normalization uses channel-wise RMSNorm without affine parameters, following
+DiTo's finding that this outperforms KL regularization.
+
+```
+Input:       x ∈ ℝ^{B×3×H×W}
+Patchify:    PixelUnshuffle(p) → Conv 1×1     →  ℝ^{B×D×h×w}
+Norm:        ChannelWise RMSNorm (affine)
+Blocks:      DiCoBlock × depth_enc              (unconditioned, learned gates)
+Bottleneck:  Conv 1×1 (D → C)
+Norm out:    ChannelWise RMSNorm (no affine)
+Output:      z ∈ ℝ^{B×C×h×w}
+```
+
+### 2.4 Decoder
+
+The decoder predicts \\(\hat{x}_0\\) from noisy input \\(x_t\\), conditioned on
+encoder latents \\(z\\) and timestep \\(t\\).
+
+```
+Patchify x_t:  PixelUnshuffle(p) → Conv 1×1   →  ℝ^{B×D×h×w}
+Norm:          ChannelWise RMSNorm (affine)
+Upsample z:    Conv 1×1 (C → D) → RMSNorm     →  ℝ^{B×D×h×w}
+Fuse:          Concat[x_feat, z_up] → Conv 1×1 →  ℝ^{B×D×h×w}
+
+Time embed:    t → sinusoidal → MLP            →  cond ∈ ℝ^{B×D}
+
+Start blocks:  DiCoBlock × 2                    (AdaLN conditioned)
+Middle blocks: DiCoBlock × (depth - 4)          (AdaLN conditioned)
+Skip fusion:   Concat[start_out, middle_out] → Conv 1×1
+End blocks:    DiCoBlock × 2                    (AdaLN conditioned)
+
+Norm:          ChannelWise RMSNorm (affine)
+Output head:   Conv 1×1 (D → 3·p²) → PixelShuffle(p)  →  x̂₀ ∈ ℝ^{B×3×H×W}
+```
+
+### 2.5 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, scale-and-gate modulation with no shift, and a
+\\(\tanh\\) on the gate.
+
+A single base projector is shared across all decoder layers, and each layer
+adds a low-rank correction:
+
+$$m_i = \text{Base}(\text{SiLU}(\text{cond})) + \Delta_i(\text{SiLU}(\text{cond}))$$
+
+where \\(\text{Base}: \mathbb{R}^D \to \mathbb{R}^{4D}\\) is a linear projection
+(zero-initialized) and \\(\Delta_i: \mathbb{R}^D \xrightarrow{\text{down}} \mathbb{R}^r \xrightarrow{\text{up}} \mathbb{R}^{4D}\\)
+is a low-rank factorization with rank \\(r\\) (zero-initialized up-projection).
+
+The packed modulation \\(m_i \in \mathbb{R}^{B \times 4D}\\) is chunked into four
+vectors \\((\text{scale}_\text{conv}, \text{gate}_\text{conv}, \text{scale}_\text{mlp}, \text{gate}_\text{mlp})\\)
+which modulate the conv and MLP paths (no shift term):
+
+$$\hat{x} = \text{RMSNorm}(x) \odot (1 + \text{scale})$$
+$$x \leftarrow x + \tanh(\text{gate}) \cdot f(\hat{x})$$
+
+### 2.6 Path-Drop Guidance (PDG)
+
+At inference, iRDiffAE supports Path-Drop Guidance — a classifier-free
+guidance analogue that does not require training with conditioning dropout.
+Instead, it exploits the decoder's skip connection:
+
+1. **Conditional pass:** run all blocks normally → \\(\hat{x}_0^\text{cond}\\)
+2. **Unconditional pass:** replace the middle block output with a learned
+   mask feature \\(m \in \mathbb{R}^{1 \times D \times 1 \times 1}\\) (initialized
+   to zero), effectively dropping the deep processing path → \\(\hat{x}_0^\text{uncond}\\)
+3. **Guided prediction:** \\(\hat{x}_0 = \hat{x}_0^\text{uncond} + s \cdot (\hat{x}_0^\text{cond} - \hat{x}_0^\text{uncond})\\)
+
+where \\(s\\) is the guidance strength.
+
+---
+
+## 3. Design Choices
+
+### 3.1 Convolutional Architecture
+
+iRDiffAE 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
+several advantages:
+
+- **Resolution generalization.** Convolutions operate on local patches and
+  generalize naturally to arbitrary image dimensions without interpolating
+  position embeddings or suffering attention distribution shift from
+  sequence length changes with global attention. Convolutions are also
+  more efficient than sliding window attention for local operations.
+- **Translation invariance.** The built-in inductive bias of weight sharing
+  across spatial positions is well matched to reconstruction, where the same
+  local patterns (edges, textures, gradients) 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 for reconstruction.
+
+Transformers are better suited for image *generation* (where global context
+and long-range dependencies are essential), but convolutions are better
+suited for autoencoders. The DiCo block provides a well-tested,
+strong building block for convolutional diffusion models, combining depthwise
+convolutions with compact channel attention in a design that has been
+validated at scale.
+
+### 3.2 Single-Stride Encoder with Final Bottleneck
+
+The encoder uses a single spatial stride (via PixelUnshuffle at the input)
+followed by a stack of DiCo blocks operating at constant spatial resolution,
+then a final 1×1 convolution to project from model dimension \\(D\\) to
+bottleneck dimension \\(C\\). 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 \\(C\\) channels to retain. Progressive compression forces early
+layers to discard information before the full feature representation has been
+computed, which is both computationally heavier and representationally
+suboptimal.
+
+### 3.3 Diffusion Decoding vs. GAN-Based Decoding
+
+Empirically, diffusion autoencoders produce a much cleaner latent space than
+patch-GAN + LPIPS-driven VAEs. The iterative diffusion process acts as a
+strong structural prior on the decoder, which in turn relaxes the pressure
+on the encoder to encode every pixel perfectly — the latent space can focus
+on semantically meaningful structure rather than adversarial reconstruction
+artifacts. This makes diffusion AE latents easier for a downstream
+latent-space diffusion model to learn.
+
+**Training efficiency.** The diffusion AE training objective is a
+straightforward weighted MSE loss with no adversarial component — no
+discriminator, no LPIPS perceptual loss, no delicate GAN balancing. At batch
+size 128, the model uses less than 30 GB of VRAM and runs at 7–10 iterations
+per second, making it trainable on a single RTX 5090 in one to two days.
+By contrast, GAN + LPIPS-based VAEs require many days of H100 time and are
+notoriously difficult to stabilize, with no publicly known working recipe
+for training from scratch at comparable quality.
+
+### 3.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. The start blocks process the fused
+input (noised image + latents) at full fidelity, the middle blocks perform
+deeper processing, and the skip connection concatenates the start block
+output with the middle block output before the end blocks.
+
+This 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,
+   preventing the middle blocks from washing out high-frequency information.
+3. **Path-Drop Guidance (PDG).** At inference, replacing the middle block
+   output with a learned zero-initialized mask feature creates an
+   "unconditional" prediction that preserves the skip path but drops the
+   deep processing. Interpolating between the conditional and unconditional
+   predictions (as in classifier-free guidance) sharpens the output
+   distribution — and hence the reconstructed image — without requiring
+   any training-time conditioning dropout.
+
+### 3.5 Half-Channel Representation Alignment (iREPA)
+
+Singh et al. ([iREPA, arXiv:2512.10794](https://arxiv.org/abs/2512.10794)) show
+that **spatial structure** of pretrained encoder representations — not global
+semantic accuracy — drives generation quality when using representation
+alignment to guide diffusion training. Their method aligns internal diffusion
+features with patch tokens from a frozen vision encoder (e.g. DINOv2) using
+patch-wise cosine similarity, with a conv-based projection and spatial
+normalization to preserve local structure.
+
+iRDiffAE adopts iREPA but aligns only the **first half** of the bottleneck
+channels (64 of 128) to a frozen DINOv3-S teacher. The rationale: models like
+DINOv3-S are trained for semantic understanding and do not preserve
+high-frequency detail. Aligning all channels biases the encoder toward dropping
+fine detail in favour of semantic structure. By aligning only half, the
+bottleneck decomposes into:
+
+- **Channels 0–63 (aligned):** semantic and spatial structure, guided by the
+  teacher's patch tokens.
+- **Channels 64–127 (free):** fine detail and high-frequency information,
+  driven purely by the reconstruction loss.
+
+The alignment operates on the encoder output **after** the final RMSNorm
+(no affine), so the teacher sees unit-RMS normalized features.
+
+**Implementation details:**
+
+```
+Encoder latents z ∈ ℝ^{B×128×h×w}  (after RMSNorm)
+                    ↓
+        z_aligned = z[:, :64, :, :]
+                    ↓
+        Conv2d 3×3 (64 → 384, padding=1)   ← iREPA conv projection
+                    ↓
+        student tokens ∈ ℝ^{B×T×384}
+                    ↓
+        patch-wise cosine similarity with DINOv3-S tokens
+```
+
+The teacher's patch tokens are spatially normalized before comparison
+(\\(\gamma = 0.7\\), removing 70% of the global mean) following iREPA's
+prescription. The alignment loss is weighted at 0.5 for most of training,
+reduced to 0.25 toward the end to improve reconstruction fidelity.
+
+**Tradeoff.** The alignment costs 2–3 dB of average reconstruction PSNR
+compared to training without it. In exchange, downstream diffusion and flow
+matching models trained on the aligned latent space converge significantly
+faster — empirically validating the iREPA finding that spatial structure of
+the latent representation matters more than raw reconstruction fidelity for
+generation quality.
+
+---
+
+## 4. Model Configuration
+
+| Parameter | Value |
+|---|---|
+| Patch size \\(p\\) | 16 |
+| Bottleneck dim \\(C\\) | 128 |
+| Compression ratio | 6× |
+| Model dim \\(D\\) | 896 |
+| Total parameters | 133.4M |
+| Encoder depth | 4 |
+| Decoder depth | 8 |
+| Decoder layout | 2 start + 4 middle + 2 end |
+| MLP ratio | 4.0 |
+| Depthwise kernel | 7×7 |
+| AdaLN rank \\(r\\) | 128 |
+| \\(\lambda_\text{min}\\) | −10 |
+| \\(\lambda_\text{max}\\) | +10 |
+| Sigmoid bias \\(b\\) | −2.0 |
+| Pixel noise std \\(s\\) | 0.558 |
+
+**Compression ratio** = \\((3 \times p^2) / C\\): the factor by which the latent
+representation is smaller than the raw pixel data. With patch size 16 and 128
+bottleneck channels, the encoder produces a \\(16\times\\) spatial downsampling
+(\\(256\times\\) area reduction) at 6× total compression.
+
+---
+
+## 5. Training
+
+### 5.1 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).
+
+### 5.2 Timestep Sampling
+
+Timesteps are drawn via **stratified uniform sampling**, a variance reduction
+technique from Monte Carlo integration. The base distribution is uniform over
+the endpoint-trimmed domain \\([\varepsilon, 1 - \varepsilon]\\). Rather than
+drawing \\(B\\) i.i.d. samples (which can cluster or leave gaps by chance),
+stratified sampling divides the domain into \\(B\\) equal-mass buckets and draws
+exactly one sample per bucket:
+
+$$t_i = u_\text{lo} + (u_\text{hi} - u_\text{lo}) \cdot \frac{i + U_i}{B}, \qquad U_i \sim \mathcal{U}(0, 1), \quad i = 0, \ldots, B-1$$
+
+where \\(u_\text{lo} = F(\varepsilon)\\), \\(u_\text{hi} = F(1 - \varepsilon)\\), and
+\\(F\\) is the CDF of the base distribution (identity for uniform). This
+guarantees that every batch covers the full timestep range evenly, reducing
+the variance of the per-batch gradient estimate without introducing bias.
+
+Endpoint trimming uses \\(\varepsilon = \sigma(-7.5) \approx 5.5 \times 10^{-4}\\),
+keeping \\(|\lambda| \leq 15\\).
+
+### 5.3 Latent Noise Synchronization (DiTo Regularization)
+
+Following DiTo, encoder latents are regularized via noise synchronization
+during training. With probability \\(p = 0.1\\), a subset of clean latents \\(z_0\\)
+are replaced with noisy versions:
+
+$$z_\tau = (1 - \tau_\text{fm}) \cdot z_0 + \tau_\text{fm} \cdot \varepsilon_z, \qquad \varepsilon_z \sim \mathcal{N}(0, I)$$
+
+where \\(\tau\\) is sampled uniformly in \\([0, t]\\) (ensuring the latent is never
+noisier than the pixel-space input) and converted to a flow-matching time
+via the logSNR mapping, since downstream latent-space models are expected
+to use flow matching:
+
+$$\tau_\text{fm} = \sigma(-\tfrac{1}{2} \, \lambda(\tau))$$
+
+This synchronizes the noising process in latent space with pixel space,
+ensuring that the latent representation remains useful when a downstream
+latent diffusion model adds noise during its own forward process.
+
+### 5.4 Pixel vs. Latent Noise Standards
+
+The model uses different noise standard deviations in pixel space and
+latent space:
+
+- **Pixel space:** \\(s = 0.558\\), matching an estimate of the per-channel standard
+  deviation of natural images over the training dataset. This ensures that at
+  \\(t = 1\\) the noise distribution roughly matches the data distribution scale.
+- **Latent space:** \\(s = 1.0\\), because encoder latents are RMSNorm'd to unit
+  scale. Downstream latent diffusion models (which use flow matching)
+  operate with this unit-variance assumption.
+
+The conversion between pixel-space VP logSNR and latent-space flow-matching
+time uses the sigmoid mapping \\(t_\text{fm} = \sigma(-\frac{1}{2}\lambda)\\),
+which naturally accounts for the different noise scales.
+
+### 5.5 Optimizer and Hyperparameters
+
+| Hyperparameter | Value |
+|---|---|
+| Optimizer | AdamW |
+| Learning rate | \\(1 \times 10^{-4}\\) |
+| Weight decay | 0 |
+| Adam \\(\varepsilon\\) | \\(1 \times 10^{-8}\\) |
+| LR schedule | Constant (after warmup), halved for last 20% of training |
+| Warmup steps | 2,000 |
+| Batch size | 128 |
+| EMA decay | 0.9999 |
+| Precision | AMP bfloat16 (FP32 master weights, TF32 matmul) |
+| Compilation | `torch.compile` enabled |
+| Training steps | 700k |
+| Training images | ~5M |
+| Hardware | Single GPU |
+
+### 5.6 Loss
+
+$$\mathcal{L} = \mathcal{L}_\text{recon} + w_\text{repa} \cdot \mathcal{L}_\text{repa}$$
+
+\\(\mathcal{L}_\text{recon}\\) is the SiD2 sigmoid-weighted x-prediction MSE
+(Section 1.4) with bias \\(b = -2.0\\), computed in float32 for numerical
+stability.
+
+\\(\mathcal{L}_\text{repa}\\) is the iREPA half-channel alignment loss
+(Section 3.5): mean patch-wise negative cosine similarity between the first
+64 encoder channels (projected via 3×3 conv) and spatially-normalized
+DINOv3-S tokens. \\(w_\text{repa} = 0.5\\) for the majority of training,
+lowered to 0.25 toward the end to recover reconstruction fidelity.
+
+---
+
+## 6. Inference
+
+### 6.1 Sampling Pipeline
+
+Decoding proceeds by iteratively denoising from Gaussian noise
+(\\(\varepsilon \sim \mathcal{N}(0, s^2 I)\\) with \\(s = 0.558\\)). A descending
+time schedule \\(t_0 > t_1 > \cdots > t_{N-1}\\) is generated (linearly spaced
+by default), and at each step \\(t_i\\) is mapped to logSNR and then to diffusion
+coefficients:
+
+1. Compute \\(\lambda_i = \lambda(t_i)\\) via the cosine-interpolated schedule
+2. Derive \\(\alpha_i = \sqrt{\sigma(\lambda_i)}\\), \\(\sigma_i = \sqrt{\sigma(-\lambda_i)}\\)
+3. Run the DDIM or DPM++2M update step (Section 1.5)
+
+The initial state is \\(x_{t_0} = \sigma_{t_0} \cdot \varepsilon\\) (pure noise
+scaled by the first-step sigma).
+
+### 6.2 Recommended Settings
+
+**1 DDIM step** with **PDG disabled** is generally recommended — it achieves
+the best PSNR and is extremely fast (a single forward pass through the
+decoder). For images with sharp text or fine line art, 10–20 steps can
+sometimes improve edge crispness.
+
+| Setting | Recommended | Sharp text |
+|---|---|---|
+| Sampler | DDIM | DDIM or DPM++2M |
+| Steps | 1 | 10–20 |
+| Schedule | Linear | Linear |
+| PDG | Disabled | Disabled or 2.0 |
+
+**Reconstruction PSNR vs. decode steps** (N=2000 images, 2/3 photos + 1/3 book
+covers, EMA weights):
+
+| Decode steps | Avg PSNR (dB) |
+|---|---|
+| 1 | 33.71 |
+| 10 | 32.69 |
+| 20 | 32.30 |
+
+PSNR decreases slightly with more steps because the model is trained for
+single-step x-prediction; additional sampling steps introduce accumulated
+discretization error. The 128-channel bottleneck preserves enough information
+that a single decoder pass suffices for high-fidelity reconstruction.
+
+Multi-step sampling can help recover sharper edges on text and line art.
+PDG (strength 2–4) further increases perceptual sharpness but tends to
+hallucinate high-frequency detail — a direct manifestation of the
+**perception-distortion tradeoff**.
+
+**Inference latency** (batch of 4 × 256×256, bf16, NVIDIA RTX PRO 6000
+Blackwell, 100 iterations after warmup):
+
+| Operation | iRDiffAE | Flux.1 VAE | Flux.2 VAE |
+|---|---|---|---|
+| Encode | 2.1 ms | 11.6 ms | 9.1 ms |
+| Decode (1 step) | 8.3 ms | 24.9 ms | 20.0 ms |
+| Decode (10 steps) | 52.7 ms | — | — |
+| Decode (20 steps) | 100.6 ms | — | — |
+| Roundtrip (enc + 1-step dec) | 11.1 ms | 36.4 ms | 29.0 ms |
+
+Encoding is ~5× faster than Flux.1 and ~4× faster than Flux.2. Single-step
+decoding is ~3× faster than both Flux VAEs; multi-step decoding trades speed
+for perceptual sharpness.
+
+### 6.3 Usage
+
+```python
+from ir_diffae import IRDiffAE, IRDiffAEInferenceConfig
+
+model = IRDiffAE.from_pretrained("data-archetype/irdiffae-v1", device="cuda")  # bfloat16 by default
+
+# Encode
+latents = model.encode(images)  # [B, 3, H, W] → [B, 128, H/16, W/16]
+
+# Decode — PSNR-optimal (1 step, single forward pass)
+cfg = IRDiffAEInferenceConfig(num_steps=1, sampler="ddim")
+recon = model.decode(latents, height=H, width=W, inference_config=cfg)
+
+# Decode — perceptual sharpness (10 steps + PDG)
+cfg_sharp = IRDiffAEInferenceConfig(
+    num_steps=10, sampler="ddim", pdg_enabled=True, pdg_strength=2.0
+)
+recon_sharp = model.decode(latents, height=H, width=W, inference_config=cfg_sharp)
+```
+
+---
+
+## Citation
+
+```bibtex
+@misc{ir_diffae,
+  title   = {iRDiffAE: A Fast, Representation Aligned Diffusion Autoencoder with DiCo Blocks},
+  author  = {data-archetype},
+  year    = {2026},
+  month   = feb,
+  url     = {https://github.com/data-archetype/irdiffae},
+}
+```
+
+---
+
+## 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/disco-diffae-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) | 8 |
+
+> All models run in bfloat16. Timings measured on an NVIDIA RTX Pro 6000 (Blackwell).
+
+### 7.3 Global Metrics
+
+| Metric | p16_c128 | Flux.1 VAE | Flux.2 VAE |
+|--------|--------|--------|--------|
+| Avg PSNR (dB) | 31.75 | 32.66 | 34.13 |
+| Avg encode (ms/image) | 2.4 | 63.3 | 45.0 |
+| Avg decode (ms/image) | 5.4 | 135.4 | 90.5 |
+
+### 7.4 Per-Image PSNR (dB)
+
+| Image | p16_c128 | Flux.1 VAE | Flux.2 VAE |
+|-------|--------|--------|--------|
+| p640x1536:94623 | 30.99 | 31.28 | 33.50 |
+| p640x1536:94624 | 27.21 | 27.62 | 30.03 |
+| p640x1536:94625 | 30.48 | 31.65 | 33.98 |
+| p640x1536:94626 | 28.97 | 29.44 | 31.53 |
+| p640x1536:94627 | 29.17 | 28.70 | 30.53 |
+| p640x1536:94628 | 25.55 | 26.38 | 28.88 |
+| p960x1024:216264 | 40.92 | 40.87 | 45.39 |
+| p960x1024:216265 | 26.18 | 25.82 | 27.80 |
+| p960x1024:216266 | 43.61 | 47.77 | 46.20 |
+| p960x1024:216267 | 37.12 | 37.65 | 39.23 |
+| p960x1024:216268 | 35.75 | 35.27 | 36.13 |
+| p960x1024:216269 | 29.14 | 28.45 | 30.24 |
+| p960x1024:216270 | 32.06 | 31.92 | 34.18 |
+| p960x1024:216271 | 38.73 | 38.92 | 42.18 |
+| p704x1472:94699 | 40.81 | 40.43 | 41.79 |
+| p704x1472:94700 | 29.52 | 29.52 | 32.08 |
+| p704x1472:94701 | 35.01 | 35.43 | 37.90 |
+| p704x1472:94702 | 30.76 | 30.73 | 32.50 |
+| p704x1472:94703 | 28.49 | 29.08 | 31.35 |
+| p704x1472:94704 | 28.68 | 29.22 | 31.84 |
+| p704x1472:94705 | 35.91 | 36.38 | 37.44 |
+| p704x1472:94706 | 31.12 | 31.50 | 33.66 |
+| r256_p1344x704:15577 | 28.10 | 28.32 | 29.98 |
+| r256_p1344x704:15578 | 28.29 | 29.35 | 30.79 |
+| r256_p1344x704:15579 | 29.86 | 30.44 | 31.83 |
+| r256_p1344x704:15580 | 34.01 | 36.12 | 36.03 |
+| r256_p1344x704:15581 | 33.41 | 37.42 | 36.94 |
+| r256_p1344x704:15582 | 29.12 | 30.64 | 32.10 |
+| r256_p1344x704:15583 | 32.61 | 34.67 | 34.54 |
+| r256_p1344x704:15584 | 28.72 | 30.34 | 31.76 |
+| r256_p896x1152:144131 | 30.73 | 33.10 | 33.60 |
+| r256_p896x1152:144132 | 33.13 | 34.23 | 35.32 |
+| r256_p896x1152:144133 | 35.70 | 37.85 | 37.33 |
+| r256_p896x1152:144134 | 31.72 | 34.25 | 34.47 |
+| r256_p896x1152:144135 | 27.34 | 28.17 | 29.87 |
+| r256_p896x1152:144136 | 32.89 | 35.24 | 35.68 |
+| r256_p896x1152:144137 | 29.78 | 32.70 | 32.86 |
+| r256_p896x1152:144138 | 24.86 | 24.15 | 25.63 |
+