Moebius 0.22B diffusion inpainting fp16/fp32 — first diffusion pipeline in coreai-community

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+moebius-unet-fp16-b2.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
+moebius-vae-decoder-fp16-b1.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
+moebius-vae-encoder-fp32-b2.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,85 @@
+---
+license: mit
+tags:
+- coreai
+- image-inpainting
+- image-to-image
+- diffusion
+- apple-silicon
+- moebius
+library_name: coreai
+---
+
+# Moebius-CoreAI
+
+[Moebius](https://github.com/hustvl/Moebius) — the 0.22B lightweight diffusion inpainting model
+(object removal / image completion, Places2 fine-tune) — as **CoreAI `.aimodel` assets** for
+Apple silicon, exported from the original [hustvl checkpoints](https://huggingface.co/hustvl/Moebius)
+(MIT weights).
+
+To our knowledge the first diffusion pipeline in `coreai-community`.
+
+| asset | role | dtype | size | PSNR vs PyTorch golden |
+|---|---|---|---|---|
+| `moebius-unet-fp16-b2.aimodel` | denoiser (CFG batch-2) | fp16 | 452 MB | **68.3 dB** |
+| `moebius-vae-encoder-fp32-b2.aimodel` | VAE posterior mean | fp32 | 137 MB | **104.7 dB** |
+| `moebius-vae-decoder-fp16-b1.aimodel` | VAE decoder | fp16 | 99 MB | **68.5 dB** |
+| `embedding_table.npy` | 20×3072 category conditioning | fp32 | 246 KB | exact |
+
+## Numbers (measured, M5 Max, macOS 27)
+
+- **UNet forward (fp16, GPU delegate): 49.8 ms** — 19-step CFG-2 projection **0.95 s**
+  (the MLX port of the same checkpoint: 117.6 ms / 2.23 s).
+- **Accuracy**: rel 9.338e-04 vs the shared PyTorch golden — the same fp16 floor as the MLX port
+  (9.257e-04). The export folds all 124 BatchNorms into fp64-precomputed per-channel scale/shift
+  (the checkpoint carries subnormal `running_var` channels that do not survive a naive fp16 cast).
+- The exported UNet carries exact, rank-safe rewrites of the LambdaNetworks attention (einsum →
+  broadcast/batched matmul; the positional Conv3d folded to a per-slice Conv2d) — numerically
+  gated at export (fp32 pre/post rel ≤ 5e-07).
+- The VAE **encoder ships fp32**: the SD-VAE encoder exceeds fp16 activation range (45.6 dB and
+  CPU-lane NaN at fp16 — the classic `sdxl-vae-fp16-fix` problem). One encode per image makes
+  fp32's cost invisible next to the denoise loop.
+
+## Placement — GPU today, honestly
+
+These assets run on the **GPU delegate**. Full-model Neural Engine compilation is currently
+blocked by an ANECCompiler bug we filed with a validated repro —
+[apple/coreai-models#138](https://github.com/apple/coreai-models/issues/138) (two 64²-level
+transformer instances per graph break the input-channel-split pass, value-dependently). 17/18
+model components already compile for ANE individually; when the OS compiler fixes #138 these
+assets inherit the ANE by re-export, no consumer change. Note the failure mode: an ANE request
+that fails **silently falls back to GPU** — verify placement with the GPU-idle signature, never
+by "it ran".
+
+## Usage
+
+Pipeline: encode `[image, masked_image]` (fp32, `[2,3,512,512]`, `[-1,1]`) → posterior mean ×
+0.13025 → DDIM (`scaled_linear` betas 0.00085–0.012, 20 steps, strength 0.99 → 19 steps from
+t=900, CFG 2.5, noise offset 0.0357) over the UNet (`sample` `[2,9,64,64]` fp16 =
+noisy(4)+mask(1)+masked(4), `timestep` `[2]` fp32, `encoder_hidden_states` `[2,10,3072]` fp16 =
+table rows [10..19; 0..9]) → decode `latents / 0.13025` (fp16, `[1,4,64,64]`) → `(x+1)/2`.
+
+A ready-made Swift package that does exactly this — scheduler, conditioning, image I/O,
+mask compositing, MLXEngine integration, tests —
+[`xocialize/coreai-moebius-swift`](https://github.com/xocialize/coreai-moebius-swift).
+
+```swift
+import CoreAI
+let model = try await AIModel(contentsOf: unetURL,
+    options: SpecializationOptions(preferredComputeUnitKind: .gpu))
+let fn = try model.loadFunction(named: "main")!
+// first load pays E5RT specialization (~40 s for the UNet, OS-cached after)
+```
+
+## Reproducibility
+
+`export_unet.py` and `export_vae.py` (in this repo) re-create every asset from the original
+checkpoints: PyTorch → `torch.export` → `coreai-torch` `TorchConverter` → `.aimodel`, with every
+graph rewrite numerically gated in-line. No opaque binaries.
+
+## Provenance & license
+
+Model: [hustvl/Moebius](https://github.com/hustvl/Moebius) (paper:
+[arXiv:2606.19195](https://arxiv.org/abs/2606.19195)) — **MIT weights**, Apache-2.0 reference
+code. VAE: the SD KL-f8 autoencoder distributed with PixelHacker (MIT). This repo redistributes
+the weights in a converted container under MIT, with the conversion scripts included.

embedding_table.npy

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

export_unet.py

@@ -0,0 +1,360 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "coreai-core==1.0.0b2",
+#     "coreai-torch==0.4.1",
+#     "diffusers",
+#     "timm",
+#     "einops",
+#     "pyyaml",
+#     "numpy",
+# ]
+#
+# [tool.uv]
+# index-url       = "https://pypi.org/simple"
+# prerelease      = "allow"
+# index-strategy  = "unsafe-best-match"
+# ///
+"""Export the Moebius UNet to a CoreAI .aimodel.
+
+WHY UNET-ONLY: the UNet is 38 of the 40 forwards per image and it IS the hypothesis under test
+(depthwise-separable + MBConv + linear attention on ANE vs Metal — memory `mlx-no-grouped-conv3d`).
+The VAE is 2 calls and does not move the measurement; it can follow using coreai-models' existing
+VAEEncoder/VAEDecoder wrappers if the answer is favourable.
+
+STATIC SHAPES throughout — required for ANE residency, and free here: Moebius is structurally
+locked to 512² (spatially-baked `rel_pos_emb` + a √n reshape in the attention wrapper), so the
+usual static-shape constraint costs nothing.
+
+Run:  uv run coreai/export_unet.py --dtype fp16
+"""
+import argparse
+import importlib
+import shutil
+import sys
+import time
+import types
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+
+ROOT = Path(__file__).resolve().parent.parent
+REF = ROOT / "reference"
+sys.path.insert(0, str(REF))
+
+CKPT = ROOT / "weights/Moebius/ft_places2/diffusion_pytorch_model.bin"
+CFG = REF / "config/model_cfg/moebius.yaml"
+NUM_EMBEDDINGS = 20
+
+
+def load_unet():
+    """The reference UNet, without executing `model_lib/__init__.py` (it eagerly imports a GLA
+    variant needing flash-linear-attention — CUDA-first and unused by Moebius)."""
+    for name, path in [
+        ("model_lib", REF / "model_lib"),
+        ("model_lib.nets", REF / "model_lib/nets"),
+        ("model_lib.nets.layers", REF / "model_lib/nets/layers"),
+    ]:
+        m = types.ModuleType(name)
+        m.__path__ = [str(path)]
+        sys.modules[name] = m
+    mod = importlib.import_module("model_lib.nets.unet_lambda_prune_lite")
+
+    cfg = yaml.safe_load(CFG.read_text())
+    model_cfg = dict(cfg["model"])
+    model_type = model_cfg.pop("model_type")
+    model_cfg["sample_size"] = cfg["data"]["image_size"] // cfg["vae"]["downsample_ratio"]
+    model_cfg["num_embeddings"] = NUM_EMBEDDINGS
+    net = getattr(mod, model_type)(**model_cfg)
+
+    sd = torch.load(CKPT, map_location="cpu", weights_only=True)
+    # The checkpoint is the RemovalModel state dict: `diff_model.*` + `embedding_layer.weight`.
+    unet_sd = {k[len("diff_model."):]: v for k, v in sd.items() if k.startswith("diff_model.")}
+    missing, unexpected = net.load_state_dict(unet_sd, strict=True)
+    print(f"[export] unet load: missing={len(missing)} unexpected={len(unexpected)}")
+    net.eval()   # the 124 BatchNorms must use running statistics
+    embedding = sd["embedding_layer.weight"]        # [20, 3072]
+    return net, embedding
+
+
+def patch_nearest_upsample(module: torch.nn.Module) -> int:
+    """Replace nearest-neighbour interpolate with repeat_interleave in `Upsample2D`.
+
+    LIFTED FROM coreai-models (`diffusion/components.py::_patch_nearest_upsample`) — and it is
+    load-bearing, not cosmetic: MPSGraph's segmenter REJECTS `coreai.interpolate` with
+    nearest_neighbor mode and routes those ops to the BNNS (CPU) backend. That both breaks
+    single-backend execution and inserts GPU→CPU→GPU copies at every upsample boundary. Exporting
+    without this yields a graph that quietly falls off the accelerator — and then a benchmark that
+    measures the wrong thing.
+
+    `repeat_interleave` is mathematically identical for integer scale factors.
+    """
+    from diffusers.models.upsampling import Upsample2D
+
+    patched = 0
+    for mod in module.modules():
+        if isinstance(mod, Upsample2D) and not mod.use_conv_transpose:
+            def _forward(hidden_states, output_size=None, _mod=mod):
+                h = hidden_states.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
+                return _mod.conv(h)
+            mod.forward = _forward
+            patched += 1
+    return patched
+
+
+def patch_lambda_einsums() -> None:
+    """Rewrite the two λ positional einsums to rank-≤4 matmul form, for ANE eligibility.
+
+    WHY (measured 2026-08-01): requesting `neuralEngine` on the unpatched export fails to compile —
+    17× `MPS-ANEC conversion failure: mps.reshape input/output rank 6 exceeds the max rank 5`, all
+    from `vanillaλ.py:146-147`, then `_ANECompiler: ANECCompile() FAILED`. torch.export decomposes
+    `einsum('n m k u, b u v m -> b n k v')` (six distinct indices) through rank-6 reshapes, and
+    **ANE's maximum tensor rank is 5**. The GPU delegate doesn't care; the ANE hard-rejects it.
+
+    Both equations fold to plain (batched) matmuls with NO change in value — same trick the MLX
+    port's `applyPositionalLambda` uses for memory reasons. One structural quirk, two backends,
+    two different symptoms.
+
+    SEAM: `_einsum` is a module-level lambda in `layers/utils.py`, but `vanillaλ.py` binds the NAME
+    at import (`from ..utils import _einsum`), so patching utils after the fact would be a no-op.
+    Rebinding the vanillaλ module global covers all four call sites (self- and cross-lambda) in one
+    move. Dispatch on the equation string; everything else falls through to the original — the
+    remaining λ einsums are rank ≤ 4 already and drew no validation warnings.
+
+    The export flow numerically gates this patch (fp32 eager, pre- vs post-patch) before casting.
+    """
+    vλ = importlib.import_module("model_lib.nets.layers.λ.vanillaλ")
+    original = vλ._einsum
+
+    def _patched(eq, *ops):
+        if eq == 'n m k u, b u v m -> b n k v':
+            # Broadcast-matmul form: [1,N,K,MU] @ [B,1,MU,Vd] → [B,N,K,Vd]. The earlier
+            # [NK, MU]-flattened form put N·K on one axis — 65536 at the 64² level, past the
+            # ANE's per-axis limit; this keeps every axis ≤ max(N, MU, K, Vd).
+            rel, V = ops                                        # [N,M,K,U], [B,U,Vd,M]
+            N, M, K, U = rel.shape
+            B, _, Vd, _ = V.shape
+            A = rel.permute(0, 2, 1, 3).reshape(1, N, K, M * U)
+            Bm = V.permute(0, 3, 1, 2).reshape(B, 1, M * U, Vd)
+            return (A @ Bm).contiguous()                        # [B,N,K,Vd]
+        if eq == 'b h k n, b n k v -> b h v n':
+            Q, lam = ops                                        # [B,H,K,N], [B,N,K,Vd]
+            Qbn = Q.permute(0, 3, 1, 2)                         # [B,N,H,K]
+            Y = Qbn @ lam                                       # [B,N,H,Vd] — batched, rank 4
+            return Y.permute(0, 2, 3, 1).contiguous()           # [B,H,Vd,N]
+        return original(eq, *ops)
+
+    vλ._einsum = _patched
+
+
+def patch_self_lambda_forward() -> None:
+    """Replace MultiQuerySelfLambda.forward with a rank-5-free, ANE-eligible formulation.
+
+    WHY (stage-bisected, probe_ane_selflambda.py): the self-λ takes the LOCAL positional branch —
+    `pos_conv = Conv3d(u, k, (1, r, r))` over V as [b,u,v,hh,ww]. The Conv3d itself compiles for
+    ANE (s4a: OK) — but any reshape/flatten CONSUMING its rank-5 output does not (s4e/s4f: FAIL;
+    s4d, the same matmul fed rank-4 tensors: OK). The fix never materialises rank 5: with u=1 and
+    depth-kernel 1, the Conv3d IS a Conv2d over each v-slice, so fold v into the conv batch and
+    land the output directly in matmul layout. The positional application then runs as a batched
+    matmul over n (the same rewrite as the MLX port's `applyPositionalLambda` — third appearance
+    of this contraction, third backend-specific formulation).
+
+    Numerically gated by the export's fp32 pre/post-patch eager comparison, same as the einsums.
+    """
+    import torch.nn.functional as F
+
+    vλ = importlib.import_module("model_lib.nets.layers.λ.vanillaλ")
+
+    def forward(self, x):                                   # x: [b, hh, ww, c]
+        b, hh, ww, _ = x.shape
+        n = hh * ww
+        xc = x.permute(0, 3, 1, 2)                          # 'b h w c -> b c h w'
+        q = self.to_q(xc)
+        k = self.to_k(xc)
+        v = self.to_v(xc)
+        Q = self.norm_q(q)
+        V = self.norm_v(v)
+        h, u = self.heads, self.u
+        dk = q.shape[1] // h
+        dv = V.shape[1] // u
+        Q = Q.reshape(b, h, dk, n)
+        k = k.reshape(b, u, dk, n).softmax(dim=-1)
+        V = V.reshape(b, u, dv, n)
+
+        lam_c = torch.einsum('b u k m, b u v m -> b k v', k, V)
+        Yc = torch.einsum('b h k n, b k v -> b h v n', Q, lam_c)
+
+        assert self.local_contexts and u == 1 and self.pos_conv.weight.shape[2] == 1, \
+            "rank-5-free fold assumes the local branch with u=1 and depth-kernel 1"
+        w2d = self.pos_conv.weight.squeeze(2)               # [k, u, r, r]
+        Vb = V.reshape(b * dv, u, hh, ww)                   # u=1: ONE rank-4 reshape, no rank-5
+        lam = F.conv2d(Vb, w2d, self.pos_conv.bias, padding=self.pos_conv.padding[1])
+        lam = lam.reshape(b, dv, dk, n).permute(0, 3, 2, 1)  # [b,n,k,v]
+        Yp = (Q.permute(0, 3, 1, 2) @ lam).permute(0, 2, 3, 1)   # [b,h,v,n]
+
+        Y = Yc + Yp
+        out = Y.reshape(b, h * dv, n).permute(0, 2, 1)      # 'b h v (hh ww) -> b (hh ww) c'
+        return out.reshape(b, hh, ww, h * dv)               # module contract: 'b h w c'
+
+    vλ.MultiQuerySelfLambda.forward = forward
+
+
+class PrecomputedBN(torch.nn.Module):
+    """BatchNorm replaced by per-channel scale/shift, constants computed at fp64 THEN cast.
+
+    WHY: the fp16 export sits at 41.4 dB vs the golden while MLX fp16 manages rel 9.3e-04 on the
+    same checkpoint. 25 of the 124 running_var tensors are below fp16's min-normal; evaluating
+    (x-mean)·rsqrt(var+eps) in fp16 arithmetic mangles those channels. The COMPOSITE constants
+    scale = γ/√(var+ε) and shift = β − mean·scale are fp16-representable even where var is not
+    (γ/√(8e-07) ≈ 3000γ ≪ 65504), so fold the four tensors into two at full precision first.
+    Numerically this is the same inference function — only the evaluation order changes.
+    """
+
+    def __init__(self, bn: torch.nn.Module, spatial: bool):
+        super().__init__()
+        var = bn.running_var.data.double()
+        mean = bn.running_mean.data.double()
+        gamma = bn.weight.data.double()
+        beta = bn.bias.data.double()
+        scale = gamma / torch.sqrt(var + bn.eps)
+        shift = beta - mean * scale
+        shape = (1, -1, 1, 1) if spatial else (1, -1, 1)
+        self.register_buffer("scale", scale.float().reshape(shape))
+        self.register_buffer("shift", shift.float().reshape(shape))
+        # ⚠️ 32 of the 124 "BatchNorms" are timm BatchNormAct2d — a subclass whose forward
+        # appends drop + activation (ReLU here). isinstance(BatchNorm2d) matches it, and a
+        # replacement that drops the activation diverges by rel ~1.0. The numeric gate caught
+        # this; carry the epilogue through.
+        self.act = getattr(bn, "act", None) or torch.nn.Identity()
+        self.drop = getattr(bn, "drop", None) or torch.nn.Identity()
+
+    def forward(self, x):
+        return self.act(self.drop(x * self.scale + self.shift))
+
+
+def patch_batchnorms_precomputed(module: torch.nn.Module) -> int:
+    replaced = 0
+    for parent in module.modules():
+        for name, child in list(parent.named_children()):
+            if isinstance(child, (torch.nn.BatchNorm2d, torch.nn.BatchNorm1d)):
+                setattr(parent, name,
+                        PrecomputedBN(child, spatial=isinstance(child, torch.nn.BatchNorm2d)))
+                replaced += 1
+    return replaced
+
+
+class MoebiusUNetWrapper(torch.nn.Module):
+    """Export surface: `(sample, timestep, encoder_hidden_states) -> noise prediction`.
+
+    The 20×3072 category table is deliberately left OUTSIDE the graph. Its lookup is a constant
+    gather (CFG always indexes rows 10–19 then 0–9), so the projected conditioning is identical on
+    every call — feeding it as an input keeps the graph free of an int64 embedding op, which is
+    friendlier to the accelerator, and lets the host hoist the lookup out of the 19-step loop
+    entirely.
+    """
+
+    def __init__(self, unet: torch.nn.Module) -> None:
+        super().__init__()
+        self.model = unet
+        n = patch_nearest_upsample(self.model)
+        print(f"[export] patched {n} Upsample2D module(s) → repeat_interleave")
+
+    def forward(self, sample, timestep, encoder_hidden_states):
+        return self.model(sample, timestep=timestep,
+                          encoder_hidden_states=encoder_hidden_states).sample
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dtype", default="fp16", choices=["fp16", "fp32"])
+    ap.add_argument("--batch", type=int, default=2, help="2 = CFG-doubled, the production shape")
+    ap.add_argument("--out", default=str(ROOT / "coreai/exports"))
+    args = ap.parse_args()
+
+    from coreai_torch import TorchConverter, get_decomp_table
+
+    net, embedding = load_unet()
+    wrapper = MoebiusUNetWrapper(net).eval()
+
+    # Gate the λ einsum rewrite numerically BEFORE any cast: fp32 eager, pre- vs post-patch.
+    # The rewrite is algebraically exact; this catches a transcription slip, not a design flaw.
+    b = args.batch
+    torch.manual_seed(0)
+    probe = (torch.randn(b, 9, 64, 64), torch.full((b,), 900, dtype=torch.float32),
+             torch.randn(b, 10, 3072))
+    with torch.no_grad():
+        pre = wrapper(*probe)
+    patch_lambda_einsums()
+    patch_self_lambda_forward()
+    n_bn = patch_batchnorms_precomputed(wrapper)
+    print(f"[export] replaced {n_bn} BatchNorms with fp64-precomputed scale/shift")
+    with torch.no_grad():
+        post = wrapper(*probe)
+    gap = (pre - post).abs().max().item() / (pre.abs().max().item() + 1e-12)
+    print(f"[export] λ einsum rewrite gate: rel {gap:.3e} (fp32 eager, pre vs post)")
+    if gap > 1e-5:
+        raise SystemExit("[export] λ rewrite diverged from the original — refusing to export.")
+
+    dtype = torch.float16 if args.dtype == "fp16" else torch.float32
+    if dtype == torch.float16:
+        # UNIFORM fp16 — including BatchNorm statistics, which DIFFERS from the MLX side.
+        #
+        # convert_weights.py pins BN running stats to fp32 as a precaution against a running_var
+        # rounding toward zero (rsqrt then explodes). That precaution is free on MLX. Here it is
+        # not: mixing fp32 BatchNorm into an fp16 graph makes the lowering fail outright —
+        #   "failed to legalize unresolved materialization from tensor<*xf32> to
+        #    tensor<2x1280x16x16xf16>" inside the λ cross-attention, because norm_q/norm_v emit
+        #   fp32 into fp16 einsums and PyTorch's silent promotion has no lowering equivalent.
+        #
+        # So the precaution was MEASURED rather than carried over: across all 124 running_var
+        # tensors the global minimum is 8.281e-07 — subnormal at fp16 but representable, and
+        # ZERO tensors round to zero. Even under flush-to-zero the result is bounded by
+        # eps (1/sqrt(1e-5) = 316), not infinite. Uniform fp16 is safe for THIS checkpoint;
+        # re-measure for any sibling before assuming it transfers.
+        wrapper = wrapper.half()
+
+    sample = torch.randn(b, 9, 64, 64, dtype=dtype)
+    timestep = torch.full((b,), 900, dtype=torch.float32)
+    context = torch.randn(b, 10, 3072, dtype=dtype)
+
+    print(f"[export] tracing — sample{tuple(sample.shape)} t{tuple(timestep.shape)} "
+          f"ctx{tuple(context.shape)} dtype={args.dtype}")
+    with torch.no_grad():
+        reference = wrapper(sample, timestep, context)
+    print(f"[export] eager forward ok → {tuple(reference.shape)}")
+
+    started = time.time()
+    ep = torch.export.export(wrapper, args=(sample, timestep, context))
+    ep = ep.run_decompositions(get_decomp_table())
+    print(f"[export] torch.export + decompositions: {time.time() - started:.1f}s")
+
+    started = time.time()
+    program = (
+        TorchConverter()
+        .add_exported_program(
+            ep,
+            input_names=["sample", "timestep", "encoder_hidden_states"],
+            output_names=["noise_pred"],
+        )
+        .to_coreai()
+    )
+    program.optimize()
+    print(f"[export] to_coreai + optimize: {time.time() - started:.1f}s")
+
+    out = Path(args.out) / f"moebius-unet-{args.dtype}-b{b}.aimodel"
+    out.parent.mkdir(parents=True, exist_ok=True)
+    if out.exists():
+        shutil.rmtree(out)
+    program.save_asset(out)   # wants a Path, not a str
+    size = sum(f.stat().st_size for f in out.rglob("*") if f.is_file()) / 1e6
+    print(f"[export] saved {out} ({size:.0f} MB)")
+
+    # The conditioning is constant — bake it next to the asset so the runtime never recomputes it.
+    np.save(Path(args.out) / "embedding_table.npy", embedding.float().numpy())
+    print(f"[export] wrote embedding_table.npy {tuple(embedding.shape)} (host-side constant gather)")
+
+
+if __name__ == "__main__":
+    main()

export_vae.py

@@ -0,0 +1,128 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "coreai-core==1.0.0b2",
+#     "coreai-torch==0.4.1",
+#     "diffusers",
+#     "numpy",
+# ]
+#
+# [tool.uv]
+# index-url       = "https://pypi.org/simple"
+# prerelease      = "allow"
+# index-strategy  = "unsafe-best-match"
+# ///
+"""Export the Moebius VAE (AutoencoderKL, KL-f8) to CoreAI .aimodel assets.
+
+Two assets, shaped for the pipeline's exact call pattern:
+  * encoder, batch 2, [2,3,512,512] -> posterior MEAN [2,4,64,64]
+    (one forward encodes image + masked_image together, as the pipeline does; the mean is the
+    deterministic moment the oracle/MLX ports gate on — no sampling in the graph)
+  * decoder, batch 1, [1,4,64,64] -> [1,3,512,512]
+
+scaling_factor stays OUT of the graph (host-side scalar), matching oracle semantics.
+
+`patch_nearest_upsample` is load-bearing here: the decoder carries 3 nearest-x2 Upsample2D
+modules, exactly the op MPSGraph's segmenter rejects (routes to BNNS/CPU) — same fix as the UNet.
+
+Run:  uv run coreai/export_vae.py
+"""
+import shutil
+import time
+from pathlib import Path
+
+import torch
+
+ROOT = Path(__file__).resolve().parent.parent
+VAE_DIR = ROOT / "weights/PixelHacker/vae"
+OUT = ROOT / "coreai/exports"
+
+
+def patch_nearest_upsample(module: torch.nn.Module) -> int:
+    from diffusers.models.upsampling import Upsample2D
+
+    patched = 0
+    for mod in module.modules():
+        if isinstance(mod, Upsample2D) and not mod.use_conv_transpose:
+            def _forward(hidden_states, output_size=None, _mod=mod):
+                h = hidden_states.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
+                return _mod.conv(h)
+            mod.forward = _forward
+            patched += 1
+    return patched
+
+
+class EncoderMean(torch.nn.Module):
+    """image [b,3,512,512] -> posterior mean [b,4,64,64] (deterministic; sf applied host-side)."""
+
+    def __init__(self, vae):
+        super().__init__()
+        self.encoder = vae.encoder
+        self.quant_conv = vae.quant_conv
+
+    def forward(self, image):
+        moments = self.quant_conv(self.encoder(image))
+        mean, _logvar = moments.chunk(2, dim=1)
+        return mean
+
+
+class Decoder(torch.nn.Module):
+    """latents [b,4,64,64] (UNSCALED — divide by sf host-side first) -> image [b,3,512,512]."""
+
+    def __init__(self, vae):
+        super().__init__()
+        self.post_quant_conv = vae.post_quant_conv
+        self.decoder = vae.decoder
+
+    def forward(self, latents):
+        return self.decoder(self.post_quant_conv(latents))
+
+
+def export(wrapper, example, name: str, dtype=torch.float16) -> None:
+    from coreai_torch import TorchConverter, get_decomp_table
+
+    # ⚠️ Eager sanity runs at fp32: torch's CPU fp16 conv path is `slow_conv2d` and a single
+    # 512² encoder forward at fp16 ground for 20+ CPU-MINUTES before being killed (the
+    # quantized-forward-on-CPU trap family). torch.export itself traces with fake tensors —
+    # no real compute — so only this sanity call ever executes kernels.
+    wrapper = wrapper.eval()
+    with torch.no_grad():
+        out = wrapper(*example)
+    print(f"[export] {name}: eager fp32 ok {tuple(example[0].shape)} -> {tuple(out.shape)}")
+    wrapper = wrapper.to(dtype)
+    example = tuple(t.to(dtype) for t in example)
+
+    started = time.time()
+    ep = torch.export.export(wrapper, args=example)
+    ep = ep.run_decompositions(get_decomp_table())
+    program = (TorchConverter()
+               .add_exported_program(ep, input_names=["x"], output_names=["out"])
+               .to_coreai())
+    program.optimize()
+    path = OUT / f"{name}.aimodel"
+    if path.exists():
+        shutil.rmtree(path)
+    program.save_asset(path)
+    size = sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) / 1e6
+    print(f"[export] saved {path.name} ({size:.0f} MB, {time.time() - started:.1f}s)")
+
+
+def main() -> None:
+    from diffusers.models import AutoencoderKL
+
+    vae = AutoencoderKL.from_pretrained(str(VAE_DIR)).eval()
+    print(f"[export] vae scaling_factor={vae.config.scaling_factor}")
+    n = patch_nearest_upsample(vae)
+    print(f"[export] patched {n} Upsample2D module(s) -> repeat_interleave")
+
+    # Encoder ships fp32: at fp16 it reads 45.6 dB (investigate) and produces NaN on the CPU
+    # lane — the classic SD-VAE fp16 activation-range problem, and mixed precision does not
+    # lower in a CoreAI graph (measured). One encode per image makes fp32's ~2x cost invisible.
+    export(EncoderMean(vae), (torch.randn(2, 3, 512, 512),), "moebius-vae-encoder-fp32-b2",
+           dtype=torch.float32)
+    # Decoder ships fp16: 68.5 dB [PASS] vs the shared golden.
+    export(Decoder(vae), (torch.randn(1, 4, 64, 64),), "moebius-vae-decoder-fp16-b1")
+
+
+if __name__ == "__main__":
+    main()

moebius-unet-fp16-b2.aimodel/main.hash

@@ -0,0 +1 @@
+¡:��b_Hrc�?�����$G,�S�����g

moebius-unet-fp16-b2.aimodel/main.mlirb

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

moebius-unet-fp16-b2.aimodel/metadata.json

@@ -0,0 +1,5 @@
+{
+  "assetVersion" : "2.0",
+  "producer" : "coreai-core 1.0.0b2",
+  "creationDate" : "20260801T200746Z"
+}

moebius-vae-decoder-fp16-b1.aimodel/main.hash

@@ -0,0 +1 @@
+'����,�lJ�{��[�%�#�̹p�}d����)

moebius-vae-decoder-fp16-b1.aimodel/main.mlirb

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

moebius-vae-decoder-fp16-b1.aimodel/metadata.json

@@ -0,0 +1,5 @@
+{
+  "assetVersion" : "2.0",
+  "creationDate" : "20260801T211110Z",
+  "producer" : "coreai-core 1.0.0b2"
+}

moebius-vae-encoder-fp32-b2.aimodel/main.hash

@@ -0,0 +1 @@
+�� !^9���~;�yy�b<qյ|`�a�fe?*�]

moebius-vae-encoder-fp32-b2.aimodel/main.mlirb

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

moebius-vae-encoder-fp32-b2.aimodel/metadata.json

@@ -0,0 +1,5 @@
+{
+  "assetVersion" : "2.0",
+  "creationDate" : "20260801T211107Z",
+  "producer" : "coreai-core 1.0.0b2"
+}