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 — the decode step itself takes ~30 ms for both TAEF1 and TAEF2 on M1 Max, versus ~0.3 s for the full VAE decode (~9–10× faster). Watch a step-by-step decode, or 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.59 GB decode memory vs ~2.8 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's decode step takes ~30 ms on M1 Max, for either TAEF2 or TAEF1 — vs ~0.3 s for the full VAE decode, with ~5× less decode 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. mflux 4-step Klein + TeaCache + TAEF2 previews measured 1.41× faster with 44% less peak memory than the same generation without TeaCache.

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

Verify the install:

mlx-taef --help

Requires Python ≥ 3.10 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
Krea2	16	Krea 2 Turbo (generates on the Qwen-Image stack, shares its Wan 2.1 VAE)	reuses 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-condition 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

TAEF1 per-step previews (left) next to the final full-VAE decode (right), from a FLUX.1-dev generation at 768x768, 26 steps, seed 20, prompt "anime key visual, a lone warrior in a tattered cloak holding an oversized greatsword, standing in an overgrown ruined city, dramatic backlight, volumetric light, highly detailed, cinematic, studio anime style", captured on an M1 Max with LivePreviewCallback registered every step. Reproduce it: adapt examples/mflux_live_preview.py (the runnable starting point for wiring up LivePreviewCallback) to those parameters, which writes a numbered preview-frame gallery plus the final image, then assemble the GIF with:

uv run python scripts/make_preview_gif.py \
    --frames-dir <run-output> \
    --frames-glob "*_step*.png" \
    --final <final.png> \
    --out docs/assets/live-preview.gif \
    --panel-width 320

The same capture as a compact, single-panel loop: the previews play through, then hold on the final frame.

Wiring the callback into a generation looks like this (FLUX.2 Klein shown; the GIFs above use FLUX.1-dev / variant="taef1" with the parameters noted):

from mflux.models.common.config.model_config import ModelConfig
from mflux.models.flux2 import Flux2Klein
from mlx_taef.integrations.mflux import LivePreviewCallback

model = Flux2Klein(quantize=4, model_config=ModelConfig.flux2_klein_base_4b())
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").

Preview failures default to on_error="disable": the callback logs one warning, stops previewing for the rest of that generation, and lets mflux finish the image. Use on_error="raise" when an integration or test needs fail-fast behavior. Calling the same callback in a later generation resets the disabled state.

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.
• v0.6.1 — released on PyPI (2026-06-30). Maintenance: converted-weights cache writes are now atomic (temp file + rename), so an interrupted download/convert can no longer leave a truncated file that later runs trust as valid. Adds Python 3.14 support, sharper public-API types (factories return their concrete subclass; the decoder/encoder role and preview variant are typed literals), and a linked live-preview demo page.
• v0.6.2 — released on PyPI (2026-07-09). Hardening and accuracy. LivePreviewCallback now rejects every < 1 and a half-set BN pair, and resets its state between generations. The converted-weights cache invalidates when a source's pinned revision or sha256 changes (Qwen-Image re-converts once on upgrade), and every conversion path enforces those pins. The decode benchmark now measures the decode step at steady state, in isolation from one-time model construction: the tiny decoders run ~30 ms per step (earlier releases reported ~180–260 ms because they timed model construction inside the decode window), a ~8–10× speedup over the full VAE decode. COMPARISON / EXAMPLES are re-measured with two Z-Image scenarios added. Also fixes the mflux quickstart earlier in this README.
• v0.7.0 — released on PyPI (2026-07-24). LivePreviewCallback gains configurable runtime failure handling: it warns once and disables previews for the current generation by default, while on_error="raise" preserves strict behavior. All built-in weight sources now have immutable revisions and role-specific sha256 pins, and mlx-taef convert uses the same verified source path as runtime loading. Conversion caches include a format version, and CI installs from the lockfile. The six-scenario benchmark requires every repetition and preview frame, exits nonzero after any scenario failure, and records source and installed-package versions separately. The measured report shows no regression against v0.6.2.
• v0.7.1 — released on PyPI (2026-07-25). Python 3.10 is now supported: the floor drops from 3.11 to 3.10 and CI runs the full offline suite on 3.10 through 3.14, so no advertised version is untested. mlx-teacache (used by the showcase and test extras, and 3.11-only) installs on 3.11+ only; the runtime dependencies all work on 3.10.
• v0.8.0 — released on PyPI (2026-08-09). Krea 2 Turbo live preview: a new Krea2 model reuses the taew2.1 weights already converted for QwenImage, with no new download. Decode quality against mflux's full Krea 2 VAE is gated by an SSIM check (measured 0.9678). mlx-teacache now installs on Python 3.10 as well, so the showcase's combined scenario runs on every supported Python version. The showcase report adds LPIPS alongside SSIM, ZImage.encode() gets its own SSIM-gated validation (measured 0.9580), and the live-preview integration is verified against mflux 0.18.1. Every model-loading showcase and benchmark subprocess now runs under the active-memory watchdog, not just the live-generation workers.

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.

Sister projects

Other MLX libraries for Apple Silicon:

• mlx-teacache — TeaCache residual caching to skip redundant FLUX denoising steps.
• mlx-model-doctor — validate an MLX / Hugging Face model repo before you load it (config, tokenizer, safetensors, memory).
• mlx-quant-fidelity — measure how much quality a quantization costs (KL divergence, top-token flips, perplexity delta).

---

By Denis Ineshin · ineshin.space