| --- |
| library_name: mlx |
| license: mit |
| pipeline_tag: image-to-image |
| base_model: microsoft/Mage-Flow-Edit-Turbo |
| tags: |
| - mlx |
| - mlx-swift |
| - image-editing |
| - instruction-based-editing |
| - diffusion |
| - rectified-flow |
| - mage-flow |
| --- |
| |
| # Mage-Flow-Edit-Turbo — MLX-Swift |
|
|
| An Apple-Silicon MLX-Swift port of |
| [microsoft/Mage-Flow-Edit-Turbo](https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo): |
| native-resolution, 4-step instruction-based image editing. |
|
|
| > **Multi-component diffusion pipeline.** Per mlx-community conventions, a |
| > pipeline like this (NR-MMDiT + MageVAE + Qwen3-VL conditioner + scheduler) has |
| > no single-`model_type` slot, so it is published under this namespace rather |
| > than a re-hosted `mlx-community/*` repo. The large component weights come from |
| > the upstream `microsoft/Mage-Flow-Edit-Turbo` repo (MIT) and are **not** |
| > re-hosted here — only the port's genuine artifact is. |
| |
| ## What's in this repo |
| |
| | file | what | |
| |---|---| |
| | `folded_adaln.safetensors` | The MageVAE adaLN constants baked at t=0 — the one weight the port adds (0.39 MB, replaces 37.7 M params of `t_embedder` + `adaLN_modulation`). | |
| | `*_config.json`, `model_index.json` | Component configs, for reference. | |
| | `sample_edit_512.png` | A validated edit ("make the background a snowy forest"). | |
|
|
| ## Use |
|
|
| Code + gates: **[github.com/xocialize/mage-flow-swift](https://github.com/xocialize/mage-flow-swift)**. |
|
|
| ```bash |
| # 1. download the upstream component weights |
| hf download microsoft/Mage-Flow-Edit-Turbo --local-dir Mage-Flow-Edit-Turbo |
| # 2. add this repo's baked artifact to the snapshot root |
| hf download xocialize/Mage-Flow-Edit-Turbo-mlx folded_adaln.safetensors \ |
| --local-dir Mage-Flow-Edit-Turbo |
| # 3. run |
| swift run -c release mage-flow-edit \ |
| --repo Mage-Flow-Edit-Turbo --ref dog.jpg \ |
| --prompt "make the background a snowy forest" --out edit.png |
| ``` |
|
|
| ## Parity (vs the PyTorch oracle, CPU stream) |
|
|
| Every component is parity-locked; relative error, since this network's |
| activations reach ~1e5 and absolute tolerances mislead: |
|
|
| | component | worst rel error | |
| |---|---| |
| | MageFlow NR-MMDiT (12-block, 4.1B) | 6.8e-6 | |
| | MageVAE (encode + decode) | 1.08e-5 | |
| | Gaussian-Shading watermark | 0 (bit-exact) | |
| | FlowMatchEuler 4-step schedule | exact | |
| | end-to-end denoise (bf16 oracle vs fp32) | 2.8e-2 | |
| | full resolution range **512–2048** | 2048² @ 34.2 dB PSNR vs oracle | |
|
|
| ## Notes |
|
|
| - **The DiT runs bf16** (upstream dtype). An earlier bf16 failure at ≥512² was |
| root-caused to mlx-swift ≤0.31.6's JIT-miscompiled NAX split-K GEMM |
| ([ml-explore/mlx#3797](https://github.com/ml-explore/mlx/issues/3797), fixed |
| upstream by [#3810](https://github.com/ml-explore/mlx/pull/3810)); the port |
| row-chunks the one affected GEMM (FFN `proj_out`) exactly, until an mlx-swift |
| release ships the fix. |
| - **The Gaussian-Shading provenance watermark is reproduced bit-exactly** — the |
| port does not strip Microsoft's watermark, so outputs remain detectable by the |
| upstream detector. |
| - **The mandatory Responsible-AI content filter is wired in** (the upstream |
| `CONTENT_FILTER_EDIT_SYSTEM` classifier, fail-closed) via the Qwen3-VL backbone. |
|
|
| Conditioning + filter run on |
| [qwen3vl-mlx-swift](https://github.com/xocialize/qwen3vl-mlx-swift). |
|
|
| ## Pre-quantized DiT (no bf16 peak) |
|
|
| | file | recipe | gate | |
| |---|---|---| |
| | `transformer-int8.safetensors` | int8 g32 affine, attn+MLP Linears only | `int8 golden-deficit 1.296e-04 vs 2x baseline 1.982e-04 -> PASS` | |
| | `transformer-int4.safetensors` | int4 g64 affine, same scope | `int4 per-pass cosine 0.991115 (threshold 0.9900) -> PASS` | |
|
|
| Modulation/io/norm layers and the trailing block 11 stay bf16 (the final norm |
| amplifies last-block quant error ~1.6× into `proj_out`). Files are |
| self-describing (metadata carries bits / group size / kept blocks) and load |
| directly with **no bf16 peak** — int8 gate: golden-relative: deficit vs the fp32 oracle ≤ 2× the bf16 baseline's own deficit (bf16 itself is ~1e-4 from fp32 at this model's ~1.2e5 activation range); int4 gate: per-pass cosine ≥ 0.99 vs bf16 + e2e-validated render (43.9 dB int8 / valid int4 vs the bf16 render). |
| Quantization is weight-only: the mandatory content filter and the bit-exact |
| Gaussian-Shading watermark are unaffected. Measured t2i peak GPU @1024²: |
| bf16 19.6 → int8 17.2 → int4 16.1 GB. |
|
|
| ```bash |
| # quant tier: skip the 7.7 GB bf16 transformer entirely |
| hf download microsoft/Mage-Flow-Edit-Turbo --include "text_encoder/*" --include "vae/*" --local-dir Mage-Flow-Edit-Turbo |
| hf download xocialize/Mage-Flow-Edit-Turbo-mlx folded_adaln.safetensors --local-dir Mage-Flow-Edit-Turbo |
| hf download xocialize/Mage-Flow-Edit-Turbo-mlx transformer-int4.safetensors --local-dir Mage-Flow-Edit-Turbo |
| swift run -c release mage-flow-edit --repo Mage-Flow-Edit-Turbo --dit-quant Mage-Flow-Edit-Turbo/transformer-int4.safetensors ... |
| ``` |
|
|
| MLXEngine: `MLXMageFlow` (same repo) wraps the family as `MageFlowT2IPackage` |
| / `MageFlowEditPackage` — quant tiers auto-materialize these files via |
| `WeightSourcing`. |
|
|
| ## Family |
|
|
| All six Mage-Flow checkpoints are supported by the same code: |
| [Mage-Flow](https://huggingface.co/xocialize/Mage-Flow-mlx) (t2i, RL) · |
| [Mage-Flow-Base](https://huggingface.co/xocialize/Mage-Flow-Base-mlx) · |
| [Mage-Flow-Turbo](https://huggingface.co/xocialize/Mage-Flow-Turbo-mlx) · |
| [Mage-Flow-Edit](https://huggingface.co/xocialize/Mage-Flow-Edit-mlx) (edit, RL) · |
| [Mage-Flow-Edit-Base](https://huggingface.co/xocialize/Mage-Flow-Edit-Base-mlx) · |
| this repo (edit, Turbo). `text_encoder/` and `vae/` are byte-identical |
| family-wide; `folded_adaln.safetensors` is the same file in every repo. |
|
|
