mlx-taef 0.6.0

Tiny AutoEncoders for diffusion on Apple Silicon — live previews + low-memory decode for FLUX, SD, and Qwen-Image.

mlx-taef

Tiny AutoEncoders for diffusion latents on Apple Silicon, in pure MLX.

mlx-taef is the first MLX port of the TAESD family — TAESD (SD1.x), TAESDXL (SDXL), TAEF1 (FLUX.1), TAEF2 (FLUX.2 Klein), and Z-Image (which reuses the TAEF1 weights for previews) — plus Qwen-Image / Qwen-Image-Edit on a port of the taew2.1 (Wan 2.1 VAE) autoencoder. All are distilled mini-autoencoders that decode diffusion latents to RGB in milliseconds using a few-MB model instead of multi-GB full VAEs.

Use it for:

• Live previews during long generations on Mac — TAEF1 decodes a 512×512 preview in ~183 ms and TAEF2 in ~258 ms on M1 Max (vs 2 s for the full VAE). See COMPARISON.md for the measured table and reproducer.
• Low-memory fallbacks when the full VAE OOMs on 16 GB Macs (TAEF2 peaks at ~0.6 GB decode memory vs ~2.6 GB for the full FLUX.2 VAE on the same latent).
• Quick latent inspection in notebooks and ML research.

import mlx.core as mx
from mlx_taef import TAEF2

taef = TAEF2.from_pretrained()              # downloads + converts on first call
img = taef.decode(latents)                  # NHWC float in [0, 1]
img_uint8 = taef.decode_image(latents)      # uint8 NHWC ready for PIL

Which library do I need?

You want live previews or low-memory FLUX decode? You're in the right place. mlx-taef decodes diffusion latents to RGB in ~260 ms (TAEF2) or ~185 ms (TAEF1) on M1 Max — vs ~2 seconds for the full VAE, with ~4× less peak memory. Drops into mflux via LivePreviewCallback.

You want FLUX generation itself to be faster on Apple Silicon? You want mlx-teacache — it skips redundant denoising steps when the schedule is cacheable (measured 1.46× on FLUX.1-dev at 25 steps).

You want both: faster generation AND live previews? Use them together — they compose cleanly. mflux 4-step Klein + TeaCache + TAEF2 previews = 1.30× wall-clock and 26% less peak memory vs vanilla.

Install

From PyPI:

pip install mlx-taef
# With the mflux preview callback:
pip install "mlx-taef[mflux]"

Or with uv:

uv add mlx-taef
# With mflux:
uv add "mlx-taef[mflux]"

Pin an exact version in a project that needs reproducibility:

pip install "mlx-taef==0.2.0"

Verify the install:

mlx-taef --help

Requires Python ≥ 3.11 and Apple Silicon (mlx itself is Apple-Silicon-only). The base mlx-taef install pulls in no PyTorch; in this repo torch is used only to generate test fixtures. The optional mflux extra is a separate case: mflux currently depends on PyTorch, so installing mlx-taef[mflux] brings it in.

Variants

Variant	latent_channels	For	HF source
TAESD	4	Stable Diffusion 1.x	madebyollin/taesd
TAESDXL	4	Stable Diffusion XL	madebyollin/taesdxl
TAEF1	16	FLUX.1	madebyollin/taef1
TAEF2	32	FLUX.2 Klein	madebyollin/taef2
ZImage	16	Z-Image / Z-Image-Turbo (shares the FLUX.1 16-ch latent contract)	reuses madebyollin/taef1
QwenImage	16	Qwen-Image / Qwen-Image-Edit (Wan 2.1 VAE 16-ch latent)	ionden/taew2.1

They all share one API.

Examples

EXAMPLES.md walks through live-preview and low-memory decode for each model with real captured frames, final images, and the measured gain. Runnable companions live in examples/.

Benchmarks

Side-by-side images + measured timings: see COMPARISON.md.

All numbers there come from scripts/run_showcase.py (subprocess-per-rep bench harness) and the committed _artifacts/showcase_report.json. Per-rep raw arrays are preserved so reviewers can see variance, not just summary stats.

The previous v0.1.x README claim — "~100 ms decode at 1024×1024, 50–100× faster than the full Flux VAE; ~1 GB peak vs ~9.6 GB" — was a same-process measurement under v0.1's tests/test_perf.py. v0.2.0 re-measures under subprocess-per-rep with per-condition memory caps; see COMPARISON.md for the honest replacement numbers.

mflux live previews

from mflux.models.flux2 import Flux2Klein
from mlx_taef.integrations.mflux import LivePreviewCallback

model = Flux2Klein.from_pretrained("4bit")
preview = LivePreviewCallback(
    flux=model,            # auto-extracts the Flux2VAE BN stats for exact color
    variant="taef2",
    every=5,
    save_to="preview.png",
    # latent_height / latent_width are auto-detected from the generation config;
    # pass both only if you want to override.
)
model.callbacks.register(preview)
model.generate_image(
    prompt="a red apple on a wooden table",
    num_inference_steps=25,
    width=512,
    height=512,
    seed=42,
)

