card: pre-quantized DiT section

8fff5a0 verified 4 days ago

3.57 kB

library_name: mlx
license: mit
pipeline_tag: image-to-image
base_model: microsoft/Mage-Flow-Edit-Base
tags:
  - mlx
  - mlx-swift
  - diffusion
  - rectified-flow
  - mage-flow

Mage-Flow-Edit-Base — MLX-Swift

Apple-Silicon MLX-Swift support for microsoft/Mage-Flow-Edit-Base (edit, defaults: 30 steps / cfg 5.0). One of six checkpoints of the Mage-Flow family, all sharing one architecture — this repo carries the port's support artifact; the component weights come from the upstream microsoft repo (MIT) and are not re-hosted (multi-component diffusion pipeline → xocialize namespace per mlx-community conventions).

Code, parity gates, and full porting notes: github.com/xocialize/mage-flow-swift. The DiT/VAE/Gaussian-Shading watermark/scheduler are parity-locked against the PyTorch oracle; T2I and CFG paths are gated per-step and validated with renders at the real defaults (sample_512.png is this checkpoint's).

Use

hf download microsoft/Mage-Flow-Edit-Base --local-dir Mage-Flow-Edit-Base
hf download microsoft/Mage-Flow-Edit-Turbo --include "text_encoder/*" --include "vae/*" --local-dir Mage-Flow-Edit-Base   # shared components (byte-identical family-wide)
hf download xocialize/Mage-Flow-Edit-Base-mlx folded_adaln.safetensors --local-dir Mage-Flow-Edit-Base
swift run -c release mage-flow-edit --repo Mage-Flow-Edit-Base --ref dog.jpg --prompt "make the background a snowy forest" --out out.png --steps 30 --cfg 5.0

Notes: the mandatory upstream RAI content filter is wired in (fail-closed); the Gaussian-Shading provenance watermark is reproduced bit-exactly; the DiT runs bf16 with one GEMM row-chunked around mlx-swift's NAX split-K JIT bug (mlx#3797) until a release ships #3810.

Pre-quantized DiT (no bf16 peak)

file	recipe	gate
`transformer-int8.safetensors`	int8 g32 affine, attn+MLP Linears only	`int8 quant-vs-bf16 0.999917 (family threshold 0.9998) -> PASS`
`transformer-int4.safetensors`	int4 g64 affine, same scope	`int4 quant-vs-bf16 0.997226 (family 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: quant-vs-bf16 per-pass, family threshold ≥ 0.9998; int4 gate: quant-vs-bf16 per-pass, family threshold ≥ 0.99. 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-Base --include "text_encoder/*" --include "vae/*" --local-dir Mage-Flow-Edit-Base
hf download xocialize/Mage-Flow-Edit-Base-mlx folded_adaln.safetensors --local-dir Mage-Flow-Edit-Base
hf download xocialize/Mage-Flow-Edit-Base-mlx transformer-int4.safetensors --local-dir Mage-Flow-Edit-Base
swift run -c release mage-flow-edit --repo Mage-Flow-Edit-Base --dit-quant Mage-Flow-Edit-Base/transformer-int4.safetensors ...

MLXEngine: MLXMageFlow (same repo) wraps the family as MageFlowT2IPackage / MageFlowEditPackage — quant tiers auto-materialize these files via WeightSourcing.