Skip to content

Run FLUX.2 Klein from DFloat11 checkpoints and fix the CAPPED tier limits - #14

Merged
IonDen merged 4 commits into
mainfrom
feature/s3-flux2-klein
Oct 8, 2026
Merged

IonDen merged 4 commits into
mainfrom
feature/s3-flux2-klein

Conversation

@IonDen

@IonDen IonDen commented Oct 8, 2026 •

Copy link
Copy Markdown
Owner

This adds the four FLUX.2 Klein models to mlx-dfloat generate and the Python API: base and distilled, 4B and 9B, each run from its DFloat11 checkpoint. It also fixes how --tier imitates a smaller Mac.

Klein reuses the per-block hook added with Z-Image support (#12): each double and single block gets its decoded weights just before it runs. The five compressed groups that are not blocks (three modulation layers, context_embedder and norm_out) are decoded once when the weights load. mflux's compiled denoising step is bypassed so the blocks can be evaluated one at a time, and weight names come from mflux's own Klein mapping.

The distilled models run at guidance 1.0 only, and the CLI refuses another value, as mflux's command does. --scheduler is refused for Klein for the same reason. Each model's default base is its own BF16 repository, and the default checkpoints are pinned to fixed revisions.

Klein's text encoder reads only hidden states 9, 18 and 27, so layers 27 to 35 are never evaluated and their weights never load. The memory estimate counts only the layers that run, read from the encoder's file headers. A warning fires if an mflux change ever makes the encode phase use more than that. A base repository missing a text-encoder or VAE tensor is now refused with the missing names; before, mflux kept the module's random initial weights for them without a word.

The --tier fix: a run that simulates a 16 or 24 GB Mac used MLX's default limits, including no wired limit. Without one, MLX allocated 9.49 GiB instead of 6.89 GiB for the same Klein VAE decode (exact MLX peak, identical output), enough to stop runs a real Mac would pass. A --tier run now installs the wired and memory caps mlx-dfloat installs on a Mac of that size (16 GB: 8 GiB wired, 10 GiB memory; 24 GB: 14 and 16 GiB), plans against the fit budget such a Mac applies (its working set minus 2 GiB: 8.67 GiB at 16 GB, 14.00 GiB at 24 GB; it used to be the watchdog ceiling, 0.5 GiB looser), and a watchdog stop records the limits in force. The committed 16 and 24 GB results were recorded under the older budget; none of their outcomes changes under the stricter one.

Measured on an M1 Max with 32 GB (macOS 27.0.1, mlx 0.32.2, mflux 0.20.0), at 1024², seed 42. Peak is the process's memory footprint as the OS counts it:

Model Recipe Peak Time Record
FLUX.2-klein-4B 4 steps, guidance 1.0 11.88 GiB 35 s bench/results/tiers/flux2-klein-4b-1024.json
FLUX.2-klein-base-4B 50 steps, guidance 4.0 12.78 GiB 8.9 min bench/results/tiers/flux2-klein-base-4b-1024.json
FLUX.2-klein-9B 4 steps, guidance 1.0 18.40 GiB 1.3 min bench/results/tiers/flux2-klein-9b-1024.json
FLUX.2-klein-base-9B 50 steps, guidance 4.0 17.60 GiB 19.9 min bench/results/tiers/flux2-klein-base-9b-1024.json

The step count of each run follows from its recorded decode count (decode_launches): 2,500 for base 4B is 25 blocks × 2 calls × 50 steps, 3,200 for base 9B is 32 × 2 × 50.

The base 4B checkpoint decodes to exactly its BF16 original: all 105 compressed matrices and 64 stored tensors (bench/results/parity/flux2-klein-base-4b/). Its final latents and pixels at 1024², 4 steps, guidance 4 are bit-identical to the same transformer streaming its BF16 weights (bench/results/identity/flux2-klein-base-4b-1024/). For the distilled 4B and both 9B checkpoints, the GPU kernel matches the CPU reference on every matrix, and sampled groups match the BF16 originals (bench/results/parity/flux2-klein-{4b,base-9b,9b}/). The same folder now also holds the Z-Image parity records the earlier release notes cite.

Smaller Macs, simulated on the 32 GB Mac under that size's caps: base 4B passes at 16 GB (8.15 GiB, the weights dropped before the VAE decode) and at 24 GB (12.29 GiB). Base 9B at 24 GB is refused by the fit check, and a forced run is stopped by the watchdog in the denoise phase at 14.64 GiB against a 14.50 GiB ceiling, so that size is marked at risk. No 16 or 24 GB Mac was measured.

The FLUX.1 and Z-Image model code is unchanged. Shared files that do change: families.py, generate.py and mflux/__init__.py (the Klein entries and CLI rules), and bench/capped.py, _memory_caps.py, integrate/memory.py and bench/table.py (the --tier fix). Schnell's and Z-Image-Turbo's 1024² images are pixel-identical to the earlier records, and schnell's latents are still bit-identical to the BF16-streamed run. Z-Image-Turbo's 16 GB stop was re-run under the new caps; it still stops, now after 4.7 s at 9.26 GiB.

The four commits land in order: the --tier fix, Klein support, the tool updates, then docs and records. The README table check (tests/test_readme_bench_block.py) passes only on the last one, because the README and the records it renders change together there.

@IonDen IonDen self-assigned this Oct 8, 2026
@IonDen IonDen added enhancement New feature or request bug Something isn't working labels Oct 8, 2026
@IonDen
IonDen force-pushed the feature/s3-flux2-klein branch from 28997da to af90b7b Compare October 8, 2026 09:20
@IonDen
IonDen force-pushed the feature/s3-flux2-klein branch from af90b7b to d42c0f5 Compare October 8, 2026 09:25
@IonDen
IonDen merged commit 89575dd into main Oct 8, 2026
6 checks passed
@IonDen
IonDen deleted the feature/s3-flux2-klein branch October 8, 2026 12:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant