Repository navigation
Run FLUX.2 Klein from DFloat11 checkpoints and fix the CAPPED tier limits - #14
Merged
Merged
Conversation
IonDen
force-pushed
the
feature/s3-flux2-klein
branch
from
October 8, 2026 09:20
28997da to
af90b7b
Compare
IonDen
force-pushed
the
feature/s3-flux2-klein
branch
from
October 8, 2026 09:25
af90b7b to
d42c0f5
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This adds the four FLUX.2 Klein models to
mlx-dfloat generateand the Python API: base and distilled, 4B and 9B, each run from its DFloat11 checkpoint. It also fixes how--tierimitates 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_embedderandnorm_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.
--scheduleris 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
--tierfix: 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--tierrun 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:
bench/results/tiers/flux2-klein-4b-1024.jsonbench/results/tiers/flux2-klein-base-4b-1024.jsonbench/results/tiers/flux2-klein-9b-1024.jsonbench/results/tiers/flux2-klein-base-9b-1024.jsonThe 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.pyandmflux/__init__.py(the Klein entries and CLI rules), andbench/capped.py,_memory_caps.py,integrate/memory.pyandbench/table.py(the--tierfix). 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
--tierfix, 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.