| --- |
| 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. |
|
|
