xocialize/Mage-Flow-Edit-Turbo-mlx

Mage-Flow-Edit-Turbo — MLX-Swift

An Apple-Silicon MLX-Swift port of 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.

# 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, fixed upstream by #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.

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.

# 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 (t2i, RL) · Mage-Flow-Base · Mage-Flow-Turbo · Mage-Flow-Edit (edit, RL) · Mage-Flow-Edit-Base · this repo (edit, Turbo). text_encoder/ and vae/ are byte-identical family-wide; folded_adaln.safetensors is the same file in every repo.