Upload folder using huggingface_hub

README.md

@@ -0,0 +1,147 @@
+---
+library_name: diffusers
+base_model: microsoft/Mage-Flow
+base_model_relation: adapter
+tags:
+- vae
+- autoencoder
+- flux.2
+- image-generation
+---
+
+# MegaFlow VAE (diffusers)
+
+A 🧨 diffusers-native `AutoencoderKLMega` port of the VAE from
+[microsoft/Mage-Flow](https://huggingface.co/microsoft/Mage-Flow)
+(DConvEncoder + DConvDenoiser/CoD decoder), usable as a drop-in adapter for
+pipelines that expect a Flux.2-style VAE.
+
+- **Encode:** one-step diffusion encoder (t=0) → `latent_dist` over `(B, 32, H/8, W/8)`
+- **Decode:** DConvDenoiser + CoD decoder (t=0) → image `(B, 3, H, W)` in `[-1, 1]`
+- **Latent space:** shaped *and* valued like the raw Flux.2 VAE latent — the native
+  128-channel `H/16` code is 2×2-unpatchified and denormalized with the Flux.2 BN
+  latent stats stored in `config.json` (anchor-latent regularization,
+  [arXiv:2607.19064](https://arxiv.org/abs/2607.19064)). No dependency on the
+  Flux.2 VAE at runtime (0.972 mean latent correlation over a 50-image test
+  set; see the [comparison](#comparison-with-the-original-flux2-vae) below).
+- **Checkpoint:** bf16, with the t=0 adaLN MLPs constant-folded at conversion time
+  (`folded: true` in the config), so it is ready to run on load.
+
+> [!NOTE]
+> The latents exposed by this model are **not** the native MegaVAE code. The
+> native `(B, 128, H/16, W/16)` code is **2×2-unpatchified** to
+> `(B, 32, H/8, W/8)` and **denormalized** with the Flux.2 BN latent statistics
+> (stored in `config.json`) so that `encode`/`decode` operate directly in the
+> original Flux.2 VAE latent space. Latents from this model can be decoded by
+> the Flux.2 VAE and vice versa.
+
+## Comparison with the original Flux.2 VAE
+
+Measured against `black-forest-labs/FLUX.2-dev` (subfolder `vae`,
+`AutoencoderKLFlux2`) on an NVIDIA L4, bf16, batch 1 (torch 2.8.0,
+diffusers 0.37.1). Speed/memory at 1024×1024, means over 5 runs after warmup,
+memory is peak CUDA allocation during the op. Quality is the mean over a
+50-image test set (06JPEG, ~2040×1524 photos, center-cropped to a multiple
+of 16, evaluated at native resolution with deterministic `.mode()` latents):
+
+| | MegaFlow VAE (this repo) | Flux.2 VAE |
+| --- | --- | --- |
+| Parameters | 100.8M | 84.0M |
+| Encode time | **23 ms** (~10× faster) | 240 ms |
+| Decode time | **71 ms** (~6× faster) | 441 ms |
+| Encode peak memory | **0.46 GiB** | 1.90 GiB |
+| Decode peak memory | **0.98 GiB** | 2.78 GiB |
+| Roundtrip PSNR (50 images) | 34.1 dB | 34.7 dB |
+
+**Quality** is on par with the original: same-VAE roundtrip reconstruction
+averages within ~0.5 dB of the Flux.2 VAE over the 50-image set. The latent
+spaces are interchangeable — mean latent correlation with raw Flux.2 latents
+is **0.972** (stable per image, ~0.97 on every one of the 50), and
+cross-decoding (`flux.decode(mega.encode(x))` at 34.0 dB,
+`mega.decode(flux.encode(x))` at 33.8 dB) stays within ~0.7 dB of the
+same-VAE roundtrip.
+
+**Repeated roundtrips** hold up slightly *better* than the original: MegaFlow
+starts ~0.5 dB below the Flux.2 VAE at one roundtrip but degrades more
+slowly, matching it by the second cycle and leading by ~1.1 dB after five —
+so it is well suited to iterative editing workflows. Alternating the two
+VAEs each cycle (MegaFlow encode → Flux.2 decode) tracks in between
+(50-image means):
+
+| PSNR vs original (dB) | k=1 | k=2 | k=3 | k=4 | k=5 |
+| --- | --- | --- | --- | --- | --- |
+| MegaFlow VAE ×k | 34.1 | 32.1 | **30.3** | **28.9** | **27.7** |
+| Flux.2 VAE ×k | **34.7** | 32.1 | 29.9 | 28.1 | 26.5 |
+| Alternating (Mega enc → Flux.2 dec) | 34.0 | 31.5 | 29.5 | 27.8 | 26.4 |
+
+**Speed** comes from the one-step architecture: both directions run a single
+t=0 forward pass of depthwise-conv (DiCo) blocks with the adaLN modulation
+constant-folded into the checkpoint, instead of the Flux.2 VAE's deep
+ResNet/attention encoder-decoder.
+
+**Memory** stays flat at high resolution: the decoder's per-patch MLP tail is
+chunked (`decode_chunk_size`, default 4096 = one 1024×1024 image worth of
+16×16 patches), so decode peak memory is roughly constant beyond 1024×1024
+instead of growing with image area.
+
+## Files
+
+| File | Purpose |
+| --- | --- |
+| `autoencoder_kl_mega.py` | Self-contained: `AutoencoderKLMega` (ModelMixin/ConfigMixin), all network blocks, and the one-time checkpoint converter |
+| `config.json` | Model config, incl. Flux.2 BN latent stats |
+| `diffusion_pytorch_model.safetensors` | Folded bf16 weights |
+
+## Usage
+
+The model class lives in this repo, so grab the code files alongside the weights:
+
+```python
+import sys
+
+import torch
+from huggingface_hub import snapshot_download
+
+repo_dir = snapshot_download("<your-username>/MegaFlow-VAE-diffusers")
+sys.path.insert(0, repo_dir)
+
+from autoencoder_kl_mega import AutoencoderKLMega
+
+vae = AutoencoderKLMega.from_pretrained(repo_dir, torch_dtype=torch.bfloat16).to("cuda")
+
+image = torch.rand(1, 3, 1024, 1024, device="cuda", dtype=torch.bfloat16) * 2 - 1
+latents = vae.encode(image).latent_dist.sample()   # (1, 32, 128, 128), Flux.2 latent space
+recon = vae.decode(latents, return_dict=False)[0]  # (1, 3, 1024, 1024) in [-1, 1]
+```
+
+Because the latents are shape- and value-compatible with the Flux.2 VAE, you can
+swap this in as the `vae` of a Flux.2 pipeline:
+
+```python
+pipe.vae = vae
+```
+
+Input `H`/`W` must be multiples of 16. `from_pretrained` accepts
+`fold_adaln=False` if you need the unfolded adaLN MLPs (only relevant for
+unfolded checkpoints).
+
+## Requirements
+
+```
+torch
+diffusers>=0.37
+safetensors
+loguru
+```
+
+## Conversion
+
+The checkpoint here was produced from the original CoD checkpoint layout with:
+
+```bash
+python autoencoder_kl_mega.py   # convert_mega_ckpt: MegaFlow/vae -> MegaFlow/vae_diffusers
+```
+
+which remaps `student.dconv_encoder.* -> encoder.*` and `pipeline.* -> decoder.*`,
+carries the Flux.2 BN stats into the config, constant-folds the t=0 adaLN MLPs
+(~74 MB smaller), and saves in bf16.

autoencoder_kl_mega.py

@@ -0,0 +1,753 @@
+"""
+Diffusers-native Mage-VAE (DConvEncoder + DConvDenoiser/CoD decoder).
+
+AutoencoderKLMega mirrors the AutoencoderKLFlux2Asym API surface used in this
+project — ModelMixin/ConfigMixin, `from_pretrained`/`save_pretrained`,
+`encode(x).latent_dist`, `decode(z, return_dict=False)[0]` — while exposing
+latents shaped AND valued like the raw Flux.2 VAE: (B, 32, H/8, W/8),
+2x2-unpatchified from the native 128ch @ H/16 code and denormalized with the
+Flux.2 BN latent stats stored in the model config (anchor-latent
+regularization, arXiv:2607.19064). No dependency on the Flux.2 VAE at runtime.
+
+Convert the original CoD checkpoint layout once:
+
+    python autoencoder_kl_mega.py            # MegaFlow/vae -> MegaFlow/vae_diffusers
+
+then load with:
+
+    vae = AutoencoderKLMega.from_pretrained("MegaFlow/vae_diffusers", torch_dtype=torch.bfloat16)
+"""
+
+from typing import List, Optional
+import math
+from functools import lru_cache
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.autoencoders.vae import DecoderOutput, DiagonalGaussianDistribution
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from diffusers.models.modeling_utils import ModelMixin
+from loguru import logger
+
+
+# ---------------------------------------------------------------------------
+# Primitive layers (vendored from GenCodec, inference subset)
+# ---------------------------------------------------------------------------
+def nonlinearity(x):
+    return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+def modulate(x, shift, scale):
+    if x.dim() == 4:
+        b, c = x.shape[:2]
+        return x * (1 + scale.view(b, c, 1, 1)) + shift.view(b, c, 1, 1)
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class LayerNorm2d(nn.LayerNorm):
+    def __init__(self, num_channels, eps=1e-6, affine=True):
+        super().__init__(num_channels, eps=eps, elementwise_affine=affine)
+
+    def forward(self, x):
+        # .contiguous() prevents a channels_last-strided NCHW view from
+        # propagating into downstream depthwise convs, which would otherwise
+        # hit a slow cuDNN path with a per-shape heuristic search.
+        x = x.permute(0, 2, 3, 1).contiguous()
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.permute(0, 3, 1, 2).contiguous()
+
+
+class _EncoderLayerNorm2d(LayerNorm2d):
+    pass
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x):
+        in_dtype = x.dtype
+        x = x.to(torch.float32)
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.variance_epsilon)
+        return self.weight * x.to(in_dtype)
+
+
+class TimestepEmbedder(nn.Module):
+    """DConv-style timestep MLP (max_period=10000, freq_size=256, hidden=384)."""
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half
+        ).to(t.device)
+        args = t[:, None].float() * freqs[None]
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
+        return emb
+
+    def forward(self, t):
+        emb = self.timestep_embedding(t, self.frequency_embedding_size)
+        return self.mlp(emb.to(self.mlp[0].weight.dtype))
+
+
+class BottleneckPatchEmbed(nn.Module):
+    """Image patch embed concatenated with a per-patch conditioning vector."""
+
+    def __init__(self, patch_size=16, in_chans=3, pca_dim=128, embed_dim=384, bias=True):
+        super().__init__()
+        self.proj1 = nn.Conv2d(in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False)
+        self.proj2 = nn.Conv2d(pca_dim + embed_dim, embed_dim, kernel_size=1, bias=bias)
+
+    def forward(self, x, cond):
+        return self.proj2(torch.cat([self.proj1(x), cond], dim=1))
+
+
+class DiCoBlock(nn.Module):
+    """DConv block with adaLN modulation."""
+
+    def __init__(self, hidden_size, mlp_ratio=4.0):
+        super().__init__()
+        self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.conv2 = nn.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True)
+        self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+
+        self.ca = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(hidden_size, hidden_size, 1, bias=True),
+            nn.Sigmoid(),
+        )
+
+        ffn = int(mlp_ratio * hidden_size)
+        self.conv4 = nn.Conv2d(hidden_size, ffn, 1, bias=True)
+        self.conv5 = nn.Conv2d(ffn, hidden_size, 1, bias=True)
+
+        self.norm1 = LayerNorm2d(hidden_size, affine=False)
+        self.norm2 = LayerNorm2d(hidden_size, affine=False)
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True),
+        )
+
+    def forward(self, inp, c):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+        x = modulate(self.norm1(inp), shift_msa, scale_msa)
+        x = F.gelu(self.conv2(self.conv1(x)))
+        x = x * self.ca(x)
+        x = self.conv3(x)
+        x = inp + gate_msa[..., None, None] * x
+        x = x + gate_mlp[..., None, None] * self.conv5(
+            F.gelu(self.conv4(modulate(self.norm2(x), shift_mlp, scale_mlp)))
+        )
+        return x
+
+
+class _EncoderDiCoBlock(nn.Module):
+    """DiCoBlock without adaLN, for the encoder pathway."""
+
+    def __init__(self, hidden_size, mlp_ratio=4.0):
+        super().__init__()
+        self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.conv2 = nn.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True)
+        self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.ca = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(hidden_size, hidden_size, 1, bias=True),
+            nn.Sigmoid(),
+        )
+        ffn = int(mlp_ratio * hidden_size)
+        self.conv4 = nn.Conv2d(hidden_size, ffn, 1, bias=True)
+        self.conv5 = nn.Conv2d(ffn, hidden_size, 1, bias=True)
+        self.norm1 = _EncoderLayerNorm2d(hidden_size)
+        self.norm2 = _EncoderLayerNorm2d(hidden_size)
+
+    def forward(self, inp):
+        x = self.norm1(inp)
+        x = F.gelu(self.conv2(self.conv1(x)))
+        x = x * self.ca(x)
+        x = self.conv3(x)
+        x = inp + x
+        return x + self.conv5(F.gelu(self.conv4(self.norm2(x))))
+
+
+class NerfEmbedder(nn.Module):
+    """Patch-position embedder used by the DConv decoder x-pathway."""
+
+    def __init__(self, in_channels, hidden_size_input, max_freqs=8):
+        super().__init__()
+        self.max_freqs = max_freqs
+        self.embedder = nn.Sequential(
+            nn.Linear(in_channels + max_freqs ** 2, hidden_size_input, bias=True),
+        )
+
+    @lru_cache
+    def fetch_pos(self, patch_size, device, dtype):
+        pos = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
+        pos_y, pos_x = torch.meshgrid(pos, pos, indexing="ij")
+        pos_x = pos_x.reshape(-1, 1, 1)
+        pos_y = pos_y.reshape(-1, 1, 1)
+        freqs = torch.linspace(0, self.max_freqs, self.max_freqs, dtype=dtype, device=device)
+        fx = freqs[None, :, None]
+        fy = freqs[None, None, :]
+        coeffs = (1 + fx * fy) ** -1
+        dct_x = torch.cos(pos_x * fx * torch.pi)
+        dct_y = torch.cos(pos_y * fy * torch.pi)
+        return (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs ** 2)
+
+    def forward(self, x):
+        B, P2, _ = x.shape
+        ps = int(P2 ** 0.5)
+        dct = self.fetch_pos(ps, x.device, x.dtype).expand(B, -1, -1)
+        return self.embedder(torch.cat([x, dct], dim=-1))
+
+
+class NerfFinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm = RMSNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, x):
+        return self.linear(self.norm(x))
+
+
+class SimpleMLPAdaLN(nn.Module):
+    """Final small MLP that maps NerfEmbedder features to per-patch RGB."""
+
+    def __init__(self, in_channels, model_channels, out_channels, z_channels, num_res_blocks, patch_size):
+        super().__init__()
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.patch_size = patch_size
+
+        self.cond_embed = nn.Linear(z_channels, patch_size ** 2 * model_channels)
+        self.input_proj = nn.Linear(in_channels, model_channels)
+
+        self.res_blocks = nn.ModuleList(_MLPResBlock(model_channels) for _ in range(num_res_blocks))
+
+    def forward(self, x, c):
+        x = self.input_proj(x)
+        c = self.cond_embed(c).reshape(c.shape[0], self.patch_size ** 2, -1)
+        for block in self.res_blocks:
+            x = block(x, c)
+        return x
+
+
+class _MLPResBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.in_ln = nn.LayerNorm(channels, eps=1e-6)
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, channels, bias=True),
+            nn.SiLU(),
+            nn.Linear(channels, channels, bias=True),
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(channels, 3 * channels, bias=True),
+        )
+
+    def forward(self, x, y):
+        shift, scale, gate = self.adaLN_modulation(y).chunk(3, dim=-1)
+        h = self.in_ln(x) * (1 + scale) + shift
+        return x + gate * self.mlp(h)
+
+
+class ResnetBlock(nn.Module):
+    """GroupNorm + Conv ResBlock used by the CoD Decoder."""
+
+    def __init__(self, *, in_channels, out_channels=None, dropout=0.0):
+        super().__init__()
+        out_channels = out_channels or in_channels
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
+        if in_channels != out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, 1)
+
+    def forward(self, x):
+        h = self.conv1(nonlinearity(self.norm1(x)))
+        h = self.conv2(self.dropout(nonlinearity(self.norm2(h))))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+
+class AttnBlock(nn.Module):
+    """Patched self-attention used at inference (eval mode of the original)."""
+
+    def __init__(self, in_channels, patch_size=32):
+        super().__init__()
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.norm = Normalize(in_channels)
+        self.q = nn.Conv2d(in_channels, in_channels, 1)
+        self.k = nn.Conv2d(in_channels, in_channels, 1)
+        self.v = nn.Conv2d(in_channels, in_channels, 1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, 1)
+
+    def forward(self, x):
+        h_ = self.norm(x)
+        Q = self.q(h_)
+        K = self.k(h_)
+        V = self.v(h_)
+
+        d = self.patch_size
+        b, c, H, W = Q.shape
+        pad_h = (d - H % d) % d
+        pad_w = (d - W % d) % d
+        if pad_h or pad_w:
+            Q = F.pad(Q, (0, pad_w, 0, pad_h), mode="replicate")
+            K = F.pad(K, (0, pad_w, 0, pad_h), mode="replicate")
+            V = F.pad(V, (0, pad_w, 0, pad_h), mode="replicate")
+        _, _, H_pad, W_pad = Q.shape
+        nph, npw = H_pad // d, W_pad // d
+        np_ = nph * npw
+
+        def to_patches(t):
+            return (t.reshape(b, c, nph, d, npw, d)
+                    .permute(0, 2, 4, 1, 3, 5)
+                    .reshape(b * np_, c, d * d))
+
+        Q = to_patches(Q)
+        K = to_patches(K)
+        V = to_patches(V)
+
+        w_ = torch.bmm(Q.permute(0, 2, 1), K) * (c ** -0.5)
+        w_ = F.softmax(w_, dim=2).permute(0, 2, 1)
+        h_ = torch.bmm(V, w_).reshape(b, nph, npw, c, d, d).permute(0, 3, 1, 4, 2, 5).reshape(b, c, H_pad, W_pad)
+        if pad_h or pad_w:
+            h_ = h_[:, :, :H, :W]
+        return x + self.proj_out(h_)
+
+
+# ---------------------------------------------------------------------------
+# adaLN constant-folding: at fixed t=0, adaLN_modulation(c) is constant.
+# Replace the MLP with a buffer so DiCoBlock.forward stays unchanged and
+# torch.compile can fuse the surrounding ops normally.
+# ---------------------------------------------------------------------------
+class _ConstAdaLN(nn.Module):
+    def __init__(self, modulation: torch.Tensor):
+        super().__init__()
+        self.register_buffer("modulation", modulation.detach().clone())
+
+    def forward(self, c):
+        b = c.shape[0]
+        if self.modulation.shape[0] != b:
+            return self.modulation.expand(b, *self.modulation.shape[1:])
+        return self.modulation
+
+
+def _replace_adaln_with_const(module: nn.Module, c: torch.Tensor) -> int:
+    # Only DiCoBlock is targeted: its adaLN is conditioned solely on t.
+    # Other adaLN_modulation submodules (e.g. _MLPResBlock in the decoder MLP)
+    # take a per-position latent and must not be folded.
+    n = 0
+    for child in module.modules():
+        if not isinstance(child, DiCoBlock):
+            continue
+        adaln = child.adaLN_modulation
+        if isinstance(adaln, _ConstAdaLN):
+            continue
+        with torch.no_grad():
+            mod = adaln(c)
+        child.adaLN_modulation = _ConstAdaLN(mod)
+        n += 1
+    return n
+
+
+# ---------------------------------------------------------------------------
+# CoD Decoder: latent → conditioning features for the denoiser
+# ---------------------------------------------------------------------------
+class _Decoder(nn.Module):
+    """ds=16, up2x=True, light=True only."""
+
+    def __init__(self, out_ch=384, z_ch=128):
+        super().__init__()
+        self.conv_in = nn.Conv2d(z_ch, out_ch, kernel_size=3, stride=1, padding=1)
+        self.block = nn.Sequential(
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+            AttnBlock(out_ch, patch_size=32),
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+            AttnBlock(out_ch, patch_size=32),
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+        )
+        self.norm_out = Normalize(out_ch)
+        self.conv_out = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1)
+        self.ada = nn.Identity()
+
+    def forward(self, z):
+        h = self.block(self.conv_in(z))
+        h = self.conv_out(nonlinearity(self.norm_out(h)))
+        return self.ada(h)
+
+
+# ---------------------------------------------------------------------------
+# DConvEncoder: image → packed (mean, logvar) latent
+# ---------------------------------------------------------------------------
+class _DConvEncoder(nn.Module):
+    def __init__(
+        self,
+        z_ch=128,
+        hidden_size=384,
+        num_blocks=21,
+        patch_size=16,
+        mlp_ratio=4.0,
+        head_size=768,
+        num_head_blocks=2,
+        out_ch_mult=2,
+    ):
+        super().__init__()
+        self.z_ch = z_ch
+        self.patch_size = patch_size
+        self.patch_cond_embed = nn.Conv2d(3, head_size, kernel_size=patch_size, stride=patch_size, bias=True)
+        self.head_blocks = nn.ModuleList([
+            _EncoderDiCoBlock(head_size, mlp_ratio=mlp_ratio) for _ in range(num_head_blocks)
+        ])
+        self.proj_down = nn.Conv2d(head_size, hidden_size, kernel_size=1, bias=True)
+        self.z_proj = nn.Conv2d(z_ch, hidden_size, kernel_size=1, bias=True)
+        self.fuse_proj = nn.Conv2d(hidden_size * 2, hidden_size, kernel_size=1, bias=True)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.blocks = nn.ModuleList([
+            DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_blocks)
+        ])
+        self.norm_out = LayerNorm2d(hidden_size)
+        self.proj_out = nn.Conv2d(hidden_size, z_ch * out_ch_mult, kernel_size=1, bias=True)
+
+    def forward_pred(self, z_t, t, y):
+        cond = self.patch_cond_embed(y)
+        for block in self.head_blocks:
+            cond = block(cond)
+        cond = self.proj_down(cond)
+
+        s = self.fuse_proj(torch.cat([cond, self.z_proj(z_t)], dim=1))
+        c = self.t_embedder(t.view(-1))
+        for block in self.blocks:
+            s = block(s, c)
+        return self.proj_out(self.norm_out(s))
+
+
+# ---------------------------------------------------------------------------
+# DConv denoiser: latent (via cond) + zero noise → reconstructed image
+# ---------------------------------------------------------------------------
+class _YEmbedder(nn.Module):
+    """Holds only the CoD decoder; the original Flux2 VAE encoder side is omitted."""
+
+    def __init__(self, ch=384, z_ch=128):
+        super().__init__()
+        self.decoder = _Decoder(out_ch=ch, z_ch=z_ch)
+
+
+class _DConvDenoiser(nn.Module):
+    def __init__(
+        self,
+        patch_size=16,
+        in_channels=3,
+        hidden_size=384,
+        hidden_size_x=32,
+        mlp_ratio=4.0,
+        num_blocks=24,
+        num_cond_blocks=21,
+        bottleneck_dim=128,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.hidden_size = hidden_size
+        self.num_cond_blocks = num_cond_blocks
+
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder_x = nn.Conv2d(hidden_size, hidden_size_x * patch_size ** 2, 1, 1, 0)
+        self.x_embedder = NerfEmbedder(in_channels + hidden_size_x, hidden_size_x, max_freqs=8)
+        self.s_embedder = BottleneckPatchEmbed(patch_size, in_channels, bottleneck_dim, hidden_size, bias=True)
+        self.blocks = nn.ModuleList([
+            DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_cond_blocks)
+        ])
+        self.dec_net = SimpleMLPAdaLN(
+            in_channels=hidden_size_x,
+            model_channels=hidden_size_x,
+            out_channels=in_channels,
+            z_channels=hidden_size,
+            num_res_blocks=num_blocks - num_cond_blocks,
+            patch_size=patch_size,
+        )
+        self.final_layer = NerfFinalLayer(hidden_size_x, in_channels)
+        self.y_embedder = _YEmbedder(ch=hidden_size, z_ch=bottleneck_dim)
+
+    def forward(self, x, t, cond, chunk_size=None):
+        b, _, h, w = x.shape
+        c = self.t_embedder(t.view(-1))
+
+        s = self.s_embedder(x, cond)
+        for block in self.blocks:
+            s = block(s, c)
+
+        length = s.shape[-2] * s.shape[-1]
+        s = s.permute(0, 2, 3, 1).reshape(b, length, self.hidden_size)
+
+        p2 = self.patch_size ** 2
+        x = torch.nn.functional.unfold(x, kernel_size=self.patch_size, stride=self.patch_size)
+
+        if chunk_size is None or chunk_size >= length:
+            x = torch.cat([x, self.y_embedder_x(cond).flatten(2)], dim=1)
+            x = x.reshape(b, -1, p2, length).permute(0, 3, 2, 1).flatten(0, 1)
+            x = self.x_embedder(x)
+            x = self.dec_net(x, s.reshape(-1, self.hidden_size))
+            x = self.final_layer(x)
+            x = x.transpose(1, 2).reshape(b, length, -1)
+            return torch.nn.functional.fold(
+                x.transpose(1, 2).contiguous(), (h, w),
+                kernel_size=self.patch_size, stride=self.patch_size,
+            )
+
+        # Chunked per-patch tail: each 16x16 output patch is independent here,
+        # so peak memory is capped at ~chunk_size patches with identical output.
+        cond_flat = cond.flatten(2)
+        out_cols = x.new_empty(b, self.in_channels * p2, length)
+        for i0 in range(0, length, chunk_size):
+            i1 = min(i0 + chunk_size, length)
+            n = i1 - i0
+            yx = self.y_embedder_x(cond_flat[:, :, i0:i1].unsqueeze(-1)).squeeze(-1)
+            xc = torch.cat([x[:, :, i0:i1], yx], dim=1)
+            xc = xc.reshape(b, -1, p2, n).permute(0, 3, 2, 1).flatten(0, 1)
+            xc = self.x_embedder(xc)
+            xc = self.dec_net(xc, s[:, i0:i1].reshape(-1, self.hidden_size))
+            xc = self.final_layer(xc)
+            out_cols[:, :, i0:i1] = xc.transpose(1, 2).reshape(b, n, -1).permute(0, 2, 1)
+        return torch.nn.functional.fold(
+            out_cols, (h, w), kernel_size=self.patch_size, stride=self.patch_size,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Wrapper
+# ---------------------------------------------------------------------------
+
+
+class AutoencoderKLMega(ModelMixin, ConfigMixin):
+    """
+    Encode: DConvEncoder (one-step diffusion, t=0) → latent_dist over (B, 32, H/8, W/8)
+    Decode: DConvDenoiser + CoD Decoder (t=0)      → image (B, 3, H, W) in [-1, 1]
+
+    With `flux_bn_mean`/`flux_bn_std` in the config, latents are emitted in and
+    accepted from the raw Flux.2 VAE latent space; without them, the normalized
+    anchor space. `latent_dist` is a DiagonalGaussianDistribution in the public
+    latent space (mean and logvar are transformed consistently), so both
+    `.mode()` and `.sample()` behave like the Flux.2 VAE's.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        latent_channels: int = 32,
+        downsample_factor: int = 8,
+        code_channels: int = 128,
+        code_downsample_factor: int = 16,
+        flux_bn_mean: Optional[List[float]] = None,
+        flux_bn_std: Optional[List[float]] = None,
+        decode_chunk_size: Optional[int] = 4096,
+        folded: bool = False,
+    ):
+        super().__init__()
+        self.encoder = _DConvEncoder()
+        self.decoder = _DConvDenoiser()
+
+        if folded:
+            # Checkpoint carries precomputed t=0 modulation buffers instead of
+            # the adaLN MLPs — install placeholder buffers so keys line up.
+            for mod in (self.encoder, self.decoder):
+                for child in mod.modules():
+                    if isinstance(child, DiCoBlock):
+                        out = child.adaLN_modulation[1].out_features
+                        child.adaLN_modulation = _ConstAdaLN(torch.zeros(1, out))
+
+        if flux_bn_mean is not None and flux_bn_std is not None:
+            if len(flux_bn_mean) != code_channels or len(flux_bn_std) != code_channels:
+                raise ValueError(
+                    f"flux_bn stats must have {code_channels} entries, "
+                    f"got {len(flux_bn_mean)}/{len(flux_bn_std)}"
+                )
+            mean = torch.tensor(flux_bn_mean, dtype=torch.float32).view(1, -1, 1, 1)
+            std = torch.tensor(flux_bn_std, dtype=torch.float32).view(1, -1, 1, 1)
+            self.register_buffer("bn_mean", mean, persistent=False)
+            self.register_buffer("bn_std", std, persistent=False)
+            self.register_buffer("bn_2logstd", 2.0 * std.log(), persistent=False)
+        else:
+            self.register_buffer("bn_mean", None, persistent=False)
+            self.register_buffer("bn_std", None, persistent=False)
+            self.register_buffer("bn_2logstd", None, persistent=False)
+
+    # -- Flux2 2x2 latent (un)packing, diffusers channel order ---------------
+    @staticmethod
+    def _patchify_latents(latents: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = latents.shape
+        latents = latents.view(b, c, h // 2, 2, w // 2, 2)
+        latents = latents.permute(0, 1, 3, 5, 2, 4)
+        return latents.reshape(b, c * 4, h // 2, w // 2)
+
+    @staticmethod
+    def _unpatchify_latents(latents: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = latents.shape
+        latents = latents.reshape(b, c // 4, 2, 2, h, w)
+        latents = latents.permute(0, 1, 4, 2, 5, 3)
+        return latents.reshape(b, c // 4, h * 2, w * 2)
+
+    def encode(self, x: torch.Tensor, return_dict: bool = True):
+        ds = self.config.code_downsample_factor
+        B, _, H, W = x.shape
+        if H % ds or W % ds:
+            raise ValueError(f"H, W must be multiples of {ds}, got ({H}, {W})")
+        z_t = torch.zeros(B, self.config.code_channels, H // ds, W // ds, device=x.device, dtype=x.dtype)
+        t = torch.zeros(B, device=x.device, dtype=x.dtype)
+        out = self.encoder.forward_pred(z_t, t, x)
+        mean = out[:, : self.config.code_channels]
+        logvar = out[:, self.config.code_channels :].clamp(min=-20.0, max=10.0)
+        if self.bn_mean is not None:
+            mean = mean * self.bn_std.to(mean.dtype) + self.bn_mean.to(mean.dtype)
+            logvar = logvar + self.bn_2logstd.to(logvar.dtype)
+        moments = torch.cat(
+            [self._unpatchify_latents(mean), self._unpatchify_latents(logvar)], dim=1
+        )
+        posterior = DiagonalGaussianDistribution(moments)
+        if not return_dict:
+            return (posterior,)
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def decode(self, z: torch.Tensor, return_dict: bool = True):
+        if z.shape[1] != self.config.latent_channels or z.shape[2] % 2 or z.shape[3] % 2:
+            raise ValueError(
+                f"expected Flux.2-shaped latent (B, {self.config.latent_channels}, H/8, W/8) "
+                f"with even spatial dims, got {tuple(z.shape)}"
+            )
+        z = self._patchify_latents(z)
+        if self.bn_mean is not None:
+            z = (z - self.bn_mean.to(z.dtype)) / self.bn_std.to(z.dtype)
+        cond = self.decoder.y_embedder.decoder(z)
+        B = z.shape[0]
+        H = z.shape[2] * self.config.code_downsample_factor
+        W = z.shape[3] * self.config.code_downsample_factor
+        noise = torch.zeros(B, 3, H, W, device=z.device, dtype=z.dtype)
+        t = torch.zeros(B, device=z.device, dtype=z.dtype)
+        sample = self.decoder.forward(noise, t, cond, chunk_size=self.config.decode_chunk_size)
+        if not return_dict:
+            return (sample,)
+        return DecoderOutput(sample=sample)
+
+    def forward(self, sample: torch.Tensor, return_dict: bool = True):
+        z = self.encode(sample, return_dict=False)[0].mode()
+        return self.decode(z, return_dict=return_dict)
+
+    @torch.no_grad()
+    def fold_adaln(self) -> int:
+        """Constant-fold the DiCoBlock adaLN MLPs at t=0 (we only run t=0).
+
+        Same optimization as mega_vae.MageVAE, folded in fp32 for numerical
+        parity with it (MageVAE folds before its bf16 cast). Mutates the module
+        structure (MLPs become buffers). No-op on checkpoints converted with
+        fold=True (already folded at conversion time).
+        """
+        if self.config.folded:
+            return 0
+        p = next(self.parameters())
+        n = 0
+        for mod in (self.encoder, self.decoder):
+            t = torch.zeros(1, device=p.device, dtype=torch.float32)
+            c = mod.t_embedder.float()(t)
+            for child in mod.modules():
+                if isinstance(child, DiCoBlock) and not isinstance(child.adaLN_modulation, _ConstAdaLN):
+                    child.adaLN_modulation = _ConstAdaLN(
+                        child.adaLN_modulation.float()(c).to(p.dtype)
+                    )
+                    n += 1
+            mod.t_embedder.to(p.dtype)
+        return n
+
+    @classmethod
+    def from_pretrained(cls, *args, fold_adaln: bool = True, **kwargs):
+        model = super().from_pretrained(*args, **kwargs)
+        if fold_adaln:
+            model.fold_adaln()
+        return model
+
+
+def convert_mega_ckpt(
+    src_ckpt: str = "MegaFlow/vae/diffusion_pytorch_model.safetensors",
+    src_config: str = "MegaFlow/vae/config.json",
+    dst_dir: str = "MegaFlow/vae_diffusers",
+    dtype: torch.dtype = torch.bfloat16,
+    fold: bool = True,
+) -> AutoencoderKLMega:
+    """One-time conversion from the CoD checkpoint layout to diffusers layout.
+
+    'student.dconv_encoder.*' -> 'encoder.*', 'pipeline.*' -> 'decoder.*'
+    (dropping the unused original-VAE 'y_embedder.encoder/bottleneck' branch);
+    BN stats are carried over from the source config.json. With fold=True the
+    t=0 adaLN MLPs are constant-folded in fp32 before the dtype cast and the
+    checkpoint stores the modulation buffers instead (~74 MB smaller,
+    ready-to-run on load with no fold step).
+    """
+    import json
+
+    from safetensors.torch import load_file
+
+    sd = load_file(src_ckpt, device="cpu")
+    new_sd = {}
+    for k, v in sd.items():
+        if k.startswith("student.dconv_encoder."):
+            new_sd["encoder." + k[len("student.dconv_encoder.") :]] = v
+        elif k.startswith("pipeline."):
+            nk = k[len("pipeline.") :]
+            if nk.startswith("y_embedder.encoder.") or nk.startswith("y_embedder.bottleneck."):
+                continue
+            new_sd["decoder." + nk] = v
+
+    with open(src_config) as f:
+        cfg = json.load(f)
+    model = AutoencoderKLMega(
+        flux_bn_mean=cfg.get("flux_bn_mean"), flux_bn_std=cfg.get("flux_bn_std")
+    )
+    missing, unexpected = model.load_state_dict(new_sd, strict=False)
+    logger.info(
+        f"convert_mega_ckpt: {len(new_sd)} keys mapped, missing={len(missing)}, "
+        f"unexpected={len(unexpected)}"
+    )
+    if missing:
+        raise RuntimeError(f"convert_mega_ckpt: missing model keys: {missing[:10]}")
+    if fold:
+        n = model.fold_adaln()
+        model.register_to_config(folded=True)
+        logger.info(f"convert_mega_ckpt: constant-folded {n} adaLN blocks at t=0")
+    model.to(dtype).save_pretrained(dst_dir)
+    logger.info(f"convert_mega_ckpt: saved to {dst_dir} ({dtype})")
+    return model
+
+
+if __name__ == "__main__":
+    convert_mega_ckpt()

config.json

@@ -0,0 +1,270 @@
+{
+  "_class_name": "AutoencoderKLMega",
+  "_diffusers_version": "0.37.1",
+  "code_channels": 128,
+  "code_downsample_factor": 16,
+  "decode_chunk_size": 4096,
+  "downsample_factor": 8,
+  "flux_bn_mean": [
+    -0.06738281,
+    -0.07128906,
+    -0.07519531,
+    -0.07470703,
+    0.02233887,
+    0.01794434,
+    0.01422119,
+    0.01831055,
+    -6.294e-05,
+    -0.0062561,
+    -0.00020981,
+    -0.00314331,
+    -0.02722168,
+    -0.02807617,
+    -0.02758789,
+    -0.02905273,
+    -0.07666016,
+    -0.06738281,
+    -0.09033203,
+    -0.08935547,
+    0.0168457,
+    0.01519775,
+    0.00787354,
+    0.00860596,
+    0.00836182,
+    0.00154114,
+    0.00025749,
+    -0.00427246,
+    -0.04394531,
+    -0.04199219,
+    -0.04370117,
+    -0.04321289,
+    -0.01025391,
+    -0.01318359,
+    -0.00662231,
+    -0.00476074,
+    -0.03100586,
+    -0.03051758,
+    -0.0279541,
+    -0.01794434,
+    0.00302124,
+    0.00150299,
+    0.01257324,
+    0.01446533,
+    0.03466797,
+    0.03369141,
+    0.03369141,
+    0.02832031,
+    0.00198364,
+    0.00473022,
+    0.00466919,
+    0.00497437,
+    0.01226807,
+    0.00811768,
+    0.00805664,
+    0.0145874,
+    0.06787109,
+    0.06787109,
+    0.07666016,
+    0.07324219,
+    -0.04614258,
+    -0.04736328,
+    -0.03930664,
+    -0.05102539,
+    -0.05273438,
+    -0.04785156,
+    -0.04711914,
+    -0.05175781,
+    -0.03173828,
+    -0.03173828,
+    -0.03442383,
+    -0.02819824,
+    0.05102539,
+    0.04443359,
+    0.05786133,
+    0.04589844,
+    -0.04125977,
+    -0.04589844,
+    -0.04882812,
+    -0.04663086,
+    -0.0088501,
+    -0.01062012,
+    -0.00878906,
+    -0.00460815,
+    -0.03759766,
+    -0.04321289,
+    -0.04345703,
+    -0.04980469,
+    0.01184082,
+    0.01660156,
+    0.02026367,
+    0.0279541,
+    0.0112915,
+    0.01287842,
+    0.0015564,
+    0.00714111,
+    -0.01177979,
+    -0.00183868,
+    -0.01416016,
+    -0.00537109,
+    -0.00909424,
+    -0.01379395,
+    -0.01446533,
+    -0.01867676,
+    0.03222656,
+    0.03051758,
+    0.02587891,
+    0.02990723,
+    0.05395508,
+    0.06152344,
+    0.04956055,
+    0.05908203,
+    -0.05102539,
+    -0.06030273,
+    -0.04785156,
+    -0.05249023,
+    -0.02270508,
+    -0.02746582,
+    -0.01538086,
+    -0.0255127,
+    -0.05712891,
+    -0.05639648,
+    -0.05175781,
+    -0.04956055,
+    0.01159668,
+    0.00543213,
+    0.01635742,
+    0.01037598
+  ],
+  "flux_bn_std": [
+    1.8028034,
+    1.77661192,
+    1.78538513,
+    1.78538513,
+    1.77220905,
+    1.75893426,
+    1.75893426,
+    1.75002849,
+    1.73207963,
+    1.73658431,
+    1.73207963,
+    1.73207963,
+    1.86248481,
+    1.8540765,
+    1.86248481,
+    1.85828543,
+    1.75893426,
+    1.75448704,
+    1.75448704,
+    1.75893426,
+    1.73658431,
+    1.74107718,
+    1.73658431,
+    1.74107718,
+    1.73207963,
+    1.7230351,
+    1.74107718,
+    1.73207963,
+    1.75448704,
+    1.75002849,
+    1.75448704,
+    1.75002849,
+    1.84562993,
+    1.84562993,
+    1.8540765,
+    1.8540765,
+    1.82434237,
+    1.78100395,
+    1.78538513,
+    1.7941153,
+    1.8028034,
+    1.78975558,
+    1.76779521,
+    1.76779521,
+    1.77661192,
+    1.76779521,
+    1.7941153,
+    1.78538513,
+    1.75893426,
+    1.74107718,
+    1.75002849,
+    1.73658431,
+    1.7941153,
+    1.79846454,
+    1.7941153,
+    1.7941153,
+    1.77661192,
+    1.77661192,
+    1.77220905,
+    1.77661192,
+    1.75448704,
+    1.74555862,
+    1.76779521,
+    1.75002849,
+    1.77661192,
+    1.77220905,
+    1.78100395,
+    1.78100395,
+    1.81575739,
+    1.8028034,
+    1.8028034,
+    1.8028034,
+    1.75448704,
+    1.73207963,
+    1.75448704,
+    1.75002849,
+    1.75893426,
+    1.75002849,
+    1.76779521,
+    1.76337028,
+    1.76779521,
+    1.75448704,
+    1.78100395,
+    1.77220905,
+    1.78975558,
+    1.79846454,
+    1.78538513,
+    1.77661192,
+    1.76337028,
+    1.76337028,
+    1.75448704,
+    1.75893426,
+    1.77661192,
+    1.76779521,
+    1.78100395,
+    1.78100395,
+    1.71849489,
+    1.70480204,
+    1.72756326,
+    1.70937848,
+    1.78100395,
+    1.75893426,
+    1.75002849,
+    1.74555862,
+    1.75893426,
+    1.75002849,
+    1.75448704,
+    1.73658431,
+    1.76337028,
+    1.75893426,
+    1.74555862,
+    1.74107718,
+    1.75002849,
+    1.75893426,
+    1.76337028,
+    1.76337028,
+    1.72756326,
+    1.71394277,
+    1.73207963,
+    1.7230351,
+    1.76779521,
+    1.76337028,
+    1.77220905,
+    1.76779521,
+    1.74555862,
+    1.74555862,
+    1.75893426,
+    1.75448704
+  ],
+  "folded": true,
+  "latent_channels": 32
+}

diffusion_pytorch_model.safetensors

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