Passing flux=model lets the callback auto-extract model.vae.bn.running_mean and running_var so TAEF2 previews are color-correct out of the box (callback.resolved_bn == "auto"). If you have a custom integration where flux= isn't convenient, pass bn_mean= and bn_var= explicitly — those take precedence (resolved_bn == "explicit"). Without either path you get identity-BN previews with correct structure but shifted colors (resolved_bn == "none").

See docs/manual-verification.md for the full verification recipe.

Status

• v0.1.0 — initial public release on PyPI (2026-05-13). All four variants, encoder + decoder, mflux integration, CI, 99 % honest coverage.
• v0.2.0 — released on PyPI (2026-05-27). Auto-bn extraction in LivePreviewCallback(flux=...); per-step gallery mode (numbered_frames=True); subprocess-per-rep showcase bench (scripts/run_showcase.py); hardware-aware memory caps via mlx_taef._memory_caps; COMPARISON.md + committed JSON report; ROADMAP.md.
• v0.2.3 — released on PyPI (2026-05-29). Weight loading is now strict: from_pretrained_local raises on an incomplete or wrong-shaped weights file instead of loading a silently-wrong model, and the HF→MLX converter checks parameter coverage and shapes at convert time (new ConversionError). The end-to-end parity tests now gate on an absolute pixel tolerance rather than cosine similarity. A bare pytest skips the network and benchmark tests by default (--run-network / --run-benchmark to opt in).
• v0.3.0 — released on PyPI (2026-06-06). Internal kernel refactor: each variant is now a composable ModelKernel (mlx_taef.kernels), so adding a model is a self-contained entry; variants.py stays a back-compat shim. Ships one user-facing fix — the mflux LivePreviewCallback FLUX.1 path fed the packed latent straight to the decoder and produced wrong previews; it now unpacks correctly.
• v0.3.1 — released on PyPI (2026-06-08). Hardening: decode()/encode() raise a clear error when weights haven't been loaded yet, or when a latent has the wrong channel count, instead of returning garbage; importing without mflux installed now raises MfluxNotInstalledError (a TaefError that is also an ImportError).
• v0.4.0 — released on PyPI (2026-06-13). Z-Image / Z-Image-Turbo live preview: a new ZImage model reuses TAEF1's FLUX.1 weights with no new download (Z-Image shares FLUX.1's 16-channel latent contract), validated by an SSIM ≥ 0.75 calibration against mflux's full Z-Image VAE (measured 0.94). Adds mlx-taef bench --variant zimage and a new top-level EXAMPLES.md with captured frames and measured decode numbers.
• v0.4.1 — released on PyPI (2026-06-14). Documentation and packaging accuracy pass, no code changes: the Z-Image SSIM calibration is described correctly (it runs as an opt-in network test, not in default CI); the "no PyTorch" note is scoped to the base install (the mflux extra pulls in PyTorch via mflux); and the source distribution is now an allowlist, so it no longer ships the test suite without its fixtures or sweeps in local tool state.
• v0.4.2 — released on PyPI (2026-06-14). Hardening patch: decode() and encode() now reject non-4-D inputs with a clear ValueError naming the expected NHWC rank, instead of failing deep in the conv stack. Also moves the release workflow's download-artifact step to a Node 24 release, and the mflux_live_preview example now prints per-step decode timing next to the full-VAE decode (thanks @ianscrivener, #20).
• v0.5.0 — released on PyPI (2026-06-18). Live-preview ergonomics: LivePreviewCallback auto-detects latent_height / latent_width from the mflux generation config, so you no longer pass them by hand. Passing both still overrides; passing exactly one now raises. The auto-extracted Flux2VAE batch-norm eps is forwarded into the TAEF2 unpack so a non-default eps previews faithfully, and auto_bn=True on a non-TAEF2 variant logs its no-op instead of staying silent. Three showcase-only exceptions (SchemaVersionError, FixtureLatentMissingError, MlxTeacacheNotInstalledError) are no longer exported from the package root; they remain importable from mlx_taef.errors.
• v0.5.1 — released on PyPI (2026-06-20). The mflux extra now installs against mflux 0.18.x as well as 0.17.x; the previous <0.18 pin excluded 0.18, so users already on it could not install the extra without downgrading. The live-preview integration is verified against mflux 0.18.0 with no API or behavior change.
• v0.6.0 — released on PyPI (2026-06-23). Qwen-Image / Qwen-Image-Edit live preview: a new QwenImage model ports madebyollin's taew2.1 tiny autoencoder (for the Wan 2.1 VAE's 16-channel latent) to pure MLX. Adds LivePreviewCallback(variant="qwen-image") and mlx-taef bench --variant qwen-image. Decode and encode match the upstream taew2.1 reference to ~3e-6 (committed parity fixtures). Live-preview quality against the full Wan VAE is community-measured — Qwen-Image is a ~20B model that won't fit a usable resolution on 32 GB.

Track future releases via the PyPI history or gh release list -R IonDen/mlx-taef.

License

MIT. Mirrors upstream madebyollin/taesd license. Pretrained weights belong to their respective authors (madebyollin).

Acknowledgements

• madebyollin for the upstream TAESD-family models and weights.
• Apple ML Explore for MLX.
• filipstrand/mflux for the MLX-native FLUX runner this library integrates with.

---

By Denis Ineshin · ineshin.space