Metadata-Version: 2.4
Name: mlx-taef
Version: 0.8.2
Summary: Tiny AutoEncoders for diffusion on Apple Silicon — live previews + low-memory decode for FLUX, SD, and Qwen-Image.
Project-URL: Homepage, https://github.com/IonDen/mlx-taef
Project-URL: Source, https://github.com/IonDen/mlx-taef
Project-URL: Issues, https://github.com/IonDen/mlx-taef/issues
Project-URL: Changelog, https://github.com/IonDen/mlx-taef/blob/main/CHANGELOG.md
Project-URL: Comparison, https://github.com/IonDen/mlx-taef/blob/main/COMPARISON.md
Project-URL: Roadmap, https://github.com/IonDen/mlx-taef/blob/main/ROADMAP.md
Author-email: Denis Ineshin <denis.ineshin@gmail.com>
License-Expression: MIT
License-File: LICENSE
Keywords: apple-silicon,autoencoder,diffusion,flux,latent-preview,mflux,mlx,qwen-image,stable-diffusion,taef,taef1,taef2,taesd,tiny-autoencoder,vae
Classifier: Development Status :: 4 - Beta
Classifier: Environment :: MacOS X
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Image Processing
Classifier: Typing :: Typed
Requires-Python: >=3.10
Requires-Dist: huggingface-hub>=0.24
Requires-Dist: mlx>=0.20
Requires-Dist: numpy>=1.26
Requires-Dist: safetensors>=0.4
Requires-Dist: typing-extensions>=4.5; python_version < '3.11'
Provides-Extra: image
Requires-Dist: pillow>=10.0; extra == 'image'
Provides-Extra: mflux
Requires-Dist: mflux<0.20,>=0.17; extra == 'mflux'
Provides-Extra: showcase
Requires-Dist: mflux<0.20,>=0.17; extra == 'showcase'
Requires-Dist: mlx-teacache<0.10,>=0.9.3; extra == 'showcase'
Requires-Dist: pillow>=10.0; extra == 'showcase'
Requires-Dist: scikit-image>=0.22; extra == 'showcase'
Description-Content-Type: text/markdown

# mlx-taef

<p align="center">
  <img src="https://raw.githubusercontent.com/IonDen/mlx-taef/main/docs/assets/mlx-taef-logo.png" alt="mlx-taef" width="100%">
</p>

[![PyPI version](https://img.shields.io/pypi/v/mlx-taef.svg)](https://pypi.org/project/mlx-taef/)
[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://pypi.org/project/mlx-taef/)
[![License: MIT](https://img.shields.io/pypi/l/mlx-taef.svg)](https://github.com/IonDen/mlx-taef/blob/main/LICENSE)

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 (~8–10× faster). [Watch a step-by-step decode](https://github.com/IonDen/mlx-taef/blob/main/PREVIEW.md), or see [COMPARISON.md](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.

```python
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`](https://github.com/IonDen/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.27× faster with 48% less peak memory than the same generation without TeaCache.

## Research notes

**[The latent in the callback is not the latent the decoder wants](https://github.com/IonDen/mlx-taef/blob/main/docs/papers/the-latent-in-the-callback-is-not-the-latent-the-decoder-wants.md)**
reads nine mflux image models at a pinned version and shows that sharing a VAE, even
byte-identical weights, does not mean sharing a latent contract: the packing, the sub-pixel fold
order, the spatial divisor and the normalization all live in the generator, and the two that matter
most, the fold order and the normalization, fail without changing the tensor's shape. It records the FLUX.2-versus-Ideogram 4 fold orders, a hash table
for the shared VAE files, and why weight sources are pinned by digest.

## Install

From PyPI:

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

Or with `uv`:

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

Pin an exact version in a project that needs reproducibility:

```bash
pip install "mlx-taef==0.7.1"
```

Verify the install:

```bash
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](https://huggingface.co/madebyollin/taesd) |
| `TAESDXL` | 4 | Stable Diffusion XL | [madebyollin/taesdxl](https://huggingface.co/madebyollin/taesdxl) |
| `TAEF1` | 16 | FLUX.1 | [madebyollin/taef1](https://huggingface.co/madebyollin/taef1) |
| `TAEF2` | 32 | FLUX.2 Klein | [madebyollin/taef2](https://huggingface.co/madebyollin/taef2) |
| `ZImage` | 16 | Z-Image / Z-Image-Turbo (shares the FLUX.1 16-ch latent contract) | reuses [madebyollin/taef1](https://huggingface.co/madebyollin/taef1) |
| `QwenImage` | 16 | Qwen-Image / Qwen-Image-Edit (Wan 2.1 VAE 16-ch latent) | [ionden/taew2.1](https://huggingface.co/ionden/taew2.1) |
| `Krea2` | 16 | Krea 2 Turbo (generates on the Qwen-Image stack, shares its Wan 2.1 VAE) | reuses [ionden/taew2.1](https://huggingface.co/ionden/taew2.1) |

They all share one API.

## Examples

[EXAMPLES.md](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/`](examples/).

## Benchmarks

Side-by-side images + measured timings: see [COMPARISON.md](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

<p align="center">
  <img src="https://raw.githubusercontent.com/IonDen/mlx-taef/main/docs/assets/live-preview.gif" alt="Side-by-side animated GIF: TAEF1 per-step live previews animating on the left against the finished full-VAE decode held static on the right, with a step-count caption." width="100%">
</p>

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`](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
```

<p align="center">
  <img src="https://raw.githubusercontent.com/IonDen/mlx-taef/main/docs/assets/taef1-live-preview.gif" alt="Looping animated GIF of the same TAEF1 live-preview frames playing through and holding on the final decode, single panel." width="40%">
</p>

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):

```python
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(
    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,
)
```

TAEF2 decodes the latent exactly as mflux hands it to the callback, so the variant and a save path are all the configuration it needs. Releases before 0.8.2 also read the Flux2VAE batch-norm statistics from `flux=model` and applied the VAE's inverse normalization before decoding. Measured against the full VAE that made previews worse (SSIM 0.59 against 0.92 on the same latent, with crushed shadows and oversaturated color), so it is off by default now. `auto_bn=True` with `flux=model`, or explicit `bn_mean=` / `bn_var=`, still turn it on and log a warning. `callback.resolved_bn` reports which path is active: `"none"`, `"auto"` or `"explicit"`.

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](COMPARISON.md) + committed JSON report; [ROADMAP.md](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](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.
- **v0.8.1 — released on PyPI** (2026-09-05). mflux 0.19.x compatibility. mflux 0.19.0 started passing a new `control_images` argument to every before-loop callback, which made any generation with a `LivePreviewCallback` registered fail before its first step; the hook now accepts it (and whichever conditioning keyword mflux adds next), and the `mflux` extra installs against 0.19.x (`>=0.17,<0.20`). Verified against mflux 0.19.1 / MLX 0.32.2: callback contract, latent layouts, and parity fixtures unchanged; a registered preview composes with mflux's new `--pid-decode` final decode.
- **v0.8.2 — released on PyPI** (2026-09-17). TAEF2 previews now match the full FLUX.2 VAE much more closely: SSIM against the full VAE decode rises from 0.616 to 0.960 on the benchmark latent (webp-scored, as in COMPARISON; 0.59 to 0.92 on lossless PNGs). Earlier releases applied the VAE's batch-norm inverse to the latent before TAEF2 whenever `flux=model` was passed, which is not the input TAEF2 wants, so previews came out dark and oversaturated. The default now feeds TAEF2 the latent as mflux produces it; `auto_bn=True` or explicit `bn_mean` / `bn_var` keep the old behavior and log a warning. COMPARISON and EXAMPLES carry the re-measured FLUX.2 numbers and frames.

Track future releases via the [PyPI history](https://pypi.org/project/mlx-taef/#history) or `gh release list -R IonDen/mlx-taef`.

## License

MIT. Mirrors upstream [madebyollin/taesd](https://github.com/madebyollin/taesd) license. Pretrained weights belong to their respective authors (madebyollin).

## Acknowledgements

- [madebyollin](https://github.com/madebyollin) for the upstream TAESD-family models and weights.
- [Apple ML Explore](https://github.com/ml-explore/mlx) for MLX.
- [filipstrand/mflux](https://github.com/filipstrand/mflux) for the MLX-native FLUX runner this library integrates with.

## Sister projects

Other MLX libraries for Apple Silicon:

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

---

By Denis Ineshin · [ineshin.space](https://ineshin.space)
