mlx-taef 0.1.0__tar.gz

mlx_taef-0.1.0/.gitignore

@@ -0,0 +1,18 @@
+__pycache__/
+*.py[cod]
+.venv/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.coverage
+.coverage.*
+htmlcov/
+dist/
+build/
+*.egg-info/
+.DS_Store
+tests/reference/*.pth
+tests/reference/*.pt
+
+# hatch-vcs auto-generated
+src/mlx_taef/_version.py

mlx_taef-0.1.0/CHANGELOG.md

@@ -0,0 +1,37 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [0.1.0] — 2026-05-13
+
+### Added
+- Public release of `mlx-taef`.
+- Pure MLX implementation of the TAESD family architecture: `Clamp`, `Block` (with optional `midblock_gn` pool branch using `pytorch_compatible=True` GroupNorm), `make_decoder`, `make_encoder`.
+- Four variants under one consistent API: `TAESD`, `TAESDXL`, `TAEF1`, `TAEF2`.
+- `Taef.from_pretrained()` — auto-download from HF Hub, convert to MLX, cache at `~/.cache/mlx-taef/`. Zero PyTorch dependency at runtime.
+- `Taef.from_pretrained_local()` — load from already-converted local safetensors.
+- `Taef.decode()` / `decode_image()` / `encode()` / `scale_latents()` / `unscale_latents()`.
+- Weight conversion handling both upstream-Sequential and Diffusers key formats (with the +1 decoder index offset for Diffusers).
+- CLI: `mlx-taef convert --variant <name> --role {decoder,encoder} --dst <path>`, `info <path>`, `bench --variant <name>`.
+- `mlx_taef.integrations.mflux.LivePreviewCallback` for live previews during FLUX.1 / FLUX.2 Klein generation in mflux.
+- `mlx_taef.integrations.mflux.unpack_flux2_latent` helper handling mflux's packed DiT latent layout + optional BN denormalization.
+- 69 tests, 100% line + branch coverage. CI on macOS 14 ARM64 (Python 3.11/3.12/3.13) plus Linux lint/typecheck.
+- Tier 1 parity tests for every MLX layer vs PyTorch reference.
+- Tier 2 parity: pre-baked reference fixtures (committed to repo); decoder cosine sim > 0.999 on 5 fixtures per variant, encoder cosine sim 1.00000000.
+- Tier 3 property tests via hypothesis (shape invariants, dtype propagation, determinism, roundtrip).
+- Tier 5 perf budget: decode 1024×1024 fp16 < 200 ms; peak memory < 1.5 GB.
+
+### Known limitations
+- `LivePreviewCallback` without `bn_mean`/`bn_var` produces structurally-correct but color-shifted previews. Pass the BN stats from `Flux2VAE.bn.running_{mean,var}` for exact recovery.
+- bf16 dtype works but isn't hardware-accelerated on M1/M2 (only M3+).
+
+## [0.0.2] — 2026-05-12
+
+Phase 2: all four variants + encoder side + property tests.
+
+## [0.0.1-alpha] — 2026-05-12
+
+Phase 1: TAEF2 decoder MVP.

mlx_taef-0.1.0/PKG-INFO

@@ -0,0 +1,125 @@
+Metadata-Version: 2.4
+Name: mlx-taef
+Version: 0.1.0
+Summary: Tiny AutoEncoders for diffusion (TAESD family) on Apple MLX.
+Project-URL: Homepage, https://github.com/ionden/mlx-taef
+Project-URL: Documentation, https://ionden.github.io/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
+Author-email: Denis Ineshin <denis.ineshin@gmail.com>
+License-Expression: MIT
+Keywords: apple-silicon,autoencoder,flux,mlx,stable-diffusion,taesd
+Classifier: Development Status :: 3 - Alpha
+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.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: huggingface-hub>=0.24
+Requires-Dist: mlx>=0.20
+Requires-Dist: numpy>=1.26
+Requires-Dist: safetensors>=0.4
+Provides-Extra: image
+Requires-Dist: pillow>=10.0; extra == 'image'
+Provides-Extra: mflux
+Requires-Dist: mflux>=0.17; extra == 'mflux'
+Description-Content-Type: text/markdown
+
+# 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) — 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 — see each step refresh in <100 ms instead of waiting 30 s for the full VAE.
+- **Low-memory fallbacks** when the full VAE OOMs on 16 GB Macs (TAEF2 peaks at ~1 GB for 1024×1024 vs ~9.6 GB for the full Flux VAE).
+- **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
+```
+
+## Install
+
+```bash
+pip install mlx-taef
+# With mflux preview callback:
+pip install "mlx-taef[mflux]"
+```
+
+Requires Python ≥ 3.11 and Apple Silicon. Runtime install has **zero PyTorch dependency**.
+
+## 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) |
+
+All four share one API.
+
+## Benchmarks (M1 Max, fp16)
+
+| Metric | TAEF2 (this library) | Full Flux VAE (reference) | Win |
+|---|---|---|---|
+| Decode latency 1024×1024 | **~100 ms** | seconds | 50–100× |
+| Peak unified memory 1024×1024 | **~1 GB** | ~9.6 GB | **9.4×** |
+| Output cosine sim vs PyTorch reference | > 0.999 | — | (parity verified) |
+
+Numbers from `tests/test_perf.py` on M1 Max 32 GB. See `notes/phase1-benchmarks.md` for details.
+
+## mflux live previews
+
+```python
+from mflux.models.flux2 import Flux2Klein
+from mlx_taef.integrations.mflux import LivePreviewCallback
+
+model = Flux2Klein.from_pretrained("4bit")
+preview = LivePreviewCallback(
+    variant="taef2",
+    every=5,
+    save_to="preview.png",
+    latent_height=32,  # 512 / 16
+    latent_width=32,
+)
+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,
+)
+```
+
+For exact value-space recovery, also pass `bn_mean=flux2_vae.bn.running_mean, bn_var=flux2_vae.bn.running_var` to the callback. Without them, previews show correct structure but colors may shift.
+
+See `docs/manual-verification.md` for the full verification recipe.
+
+## Status
+
+- v0.1.0 — initial public release. All four variants, encoder + decoder, mflux integration, CI, 100% test coverage.
+
+## 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.

mlx_taef-0.1.0/README.md

@@ -0,0 +1,92 @@
+# 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) — 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 — see each step refresh in <100 ms instead of waiting 30 s for the full VAE.
+- **Low-memory fallbacks** when the full VAE OOMs on 16 GB Macs (TAEF2 peaks at ~1 GB for 1024×1024 vs ~9.6 GB for the full Flux VAE).
+- **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
+```
+
+## Install
+
+```bash
+pip install mlx-taef
+# With mflux preview callback:
+pip install "mlx-taef[mflux]"
+```
+
+Requires Python ≥ 3.11 and Apple Silicon. Runtime install has **zero PyTorch dependency**.
+
+## 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) |
+
+All four share one API.
+
+## Benchmarks (M1 Max, fp16)
+
+| Metric | TAEF2 (this library) | Full Flux VAE (reference) | Win |
+|---|---|---|---|
+| Decode latency 1024×1024 | **~100 ms** | seconds | 50–100× |
+| Peak unified memory 1024×1024 | **~1 GB** | ~9.6 GB | **9.4×** |
+| Output cosine sim vs PyTorch reference | > 0.999 | — | (parity verified) |
+
+Numbers from `tests/test_perf.py` on M1 Max 32 GB. See `notes/phase1-benchmarks.md` for details.
+
+## mflux live previews
+
+```python
+from mflux.models.flux2 import Flux2Klein
+from mlx_taef.integrations.mflux import LivePreviewCallback
+
+model = Flux2Klein.from_pretrained("4bit")
+preview = LivePreviewCallback(
+    variant="taef2",
+    every=5,
+    save_to="preview.png",
+    latent_height=32,  # 512 / 16
+    latent_width=32,
+)
+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,
+)
+```
+
+For exact value-space recovery, also pass `bn_mean=flux2_vae.bn.running_mean, bn_var=flux2_vae.bn.running_var` to the callback. Without them, previews show correct structure but colors may shift.
+
+See `docs/manual-verification.md` for the full verification recipe.
+
+## Status
+
+- v0.1.0 — initial public release. All four variants, encoder + decoder, mflux integration, CI, 100% test coverage.
+
+## 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.

mlx_taef-0.1.0/docs/manual-verification.md

@@ -0,0 +1,54 @@
+# Manual verification: mflux LivePreviewCallback
+
+These steps require:
+- mflux installed (`pip install mflux`)
+- FLUX.2 Klein 4-bit weights downloaded (~5 GB)
+- M-series Mac with 16 GB+ unified memory
+
+This validates that `LivePreviewCallback` produces recognizable previews during a real generation. Automated tests use a fake-shaped latent and can't catch value-space drift.
+
+## Steps
+
+1. Install mflux (in the mlx-taef venv):
+   ```bash
+   uv pip install mflux
+   ```
+
+2. Run a generation with the callback:
+   ```python
+   from pathlib import Path
+   from mflux.models.flux2 import Flux2Klein
+   from mlx_taef.integrations.mflux import LivePreviewCallback
+
+   model = Flux2Klein.from_pretrained("4bit")
+   callback = LivePreviewCallback(
+       variant="taef2",
+       every=5,
+       save_to=Path("preview.png"),
+       latent_height=32,  # 512 / 16
+       latent_width=32,
+   )
+   model.callbacks.register(callback)
+   model.generate_image(
+       prompt="a red apple on a wooden table",
+       num_inference_steps=25,
+       width=512,
+       height=512,
+       seed=42,
+   )
+   ```
+
+3. Verify `preview.png` updates every 5 steps with progressively clearer images. The final preview should be recognizable as "red apple on a wooden table."
+
+## Known limitation
+
+Without passing `bn_mean` and `bn_var` from `Flux2VAE.bn.running_{mean,var}`, the preview uses identity BN — structure is correct but colors may shift. For exact previews:
+
+```python
+flux2_vae = model.vae
+callback = LivePreviewCallback(
+    ...,
+    bn_mean=flux2_vae.bn.running_mean,
+    bn_var=flux2_vae.bn.running_var,
+)
+```

mlx_taef-0.1.0/pyproject.toml

@@ -0,0 +1,176 @@
+[build-system]
+requires = ["hatchling>=1.27", "hatch-vcs>=0.4"]
+build-backend = "hatchling.build"
+
+[project]
+name = "mlx-taef"
+description = "Tiny AutoEncoders for diffusion (TAESD family) on Apple MLX."
+readme = "README.md"
+requires-python = ">=3.11"
+license = "MIT"
+authors = [{ name = "Denis Ineshin", email = "denis.ineshin@gmail.com" }]
+keywords = ["mlx", "stable-diffusion", "flux", "autoencoder", "taesd", "apple-silicon"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: MacOS :: MacOS X",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "Typing :: Typed",
+]
+dynamic = ["version"]
+dependencies = [
+  "mlx>=0.20",
+  "numpy>=1.26",
+  "huggingface-hub>=0.24",
+  "safetensors>=0.4",
+]
+
+[project.optional-dependencies]
+mflux = ["mflux>=0.17"]
+image = ["pillow>=10.0"]
+
+[project.urls]
+Homepage = "https://github.com/ionden/mlx-taef"
+Documentation = "https://ionden.github.io/mlx-taef"
+Issues = "https://github.com/ionden/mlx-taef/issues"
+Changelog = "https://github.com/ionden/mlx-taef/blob/main/CHANGELOG.md"
+
+[project.scripts]
+mlx-taef = "mlx_taef.cli:main"
+
+[dependency-groups]
+test = [
+  "pytest>=8.3",
+  "pytest-cov>=5.0",
+  "hypothesis>=6.130",
+  "pytest-benchmark>=4.0",
+]
+typecheck = [
+  "mypy>=1.13",
+]
+lint = [
+  "ruff>=0.7",
+]
+docs = [
+  "mkdocs-material>=9.5",
+  "mkdocstrings[python]>=0.27",
+]
+# Reference-fixture generation only. Not for runtime.
+fixtures = [
+  "torch>=2.4",
+  "pillow>=10.0",
+]
+dev = [
+  { include-group = "test" },
+  { include-group = "typecheck" },
+  { include-group = "lint" },
+  { include-group = "docs" },
+  { include-group = "fixtures" },
+  "pre-commit>=3.8",
+]
+
+[tool.hatch.version]
+source = "vcs"
+
+[tool.hatch.build.hooks.vcs]
+version-file = "src/mlx_taef/_version.py"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/mlx_taef"]
+
+[tool.hatch.build.targets.sdist]
+exclude = [
+  ".github/",
+  ".hypothesis/",
+  ".pre-commit-config.yaml",
+  "notes/",
+  "tests/_third_party/",
+  "tests/converted/",
+  "tests/reference/",
+]
+
+[tool.ruff]
+src = ["src"]
+line-length = 100
+target-version = "py311"
+
+[tool.ruff.lint]
+select = [
+  "E", "W", "F",
+  "I",
+  "B",
+  "UP",
+  "RUF",
+  "SIM",
+  "C4",
+  "PIE", "PT", "RET",
+  "PTH",
+  "TID",
+  "TCH",
+  "N",
+  "D",
+  "ANN",
+  "PYI",
+  "PERF",
+  "FURB",
+  "NPY",
+]
+ignore = [
+  "ANN401",
+  "E501",
+]
+
+[tool.ruff.lint.pydocstyle]
+convention = "google"
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**/*.py" = ["D", "ANN", "S101"]
+"tests/_third_party/**/*.py" = ["ALL"]
+"examples/**/*.py" = ["D", "ANN"]
+"notes/**/*.py" = ["D", "ANN", "RUF", "RET"]
+
+[tool.ruff.format]
+docstring-code-format = true
+exclude = ["tests/_third_party/**/*.py"]  # vendored; preserve original formatting
+
+[tool.mypy]
+python_version = "3.11"
+strict = true
+warn_unreachable = true
+enable_error_code = ["redundant-expr", "truthy-bool", "ignore-without-code"]
+files = ["src/mlx_taef"]
+
+[tool.pytest.ini_options]
+minversion = "8.0"
+addopts = "-ra --strict-markers --strict-config"
+testpaths = ["tests"]
+xfail_strict = true
+markers = [
+  "slow: tests >1s",
+  "gpu: requires Metal GPU",
+  "network: requires network access (skipped in CI)",
+  "benchmark: perf test (skipped in CI)",
+  "integration: needs real model checkpoint",
+]
+
+[tool.coverage.run]
+branch = true
+source = ["mlx_taef"]
+parallel = true
+omit = ["src/mlx_taef/_version.py"]  # generated by hatch-vcs; no user code
+
+[tool.coverage.report]
+exclude_lines = [
+  "pragma: no cover",
+  "if TYPE_CHECKING:",
+  "@overload",
+  "raise NotImplementedError",
+]
+show_missing = true
+fail_under = 85

mlx_taef-0.1.0/src/mlx_taef/__init__.py

@@ -0,0 +1,9 @@
+"""mlx-taef: Tiny AutoEncoder family ported to Apple MLX."""
+
+import logging
+
+from mlx_taef.api import TAEF1, TAEF2, TAESD, TAESDXL, Taef
+
+__all__ = ["TAEF1", "TAEF2", "TAESD", "TAESDXL", "Taef"]
+
+logging.getLogger("mlx_taef").addHandler(logging.NullHandler())

mlx_taef-0.1.0/src/mlx_taef/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None