Upload folder using huggingface_hub

README.md

@@ -0,0 +1,144 @@
+---
+license: apache-2.0
+tags:
+  - diffusion
+  - autoencoder
+  - image-reconstruction
+  - pytorch
+  - masked-autoencoder
+library_name: mdiffae
+---
+
+# mdiffae_v1
+
+**mDiffAE** — **M**asked **Diff**usion **A**uto**E**ncoder.
+A fast, single-GPU-trainable diffusion autoencoder with a **64-channel**
+spatial bottleneck. Uses decoder token masking as an implicit regularizer instead of
+REPA alignment, achieving Flux.2-level conditioning quality with a simpler
+flat decoder architecture (4 blocks, no skip connections).
+
+This variant (mdiffae_v1): 81.4M parameters, 310.6 MB.
+Bottleneck: **64 channels** at patch size 16
+(compression ratio 12x).
+
+## Documentation
+
+- [Technical Report](technical_report_mdiffae.md) — architecture, masking strategy, and results
+- [iRDiffAE Technical Report](https://huggingface.co/data-archetype/irdiffae-v1/blob/main/technical_report.md) — full background on VP diffusion, DiCo blocks, patchify encoder, AdaLN
+- [Results — interactive viewer](https://huggingface.co/spaces/data-archetype/mdiffae-results) — full-resolution side-by-side comparison
+
+## Quick Start
+
+```python
+import torch
+from m_diffae import MDiffAE
+
+# Load from HuggingFace Hub (or a local path)
+model = MDiffAE.from_pretrained("data-archetype/mdiffae_v1", device="cuda")
+
+# Encode
+images = ...  # [B, 3, H, W] in [-1, 1], H and W divisible by 16
+latents = model.encode(images)
+
+# Decode (1 step by default — PSNR-optimal)
+recon = model.decode(latents, height=H, width=W)
+
+# Reconstruct (encode + 1-step 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 | 81,410,624 |
+| File size | 310.6 MB |
+| Patch size | 16 |
+| Model dim | 896 |
+| Encoder depth | 4 |
+| Decoder depth | 4 |
+| Decoder topology | Flat sequential (no skip connections) |
+| Bottleneck dim | 64 |
+| MLP ratio | 4.0 |
+| Depthwise kernel | 7 |
+| AdaLN rank | 128 |
+| PDG mechanism | Token-level masking (ratio 0.75) |
+| Training regularizer | Decoder token masking (75% ratio, 50% apply prob) |
+
+**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. 4 flat
+sequential blocks (no skip connections). Supports token-level Path-Drop
+Guidance (PDG) at inference — very sensitive, use small strengths only.
+
+**Compared to iRDiffAE's decoder**: iRDiffAE uses an 8-block decoder split into
+start (2), middle (4), and end (2) groups with a skip connection that concatenates
+start-block output with middle-block output and fuses them through a Conv1x1 before
+the end blocks. PDG works by dropping the entire middle block computation and
+replacing it with a learned mask feature. In contrast, mDiffAE uses a simple flat
+stack of 4 blocks with no skip connections or block groups.
+PDG instead works at the token level: 75% of spatial tokens in the fused decoder
+input are replaced with a learned mask feature, providing a much finer-grained
+guidance signal. The bottleneck is also halved from 128 to 64
+channels, giving a 12x
+compression ratio vs iRDiffAE's 6x.
+
+### Key Differences from iRDiffAE
+
+| Aspect | iRDiffAE v1 | mDiffAE v1 |
+|---|---|---|
+| Bottleneck dim | 128 | **64** |
+| Decoder depth | 8 (2+4+2 skip-concat) | **4 (flat sequential)** |
+| PDG mechanism | Block dropping | **Token masking** |
+| Training regularizer | REPA + covariance reg | **Decoder token masking** |
+| PDG sensitivity | Moderate (1.5–3.0) | **Very sensitive (1.05–1.2)** |
+
+## Recommended Settings
+
+Best quality is achieved with just **1 DDIM step** and PDG disabled,
+making inference extremely fast.
+
+PDG in mDiffAE is **very sensitive** — use tiny strengths (1.05–1.2)
+if enabled. Higher values will cause artifacts.
+
+| Setting | Default |
+|---|---|
+| Sampler | DDIM |
+| Steps | 1 |
+| PDG | Disabled |
+| PDG strength (if enabled) | 1.1 |
+
+```python
+from m_diffae import MDiffAEInferenceConfig
+
+# PSNR-optimal (fast, 1 step)
+cfg = MDiffAEInferenceConfig(num_steps=1, sampler="ddim")
+recon = model.decode(latents, height=H, width=W, inference_config=cfg)
+```
+
+## Citation
+
+```bibtex
+@misc{m_diffae,
+  title   = {mDiffAE: A Masked Diffusion Autoencoder with Flat Decoder and Token-Level Guidance},
+  author  = {data-archetype},
+  year    = {2026},
+  month   = mar,
+  url     = {https://huggingface.co/data-archetype/mdiffae_v1},
+}
+```
+
+## Dependencies
+
+- PyTorch >= 2.0
+- safetensors (for loading weights)
+
+## License
+
+Apache 2.0

config.json

@@ -0,0 +1,15 @@
+{
+  "in_channels": 3,
+  "patch_size": 16,
+  "model_dim": 896,
+  "encoder_depth": 4,
+  "decoder_depth": 4,
+  "bottleneck_dim": 64,
+  "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,
+  "pdg_mask_ratio": 0.75
+}

m_diffae/__init__.py

@@ -0,0 +1,28 @@
+"""mDiffAE: Standalone diffusion autoencoder for HuggingFace distribution.
+
+Masked Diffusion AutoEncoder — a compact diffusion autoencoder
+that encodes images to spatial latents and decodes via iterative VP diffusion.
+Uses decoder token masking as an implicit regularizer instead of REPA alignment,
+with a flat 4-block decoder (no skip connections).
+
+Usage::
+
+    from m_diffae import MDiffAE, MDiffAEInferenceConfig
+
+    model = MDiffAE.from_pretrained("path/to/weights", device="cuda")
+
+    # Encode
+    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 + PDG)
+    cfg = MDiffAEInferenceConfig(num_steps=10, sampler="ddim", pdg_enabled=True)
+    recon = model.decode(latents, height=H, width=W, inference_config=cfg)
+"""
+
+from .config import MDiffAEConfig, MDiffAEInferenceConfig
+from .model import MDiffAE
+
+__all__ = ["MDiffAE", "MDiffAEConfig", "MDiffAEInferenceConfig"]

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

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

m_diffae/config.py

@@ -0,0 +1,56 @@
+"""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 MDiffAEConfig:
+    """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 = 4
+    bottleneck_dim: int = 64
+    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
+    # Token mask ratio for PDG (fraction of spatial tokens replaced with mask_feature)
+    pdg_mask_ratio: float = 0.75
+
+    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) -> MDiffAEConfig:
+        """Load config from JSON."""
+        data = json.loads(Path(path).read_text())
+        return cls(**data)
+
+
+@dataclass
+class MDiffAEInferenceConfig:
+    """User-tunable inference parameters with sensible defaults.
+
+    PDG is very sensitive in mDiffAE — use small strengths (1.05–1.2).
+    """
+
+    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 = 1.1
+    seed: int | None = None

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

m_diffae/decoder.py

@@ -0,0 +1,198 @@
+"""mDiffAE decoder: flat sequential DiCoBlocks with token-level PDG masking."""
+
+from __future__ import annotations
+
+import math
+
+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
+        -> Blocks (flat sequential, depth blocks) -> Norm -> Conv1x1 -> PixelShuffle
+
+    Token-level PDG: at inference, a fraction of spatial tokens in the fused input
+    are replaced with a learned mask_feature before the decoder blocks. Comparing
+    the masked vs unmasked outputs provides guidance signal.
+    """
+
+    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,
+        pdg_mask_ratio: float = 0.75,
+    ) -> None:
+        super().__init__()
+        self.patch_size = int(patch_size)
+        self.model_dim = int(model_dim)
+        self.pdg_mask_ratio = float(pdg_mask_ratio)
+
+        # 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)
+            ]
+        )
+
+        # Flat sequential blocks (no start/middle/end split, no skip connections)
+        self.blocks = nn.ModuleList(
+            [
+                DiCoBlock(
+                    model_dim,
+                    mlp_ratio,
+                    depthwise_kernel_size=depthwise_kernel_size,
+                    use_external_adaln=True,
+                )
+                for _ in range(depth)
+            ]
+        )
+
+        # Learned mask feature for token-level PDG
+        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 _apply_token_mask(self, fused: Tensor) -> Tensor:
+        """Replace a fraction of spatial tokens with mask_feature (2x2 groupwise).
+
+        Divides the spatial grid into 2x2 groups. Within each group, masks
+        floor(ratio * 4) tokens deterministically (lowest random scores).
+
+        Args:
+            fused: [B, C, H, W] fused decoder input.
+
+        Returns:
+            Masked tensor with same shape, where masked positions contain mask_feature.
+        """
+        b, c, h, w = fused.shape
+        # Pad to even dims if needed
+        h_pad = (2 - h % 2) % 2
+        w_pad = (2 - w % 2) % 2
+        if h_pad > 0 or w_pad > 0:
+            fused = torch.nn.functional.pad(fused, (0, w_pad, 0, h_pad))
+            _, _, h, w = fused.shape
+
+        # Reshape into 2x2 groups: [B, C, H/2, 2, W/2, 2] -> [B, C, H/2, W/2, 4]
+        x = fused.reshape(b, c, h // 2, 2, w // 2, 2)
+        x = x.permute(0, 1, 2, 4, 3, 5).reshape(b, c, h // 2, w // 2, 4)
+
+        # Random scores for each token in each group
+        scores = torch.rand(b, 1, h // 2, w // 2, 4, device=fused.device)
+
+        # Mask the floor(ratio * 4) lowest-scoring tokens per group
+        num_mask = math.floor(self.pdg_mask_ratio * 4)
+        if num_mask > 0:
+            # argsort ascending, mask the first num_mask
+            _, indices = scores.sort(dim=-1)
+            mask = torch.zeros_like(scores, dtype=torch.bool)
+            mask.scatter_(-1, indices[..., :num_mask], True)
+        else:
+            mask = torch.zeros_like(scores, dtype=torch.bool)
+
+        # Apply mask: replace masked tokens with mask_feature
+        mask_feat = self.mask_feature.to(device=fused.device, dtype=fused.dtype)
+        mask_feat = mask_feat.squeeze(-1).squeeze(-1)  # [1, C]
+        mask_feat = mask_feat.view(1, c, 1, 1, 1).expand_as(x)
+        mask_expanded = mask.expand_as(x)
+        x = torch.where(mask_expanded, mask_feat, x)
+
+        # Reshape back to [B, C, H, W]
+        x = x.reshape(b, c, h // 2, w // 2, 2, 2)
+        x = x.permute(0, 1, 2, 4, 3, 5).reshape(b, c, h, w)
+
+        # Remove padding if applied
+        if h_pad > 0 or w_pad > 0:
+            x = x[:, :, : h - h_pad, : w - w_pad]
+
+        return x
+
+    def forward(
+        self,
+        x_t: Tensor,
+        t: Tensor,
+        latents: Tensor,
+        *,
+        mask_tokens: 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].
+            mask_tokens: If True, apply token-level masking to decoder input (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)
+
+        # Token masking for PDG (replaces tokens with mask_feature)
+        if mask_tokens:
+            fused = self._apply_token_mask(fused)
+
+        # Time conditioning
+        cond = self.time_embed(t.to(torch.float32).to(device=x_t.device))
+
+        # Run all blocks sequentially
+        x = fused
+        for layer_idx, block in enumerate(self.blocks):
+            adaln_m = self._adaln_m_for_layer(cond, layer_idx=layer_idx)
+            x = block(x, adaln_m=adaln_m)
+
+        # Output head
+        x = self.norm_out(x)
+        patches = self.out_proj(x)
+        return self.unpatchify(patches)

m_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

m_diffae/encoder.py

@@ -0,0 +1,55 @@
+"""mDiffAE 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

m_diffae/model.py

@@ -0,0 +1,294 @@
+"""MDiffAE: standalone HuggingFace-compatible mDiffAE model."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from torch import Tensor, nn
+
+from .config import MDiffAEConfig, MDiffAEInferenceConfig
+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 MDiffAE(nn.Module):
+    """Standalone mDiffAE model for HuggingFace distribution.
+
+    A masked diffusion autoencoder that encodes images to compact latents and
+    decodes them back via iterative VP diffusion. Uses a flat 4-block decoder
+    with token-level masking for PDG instead of the skip-concat + block-drop
+    approach of iRDiffAE.
+
+    Usage::
+
+        model = MDiffAE.from_pretrained("data-archetype/mdiffae-v1")
+        model = model.to("cuda", dtype=torch.bfloat16)
+
+        # Encode
+        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: MDiffAEConfig) -> 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,
+            pdg_mask_ratio=config.pdg_mask_ratio,
+        )
+
+    @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,
+    ) -> MDiffAE:
+        """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/mdiffae-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 = MDiffAEConfig.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: MDiffAEInferenceConfig | 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 MDiffAEInferenceConfig()
+        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,
+                *,
+                mask_tokens: bool = False,
+            ) -> Tensor:
+                return self.decoder(
+                    x_t.to(dtype=model_dtype),
+                    t,
+                    latents.to(dtype=model_dtype),
+                    mask_tokens=mask_tokens,
+                )
+
+            # 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: MDiffAEInferenceConfig | 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
+        )

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

m_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,
+        *,
+        mask_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 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.1,
+    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 token-level Path-Drop Guidance.
+        pdg_strength: CFG-like strength for PDG (use small values: 1.05–1.2).
+        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, mask_tokens=True).to(
+                torch.float32
+            )
+            x0_cond = forward_fn(state, t_vec, latents, mask_tokens=False).to(
+                torch.float32
+            )
+            x0_hat = x0_uncond + pdg_strength * (x0_cond - x0_uncond)
+        else:
+            x0_hat = forward_fn(state, t_vec, latents, mask_tokens=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.1,
+    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, mask_tokens=True).to(
+                torch.float32
+            )
+            x0_cond = forward_fn(state, t_vec, latents, mask_tokens=False).to(
+                torch.float32
+            )
+            x0_hat = x0_uncond + pdg_strength * (x0_cond - x0_uncond)
+        else:
+            x0_hat = forward_fn(state, t_vec, latents, mask_tokens=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

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

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

m_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:196bd7a0607c495c0b72bb8c306c66d871792ea3d3f120afa14115ba3ca8e7ae
+size 325656824

technical_report_mdiffae.md

@@ -0,0 +1,211 @@
+# mDiffAE: Masked Diffusion AutoEncoder — Technical Report
+
+**Version 1** — March 2026
+
+## 1. Introduction
+
+mDiffAE (**M**asked **Diff**usion **A**uto**E**ncoder) builds on the [iRDiffAE](https://huggingface.co/data-archetype/irdiffae-v1/blob/main/technical_report.md) model family, which provides a fast, single-GPU-trainable diffusion autoencoder platform capable of high-quality image reconstruction. See that report for full background on the shared components: VP diffusion math (logSNR schedules, alpha/sigma, x-prediction), DiCo block architecture (depthwise conv + compact channel attention + GELU MLP), patchify encoder (PixelUnshuffle + 1×1 conv), shared-base + low-rank AdaLN-Zero conditioning, and the Path-Drop Guidance (PDG) concept.
+
+The iRDiffAE platform is designed to make it easy to experiment with different ways of regularizing the latent space. iRDiffAE v1 used REPA — aligning encoder features with a frozen DINOv2 teacher — which produces well-structured latents but tends toward overly smooth representations that are hard to reconcile with fine detail. Here we take a different approach entirely: **decoder token masking**.
+
+### 1.1 Token Masking as Regularizer
+
+A fraction of the time (50% of samples per batch), the decoder only sees **25% of tokens** in the fused conditioning input. The spatial token grid is divided into non-overlapping 2×2 groups, and within each group a single token is randomly kept while the other three are replaced with a learned mask feature. Hiding such a large fraction (75%) pushes the encoder to learn a form of representation consistency — each spatial token must carry enough information to support reconstruction even when most of its neighbors are absent. A smaller masking fraction helps downstream models learn sharp details quickly, but they fail to learn spatial coherence nearly as well — the task becomes too close to local inpainting, and the encoder is not pressured into globally consistent representations. The importance of a high masking ratio echoes findings in the masked autoencoder literature (He et al., 2022); we tested lower ratios and confirmed this tradeoff empirically.
+
+The 50% per-sample application probability is the knob that controls the compromise between reconstruction quality and latent space quality: samples that receive masking push the encoder toward consistent representations, while unmasked samples maintain reconstruction fidelity.
+
+### 1.2 Latent Noise Regularization
+
+To further regularize the latent space, we retain the random latent noising mechanism 10% of the time. However, unlike the pixel-space diffusion noise, the latent noise level is sampled independently using a **Beta(2,2)** distribution (stratified), with a **logSNR shift of +1.0** that biases it toward low noise levels (low *t* in our convention). This decouples the latent regularization schedule from the decoder's diffusion schedule, providing a gentle push toward noise-robust representations without disrupting reconstruction training.
+
+### 1.3 Simplified Decoder
+
+To keep the representational pressure on the encoder, we restrict the decoder to only **4 blocks** (down from 8 in iRDiffAE v1) and simplify it to a flat sequential architecture — no start/middle/end block groups, no skip connections. This halves the decoder's parameter count and makes it roughly 2× faster, while forcing the encoder to compensate by producing more informative latents.
+
+### 1.4 Empirical Results
+
+Compared to the REPA-regularized iRDiffAE v1, mDiffAE achieves slightly higher PSNR (to be confirmed with final benchmarks, but initial results were quite decisive) and produces a less oversmoothed but very locally consistent latent space PCA. In downstream diffusion model training, mDiffAE's latent space does not exhibit the very steep initial loss descent seen with iRDiffAE, but it quickly catches up after 50k–100k training steps, producing more spatially coherent images earlier with better high-frequency detail.
+
+### 1.5 References
+
+- He, K., Chen, X., Xie, S., Li, Y., Dollár, P., & Girshick, R. (2022). *Masked Autoencoders Are Scalable Vision Learners*. CVPR 2022.
+- Li, T., Chang, H., Mishra, S.K., Zhang, H., Katabi, D., & Krishnan, D. (2023). *MAGE: MAsked Generative Encoder to Unify Representation Learning and Image Synthesis*. CVPR 2023.
+
+## 2. Architecture Differences from iRDiffAE
+
+| Aspect | iRDiffAE v1 (halfrepa c128) | mDiffAE v1 (masked c64) |
+|--------|------------------------------|--------------------------|
+| Bottleneck dim | 128 | **64** |
+| Decoder depth | 8 (2 start + 4 middle + 2 end) | **4 (flat sequential)** |
+| Decoder topology | START_MIDDLE_END_SKIP_CONCAT | **FLAT (no skip concat)** |
+| Skip fusion | Yes (`fuse_skip` Conv1×1) | **No** |
+| PDG mechanism | Drop middle blocks → mask_feature | **Token-level masking** (75% spatial tokens → mask_feature) |
+| PDG sensitivity | Moderate (strength 1.5–3.0) | **Very sensitive** (strength 1.05–1.2 only) |
+| Training regularizer | REPA (half-channel DINOv2 alignment) + covreg | **Decoder token masking** (75% ratio, 50% apply prob) |
+| Latent noise reg | Same mechanism | **Independent Beta(2,2), logSNR shift +1.0, 10% prob** |
+| Depthwise kernel | 7×7 | 7×7 (same) |
+| Model dim | 896 | 896 (same) |
+| Encoder depth | 4 | 4 (same) |
+| Best decode | 1 step DDIM | 1 step DDIM (same) |
+
+## 3. Training-Time Masking Details
+
+### 3.1 Token Masking Procedure
+
+During training, with 50% probability per sample:
+1. The fused decoder input (patchified x_t + upsampled encoder latents) is divided into non-overlapping 2×2 spatial groups
+2. Within each group, 3 of 4 tokens (75%) are selected for masking using random scoring
+3. Masked tokens are replaced with a learned `mask_feature` parameter (same dimensionality as model_dim)
+4. The decoder processes the partially-masked input normally through all blocks
+
+### 3.2 Inference-Time PDG via Token Masking
+
+At inference, the trained mask_feature enables Path-Drop Guidance (PDG) through token-level masking rather than block-level dropping:
+
+- **Conditional pass**: Full decoder input (no masking)
+- **Unconditional pass**: Apply 2×2 groupwise token masking at the trained ratio (75%)
+- **Guided output**: `x0 = x0_uncond + strength × (x0_cond − x0_uncond)`
+
+Because the decoder has only 4 blocks and no skip connections, the guidance signal from token masking is very concentrated. This makes PDG extremely sensitive — even a strength of 1.2 produces noticeable sharpening, and values above 1.5 cause severe artifacts.
+
+## 4. Flat Decoder Architecture
+
+### 4.1 iRDiffAE v1 Decoder (for comparison)
+
+The iRDiffAE v1 decoder uses an 8-block layout split into three groups with a skip connection:
+
+```
+Fused input → Start blocks (2) → [save for skip] →
+  Middle blocks (4) → [cat with saved skip] → FuseSkip Conv1×1 →
+  End blocks (2) → Output head
+```
+
+The skip connection concatenates the start-block output with the middle-block output and fuses them through a learned Conv1×1 before feeding into the end blocks. For PDG, the entire middle block computation is dropped and replaced with a broadcasted learned `mask_feature`, effectively removing all 4 middle blocks from the forward pass. This produces a coarse "unconditional" signal for classifier-free guidance.
+
+### 4.2 mDiffAE v1 Decoder
+
+The mDiffAE decoder replaces this with a flat sequential architecture — no block groups, no skip connection:
+
+```
+Input: x_t [B, 3, H, W], t [B], z [B, 64, h, w]
+
+Patchify(x_t) → RMSNorm → x_feat [B, 896, h, w]
+LatentUp(z) → RMSNorm → z_up [B, 896, h, w]
+FuseIn(cat(x_feat, z_up)) → fused [B, 896, h, w]
+[Optional: token masking for PDG]
+TimeEmbed(t) → cond [B, 896]
+Block_0(fused, AdaLN(cond)) → ...
+Block_1(..., AdaLN(cond)) → ...
+Block_2(..., AdaLN(cond)) → ...
+Block_3(..., AdaLN(cond)) → out [B, 896, h, w]
+RMSNorm → Conv1x1 → PixelShuffle → x0_hat [B, 3, H, W]
+```
+
+With only 4 blocks and no skip fusion layer, the decoder has roughly half the parameters of iRDiffAE's decoder. The `fuse_skip` Conv1×1 layer is eliminated entirely. For PDG, instead of dropping blocks, 75% of spatial tokens in the fused input are replaced with a learned `mask_feature` before the blocks run. This token-level masking provides a finer-grained guidance signal but is much more sensitive to strength — the decoder sees the full block computation in both the conditional and unconditional paths, so the difference between them is subtle.
+
+### 4.3 Bottleneck
+
+The bottleneck dimension is halved from 128 channels (iRDiffAE) to 64 channels, giving a 12x compression ratio at patch size 16 (vs 6x for iRDiffAE). Despite the higher compression, the masking regularizer forces the encoder to produce informative per-token representations, maintaining reconstruction quality.
+
+## 5. Model Configuration
+
+| Parameter | Value |
+|-----------|-------|
+| `in_channels` | 3 |
+| `patch_size` | 16 |
+| `model_dim` | 896 |
+| `encoder_depth` | 4 |
+| `decoder_depth` | 4 |
+| `bottleneck_dim` | 64 |
+| `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 |
+| `pdg_mask_ratio` | 0.75 |
+
+Training checkpoint: step 708,000 (EMA weights).
+
+## 6. Inference Recommendations
+
+| Setting | Value | Notes |
+|---------|-------|-------|
+| Sampler | DDIM | Best for 1-step |
+| Steps | 1 | PSNR-optimal |
+| PDG | Disabled | Default, safest |
+| PDG strength | 1.05–1.2 | If enabled, very sensitive |
+
+## 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/mdiffae-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 | mdiffae_v1 (1 step) | Flux.1 VAE | Flux.2 VAE |
+|--------|--------|--------|--------|
+| Avg PSNR (dB) | 31.89 | 32.76 | 34.16 |
+| Avg encode (ms/image) | 2.4 | 63.9 | 45.7 |
+| Avg decode (ms/image) | 3.0 | 138.2 | 92.8 |
+
+### 7.4 Per-Image PSNR (dB)
+
+| Image | mdiffae_v1 (1 step) | Flux.1 VAE | Flux.2 VAE |
+|-------|--------|--------|--------|
+| p640x1536:94623 | 31.20 | 31.28 | 33.50 |
+| p640x1536:94624 | 27.32 | 27.62 | 30.03 |
+| p640x1536:94625 | 30.68 | 31.65 | 33.98 |
+| p640x1536:94626 | 29.14 | 29.44 | 31.53 |
+| p640x1536:94627 | 29.63 | 28.70 | 30.53 |
+| p640x1536:94628 | 25.60 | 26.38 | 28.88 |
+| p960x1024:216264 | 44.50 | 40.87 | 45.39 |
+| p960x1024:216265 | 26.42 | 25.82 | 27.80 |
+| p960x1024:216266 | 44.90 | 47.77 | 46.20 |
+| p960x1024:216267 | 37.78 | 37.65 | 39.23 |
+| p960x1024:216268 | 36.15 | 35.27 | 36.13 |
+| p960x1024:216269 | 29.37 | 28.45 | 30.24 |
+| p960x1024:216270 | 32.43 | 31.92 | 34.18 |
+| p960x1024:216271 | 41.23 | 38.92 | 42.18 |
+| p704x1472:94699 | 41.88 | 40.43 | 41.79 |
+| p704x1472:94700 | 29.66 | 29.52 | 32.08 |
+| p704x1472:94701 | 35.14 | 35.43 | 37.90 |
+| p704x1472:94702 | 30.90 | 30.73 | 32.50 |
+| p704x1472:94703 | 28.65 | 29.08 | 31.35 |
+| p704x1472:94704 | 28.98 | 29.22 | 31.84 |
+| p704x1472:94705 | 36.09 | 36.38 | 37.44 |
+| p704x1472:94706 | 31.53 | 31.50 | 33.66 |
+| r256_p1344x704:15577 | 27.89 | 28.32 | 29.98 |
+| r256_p1344x704:15578 | 28.07 | 29.35 | 30.79 |
+| r256_p1344x704:15579 | 29.56 | 30.44 | 31.83 |
+| r256_p1344x704:15580 | 32.89 | 36.12 | 36.03 |
+| r256_p1344x704:15581 | 32.26 | 37.42 | 36.94 |
+| r256_p1344x704:15582 | 28.74 | 30.64 | 32.10 |
+| r256_p1344x704:15583 | 31.99 | 34.67 | 34.54 |
+| r256_p1344x704:15584 | 28.42 | 30.34 | 31.76 |
+| r256_p896x1152:144131 | 30.02 | 33.10 | 33.60 |
+| r256_p896x1152:144132 | 33.19 | 34.23 | 35.32 |
+| r256_p896x1152:144133 | 35.42 | 37.85 | 37.33 |
+| r256_p896x1152:144134 | 31.41 | 34.25 | 34.47 |
+| r256_p896x1152:144135 | 27.13 | 28.17 | 29.87 |
+| r256_p896x1152:144136 | 32.75 | 35.24 | 35.68 |
+| r256_p896x1152:144137 | 28.60 | 32.70 | 32.86 |
+| r256_p896x1152:144138 | 24.76 | 24.15 | 25.63 |
+| VAE_accuracy_test_image | 31.52 | 36.69 | 35.25 |
+