TeaCache step-skipping for FLUX, Qwen-Image & Z-Image diffusion on Apple Silicon, in pure MLX.
TL;DR — Wrap an mflux generation in apply_teacache(flux) and non-distilled diffusion runs faster with a visually indistinguishable image: one line, no retraining, prompt and seed unchanged. Measured on an M1 Max 32 GB, Qwen-Image runs 2.68× faster, with FLUX.1 Krea at 1.62×, FLUX.1-dev at 1.57×, FLUX.2 Klein base at 1.22–1.37× (canonical CFG recipe), and Z-Image at 1.31×, each matching vanilla under its SSIM gate (numbers in the benchmark table). Distilled schedules — FLUX.1 schnell, FLUX.2 Klein 4B/9B at their 4–8 step defaults — skip nothing and gain nothing from the gate, so run those through vanilla mflux.
mlx-teacache is the first MLX port of TeaCache, a training-free inference optimization that predicts which denoising steps add little to the final image and reuses the previous step's output instead of running the full transformer. On FLUX.1-dev at 25 steps the polynomial gate skips 6 of 25 steps for a measured 1.57× wall-clock speedup (median of three cold reps; 1.50× fastest-to-fastest), with output visually indistinguishable from vanilla (SSIM ≥ 0.80 across a 5-prompt suite, ≥ 0.90 on the PR-gate prompt). Qwen-Image gains the most, at 2.68×, skipping about two-thirds of its steps.
FLUX.2 Klein carries a second, smaller wall-clock effect alongside step-skipping: the eager wrapper sidesteps mflux's mx.compile of _predict, which on Max and Ultra chips runs slower than eager. That is a compile-avoidance win, not a caching one, and the docs keep the two apart — it is also why distilled Klein 4B/9B can run faster under the wrapper on Max/Ultra chips even though the gate skips nothing. That effect is thermally variable and chip-dependent, though, and brings no gate benefit, so distilled schedules are still best run through vanilla mflux. See How the speedup happens for the breakdown.
You want FLUX generation to be faster on Apple Silicon? You're in the right place. mlx-teacache skips redundant denoising steps on FLUX.1 and non-distilled FLUX.2 Klein — measured 1.57× on FLUX.1-dev at 25 steps. Drops into mflux via one line.
You want live previews while generating, or low-memory latent decode? You want mlx-taef — tiny TAESD-family decoders in MLX.
You want both? They compose cleanly. mflux 4-step Klein + TeaCache + TAEF2 previews = 1.30× wall-clock and 26% less peak memory vs vanilla.
Qwen-Image mixed precision on a 32 GB Mac documents a uniform-q4 portrait artifact and the mixed q8/q4/bf16 recipe that removed it in one controlled pair for about 1.9 GiB of additional peak MLX allocation. The recipe changes model construction, not TeaCache.
Why the TeaCache gate did not engage on short distilled FLUX schedules documents the zero-skip negative result on the 8-step distilled Klein schedules and the measurement practice it forced: skip counts published next to every wall-clock number, so a speedup cannot be mis-attributed to caching.
Why byte-exact parity is a poor MLX integration oracle uses a pre-release FLUX.1 fixture failure to separate artifact drift, no-op parity, numerical equivalence, and intentionally approximate output. It explains why paired same-process tests and measured per-variant gates are more informative than committed real-model bytes.
Diffusion models run the same big transformer 20-50 times in a loop. Between consecutive steps the output changes very little, and TeaCache uses a tiny polynomial fit to predict which steps can reuse the previous step's output. On M1 Max with FLUX.1-dev at 25 steps the default threshold (rel_l1_thresh=0.20) skips 6 of 25 steps and produces a 1.57× speedup.
from mflux.models.flux.variants.txt2img.flux import Flux1
from mlx_teacache import apply_teacache
flux = Flux1.from_name("dev")
with apply_teacache(flux): # default rel_l1_thresh=0.20
flux.generate_image(prompt="a red apple", seed=42, num_inference_steps=25)pip install "mlx-teacache[mflux]"
# or with uv:
uv add "mlx-teacache[mflux]"Requires Python ≥ 3.10 and Apple Silicon. The [mflux] extra pulls in mflux>=0.17.5,<0.20. mflux 0.19 needs MLX 0.32 and torch 2.13 or newer (torch and opencv have been mflux dependencies since before 0.18); if you pair it with mlx-taef's live preview, use mlx-taef 0.8.1 or later. One caveat on 0.19: its qwen-image alias loads Qwen/Qwen-Image-2512, a checkpoint this library's Qwen coefficients were not calibrated on, and apply_teacache warns about it (see the Qwen-Image section).
pip install "mlx-teacache[mflux]==0.11.0" # pin for reproducibilityFLUX.1 dev:
from mflux.models.flux.variants.txt2img.flux import Flux1
from mlx_teacache import apply_teacache
flux = Flux1.from_name("dev", quantize=4)
with apply_teacache(flux) as handle: # default rel_l1_thresh=0.20
flux.generate_image(prompt="...", seed=42, num_inference_steps=25, guidance=3.5)
print(f"Speedup: {handle.stats.speedup_estimate:.2f}×")mlx-teacache supports mflux's image-to-image generation starting with v0.2.0. Pass image_path and image_strength > 0 to flux.generate_image() with TeaCache active:
from mlx_teacache import apply_teacache
with apply_teacache(flux):
flux.generate_image(
prompt="a red apple on a wooden table",
image_path="/path/to/init.png",
image_strength=0.7,
num_inference_steps=25,
seed=42,
height=512,
width=512,
)Caching engages on the active denoising window only, which mflux computes as num_inference_steps - init_time_step. So image_strength=0.7 with num_inference_steps=25 gives 8 predict calls available for caching. At image_strength=1.0, mflux skips denoising entirely and runs only VAE reconstruction, so TeaCache becomes a no-op.
Txt2img and img2img use the same polynomial coefficients. Image-quality is verified by SSIM gates over a fixed init-image suite in tests/test_image_quality_*.py.
Measured on M1 Max 32GB, FLUX.1-dev @ 25 steps, bf16, seed=42, guidance=3.5, red-apple prompt:
rel_l1_thresh |
Skipped steps | Speedup | SSIM vs vanilla | Recommended use |
|---|---|---|---|---|
| 0.10 | 0 / 25 | 1.07× | 1.0000 | Cache never engages |
| 0.15 | 0 / 25 | 1.13× | 1.0000 | Cache never engages |
| 0.20 (default) | 6 / 25 | 1.57× | ≥ 0.80 (5-prompt suite) | Visually-lossless sweet spot |
| 0.25 | 11 / 25 | 1.96× | 0.57-0.93 | Visible style changes on text/synthetic prompts |
The committed benchmark _artifacts/v0.10.0_bench_flux1_dev.json validates the 0.20 default row (median of three cold reps, 2026-08-15); the other rows are single-run measurements from an earlier session with a faster vanilla baseline (~104 s), so read them for the trend, not as a committed multi-rep sweep.
0.20 was picked after side-by-side visual comparison. At 0.25, text prompts that vanilla renders as neon tubes can come out as dot-matrix. At 0.20, the output is indistinguishable from vanilla and the cache still skips around 25% of steps. SSIM is conservative on high-frequency-detail prompts like text and synthetic patterns, which is why the suite floor (0.80) is lower than the PR-gate floor (0.90) on the red-apple prompt.
The table below is generated from the variant registry — see docs/_generate_supported_models.py. Per-variant detail (mflux constructor, coefficient provenance, quirks) lives in docs/variants/<id>.md.
| Variant id | Display name | Distilled? | Default recipe | License |
|---|---|---|---|---|
flux1-dev |
FLUX.1 dev | no | 25 steps, g=3.5 | FLUX.1-dev Non-Commercial License |
flux1-krea-dev |
FLUX.1 Krea [dev] | no | 28 steps, g=4.5 | FLUX.1-dev Non-Commercial License |
flux1-schnell |
FLUX.1 schnell | yes | 4 steps, g=1.0 | Apache-2.0 |
flux2-klein-4b |
FLUX.2 Klein 4B | yes | 8 steps, g=1.0 | Apache-2.0 |
flux2-klein-9b |
FLUX.2 Klein 9B | yes | 8 steps, g=1.0 | FLUX Non-Commercial |
flux2-klein-base-4b |
FLUX.2 Klein base 4B | no | 50 steps, g=4.0 | Apache-2.0 |
flux2-klein-base-9b |
FLUX.2 Klein base 9B | no | 50 steps, g=4.0 | FLUX Non-Commercial |
qwen-image |
Qwen-Image | no | 50 steps, g=4.0 | Apache-2.0 |
z-image-base |
Z-Image base | no | 50 steps, g=4.0 | Apache-2.0 |
Each variant has its own page under docs/variants/ — mflux constructor, recipe, license obligations, coefficient provenance, quirks. The highlights for the variants where behavior diverges from the default story:
flux2-klein-9b — coefficients are calibrated at num_inference_steps=8, origin-constrained polyfit. At the default threshold, the gate produces 0 step-skips on Klein 9B's 8-step schedule (the empirical adjacent-step body-output rel-L1 starts at 0.25 — above the 0.20 threshold). The library still helps via mx.compile-path avoidance (measured ~1.5-2.0× wall-clock improvement), and output quality is preserved (SSIM ≥ 0.85 PR-gate). See Benchmarks → "How the speedup happens".
flux2-klein-base-4b — non-distilled FLUX.2 Klein 4B (Apache-2.0), the primary FLUX.2 use case. At the canonical upstream CFG recipe (guidance=4.0, num_inference_steps=50) with its per-variant default rel_l1_thresh=0.17, the gate skips 9/48 active steps for a 1.22× combined speedup on M1 Max — almost all of it step-skipping, with compile-avoidance at noise level (1.01×). At the alternative guidance=1.0 / 25-step recipe it skips 3/25 for 1.41× at SSIM > 0.99. CFG runs through a per-branch gated path (each branch keeps its own cached residual). Measurement history, including v0.4.1's superseded decomposition, is in footnotes ³ and ⁴.
flux2-klein-base-9b — non-distilled FLUX.2 Klein 9B (FLUX Non-Commercial — see License obligations and accept on the Hugging Face model page before downloading). Reuses base-4b's coefficients and rel_l1_thresh=0.17 default (shared FLUX.2-Klein architecture family and calibration recipe). At the canonical 50-step + g=4.0 CFG recipe it skips 13/48 active steps for 1.37× combined wall-clock (520.6 s → 379.1 s median; 1.34× gating + 1.02× compile-avoidance), at SSIM 0.986 and a peak-memory drop from ~22 GB to ~9.5 GB. Full history in footnote ⁵.
Correction. v0.5.0 advertised 2.68× on this variant. Same-process MLX state leakage inflated it — vanilla ran cold while the wrapper inherited a warm allocator. Under cold subprocess-per-rep isolation the honest number is 1.37× (footnote ⁵).
See _artifacts/v0.10.0_bench_klein_base_9b.json for the full report (_artifacts/v0.6.0_bench_klein_base_9b.json is the v0.6.0 measurement); regenerate side-by-side images with scripts/bench_comparison.py.
z-image-base — Z-Image base (Tongyi-MAI, Apache-2.0), a single-stream DiT and the first non-FLUX model with a TeaCache mini-kernel. Its adaLN modulation is timestep-only, so there is no cheap caption-independent modulation input to gate on; the gate signal is the first-main-layer residual (calibrated in-repo as "Signal B", R² 0.400 — the threshold sweep, not the fit R², sets the quality bar here; the caption-independent noise-refiner tap was tried and rejected at R² 0.069). Per-variant default rel_l1_thresh=0.12, set at the SSIM knee. At the 512×512 red-apple recipe it skips 15/48 active steps for 1.31× combined wall-clock (227.4 s → 174.2 s), entirely from gating — the no-gate wrapper runs at vanilla speed (0.99×), so compile-avoidance is not a tailwind here. Peak memory drops 17.2 GB → 11.5 GB (the eager wrapper bypassing mflux's compiled _predict, not the gate). SSIM 0.991 vs vanilla. Footnote ⁶ has the measurement history; the COMPARISON.md portrait row is a separate 640×896 generation.
See scripts/_bench_z_image_v0_7_0.json for the full report.
qwen-image — Qwen-Image base (Alibaba, Apache-2.0), a ~20B dual-stream MMDiT and the first variant that proxies flux.transformer (the FLUX.1 pattern) and runs true two-pass CFG. It's FLUX-shaped, so the gate taps the FLUX-canonical modulated block-0 input, calibrated in-repo at R² 0.849 (well above Z-Image's 0.400 and the FLUX.2 family's 0.11–0.47). Per-variant default rel_l1_thresh=0.30, set from the threshold sweep. At the red-apple 768×768 recipe it skips 33 of 48 active steps for 2.68× combined (642.6 s → 239.4 s), SSIM 0.967 vs vanilla — the largest gain of any variant, all of it step-skipping (Qwen has no mx.compile path). Lowering the threshold to 0.20 gives back the earlier SSIM-0.978 operating point; the variant page has the whole sweep and footnote ⁷ has the measurement history. On mflux 0.19 the qwen-image alias loads Qwen/Qwen-Image-2512, a checkpoint the coefficients were not calibrated on: the variant still applies, apply_teacache raises TeaCacheUncalibratedCheckpointWarning once, and that checkpoint's numbers stay unverified until it is calibrated with scripts/calibrate_qwen.py. On a 32 GB Mac stock 4-bit Qwen is grainy — the showcase portraits use a mixed-precision build (see the variant page); mlx-teacache stays quantization-agnostic.
flux1-krea-dev — FLUX.1 Krea [dev] (Black Forest Labs with Krea; FLUX.1-dev Non-Commercial, gated), a FLUX.1-dev-architecture finetune that reuses the FLUX.1 proxy strategy unchanged. It does not reuse FLUX.1-dev's coefficients: scored on Krea's own calibration pairs, dev's tuple gives R² −496, because Krea changes roughly three times more per step, so Krea ships its own fit (R² 0.68, calibrated in-repo at the model card's 28-step, guidance 4.5 recipe) and a per-variant default rel_l1_thresh=0.30, the sharp knee of its threshold sweep (SSIM 0.990 at 0.30, 0.890 at 0.35). The package fallback 0.20 would skip nothing on this model. The v0.11.0 three-way bench at the red-apple 512×512 recipe (subprocess-per-rep, three cold reps, q4, 28 steps, g=4.5): 1.62× combined (125.3 s → 77.2 s), 10 of 26 active steps skipped in every rep and never two in a row; 1.57× of it is gating, and the 1.03× the no-gate wrapper shows is noise, since FLUX.1 has no mx.compile path to avoid. Reproduce with uv run python scripts/bench_speedup.py --variant krea-dev --three-way --reps 3.
The wrapper helps when the underlying schedule actually has cacheable redundancy. That is the case for non-distilled FLUX schedules with enough denoising steps for adjacent transformer outputs to look similar, which is what TeaCache's gate exploits.
In practice, that means:
- Use mlx-teacache for
flux1-devat 20-50 steps, the non-distilled FLUX.2 Klein family (flux2-klein-base-4b,flux2-klein-base-9b) at 20-50 steps with or without CFG,z-image-baseat 50 steps with CFG, andqwen-imageat 50 steps with CFG. These are the variants featured in COMPARISON.md, and the wrapper measurably skips steps and produces visually equivalent output. - Do not reach for it on the distilled variants —
flux1-schnell(4 steps),flux2-klein-4bandflux2-klein-9bat their distilled defaults (4-8 steps). The residual between adjacent steps is too large for the gate to engage at any reasonable threshold, so it skips zero steps and adds about 1-2% gating overhead. Run those through vanilla mflux. The full story of that zero-skip result is in the distilled-schedules research note.
There is a separate, incidental benefit on FLUX.2 variants regardless of whether the gate engages: the wrapper sidesteps mflux's compiled _predict path, which on Max and Ultra chips happens to be slower than the uncompiled path on the current MLX release. That is a wall-clock effect from compile avoidance, not from step-skipping, and we keep the two attributions separate in the docs.
Distilled Klein or Base plus the wrapper? scripts/bench_klein_base_vs_distilled.py runs distilled, Base, and Base + TeaCache on one prompt and seed and prints a table. On an M1 Max the distilled model is roughly 20× faster than Base, and the wrapper trims 15 to 30 percent off Base rather than closing that gap. So the distilled model stays the everyday choice, and Base is for when you need guidance, negative prompts, or a fine-tune. Run it with uv run python scripts/bench_klein_base_vs_distilled.py --size 4b --quantize 4 --reps 3 (or --size 9b); on a 64 GB or larger Mac, run it at --quantize 8 and post your table.
See COMPARISON.md for side-by-side images and warm-median wall-clock numbers on an M1 Max.
from mlx_taef.integrations.mflux import LivePreviewCallback
from mlx_teacache import apply_teacache
preview = LivePreviewCallback(variant="taef2", every=5, save_to="preview.png",
latent_height=32, latent_width=32)
flux.callbacks.register(preview)
with apply_teacache(flux): # default rel_l1_thresh=0.20
flux.generate_image(prompt="...", seed=42, num_inference_steps=25)handle = apply_teacache(flux)
flux.generate_image(prompt="apple", seed=42, num_inference_steps=25)
print(handle.stats.computed_count, handle.stats.skipped_count)
print(handle.stats.last_generation.decisions[5]) # per-step recordcustom = [...] # length 5, all finite
apply_teacache(flux, coefficients=custom)TeaCache observes that in diffusion denoising, consecutive transformer outputs change very little between most pairs of adjacent steps. The expensive transformer body (all the joint and single attention blocks) produces a residual that's added to the input, and that residual stays roughly stable for stretches of the denoising trajectory.
TeaCache trains a tiny polynomial that predicts how much the output will change given how much the input has changed, measured as the relative L1 distance between each step's modulated block-0 input and the immediately previous step's — matching how the polynomials are calibrated, and the comparison anchor advances on every step, including skipped ones. TeaCache accumulates the predicted change across consecutive skips; once that running total would cross rel_l1_thresh, TeaCache recomputes the transformer body instead of reusing the cached residual, and resets the running total. Only the cheap prelude (embeddings) and tail (norm + projection) still run on a skipped step.
mlx-teacache implements this for mflux on Apple Silicon. For FLUX.1 we replace flux.transformer with a per-instance proxy; for FLUX.2 we replace flux._predict with an instance-level closure, which keeps gating live on chips where mflux would otherwise wrap _predict in mx.compile. The FLUX.1 polynomial coefficients are vendored from upstream. The FLUX.2 coefficients are derived in-repo by scripts/calibrate_flux2.py (distilled klein-4b and klein-9b, base-4b each have their own calibrated tuple; klein-base-9b cross-imports base-4b's tuple because the architectures share the calibration recipe). See docs/calibration.md for procedure and provenance, and the per-variant docs under docs/variants/ for each tuple's R² and origin date. The original method is described in the TeaCache paper at https://liewfeng.github.io/TeaCache/.
For the full contributor's map — the variant registry, the mflux-free import contract, the per-family patch strategies, the gate internals, and the teardown contract — see docs/architecture.md.
All numbers are reproducible via scripts/bench_speedup.py. M1 Max 32GB, macOS 26.x, bf16, quantize=4 (q8 on the z-image-base row — its pinned recipe), 512×512 (768×768 on the qwen-image row — its pinned recipe), seed=42, red-apple prompt; default rel_l1_thresh=0.20 (per-variant defaults 0.17 on base-4b rows, 0.12 on z-image-base, 0.30 on qwen-image). The flux1-dev, base-4b CFG, base-9b CFG, and z-image-base rows were re-measured on 2026-08-15 with the v0.10.0 gate on mflux 0.18.0, three-way, three cold reps per condition (footnotes ¹ ⁴ ⁵ ⁶), qwen-image the same way on 2026-09-06 (footnote ⁷), and flux1-krea-dev on 2026-09-05 (footnote ⁸); the distilled klein rows and the base-4b 25-step row are older measurements (2026-05-16 and 2026-05-17, mflux 0.17.x, same-process harness) flagged with their era. Subprocess-per-rep means every (variant, condition, rep) gets a fresh Python interpreter so each timing starts from a cold MLX allocator.
| Variant | Steps | Vanilla | Wrapper | Speedup | Skipped | Mechanism |
|---|---|---|---|---|---|---|
flux1-dev |
25 | 113.1s | 71.9s | 1.57× | 6 / 25 | TeaCache step-skipping¹ |
flux1-schnell |
— | — | — | — | — | shares dev's coefficients; gate behaves like dev at long schedules, like Klein at the 4-step distilled default (no benefit) |
flux2-klein-4b† |
8 | 28.1s | 22.3s | 1.26× | 0 / 8 | mx.compile avoidance only |
flux2-klein-9b† |
8 | 119.0s | 61.8s | 1.93× | 0 / 8 | mx.compile avoidance only |
flux2-klein-base-4b³ |
25 | 77.5s | 55.1s | 1.41× | 3 / 25 | step-skipping + mx.compile avoidance |
flux2-klein-base-4b (CFG)⁴ |
50 | 233.9s | 192.4s | 1.22× | 9 / 50 | step-skipping (compile-avoidance ≈ noise) |
flux2-klein-base-9b (CFG)⁵ |
50 | 520.6s | 379.1s | 1.37× | 13 / 50 | step-skipping + small compile-avoidance |
z-image-base (CFG)⁶ |
50 | 227.4s | 174.2s | 1.31× | 15 / 50 | step-skipping (q8; compile-avoidance not a tailwind) |
qwen-image (CFG)⁷ |
50 | 642.6s | 239.4s | 2.68× | 33 / 50 | step-skipping (2.57× against its own no-gate run; the no-gate wrapper's 1.04× is inside the rep spread, see ⁷) |
flux1-krea-dev⁸ |
28 | 125.3s | 77.2s | 1.62× | 10 / 28 | step-skipping (1.57× gating; the no-gate wrapper's 1.03× is noise, FLUX.1 has no compile path) |
¹ flux1-dev at 25 steps, guidance=3.5, 512×512, default rel_l1_thresh=0.20. Measured 2026-08-15 under the subprocess-per-rep harness (mflux 0.18.0, v0.10.0 gate): 1.57× combined as the median of three cold reps (vanilla 113.1 s → gated 71.9 s), decomposed as 1.41× from step-skipping and 1.11× from mx.compile-path avoidance. One caveat on the split: the vanilla reps were spread wide this session (104.7 / 113.1 / 119.8 s) while the gated wrapper held steady (69.8 / 71.9 / 73.1 s), the same 71 s v0.6.3 measured. Comparing fastest to fastest gives 1.50×, so 1.5× is the conservative headline. The earlier 1.46× (v0.6.3, vanilla 103.8 s) came from a session with a quicker vanilla baseline, not from a slower wrapper. 6/25 skips in every rep, never two consecutive. Full report: _artifacts/v0.10.0_bench_flux1_dev.json (v0.6.3's is kept alongside as _artifacts/v0.6.3_bench_flux1_dev.json). Reproduce with uv run python scripts/bench_speedup.py --variant flux1-dev --three-way --reps 3 --report out.json.
† Distilled klein rows are v0.4-era same-process measurements with high thermal variance. The klein-9b 1.93× median combined a thermally-throttled vanilla rep (227s) with a recovered wrapper rep (46s); the steady-state range across reps is roughly 1.5-2.0× depending on system load. The 0/8 skip count is stable across all reps. These rows are pending a re-bench under the v0.6.0 subprocess-per-rep harness; v0.6.0's measurement on klein-base CFG showed the compile-avoidance contribution at 1.01-1.02× on 50-step schedules, so the distilled 1.5-2.0× figure is specific to the 8-step distilled path and may be smaller under cold-isolation conditions.
³ flux2-klein-base-4b at guidance=1.0, per-variant default rel_l1_thresh=0.17. The 1.41× was measured under the v0.4.0 same-process harness; not yet re-bench'd under subprocess-per-rep. The combined number historically credited step-skipping (3/25 skips save ~12% directly) and mx.compile-path avoidance. Under v0.6.0's subprocess-per-rep harness on the related 50-step CFG recipe, compile-avoidance came out at 1.01× — so the 1.41× decomposition may shift toward "almost entirely step-skipping" when re-measured. CFG (guidance > 1.0) is gated end-to-end as of v0.4.1; see footnote ⁴.
⁴ flux2-klein-base-4b under CFG at the canonical upstream BFL recipe (guidance=4.0, 50 steps), per-variant default rel_l1_thresh=0.17. Measured 2026-08-15 under the subprocess-per-rep harness (mflux 0.18.0, v0.10.0 gate): 1.22× combined = 1.20× from step-skipping (wrapped-no-gate vs wrapped-gated) × 1.01× from mx.compile-path avoidance (vanilla vs wrapped-no-gate); vanilla 233.9 s, no-gate 231.2 s, gated 192.4 s medians. v0.6.0 measured 1.23× (236.2 s → 191.8 s) at the same recipe on 2026-05-26 (_artifacts/v0.6.0_bench_klein_base_4b.json), the same number to within noise. Compile-avoidance is effectively noise on this recipe — the v0.4.1-era 1.09× attribution was inflated by same-process MLX state leakage (vanilla ran cold, wrapper inherited warm allocator state). Skip count is byte-identical to v0.4.1 at 9/50 across three reps; the longest run of consecutive skips is 2. Reproduce with uv run python scripts/bench_speedup.py --variant klein-base-4b --three-way --reps 3 --report out.json. Full report: _artifacts/v0.10.0_bench_klein_base_4b.json.
⁵ flux2-klein-base-9b under CFG at the canonical 50-step + g=4.0 recipe. Measured 2026-08-15 under the subprocess-per-rep harness (mflux 0.18.0, v0.10.0 gate): 1.37× combined = 1.34× gating × 1.02× compile-avoidance (vanilla 520.6 s, no-gate 509.8 s, gated 379.1 s medians; every rep within 2 s of its median). 13/50 skips in every rep, never two consecutive; SSIM 0.986 vs vanilla. v0.6.0 measured 1.36× at the same recipe (_artifacts/v0.6.0_bench_klein_base_9b.json); both replace v0.5.0's advertised 2.68×, which was inflated by same-process MLX state leakage in the v0.5.x harness (see Per-variant notes → klein-base-9b correction blockquote). Full report: _artifacts/v0.10.0_bench_klein_base_9b.json. Wrapper peak memory ~9.5 GB vs vanilla's ~22 GB. Reproduce with uv run python scripts/bench_speedup.py --variant klein-base-9b --three-way --reps 3 --report out.json.
⁶ z-image-base under CFG at 50 steps, guidance=4.0, 512×512, q8 (its pinned recipe — the rest of the table is q4), per-variant default rel_l1_thresh=0.12. Measured 2026-08-15 under the subprocess-per-rep harness (mflux 0.18.0, v0.10.0 gate): 1.31× combined (vanilla 227.4 s → wrapper 174.2 s median), entirely from step-skipping at 15 of the 48 active steps (the table's 15 / 50 is over nominal steps; 48 active = 50 minus the skip-first/skip-last windows), the same 15 in every rep and never two consecutive. mx.compile-path avoidance is not a tailwind here: the no-gate wrapper ran at 228.9 s median against vanilla's 227.4 s (0.99×), so the whole 1.31× is gating. Peak memory 17.2 GB → 11.5 GB, from the eager wrapper bypassing mflux's compiled _predict (the no-gate wrapper shows the same ~11.5 GB), not from gating. SSIM 0.991 vs vanilla. v0.7.0 reported 1.17× at this recipe (scripts/_bench_z_image_v0_7_0.json, wrapper 209.4 s) with a thermally confounded three-way split; the skip count and pattern are identical, so the difference is host state, not the gate. Full report: _artifacts/v0.10.0_bench_z_image.json. Reproduce with uv run python scripts/bench_speedup.py --variant z-image --three-way --reps 3 --report out.json.
⁷ qwen-image under CFG at 50 steps, guidance=4.0, 768×768 (its pinned recipe — the rest of the table is 512×512), q4, per-variant default rel_l1_thresh=0.30. Measured 2026-09-06 under the subprocess-per-rep harness (mflux 0.18.0, chunks interleaved rep-outer, medians of three cold reps): 2.68× combined against vanilla (642.6 s → gated 239.4 s), the largest gain in this table, and 2.57× against the wrapper's own no-gate run (615.6 s → 239.4 s). The no-gate wrapper's 1.04× over vanilla is inside the spread of the vanilla reps themselves (610.9 / 642.6 / 715.9 s), so what 0.10.0 reported as an unexplained 1.10× reads as run-to-run noise on this harness; mflux does not mx.compile Qwen's _predict, so there is no compiled path to avoid. 33 of the 48 active steps skipped in every rep (the table's 33 / 50 counts nominal steps), longest run of consecutive skips 4, the same per-step pattern in all three reps and the same one the threshold sweep and the committed calibration replay produce. SSIM 0.967 against vanilla at this recipe, from the sweep on the variant page; 0.9.x skipped 24 here at SSIM 0.978, and rel_l1_thresh=0.20 returns to that point. The bench and the sweep now free Qwen's text encoders once the prompt is encoded (mflux's MemorySaver; the report records memory_saver), which brings the loop peak from 26.2 GB to 24.1 GB in every condition, under this machine's 25.0 GB recommended working set, with the cache pool held at 1.1 GB; the worker still aborts with an artifact if active plus cached MLX memory exceeds physical memory minus 4 GiB. The 0.10.0 numbers (850.3 s vanilla, 3.02×) were taken with the loop peak above that working set; this run stays under it, and the vanilla time fell more than the gated one, which is where the ratio moved. The 0.10.0 report stays committed as _artifacts/v0.10.0_bench_qwen_image.json. Full report: _artifacts/v0.11.0_bench_qwen_image.json. Reproduce with uv run python scripts/bench_speedup.py --variant qwen --three-way --reps 3 --report out.json.
⁸ flux1-krea-dev at 28 steps, guidance=4.5, 512×512, q4, per-variant default rel_l1_thresh=0.30. Measured 2026-09-05 under the subprocess-per-rep harness (mflux 0.18.0): 1.62× combined as the median of three cold reps (vanilla 125.3 s → gated 77.2 s), 1.57× from gating; the no-gate wrapper timed 1.03× ahead of vanilla, and since mflux does not mx.compile the FLUX.1 predict step that residual is run-to-run noise, not compile avoidance. 10 of the 26 active steps skipped in every rep (the table's 10 / 28 counts nominal steps), never two in a row. SSIM 0.990 against vanilla at this recipe, from the threshold sweep documented on the variant page; 0.35 already falls to 0.890, which is why the default sits exactly at the knee. Peak memory 11.1 GB. Full report: _artifacts/v0.11.0_bench_krea_dev.json. Reproduce with uv run python scripts/bench_speedup.py --variant krea-dev --three-way --reps 3 --report out.json.
Reproduce any row:
uv run python scripts/bench_speedup.py --variant flux1-dev # 25-step dev
uv run python scripts/bench_speedup.py --variant klein-4b # 8-step Klein 4B
uv run python scripts/bench_speedup.py --variant klein-9b # 8-step Klein 9B
uv run python scripts/bench_speedup.py --variant klein-base-4b # 50-step base-4B under CFG (g=4.0, v0.4.1+ default)
uv run python scripts/bench_speedup.py --variant klein-base-4b --guidance 1.0 --num-inference-steps 25 # v0.4.0 row
uv run python scripts/bench_speedup.py --variant klein-base-9b # 50-step base-9B under CFG (g=4.0, v0.5.0+ default)
uv run python scripts/bench_speedup.py --variant z-image # 50-step Z-Image base under CFG (g=4.0, q8)
uv run python scripts/bench_speedup.py --variant krea-dev # 28-step FLUX.1 Krea [dev] (g=4.5)For the three-way decomposition (vanilla / wrapped-no-gate / wrapped-gated), add --three-way --reps 3 --report out.json. This is how the v0.10.0 rows above were produced.
The wall-clock improvement above comes from two distinct mechanisms; they fire independently depending on variant and schedule.
1. TeaCache step-skipping. This is the headline feature. The polynomial gate predicts how much the transformer body output will change from the immediately previous step, whether or not that step was itself skipped. The wrapper accumulates the predicted change across consecutive skips, and once that running total would cross rel_l1_thresh, it recomputes the transformer body instead of reusing the cached residual and resets the running total. On FLUX.1-dev at 25 steps, 6 of 25 steps are skippable and this is where the 1.57× speedup on FLUX.1-dev comes from (1.41× of it is gating in the three-way split). On non-distilled FLUX.2 Klein at 50 steps + CFG, the same mechanism produces 9-13 skips per generation and dominates the wall-clock win (the v0.10.0 three-way bench attributes 1.34× to gating on base-9b).
2. mx.compile avoidance on FLUX.2. mflux wraps Flux2Klein._predict in mx.compile on every chip except base + Pro M1/M2 — i.e., compilation is active on M1/M2 Max + Ultra and on every M3, M4, M5 chip. mlx-teacache replaces the compiled _predict with an eager Python closure so the gate can run live per step. The magnitude of this effect varies by schedule and chip. v0.4-era same-process benches on M1 Max at quantize=4 showed the eager closure running 1.5-2.0× faster than the compiled path on distilled klein's 8-step schedule even with zero gate-engagements. v0.6.0's subprocess-per-rep harness re-measured the same mechanism on the 50-step klein-base CFG recipe and found it contributing only 1.01-1.02× — kernel-dispatch round-trips drop slightly, but with 50 steps per generation that gain is small relative to the per-step compute. The distilled schedules likely benefit more because each step is a larger fraction of the wall-clock, so per-step dispatch overhead matters more. The distilled klein rows are still measured under same-process conditions and pending a re-bench. On chips where mflux is already eager (base + Pro M1/M2), this mechanism does not fire: the wrapper just adds per-step gate overhead, and Klein with mlx-teacache on those chips is approximately neutral or slightly slower than vanilla.
On FLUX.2 Klein 4B and 9B at the distilled 4-8 step defaults, mechanism (1) does not engage: the empirical adjacent-step rel-L1 between consecutive transformer outputs is ≥ 0.25, so every step's predicted change exceeds the default 0.20 threshold and the gate signals "compute" every time. This is expected — distilled schedules collapse the entire denoising trajectory into a handful of consequential steps, so adjacent steps are not similar enough to skip. Klein's wall-clock improvement on these variants is real and reproducible, but it comes entirely from mechanism (2).
For algorithmic step-skipping on FLUX.2, use the non-distilled flux2-klein-base-4b (Apache-2.0) or flux2-klein-base-9b (FLUX Non-Commercial) variants. Both ship with a per-variant default rel_l1_thresh=0.17. At the canonical upstream 50-step + g=4.0 CFG recipe, base-4b skips 9/50 steps for a measured 1.22× combined speedup (subprocess-per-rep), and base-9b skips 13/50 steps for 1.37×. Pushing the threshold higher on distilled Klein is not recommended: the gate's prediction quality at thresholds > 0.25 is uncalibrated on a 4-8 step trajectory and image quality is not characterized there.
Quality gates use a 5-prompt SSIM suite defined at tests/test_image_quality_flux1.py:45 and reused at tests/test_image_quality_flux2.py:28:
- "a red apple on a wooden table"
- "mountain landscape at sunset"
- "portrait of a woman"
- "abstract pattern with circles"
- "text saying HELLO"
The PR-gate prompt is the red-apple one; SSIM ≥ 0.90 on FLUX.1-dev and ≥ 0.85 on Klein 4B / 9B at the default threshold. Full suite floor is 0.80 to absorb high-frequency-detail variance (text, synthetic patterns). Run uv run pytest tests/test_image_quality_flux1.py tests/test_image_quality_flux2.py -m parity with real model weights to reproduce.
mflux wraps _predict in mx.compile on every Apple Silicon chip except base + Pro M1/M2 (the behaviour is unchanged from mflux 0.17.5 through 0.19). The is_m1_or_m2() predicate returns true (eager path) when the chip brand contains "Apple M1" or "Apple M2" and does not contain "Max" or "Ultra" — so M1 Pro and M2 Pro are eager too, while M1/M2 Max + Ultra and every M3/M4/M5 chip get the compiled path. mlx-teacache replaces _predict with an eager closure so per-step gating stays live, trading the compile gain for the skip gain on compiled chips. That compile-avoidance effect is FLUX.2 Klein only — FLUX.1, FLUX.1 Krea, and Qwen-Image have no mx.compile path, so their step-skipping speedup is the same on every chip. See docs/m3-plus-tradeoff.md for a benchmark recipe.
| Chip | Vanilla _predict in mflux |
Expected speedup |
|---|---|---|
| Apple M1 / M2 (base) | eager | the model's skip fraction, chip-independent (1.57× FLUX.1-dev, 1.62× Krea, 2.68× Qwen); FLUX.2 loses its small compile-avoidance share here |
| M1 Pro / M2 Pro | eager | same as base — the skip fraction, chip-independent |
| M1 Max / Ultra, M2 Max / Ultra | compiled | measured on M1 Max: FLUX.1-dev 1.57×, FLUX.1 Krea 1.62×, Qwen-Image 2.68×, klein-base-4b 1.22×, klein-base-9b 1.37×, z-image-base 1.31× (see Benchmarks) |
| M3 / M3 Pro / M3 Max / Ultra | compiled | FLUX.1 / Qwen keep their skip fraction; FLUX.2 likely 1.1–1.3× — untested |
| M4 / M4 Pro / M4 Max | compiled | FLUX.1 / Qwen keep their skip fraction; FLUX.2 likely 1.1–1.3× — untested |
| M5+ (Neural Accelerators / TensorOps) | compiled + accelerator | FLUX.2 may approach 1.0× — the eager wrapper can lose some or all of the M5 TensorOps advantage; confirm with a profiler. FLUX.1 / Qwen stay skip-limited |
img2img reuses the txt2img calibration. A dedicated img2img calibration may follow in a future release if SSIM gates flag drift on specific schedules.
FLUX.2 with CFG (guidance > 1.0) runs through the gated path as of v0.4.1. The wrapper keeps two cached residuals (positive and negative branch) and shares one gate decision per step across both. The canonical base-4b recipe (guidance_scale=4.0, num_inference_steps=50) measures 1.22× combined on M1 Max under the subprocess-per-rep harness (v0.10.0; v0.6.0 measured 1.23×, v0.4.1 advertised 1.26× combined under same-process measurement; combined within day-to-day noise, but the decomposition shifted — see Benchmarks footnote ⁴).
Distilled schedules are out of scope for algorithmic step-skipping by design. This includes FLUX.2 Klein 4B / 9B at their 4-8 step defaults and FLUX.1 schnell at its 4-step default. The polynomial gate's premise — that consecutive transformer outputs are similar enough that the residual can be reused — does not hold on distilled trajectories where each step does a much larger share of the denoising work. On the v0.3.0 bench (M1 Max, quantize=4) the gate signals "compute" on every Klein step at the package default rel_l1_thresh=0.20 (0 skips across 3 reps on both Klein 4B and 9B); empirical adjacent-step body-output rel-L1 on Klein is ≥ 0.25. Klein still gets a real wall-clock improvement (~1.2-1.9×) from mx.compile-path avoidance, but the headline TeaCache step-skipping feature only fires on non-distilled schedules.
flux2-klein-base-4b runs TeaCache at both guidance=1.0 (single-branch path) and guidance > 1.0 (per-branch path, v0.4.1+). The upstream BFL base-4b model card recommends guidance_scale=4.0, num_inference_steps=50; the subprocess-per-rep harness measures 1.22× combined wall-clock vs vanilla on M1 Max at that recipe (v0.10.0: 1.20× gating + 1.01× compile-avoidance; 9/50 skips, SSIM PR-gate passed; v0.6.0 measured 1.23×). v0.4.1's same-process bench reported 1.26× combined and decomposed it 1.16× gating / 1.09× compile-avoidance; the combined was honest within day-to-day noise but the decomposition over-attributed to compile-avoidance.
flux2-klein-base-9b reuses base-4b's polynomial coefficients verbatim (same architecture family, same calibration recipe). The v0.10.0 subprocess-per-rep bench at the canonical 50-step / guidance=4.0 recipe measures 1.37× combined wall-clock (1.34× gating + 1.02× compile-avoidance; v0.6.0 measured 1.36×), 13/48 active steps skipped at rel_l1_thresh=0.17, SSIM 0.986 vs vanilla. v0.5.0 advertised 2.68× on this variant — that number was inflated by same-process MLX state leakage and is corrected here. Same FLUX Non-Commercial license + BFL safety-filter obligations as flux2-klein-9b — see License obligations.
The wrapper runs eager, which gives up mflux's mx.compile of _predict in exchange for live per-step gating. Vanilla mflux compiles _predict on every chip except base + Pro M1/M2 (the is_m1_or_m2() predicate only excludes Max + Ultra). The 1.57× measurement is from M1 Max / FLUX.1-dev / 25 steps; speedup on M3 and newer is plausible but untested locally. On M5, the GPU Neural Accelerators (Metal 4 TensorOps) are only reachable through the compiled path, so the eager wrapper can lose some or all of that advantage. Output stays correct either way. See docs/m3-plus-tradeoff.md for the per-chip recipe; PRs with measurements welcome.
FLUX.2 parity is numerical, not bit-exact. Replacing a function that mflux wraps in mx.compile produces about 1 ULP per element of divergence from Metal kernel-dispatch noise, which compounds across steps but keeps cosine similarity ≥ 0.97 on Klein 4B, Klein 9B, and base-4b under CFG at threshold 0. The user-facing guarantee is end-to-end image quality (SSIM ≥ 0.85 on all supported FLUX.2 variants at the package default threshold).
The mflux pin is strict at >=0.17.5,<0.20. Bumping it is a deliberate release: each new mflux minor is verified on real weights before the range widens.
Calling flux.parameters() at the parent level can miss transformer parameters while the wrapper is active. Use flux.transformer.parameters() directly, or call handle.restore() first.
The FLUX.1 variants (flux1-dev, flux1-schnell) and flux2-klein-4b come with their own upstream weight licenses; the wrapper this library applies does not change those terms.
flux2-klein-9b and flux2-klein-base-9b are both distributed under the FLUX.2 Klein license (non-commercial use + BFL safety-filter obligations). These terms flow with the weights, not with mlx-teacache. If you call apply_teacache on either variant — Flux2Klein(model_config=ModelConfig.flux2_klein_9b()) or Flux2Klein(model_config=ModelConfig.flux2_klein_base_9b()) — you are responsible for ensuring your use complies with the upstream license, including the safety-filter requirements the BFL model cards describe. See the official model cards at https://huggingface.co/black-forest-labs/FLUX.2-klein-9B and https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B for the full terms.
Open an issue at https://github.com/IonDen/mlx-teacache/issues. If you're adding a variant or changing the gate, docs/architecture.md is the contributor's map and docs/calibration.md covers coefficient calibration.
Apache-2.0. See LICENSE and NOTICE.
- ali-vilab/TeaCache — upstream method and FLUX.1 coefficients.
- mflux-community/mflux — the MLX diffusion runner this library integrates with (moved from
filipstrand/mfluxin August 2026; the old URL redirects). - Apple ML Explore — MLX.
Other MLX libraries for Apple Silicon:
- mlx-taef — tiny autoencoders for fast diffusion-latent previews and low-memory decode (FLUX / SD).
- 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