Skip to content

Add Krea 2 Turbo and Krea 2 Raw, and prepare the 0.2.0 release - #17

Merged
IonDen merged 5 commits into
mainfrom
feature/s6-krea-2
Oct 9, 2026
Merged

IonDen merged 5 commits into
mainfrom
feature/s6-krea-2

Conversation

@IonDen

@IonDen IonDen commented Oct 8, 2026

Copy link
Copy Markdown
Owner

This pull request adds Krea 2 Turbo and Krea 2 Raw to mlx-dfloat generate and to the Python API, and it carries the
release text for 0.2.0. With it, 0.2.0 runs five image-model families besides FLUX.1: Z-Image, FLUX.2 Klein,
Qwen-Image-2.1, ERNIE-Image and Krea 2. As before, the weights stay compressed in memory and the GPU turns each block
back into the exact bits the checkpoint stores just before the block runs.

Krea 2

mlx-dfloat generate --model krea-2 runs Krea 2 Turbo and --model krea-2-raw runs Krea 2 Raw; from Python,
DFloatModel("krea-2") and DFloatModel("krea-2-raw") return a DFloatKrea2. The DFloat11 checkpoints are the
ComfyUI single files mingyi456/Krea-2-Turbo-DF11-ComfyUI and mingyi456/Krea-2-Raw-DF11-ComfyUI, which have no
config.json. The reader opens each through a layout pinned to the file's revision, and building one model from the
other's file is refused with the --model the file belongs to. The text encoder, VAE and tokenizer come from
krea/Krea-2-Turbo and krea/Krea-2-Raw, pinned to the revisions the measured runs used. Both repositories are gated
on the Hub: a user logs in once and accepts the Krea 2 Community License, and the Hub grants access on the spot. Their
26 GB transformers are never downloaded.

The defaults follow mflux: 8 steps for Turbo and 25 for Raw, both at guidance 1.0 with the er_sde scheduler, so each
step runs the transformer once. Raw's model card uses 52 steps at guidance 3.5, and the README gives that recipe as
flags. As in mflux, any guidance other than 1.0 turns on classifier-free guidance, which runs the transformer twice
per step; the command says so in one line, warns about a --negative-prompt that has no effect at 1.0, and refuses a
scheduler other than er_sde, euler or linear. An empty or blank prompt is refused, because Krea 2's tokenizer
gives an empty prompt no tokens.

Both checkpoints were compared in full with their originals, read from the Hub by HTTP range reads. In each, the 256
matrices of the 28 transformer blocks and the four text-fusion blocks equal the BF16 originals, and the 5 other
matrices and 169 other tensors equal the FP32 originals rounded to BF16, nearest even. The GPU decoder matches the CPU
reference on all 261 matrices of both models. The records are in bench/results/parity/krea-2/ and
bench/results/parity/krea-2-raw/.

Krea 2 Raw's final latents, after 4 steps at 1024² and guidance 3.5, match bit for bit those of
the same transformer streaming its original weights one block at a time, and the two images match pixel for pixel
(bench/results/identity/krea-2-raw-1024/compare.json). Turbo's latents were not compared: that check needs its 26 GB
original transformer on disk.

Memory on a 32 GB Mac

Krea 2 is the tightest fit so far. Its compressed transformer alone is 16.34 GiB. Besides its 28 blocks it has seven
smaller groups (four text-fusion blocks and three small layers) that decode to 1.23 GiB. Kept decoded for the whole
run, they pushed a guided Raw step to 22.56 GiB, close to the 22.96 GiB budget. So mlx-dfloat decodes these groups at
the start of every transformer call and frees them before the first block, and it empties MLX's cache at the start of
each call. The same step then peaked at 21.14 GiB (both runs in bench/results/calibration/krea-2-raw-1024.json).

Measured at 1024² with seed 42 on an M1 Max (macOS 27.0.1, mlx 0.32.2, mflux 0.20.0), Krea 2 Turbo peaked at
21.17 GiB in 5.1 minutes, and Krea 2
Raw at its default 25 steps peaked at 21.15 GiB in
14.1 minutes. With the card's 52 steps at guidance 3.5, Raw peaked at
21.15 GiB in
56.2 minutes. That run needs --wall-budget 7200,
because the command's watchdog otherwise stops a run after an hour. Its file sits in bench/results/recipes/, outside
the tier table, which shows no step count.

The fit estimate puts the denoising step at 22.58 GiB against the 22.96 GiB budget. Every call therefore drops the
compressed transformer before the VAE decode and loads it again for the next one. The estimate also grows with the
prompt, so a prompt of more than about 500 tokens is refused on a 32 GB Mac unless --no-fit-check is passed; the
lighthouse prompt used for every measurement is 30 tokens.

