Repository navigation
Add Krea 2 Turbo and Krea 2 Raw, and prepare the 0.2.0 release - #17
Merged
Merged
Conversation
IonDen
force-pushed
the
feature/s6-krea-2
branch
from
October 8, 2026 23:27
ca8f35d to
a09d436
Compare
IonDen
force-pushed
the
feature/s6-krea-2
branch
from
October 8, 2026 23:43
a09d436 to
5c03d48
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 pull request adds Krea 2 Turbo and Krea 2 Raw to
mlx-dfloat generateand to the Python API, and it carries therelease 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-2runs Krea 2 Turbo and--model krea-2-rawruns Krea 2 Raw; from Python,DFloatModel("krea-2")andDFloatModel("krea-2-raw")return aDFloatKrea2. The DFloat11 checkpoints are theComfyUI single files
mingyi456/Krea-2-Turbo-DF11-ComfyUIandmingyi456/Krea-2-Raw-DF11-ComfyUI, which have noconfig.json. The reader opens each through a layout pinned to the file's revision, and building one model from theother's file is refused with the
--modelthe file belongs to. The text encoder, VAE and tokenizer come fromkrea/Krea-2-Turboandkrea/Krea-2-Raw, pinned to the revisions the measured runs used. Both repositories are gatedon 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_sdescheduler, so eachstep 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-promptthat has no effect at 1.0, and refuses ascheduler other than
er_sde,eulerorlinear. An empty or blank prompt is refused, because Krea 2's tokenizergives 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/andbench/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 GBoriginal 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/, outsidethe 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-checkis passed; thelighthouse 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 watchdogwhile 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 generaterun per model at 1024², seed 42, on the 32 GB M1 Max (watched_peak_bytesofthe file in
bench/results/tiers/, except the last row):z-image-turbo-1024.jsonz-image-1024.jsonflux2-klein-4b-1024.jsonflux2-klein-base-4b-1024.jsonflux2-klein-9b-1024.jsonflux2-klein-base-9b-1024.jsonqwen-image-2.1-1024.jsonernie-image-1024.jsonernie-image-turbo-1024.jsonkrea-2-1024.jsonkrea-2-raw-1024.jsonbench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.jsonThe 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_bytesin each file). The same command does notpeak 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_imageorencode()now runs those calls under mlx-dfloat's memory caps, 20 GiBwired 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 16and
--tier 24now install that Mac's caps and check the fit against that Mac's budget; the stricter budget can refusea smaller-Mac run that 0.1.0 accepted, and
--no-fit-checkstill forces it.The FLUX.1 commands and defaults are unchanged,
--modelstill defaults toschnell, and themfluxextra stillinstalls mflux 0.20. The command's module moved from
mlx_dfloat.mflux.flux1.clitomlx_dfloat.mflux.generate; themlx-dfloatconsole script is the same, so only code that imported the old module directly is affected.After the merge, the
v0.2.0tag publishes the release to PyPI.