No 16 GB or 24 GB Mac was measured. Under those Macs' caps on the 32 GB Mac, the fit check refuses both models at both
sizes (bench/results/refusals/krea-2{,-raw}-1024-tier{16,24}.json), and forced runs were stopped by the watchdog
while the compressed transformer was loading (bench/results/tiers/aborts/), so both sizes are at risk.

This pull request makes no speed claim. Like the other new families, Krea 2 runs its denoising step without mflux's step compilation,
because the step has to pause between blocks to decode the next one, and that cost against stock mflux was not
measured.

The 0.2.0 release text

The changelog's unreleased entries move under a 0.2.0 heading with a short summary. The README's Status section, its
model table and its install text say "released in 0.2.0" where they said "next release". The 20B Qwen-Image models
are listed for a later release, for Macs with more than 32 GB. The one-line summary PyPI shows names the image models
instead of "today that runs FLUX.1". The earlier families came in
#12 (Z-Image), #13 (single-file
checkpoints), #14 (FLUX.2 Klein), #15
(Qwen-Image-2.1) and #16 (ERNIE-Image), and the GPU decoder self-test in
#11.

Peak memory of one mlx-dfloat generate run per model at 1024², seed 42, on the 32 GB M1 Max (watched_peak_bytes of
the file in bench/results/tiers/, except the last row):

Model Steps Peak File
Z-Image-Turbo 9 13.11 GiB z-image-turbo-1024.json
Z-Image 50, guidance 4 13.15 GiB z-image-1024.json
FLUX.2-klein-4B 4 11.88 GiB flux2-klein-4b-1024.json
FLUX.2-klein-base-4B 50, guidance 4 12.78 GiB flux2-klein-base-4b-1024.json
FLUX.2-klein-9B 4 18.40 GiB flux2-klein-9b-1024.json
FLUX.2-klein-base-9B 50, guidance 4 17.60 GiB flux2-klein-base-9b-1024.json
Qwen-Image-2.1 40 20.42 GiB qwen-image-2.1-1024.json
ERNIE-Image 50, guidance 4 16.79 GiB ernie-image-1024.json
ERNIE-Image-Turbo 8 17.34 GiB ernie-image-turbo-1024.json
Krea 2 Turbo 8 21.17 GiB krea-2-1024.json
Krea 2 Raw 25 21.15 GiB krea-2-raw-1024.json
Krea 2 Raw 52, guidance 3.5 21.15 GiB bench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.json

The FLUX.1 rows 0.1.0 shipped are unchanged: 19.78, 19.96 and 20.04 GiB for schnell, dev and Krea-dev. Every row stays
under the 22.96 GiB budget mlx-dfloat checks on this Mac (fit.budget_bytes in each file). The same command does not
peak at the same value every time: six runs of the Z-Image-Turbo command peaked between 13.11 and 14.26 GiB
(bench/results/repeats/z-image-turbo-1024/).

Upgrade risk

A Python program that calls generate_image or encode() now runs those calls under mlx-dfloat's memory caps, 20 GiB
wired and 22 GiB memory on a 32 GB Mac, where 0.1.0 left MLX's defaults in place. A program that sets its own wired
limit before the call keeps it. Tier reports made with 0.1.0 are not comparable with 0.2.0 ones, because --tier 16
and --tier 24 now install that Mac's caps and check the fit against that Mac's budget; the stricter budget can refuse
a smaller-Mac run that 0.1.0 accepted, and --no-fit-check still forces it.

The FLUX.1 commands and defaults are unchanged, --model still defaults to schnell, and the mflux extra still
installs mflux 0.20. The command's module moved from mlx_dfloat.mflux.flux1.cli to mlx_dfloat.mflux.generate; the
mlx-dfloat console script is the same, so only code that imported the old module directly is affected.

After the merge, the v0.2.0 tag publishes the release to PyPI.

@IonDen IonDen self-assigned this Oct 8, 2026
@IonDen IonDen added enhancement New feature or request release Release preparation labels Oct 8, 2026
@IonDen
IonDen force-pushed the feature/s6-krea-2 branch from ca8f35d to a09d436 Compare October 8, 2026 23:27
@IonDen
IonDen force-pushed the feature/s6-krea-2 branch from a09d436 to 5c03d48 Compare October 8, 2026 23:43
@IonDen
IonDen merged commit 620c99a into main Oct 9, 2026
6 checks passed
@IonDen
IonDen deleted the feature/s6-krea-2 branch October 9, 2026 07:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request release Release preparation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant