diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 90763bc..39e0e79 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,7 +54,7 @@ jobs: - uses: astral-sh/setup-uv@v5 - run: uv sync --locked --group dev --extra mflux - run: uv run --extra mflux --group dev python -c "import mflux, mlx.core, importlib.metadata as m; print('mflux', m.version('mflux'))" - # No combined threshold: init.py's tests run in the default lane (its live ones skip without a token); the next five steps gate each model.py on its own. + # No combined threshold: the steps after this one gate each model.py, and each init.py and transformer.py the required job omits, on its own. - name: Integration tests (mflux lane) env: # Every FLUX.1 base repository is gated on the Hub (schnell auto-approved), so the three @@ -62,7 +62,7 @@ jobs: # Only this step gets it: the actions, the sync and the collection and coverage steps have # no use for it. The secret is a read-only Hub token; the default lane runs without one. HF_TOKEN: ${{ secrets.HF_TOKEN }} - run: uv run --extra mflux --group dev pytest -m mflux --run-network --cov=mlx_dfloat.mflux.flux1.model --cov=mlx_dfloat.mflux.flux1.init --cov=mlx_dfloat.mflux.zimage.model --cov=mlx_dfloat.mflux.flux2.model --cov=mlx_dfloat.mflux.qwen21.model --cov=mlx_dfloat.mflux.qwen21.init --cov=mlx_dfloat.mflux.qwen21.transformer --cov=mlx_dfloat.mflux.ernie.model --cov=mlx_dfloat.mflux.ernie.init --cov=mlx_dfloat.mflux.ernie.transformer --cov-config=.coveragerc-integration --cov-report=term-missing -q -p no:cacheprovider -rs + run: uv run --extra mflux --group dev pytest -m mflux --run-network --cov=mlx_dfloat.mflux.flux1.model --cov=mlx_dfloat.mflux.flux1.init --cov=mlx_dfloat.mflux.zimage.model --cov=mlx_dfloat.mflux.zimage.init --cov=mlx_dfloat.mflux.flux2.model --cov=mlx_dfloat.mflux.flux2.init --cov=mlx_dfloat.mflux.flux2.transformer --cov=mlx_dfloat.mflux.qwen21.model --cov=mlx_dfloat.mflux.qwen21.init --cov=mlx_dfloat.mflux.qwen21.transformer --cov=mlx_dfloat.mflux.ernie.model --cov=mlx_dfloat.mflux.ernie.init --cov=mlx_dfloat.mflux.ernie.transformer --cov=mlx_dfloat.mflux.krea2.model --cov=mlx_dfloat.mflux.krea2.init --cov=mlx_dfloat.mflux.krea2.transformer --cov-config=.coveragerc-integration --cov-report=term-missing -q -p no:cacheprovider -rs - name: model.py coverage on its own (mflux lane) run: uv run --extra mflux --group dev coverage report --include='*/mflux/flux1/model.py' --fail-under=80 - name: zimage/model.py coverage on its own (mflux lane) @@ -73,6 +73,28 @@ jobs: run: uv run --extra mflux --group dev coverage report --include='*/mflux/qwen21/model.py' --fail-under=80 - name: ernie/model.py coverage on its own (mflux lane) run: uv run --extra mflux --group dev coverage report --include='*/mflux/ernie/model.py' --fail-under=80 + - name: krea2/model.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/krea2/model.py' --fail-under=80 + - name: flux1/init.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/flux1/init.py' --fail-under=80 + - name: zimage/init.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/zimage/init.py' --fail-under=80 + - name: flux2/init.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/flux2/init.py' --fail-under=80 + - name: flux2/transformer.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/flux2/transformer.py' --fail-under=80 + - name: qwen21/init.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/qwen21/init.py' --fail-under=80 + - name: qwen21/transformer.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/qwen21/transformer.py' --fail-under=80 + - name: ernie/init.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/ernie/init.py' --fail-under=80 + - name: ernie/transformer.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/ernie/transformer.py' --fail-under=80 + - name: krea2/init.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/krea2/init.py' --fail-under=80 + - name: krea2/transformer.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/krea2/transformer.py' --fail-under=80 - name: The lane collected mflux tests run: test "$(uv run --extra mflux --group dev pytest -m mflux --co -q -p no:cacheprovider | grep -c '::')" -ge 2 build: diff --git a/CHANGELOG.md b/CHANGELOG.md index df4560c..5788d74 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,26 @@ All notable changes to this project are documented here. The format follows ## [Unreleased] +## [0.2.0] - 2026-10-08 + +0.2.0 adds five image-model families to the FLUX.1 support of 0.1.0. Each runs from its DFloat11 checkpoint through +mflux, and the GPU decodes every compressed weight to exactly the bits the checkpoint stores. The peaks below are from +one 1024² run per model, six for Z-Image-Turbo, on a 32 GB M1 Max (macOS 27.0.1, mlx 0.32.2, mflux 0.20.0; the files +in `bench/results/tiers/` and `bench/results/repeats/`): + +- Z-Image-Turbo and Z-Image: 13.11 and 13.15 GiB. Six runs of the Turbo command peaked between 13.11 and 14.26 GiB. +- FLUX.2 Klein, 4B and 9B, base and distilled: 11.88 GiB (FLUX.2-klein-4B) to 18.40 GiB (FLUX.2-klein-9B). +- Qwen-Image-2.1: 20.42 GiB. +- ERNIE-Image and ERNIE-Image-Turbo: 16.79 and 17.34 GiB. +- Krea 2 Turbo and Krea 2 Raw: 21.17 and 21.15 GiB. + +None of these models was measured on a 16 GB or 24 GB Mac. Under those Macs' memory caps on the 32 GB Mac, only +FLUX.2-klein-base-4B passed at both sizes, and ERNIE-Image-Turbo passed one forced run at 24 GB. + +From Python, every model runs under the same memory caps as the command, unless the process already set a wired +limit. `mlx-dfloat selftest` checks the GPU decoder +on your Mac. + ### Added - `mlx-dfloat selftest` and `mlx_dfloat.decode.selftest()` decode two packaged groups, one cut from a real Qwen3-4B @@ -80,6 +100,38 @@ All notable changes to this project are documented here. The format follows - A FLUX.2 Klein call warns when the prompt encode holds more memory than the fit estimate allows for it. The estimate counts only the text-encoder layers whose output FLUX.2 Klein reads, so the warning fires if a future mflux runs them all. +- Krea 2 Turbo and Krea 2 Raw generation through mflux: `mlx-dfloat generate --model krea-2|krea-2-raw` and the + `DFloatKrea2` class, also returned by `DFloatModel`. 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; their transformers are not downloaded. Both + repositories are gated on the Hub (the Hub grants access once you accept the Krea 2 Community License). The defaults + are 8 steps for Turbo and 25 for Raw, both at guidance 1.0 with the `er_sde` scheduler, so one transformer call per + step. Turbo's come from mflux's own Krea 2 command, which does not run Raw; Raw's 25 steps are mflux's fallback for a + model without a step default of its own. Any other guidance runs classifier-free guidance as two transformer calls per step; Raw's model card + uses 52 steps at guidance 3.5. An empty or blank prompt is refused as a user error. +- The default Krea 2 checkpoints, `mingyi456/Krea-2-Turbo-DF11-ComfyUI` and `mingyi456/Krea-2-Raw-DF11-ComfyUI`, are + ComfyUI files with no `config.json`. The reader opens each through a layout pinned to that file's revision, and a + Krea 2 model built from the other model's file is refused with the `--model` the file belongs to. +- Both Krea 2 checkpoints were compared in full with their originals by HTTP range reads. In each, the 256 matrices of + the transformer and text-fusion blocks equal the BF16 originals, and the other 5 matrices and 169 stored tensors + equal the FP32 originals rounded to BF16, nearest even. The GPU decoder matches the CPU reference on all 261 matrices + of each (records in `bench/results/parity/krea-2/` and `bench/results/parity/krea-2-raw/`). In a 4-step, 1024² run + at guidance 3.5, Krea 2 Raw's final latents match bit for bit those of the same transformer streaming its original + weights one block at a time (`bench/results/identity/krea-2-raw-1024/compare.json`); Turbo's latents have not been + compared with a reference. +- Krea 2 decodes its seven non-block groups (the four text-fusion blocks and three small layers, 1.23 GiB decoded) at + the start of every transformer call and frees them before the first block, and it empties MLX's cache at each call's + start and again before the first block. On one guided Raw step the denoising peak was 21.14 GiB, against 22.56 GiB + with the groups kept decoded (one run each, `bench/results/calibration/krea-2-raw-1024.json`); most of that comes + from emptying the cache, and decoding per call on its own saves about 0.14 GiB. +- The measured 32 GB runs peak at 1024² at 21.17 GiB for Krea 2 Turbo (8 steps) and at 21.15 GiB for Krea 2 Raw (25 + steps) and for Raw with its model card's 52 steps at guidance 3.5 (`bench/results/tiers/krea-2-1024.json`, + `krea-2-raw-1024.json` and `bench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.json`). Every call drops the compressed transformer + before the VAE decode. The fit estimate puts the denoising step at 22.58 GiB against the 22.96 GiB budget, and it + predicts a refusal on a 32 GB Mac at about 500 prompt tokens (only a 30-token prompt was measured). Under a 16 GB and a 24 GB Mac's caps the fit check + refuses both models (`bench/results/refusals/krea-2-1024-tier16.json` and the other three), and forced runs were + stopped by the watchdog while the compressed transformer was loading, so both sizes are at risk. +- `scripts.verify_remote_group --extras` also compares every uncompressed tensor with its original of the same name, + BF16 as is or, with `--cast-fp32-to-bf16`, FP32 rounded to BF16; `--groups none --extras` compares the extras alone. - `scripts.verify_remote_group --cast-fp32-to-bf16` compares a DFloat11 repository with an FP32 original rounded to BF16, nearest even. @@ -95,6 +147,10 @@ All notable changes to this project are documented here. The format follows `--guidance` other than 1.0 and ignores `--negative-prompt` with a warning, as mflux's Turbo command does. The image identity check covers ERNIE-Image and refuses `--negative-prompt` and a guidance other than 1.0 for ERNIE-Image-Turbo. +- `mlx-dfloat generate --model` also accepts `krea-2` and `krea-2-raw`. For them, the `--negative-prompt` warning + follows mflux's Krea 2 rule (classifier-free guidance runs at any guidance other than 1.0), the command prints a line + when a call runs two transformer calls per step, and `--scheduler` takes `er_sde`, `euler` or `linear` and refuses + any other name. The image identity check accepts Krea 2 Raw, and Krea 2 Turbo at guidance 1.0. - `mlx-dfloat generate --report` records `elapsed_seconds`, and the watchdog's `abort.json` names the phase that was running (`build`, `encode`, `set_load`, `denoise` or `vae`). A checkpoint or base taken as a Hub id (the defaults included) prints one line saying it is resolved on the Hub, from the local cache or by download. On the Z-Image base @@ -115,8 +171,8 @@ All notable changes to this project are documented here. The format follows - When a call drops the compressed transformer before the VAE decode, the fit estimate for that step now counts the transformer's uncompressed BF16 tensors, which stay loaded. The estimate grows by 6 MB to 137 MB depending on the model; at 1024² none of the 16, 24 or 32 GB verdicts changes. -- The package summary and keywords name the FLUX.1 use case; `mlx-lm` is no longer a keyword, since no mlx-lm path - exists yet. +- The package summary names the image models mlx-dfloat runs, and `mlx-lm` is no longer a keyword, since no mlx-lm + path exists yet. - The `mlx-dfloat generate --help` text no longer prints reST markup, explains the `--eval-policy` choices, and says "the host's own limits" where it said "the host caps". @@ -158,8 +214,8 @@ The first release. (`scripts/bench_decode_kernel.py`, median of 5 repetitions after a warm-up); staged is the default, and the direct path is the fallback for blocks larger than the staging buffer. `verify_checkpoint.py --decoder metal` checks a whole checkpoint through the kernel, against the BF16 original when one is given and against the reference - otherwise: Qwen3-4B decodes bit-identically in 12.8 s where the reference took 223 s, and all 57 groups of - FLUX.1-schnell match the reference. + otherwise: Qwen3-4B decodes bit-identically against its BF16 original, and all 57 groups of FLUX.1-schnell match the + reference. - A measurement rig for the question the project hinges on: how much a FLUX.1 denoise step slows down when every transformer block's weights are decoded just in time. `scripts/bench_flux_step.py` runs one mode per process, either decoding per block or a control that swaps pre-decoded weights into the same graph, and reports the paired overhead; @@ -182,7 +238,9 @@ The first release. isolated prediction. Decoding the next block on a second GPU stream (`df11-prefetch`, submitted after the previous block's evaluation returns) hides about a fifth of the kernel's time and, at that submission point, none of the allocation; since the cache limit removes that cost anyway, the integration will evaluate per block and size the - cache limit instead. A timed process peaks at about 19 GiB of memory, 20.3 GiB with the look-ahead. + cache limit instead. A timed process peaks at about 19 GiB of memory, 20.3 GiB with the look-ahead. (Note, + 2026-10-08: the overhead figures in this entry were measured at a 1.4 GB cache limit and are superseded by the + committed overhead block in the README.) - The scripts' memory watchdog now enforces its ceiling on the process footprint the OS reports rather than on RSS plus MLX memory, which counted loaded arrays twice. - A block-boundary integration layer, `mlx_dfloat.integrate`: zero-size placeholders for a module's matrices, weight @@ -280,5 +338,6 @@ The first release. - CI's mflux lane gets the Hub token for its test step only, so the three tokenizer tests against the gated FLUX.1 bases run there instead of skipping. -[Unreleased]: https://github.com/IonDen/mlx-dfloat/compare/v0.1.0...HEAD +[Unreleased]: https://github.com/IonDen/mlx-dfloat/compare/v0.2.0...HEAD +[0.2.0]: https://github.com/IonDen/mlx-dfloat/compare/v0.1.0...v0.2.0 [0.1.0]: https://github.com/IonDen/mlx-dfloat/releases/tag/v0.1.0 diff --git a/README.md b/README.md index fb471be..fd92ff2 100644 --- a/README.md +++ b/README.md @@ -6,16 +6,15 @@ Run [DFloat11](https://github.com/LeanModels/DFloat11) checkpoints on a Mac with [MLX](https://github.com/ml-explore/mlx). The weights stay about 30 % smaller than BF16 in memory, and the GPU decodes -each block back to exactly the same bits just before it runs. Today that means FLUX.1 image generation on a 32 GB Mac; -Z-Image, FLUX.2 Klein, Qwen-Image-2.1 and ERNIE-Image arrive in the next release. +each block back to exactly the same bits just before it runs. It generates images on a 32 GB Mac with FLUX.1, Z-Image, +FLUX.2 Klein, Qwen-Image-2.1, ERNIE-Image and Krea 2. DFloat11 is lossless compression for BF16 model weights. Each weight is 16 bits. DFloat11 stores the 8 exponent bits as a short variable-length code, the same idea a zip file uses, and keeps the other 8 bits (the sign and the fraction) exactly as they are. The DFloat11 authors report models at about 70% of their BF16 size with output that is bit for bit the same as the original.[^size] Their decoder runs on NVIDIA GPUs only. This project is an independent reader and decoder for Apple Silicon. The weights stay compressed in memory, and a Metal kernel decodes each block -on the GPU as it runs. That is how it generates FLUX.1 images through mflux; Z-Image, FLUX.2 Klein, Qwen-Image-2.1 and -ERNIE-Image follow in the next release. +on the GPU as it runs. ``` one BF16 weight, 16 bits: s eeeeeeee mmmmmmm @@ -30,10 +29,11 @@ DFloat11 FLUX.1 transformer stays compressed in a Mac's memory, and the GPU unpa before it runs." width="720">

-Why that matters on a Mac: unified memory is the limit. The BF16 FLUX.1-dev transformer is about 24 GB,[^flux] more +On a Mac, unified memory is the limit. The BF16 FLUX.1-dev transformer is about 24 GB,[^flux] more than the GPU on a 32 GB machine can comfortably hold. At 70% it comes in under that line, and unlike 4-bit or 8-bit quantisation it changes nothing in the weights. On a 32 GB M1 Max a 1024² generation peaked at about 20 GiB for each -of the three FLUX.1 models; no other Mac has been measured. The numbers are under "Measured numbers" below. +of the three FLUX.1 models, and between 11.88 GiB (FLUX.2-klein-4B) and 21.17 GiB (Krea 2 Turbo) for the others; no +other Mac has been measured. The numbers are under "Measured numbers" below. ## Does this help me? @@ -41,24 +41,26 @@ of the three FLUX.1 models; no other Mac has been measured. The numbers are unde |---|---| | A 32 GB Mac, FLUX.1-schnell, FLUX.1-dev or FLUX.1-Krea-dev, and you want the BF16 weights unchanged | mlx-dfloat. Measured peak about 20 GiB at 1024² for all three. The final latents were checked bit for bit against BF16 on FLUX.1-schnell, not yet on the other two. Unpacking makes a denoising step 3.5 % slower on dev and 4.0 % slower on schnell than a control that runs the same graph on weights decoded in advance (Krea-dev's step was not timed). | | Less memory or a faster step, and a different image is fine | mflux's own quantised transformer. Its 8-bit mode, the only one measured here, ran a step roughly 1.3 times faster and peaked about 5 GiB lower, under the same 2.5 GB cache limit; mflux's own generate sets no cache limit, so this is not quite its default setup (see "Measured numbers" for why the ratio is rough). It changes the weights and the output; by how much was not measured. | -| A 32 GB Mac, Z-Image-Turbo or Z-Image, with the weights unchanged | mlx-dfloat, from the next release. Peak at 1024²: 13.1 to 14.3 GiB for Z-Image-Turbo over six runs, 13.15 GiB for Z-Image. See "Generate a Z-Image image". | -| A 32 GB Mac, FLUX.2 Klein (4B or 9B, base or distilled), with the weights unchanged | mlx-dfloat, from the next release. Peak at 1024²: 11.88 to 12.78 GiB for the 4B models, 17.60 to 18.40 GiB for the 9B models. See "Generate a FLUX.2 Klein image". | +| A 32 GB Mac, Z-Image-Turbo or Z-Image, with the weights unchanged | mlx-dfloat. Peak at 1024²: 13.1 to 14.3 GiB for Z-Image-Turbo over six runs, 13.15 GiB for Z-Image. See "Generate a Z-Image image". | +| A 32 GB Mac, FLUX.2 Klein (4B or 9B, base or distilled), with the weights unchanged | mlx-dfloat. Peak at 1024²: 11.88 to 12.78 GiB for the 4B models, 17.60 to 18.40 GiB for the 9B models. See "Generate a FLUX.2 Klein image". | | A 24 GB Mac, FLUX.2 Klein | CAPPED: FLUX.2-klein-base-4B passed, peak 12.29 GiB; FLUX.2-klein-base-9B is at risk (the fit check refuses it, and a forced run was stopped by the watchdog). See "Generate a FLUX.2 Klein image". | | A 16 GB Mac, FLUX.2-klein-base-4B | CAPPED: passed, peak 8.15 GiB. See "Generate a FLUX.2 Klein image". | | A 16 GB Mac, Z-Image-Turbo | CAPPED: at risk. The fit check refuses it, and a forced run was stopped by the watchdog at 9.26 GiB before the transformer ran. See "Generate a Z-Image image". | -| A 32 GB Mac, Qwen-Image-2.1, with the weights unchanged | mlx-dfloat, from the next release. Peak at 1024²: 20.42 GiB. See "Generate a Qwen-Image-2.1 image". | +| A 32 GB Mac, Qwen-Image-2.1, with the weights unchanged | mlx-dfloat. Peak at 1024²: 20.42 GiB. See "Generate a Qwen-Image-2.1 image". | | A 16 GB or 24 GB Mac, Qwen-Image-2.1 | CAPPED: at risk. The fit check refuses it, and forced runs were stopped by the watchdog while the prompt was being encoded, before the transformer ran. See "Generate a Qwen-Image-2.1 image". | -| A 32 GB Mac, ERNIE-Image or ERNIE-Image-Turbo, with the weights unchanged | mlx-dfloat, from the next release. Peak at 1024²: 16.79 GiB for ERNIE-Image (50 steps), 17.34 GiB for ERNIE-Image-Turbo (8 steps). See "Generate an ERNIE-Image image". | +| A 32 GB Mac, ERNIE-Image or ERNIE-Image-Turbo, with the weights unchanged | mlx-dfloat. Peak at 1024²: 16.79 GiB for ERNIE-Image (50 steps), 17.34 GiB for ERNIE-Image-Turbo (8 steps). See "Generate an ERNIE-Image image". | | A 24 GB Mac, ERNIE-Image-Turbo | CAPPED: at risk. The fit check refuses it, 0.39 GiB over the budget; a forced run passed with a 14.39 GiB peak, 0.11 GiB under the watchdog ceiling. See "Generate an ERNIE-Image image". | | A 24 GB Mac, ERNIE-Image | CAPPED: at risk. The fit check refuses it, and a forced run was stopped by the watchdog. See "Generate an ERNIE-Image image". | | A 16 GB Mac, ERNIE-Image or ERNIE-Image-Turbo | CAPPED: at risk. The fit check refuses both, and forced runs were stopped by the watchdog. See "Generate an ERNIE-Image image". | +| A 32 GB Mac, Krea 2 Turbo or Krea 2 Raw, with the weights unchanged | mlx-dfloat. Peak at 1024²: 21.17 GiB for Turbo at 8 steps, 21.15 GiB for Raw at 25 steps and for Raw with its model card's 52 steps at guidance 3.5. The fit estimate predicts a refusal at about 500 prompt tokens; only a 30-token prompt was measured. See "Generate a Krea 2 image". | +| A 16 GB or 24 GB Mac, Krea 2 Turbo or Krea 2 Raw | CAPPED: at risk. The fit check refuses both models, and forced runs were stopped by the watchdog while the compressed weights were loading, before the transformer ran. See "Generate a Krea 2 image". | | A Mac with room for the BF16 transformer next to everything else | Plain mflux in BF16. There is nothing to decode. | | A Mac with less than 32 GB | Not measured. The rows marked CAPPED are runs on the 32 GB Mac under the caps mlx-dfloat installs on a 16 GB or 24 GB Mac; the tables below make no other claim for those Macs. | -| Another model: an LLM, the original Qwen-Image or Qwen-Image-Edit, FLUX.2-dev, Krea-2 | Not wired up yet. The reader and the decoder handle all four published versions of the DFloat11 format, but generation in 0.1.0 covers FLUX.1 only. The next release adds Z-Image-Turbo, Z-Image, the four FLUX.2 Klein models, Qwen-Image-2.1, ERNIE-Image and ERNIE-Image-Turbo. | +| Another model: an LLM, the 20B Qwen-Image models (the original Qwen-Image, Qwen-Image-2512, Qwen-Image-Edit) or FLUX.2-dev | Not wired up yet. The reader and the decoder handle all four published versions of the DFloat11 format, but generation covers only the models in the rows above. The 20B Qwen-Image models are planned for a later release, for Macs with more than 32 GB; none of them has been run for generation. | ## Status -Pre-alpha: version 0.1.0 is on PyPI. +Pre-alpha: version 0.2.0 is on PyPI. Released in 0.1.0: @@ -67,7 +69,9 @@ Released in 0.1.0: Qwen-Image-Edit-2509. - A Metal kernel decodes to the same bits as the CPU reference decoder, at about 50 GB/s on an M1 Max on its default path, timed on Qwen3-4B and FLUX.1-schnell groups with `uv run python -m scripts.bench_decode_kernel --df11 - --groups --out decode.json`. + --groups --out decode.json`. The records, `bench/results/decode/qwen3-4b-layer0.json` + and `bench/results/decode/flux1-schnell.json`, give 49.9 to 54.2 GB/s, each the median of five timed repetitions + (M1 Max, macOS 27.0, mlx 0.32.2). - FLUX.1-schnell, FLUX.1-dev and FLUX.1-Krea-dev make 1024² images through mflux, from the command line or from Python. For FLUX.1-schnell, the final latents match the same transformer run block by block from its BF16 weights bit for bit, and the saved images match pixel for pixel. @@ -76,7 +80,7 @@ Released in 0.1.0: - `mlx-dfloat generate`, the parity scripts and the benchmark run under a memory watchdog that stops them when they cross its ceiling. -In the next release (not yet on PyPI): +Released in 0.2.0: - Z-Image-Turbo and Z-Image make 1024² images through mflux, from the command line or from Python. For Z-Image, every one of the 271 compressed matrices and 250 other tensors equals the BF16 original @@ -122,6 +126,20 @@ In the next release (not yet on PyPI): (`bench/results/parity/ernie-image-turbo/full-vs-bf16.json`), and the GPU decoder matches the CPU reference on all 256 (`bench/results/parity/ernie-image-turbo/kernel-vs-reference.json`). Turbo's latents were not compared with a reference: that check needs its 16 GB BF16 transformer on disk. +- Krea 2 Turbo and Krea 2 Raw make 1024² images through mflux, from the command line or from Python. Their DFloat11 + checkpoints (`mingyi456/Krea-2-Turbo-DF11-ComfyUI` and `mingyi456/Krea-2-Raw-DF11-ComfyUI`) are single files made for + ComfyUI, with no `config.json`; mlx-dfloat carries the layout of each, pinned to one revision on the Hub. Both were + checked in full against their originals in `krea/Krea-2-Turbo` and `krea/Krea-2-Raw`, read from the Hub one piece at + a time. The 256 matrices of the 28 transformer blocks and the four text-fusion blocks equal the BF16 originals. The + originals store the other 5 matrices and the 169 other tensors in FP32, and the DFloat11 values equal them rounded to + BF16, nearest even (`bench/results/parity/krea-2/full-vs-original.json` and + `bench/results/parity/krea-2-raw/full-vs-original.json`). The GPU decoder gives the same bits as the CPU reference on + all 261 matrices of each model (`bench/results/parity/krea-2/kernel-vs-reference.json` and + `bench/results/parity/krea-2-raw/kernel-vs-reference.json`). Krea 2 Raw's final latents match, bit for bit, those of + the same transformer streaming its original weights one block at a time, and the images match pixel for pixel. That + check ran at 1024² for 4 steps at guidance 3.5, so classifier-free guidance ran (record: + `bench/results/identity/krea-2-raw-1024/compare.json`). Turbo's latents were not compared with a reference: that + check needs its 26 GB original transformer on disk. - `mlx-dfloat generate --tier 16` and `--tier 24` run under the memory caps mlx-dfloat installs on a Mac of that size. The 16 GB and 24 GB rows under "Measured numbers" are CAPPED runs on the 32 GB Mac under those caps; no smaller Mac was measured. @@ -138,7 +156,8 @@ pip install "mlx-dfloat[mflux]" ``` Without the extra you get the checkpoint reader and the decoder. The `mflux` extra installs mflux 0.20, which the -`mlx-dfloat generate` command needs for FLUX.1, Z-Image, FLUX.2 Klein, Qwen-Image-2.1 and ERNIE-Image generation. +`mlx-dfloat generate` command needs for FLUX.1, Z-Image, FLUX.2 Klein, Qwen-Image-2.1, ERNIE-Image and Krea 2 +generation. The parity scripts and the benchmark live in the repository, not in the package, so run them from a checkout with [uv](https://docs.astral.sh/uv/): @@ -153,7 +172,8 @@ A plain `uv sync` installs the reader, the decoder and the parity scripts, and ` benchmark runs with `uv run --group bench`, which installs mflux as well. Every FLUX.1 base repository on the Hugging Face Hub is gated, so log in once with `hf auth login` and accept the -model's licence on its Hub page (details under "Generate a FLUX.1 image"). On disk, a DFloat11 FLUX.1 transformer +model's licence on its Hub page (details under "Generate a FLUX.1 image"). The Krea 2 bases are gated too (see +"Generate a Krea 2 image"). On disk, a DFloat11 FLUX.1 transformer takes about 16 GB, and the base's text encoders, VAE and tokenizers about 10 GB more. The benchmark's `q8` condition also needs the base's BF16 transformer, about 24 GB, which generation itself never downloads. @@ -190,10 +210,11 @@ result is compared with bits known in advance. It exits 0 when every check passe GPU decoder cannot run at all, for example because there is no Metal device. `--json` prints the same report as JSON. The same check runs automatically the first time a process decodes on the GPU, once for each write path, and when a -FLUX.1, Z-Image, FLUX.2 Klein, Qwen-Image-2.1 or ERNIE-Image model or the `generate` command sets up its decoder, so a -broken GPU decoder refuses before any model loads. On an M1 Max the check itself took about 13 ms, measured after the -kernel pipelines were compiled; the one-time compile took about 0.4 s with a cold shader cache and 0.03 s with a warm -one. It does not check whether a real model fits in your memory or how fast it runs; the numbers below cover that. +FLUX.1, Z-Image, FLUX.2 Klein, Qwen-Image-2.1, ERNIE-Image or Krea 2 model or the `generate` command sets up its +decoder, so a broken GPU decoder refuses before any model loads. In one measurement on an M1 Max, not committed, the +check itself took about 13 ms after the kernel pipelines were compiled, and the one-time compile took about 0.4 s with +a cold shader cache and 0.03 s with a warm one. It does not check whether a real model fits in your memory or how fast it runs; the numbers +below cover that. ### Generate a FLUX.1 image @@ -239,24 +260,29 @@ FLUX.1-dev and FLUX.1-Krea-dev want more steps; the measured runs below used 20 returns; a wired limit your process already set is left alone. Only the command adds a memory watchdog, so prefer it on a Mac near its memory limit. -Measured on an M1 Max (32 GB, macOS 27.0, mlx 0.32.2, mflux 0.20.0, 2026-09-28/29) at 1024², one image per model: -FLUX.1-schnell (4 steps) peaked at 19.94 GiB and took 1 minute 54 seconds including imports; FLUX.1-dev (20 steps, -guidance 3.5) peaked at 20.10 GiB in 7 minutes 6 seconds; FLUX.1-Krea-dev (20 steps, guidance 3.5) peaked at 20.09 -GiB in 7 minutes 2 seconds. The machine's budget for this, its recommended working set minus a 2 GiB reserve, is -22.96 GiB, so all three stay under it. `scripts/verify_image.py` checked FLUX.1-schnell's output against the same -transformer streaming its BF16 shards block by block instead of the compressed set: the final latents matched bit -for bit and the two images matched pixel for pixel. FLUX.1-dev and FLUX.1-Krea-dev were not checked this way, -because their BF16 transformers are gated and were not on disk to compare against. A later run of each model, the -one recorded under "Measured numbers", peaked within 0.2 GiB of these figures. - -Generating a second image in the same process pays a reload. On a 32 GB Mac every call drops the compressed transformer -before the VAE decode, whatever the image size: at 1024² holding it next to the decode measured 23.29 GiB, over the -budget, and smaller sizes have not been measured, so the estimate assumes the decode needs at least as much there. The -next call reloads it, about 26 seconds. A Mac with a larger budget keeps the transformer loaded between calls. Even a -repeated prompt pays the reload: three consecutive 4-step FLUX.1-schnell calls with the same prompt each took about 103 -seconds and peaked between 19.77 and 19.98 GiB. A different prompt adds a further reload, this time of the text encoders -instead of the transformer, since the two are never resident together. `model.encode(*prompts)` pays that encoder reload -once for several prompts, by encoding all of them while the compressed transformer is not yet loaded. +Measured on an M1 Max at 1024², one image per model, with the commands under "Measured numbers": FLUX.1-schnell (4 +steps) peaked at 19.78 GiB, FLUX.1-dev (20 steps, guidance 3.5) at 19.96 GiB and FLUX.1-Krea-dev (20 steps, guidance +3.5) at 20.04 GiB (`bench/results/tiers/schnell-1024.json`, `dev-1024.json` and `krea-dev-1024.json`). The machine's +budget, its recommended working set minus a 2 GiB reserve, is 22.96 GiB, so all three stay under it. Those records +hold no wall time. Earlier single runs on 2026-09-28 and 29, not committed, took 1 minute 54 seconds for +FLUX.1-schnell, imports included, and about 7 minutes each for FLUX.1-dev and FLUX.1-Krea-dev. + +`scripts/verify_image.py` checked FLUX.1-schnell's output against the same transformer streaming its BF16 shards block +by block instead of the compressed set. After 4 steps at 1024², the final latents matched bit for bit and the two +images matched pixel for pixel (record: `bench/results/identity/flux1-schnell-1024/compare.json`). FLUX.1-dev and +FLUX.1-Krea-dev were not checked this way, because their BF16 transformers are gated and were not on disk to compare +against. + +Generating a second image in the same process pays a reload. On a 32 GB Mac every call drops the compressed +transformer before the VAE decode, whatever the image size. At 1024², holding it next to the decode measured 23.29 GiB, +over the budget (a run on 2026-09-28, not committed). Smaller sizes have not been measured, so the estimate assumes the +decode needs at least as much there. The next call loads the transformer again, which took 27 to 28 seconds in the +three committed runs (`lifecycle.set_seconds`). A Mac with a larger budget keeps the transformer loaded between calls. + +Even a repeated prompt pays the reload. In a test on 2026-09-28, not committed, three consecutive 4-step FLUX.1-schnell +calls with the same prompt each took about 103 seconds and peaked between 19.77 and 19.98 GiB. A different prompt also +reloads the text encoders, because they and the transformer are never resident together. `model.encode(*prompts)` pays +that encoder reload once for several prompts: it encodes all of them before the compressed transformer loads. Sizes above 1024² are refused for now, because no run above 1024² has been measured on this path. `--no-fit-check` (`fit_check=False` in Python) runs them anyway, on a memory estimate that is then an extrapolation. @@ -323,11 +349,11 @@ Z-Image-Turbo has no BF16 original to compare with. On a 16 GB Mac, Z-Image-Turbo is at risk. No 16 GB Mac was measured: the run behind that row is CAPPED, made on the 32 GB M1 Max under the caps mlx-dfloat installs on a 16 GB Mac and the tier's watchdog ceiling (`--tier 16`). Without `--no-fit-check` the fit check refuses the run: it predicts a 12.2 GiB peak against a 9.17 GiB budget -(`bench/results/refusals/z-image-turbo-1024-tier16.json`). That run used the older, looser budget for a 16 GB Mac; -`generate` now uses 8.67 GiB there, which refuses it too. With the flag, the watchdog stopped the run after 4.7 s, at -the start of loading the transformer weights (the abort file names the phase), with a footprint of 9.26 GiB against the -9.17 GiB ceiling, while MLX held almost nothing (9 MB). The transformer never ran, so this run says nothing about -whether it would fit. +(`bench/results/refusals/z-image-turbo-1024-tier16.json`). That record was made with a looser 16 GB budget; the budget +`generate` applies on a 16 GB Mac is 8.67 GiB, which refuses the run too. With the flag, the watchdog stopped the run +after 4.7 s, at the start of loading the transformer weights (the abort file names the phase), with a footprint of 9.26 +GiB against the 9.17 GiB ceiling, while MLX held almost nothing (9 MB). The transformer never ran, so this run says +nothing about whether it would fit. Quantisation, LoRA, img2img, mflux's alternative image decoder (`--pid-decode`) and ControlNet are refused, as for FLUX.1. Sizes above 1024² are refused unless you pass `--no-fit-check`, because none has been measured. @@ -404,9 +430,9 @@ installs on those Macs. FLUX.2-klein-base-4B (50 steps, guidance 4.0) passed und 16 GB Mac's caps, where it dropped the compressed transformer before the VAE decode to stay under the budget, and at 12.29 GiB under a 24 GB Mac's. FLUX.2-klein-base-9B is at risk on a 24 GB Mac. The fit check refuses it, predicting 16.3 GiB while denoising against a 14.50 GiB budget (`bench/results/refusals/flux2-klein-base-9b-1024-tier24.json`). -That run used the older, looser budget for a 24 GB Mac; `generate` now uses 14.00 GiB there, which refuses it too. A -forced run (`--no-fit-check`) was stopped by the watchdog while denoising, at 14.64 GiB against the 14.50 GiB -ceiling. Under a 16 GB Mac's caps only FLUX.2-klein-base-4B was run, and under a 24 GB Mac's only the two base models. +That record was made with a looser 24 GB budget; the budget `generate` applies on a 24 GB Mac is 14.00 GiB, which +refuses the run too. A forced run (`--no-fit-check`) was stopped by the watchdog while denoising, at 14.64 GiB against +the 14.50 GiB ceiling. Under a 16 GB Mac's caps only FLUX.2-klein-base-4B was run, and under a 24 GB Mac's only the two base models. Quantisation, LoRA, img2img, mflux's alternative image decoder (`--pid-decode`) and ControlNet are refused, as for FLUX.1. FLUX.2 Klein's image-editing and KV-cache variants are not on this path. `--negative-prompt` is ignored with @@ -549,8 +575,9 @@ Under a 24 GB Mac's caps (`--tier 24`, run on the 32 GB M1 Max), the fit check r GiB peak in 2.8 minutes, the compressed transformer dropped before the VAE decode (`bench/results/tiers/ernie-image-turbo-1024-tier24.json`). That peak is only 0.11 GiB under the 14.50 GiB watchdog ceiling. The run's report still shows the earlier 14.03 GiB estimate in its `fit` block, because the estimate was raised -from this run's peak. Two identical runs under these caps peaked 0.35 GiB apart, at 14.03 and 14.39 GiB, which is why -the estimate takes the higher one. ERNIE-Image-Turbo on a 24 GB Mac is at risk. No 24 GB Mac was measured. +from this run's peak. Two identical runs under these caps peaked 0.35 GiB apart, at 14.03 and 14.39 GiB +(`bench/results/repeats/ernie-image-turbo-1024-tier24/run-1.json` and the tier file above), which is why the estimate +takes the higher one. ERNIE-Image-Turbo on a 24 GB Mac is at risk. No 24 GB Mac was measured. ERNIE-Image under a 24 GB Mac's caps, and both models under a 16 GB Mac's, are at risk too. The fit check refuses them: it predicts 15.2 GiB for ERNIE-Image's denoising step and 14.39 GiB for Turbo's, against budgets of 14.00 GiB at 24 GB @@ -562,6 +589,119 @@ after 23.3 s at 9.44 GiB. Quantisation, LoRA, img2img, mflux's alternative image decoder (`--pid-decode`) and ControlNet are refused, as for FLUX.1. Sizes above 1024² are refused unless you pass `--no-fit-check`, because none has been measured. +### Generate a Krea 2 image + +Krea 2 Turbo and Krea 2 Raw run through the same command. Turbo is the post-trained, distilled release for image +generation; Raw is the pre-training base. Raw's model card calls it "not recommended for inference use" and a good base +for fine-tuning ([Krea 2 Raw](https://huggingface.co/krea/Krea-2-Raw)). + +The command takes each model's text encoder, VAE and tokenizer from its BF16 repository, `krea/Krea-2-Turbo` or +`krea/Krea-2-Raw`, and both are gated on the Hub. Log in once (`hf auth login`) and accept the Krea 2 Community License +([licence](https://huggingface.co/krea/Krea-2-Turbo/blob/main/LICENSE.pdf)) on the Hub page of each model you run; +the Hub grants access as soon as you accept. The DFloat11 checkpoints (`mingyi456/Krea-2-Turbo-DF11-ComfyUI` and +`mingyi456/Krea-2-Raw-DF11-ComfyUI`) are not gated and list the same licence. + +``` +mlx-dfloat generate --model krea-2 \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 --report report.json +mlx-dfloat generate --model krea-2-raw \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 --report report.json +``` + +The defaults are 8 steps for Turbo and 25 for Raw, both at guidance 1.0 with the `er_sde` scheduler. Turbo's are those +of mflux's own Krea 2 command. That command does not run Raw, so Raw's 25 steps are mflux's fallback for a model +without a step default of its own.[^krea-defaults] At guidance 1.0 each step calls the transformer once. Raw's model card runs 52 steps at guidance 3.5 in every example. +To use that recipe, pass `--steps 52 --guidance 3.5` and raise the time limit. A single run took 56.2 minutes here, +close to the one hour after which the command's watchdog stops a run unless you pass `--wall-budget`: + +``` +mlx-dfloat generate --model krea-2-raw --steps 52 --guidance 3.5 --wall-budget 7200 \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 --report report.json +``` + +Any guidance other than 1.0, below 1.0 too, turns on classifier-free guidance, as in mflux. Each step then runs the +transformer twice, once on the prompt and once on the negative prompt, and the command prints a line saying so. Without +`--negative-prompt` the second call uses a single space, as mflux does. At guidance 1.0 a `--negative-prompt` has no +effect, and the command warns. `--scheduler` takes `er_sde` (the default), `euler` or `linear`, which mflux runs as +`er_sde`; any other name is refused. An empty or blank `--prompt` is refused as a user error: Krea 2's tokenizer gives an +empty prompt no tokens at all. + +On a 32 GB Mac the fit estimate predicts a refusal for a prompt of about 500 tokens or more. Only the 30-token +lighthouse prompt above was measured. The estimate grows with the prompt's length, and at 1024² it reaches the 22.96 +GiB budget near 500 tokens.[^krea-tokens] The command refuses such a prompt before the text encoder or the transformer +loads. `--no-fit-check` runs it anyway. + +The command takes the DFloat11 transformer from that checkpoint, one 17.5 GB file for either model, and the text +encoder, VAE and tokenizer from the model's BF16 repository. It does not download that repository's own transformer +(26.3 GB). The text-encoder file is 8.9 GB and also holds a vision model that is never loaded; the VAE is 0.5 GB. All +these sizes are as listed on the Hub. Both BF16 repositories publish the same text-encoder, VAE and tokenizer files +(the Hub lists the same hashes for them). Each DFloat11 checkpoint is a single file made for ComfyUI, with no +`config.json`. mlx-dfloat carries the layout of each, pinned to one revision, and refuses to build one Krea 2 model +from the other's file, naming the `--model` it belongs to. The default checkpoints and BF16 repositories are pinned to +the revisions the measured runs used; a `--df11` or `--base` you pass is used as given. + +The same from Python: + +```python +from mlx_dfloat.mflux import DFloatModel + +model = DFloatModel("krea-2") # or "krea-2-raw" +image = model.generate_image( + seed=42, + prompt="A stone lighthouse on a rocky shore at dawn", + height=1024, + width=1024, +) +image.save("lighthouse.png") +``` + +Left out, `num_inference_steps`, `guidance` and `scheduler` take the same per-model defaults as the command; mflux's own +`Krea2` class defaults to 8 steps for both models. For Raw's card recipe pass `num_inference_steps=52, guidance=3.5`. +An empty or blank prompt is refused here too. Each call runs under the memory caps the command installs, unless the +process has already set a wired limit. + +mflux compiles Krea 2's denoising step on M1 and M2 Max and Ultra chips and on M3 and later. This path runs it without +that compilation on every chip (see "Generate a Z-Image image" for why). mflux also passes Krea 2's hidden state from +block to block in float32, not BF16, because the latents it starts from are float32; this path runs mflux's own block +code, so it does the same. No Krea 2 step time is reported, against stock mflux or otherwise. + +Besides its 28 transformer blocks, Krea 2 has seven smaller compressed groups: four blocks that merge the text +encoder's layers into the prompt input, and three small layers for the timestep and the text.[^krea-groups] Decoded, +they take 1.23 GiB. mlx-dfloat decodes them at the start of every transformer call and frees them before the first +block runs. It also empties MLX's buffer cache at the start of each call and again before the first block. On one Raw +step at guidance 3.5, the denoising peak was 21.14 GiB this way, against 22.56 GiB with the seven groups kept decoded +for the whole run (two single-step runs in `bench/results/calibration/krea-2-raw-1024.json`, one per setup, made at +different development commits). Most of that 1.4 GiB comes from emptying the cache; decoding per call saves about +0.14 GiB on its own.[^krea-nonblock] The price is about 1.3 GB more decoded weights resident at the start of each +call, decoded again every call; that decode time was not measured on its own. + +Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0) at 1024², seed 42, with the commands under +"Measured numbers", one run each: Krea 2 Turbo (8 steps at guidance 1.0) peaked at 21.17 GiB and took 5.1 minutes, +and Krea 2 Raw (25 steps at guidance 1.0) peaked at 21.15 GiB in 14.1 minutes. With the card's recipe (52 steps at +guidance 3.5) Raw peaked at 21.15 GiB in 56.2 minutes +(`bench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.json`; it is not a row of the table below, which shows no +step count). In all three runs the peak came while denoising. + +The fit estimate is mlx-dfloat's prediction of each step's peak memory, checked before a call. For this prompt it puts +the denoising step at 22.58 GiB, under this Mac's 22.96 GiB budget. That is about 1.4 GiB above the measured peaks, +because the estimate counts a full buffer cache that the emptied cache never reaches. The estimate implies that +holding the compressed transformer through the VAE decode would go over the budget. So on a 32 GB Mac every call +drops it before the decode and loads it again at the next call (31 s in the Turbo run). + +On a 16 GB or 24 GB Mac, both models are at risk. No such Mac was measured: the runs behind those rows are CAPPED, made +on the 32 GB M1 Max under the caps mlx-dfloat installs on a Mac of that size (`--tier 16` and `--tier 24`). The fit +check refuses both models at both sizes. It predicts 22.6 GiB for the denoising step, against a budget of 8.67 GiB at +16 GB and 14.00 GiB at 24 GB; the compressed transformer alone is 16.34 GiB. Forced runs (`--no-fit-check`) were +stopped by the watchdog while the compressed transformer was loading. At 16 GB, against the 9.17 GiB ceiling, Turbo +was stopped after 17.1 s at 9.50 GiB and Raw after 17.7 s at 9.49 GiB. At 24 GB, against the 14.50 GiB ceiling, Turbo +was stopped after 29.7 s at 15.05 GiB and Raw after 32.0 s at 14.91 GiB. + +Quantisation, LoRA, img2img, mflux's alternative image decoder (`--pid-decode`) and ControlNet are refused, as for +FLUX.1. Sizes above 1024² are refused unless you pass `--no-fit-check`, because none has been measured. + ## Measured numbers The blocks below are generated from the files under `bench/results/` by `scripts/bench_table.py`, and a test fails when @@ -575,23 +715,24 @@ ceiling, made only to show that the watchdog works; it never appears as a tier r The 32 GB rows are MEASURED, on one M1 Max. The 16 GB and 24 GB rows are CAPPED, run on that same Mac: FLUX.2-klein-base-4B passed under both and ERNIE-Image-Turbo under a 24 GB Mac's caps (a forced run: the fit check refuses it), while FLUX.2-klein-base-9B under a 24 GB Mac's caps, Z-Image-Turbo under a 16 GB Mac's, ERNIE-Image-Turbo -under a 16 GB Mac's, and Qwen-Image-2.1 and ERNIE-Image under both were stopped by the watchdog. The table makes no -claim about any Mac it does not list. Each row is one `mlx-dfloat generate` run at 1024² with seed 42 (4 steps for -FLUX.1-schnell, FLUX.2-klein-4B and FLUX.2-klein-9B, 20 for FLUX.1-dev and FLUX.1-Krea-dev, 9 for Z-Image-Turbo, 40 at -guidance 1.0 for Qwen-Image-2.1, 8 at guidance 1.0 for ERNIE-Image-Turbo, 50 with guidance 4.0 for Z-Image, ERNIE-Image -and the two FLUX.2 Klein base models), with `--tier` and `--report` writing the file in the last column. +under a 16 GB Mac's, and Qwen-Image-2.1, ERNIE-Image, Krea 2 Turbo and Krea 2 Raw under both were stopped by the +watchdog. The table makes no claim about any Mac it does not list. Each row is one `mlx-dfloat generate` run at 1024² +with seed 42 (4 steps for FLUX.1-schnell, FLUX.2-klein-4B and FLUX.2-klein-9B, 20 for FLUX.1-dev and FLUX.1-Krea-dev, 9 +for Z-Image-Turbo, 40 at guidance 1.0 for Qwen-Image-2.1, 8 at guidance 1.0 for ERNIE-Image-Turbo and Krea 2 Turbo, 25 +at guidance 1.0 for Krea 2 Raw, 50 with guidance 4.0 for Z-Image, ERNIE-Image and the two FLUX.2 Klein base models), +with `--tier` and `--report` writing the file in the last column. The second column is the Mac's recommended GPU working set minus a reserve. On the 16 GB and 24 GB rows the working set is taken as two thirds of that Mac's memory. Two limits come off that working set, with different reserves. The fit budget, which `generate` checks before a run, is the working set minus 2 GiB on every Mac: 22.96 GiB on this one, 8.67 GiB for a 16 GB Mac and 14.00 GiB for a 24 GB Mac. The watchdog ceiling, where a CAPPED run is stopped, is the working set minus 1.5 GiB at 16 GB and 24 GB: 9.17 and 14.50 GiB. So the second column is the fit budget on the 32 GB -rows and the watchdog ceiling on the CAPPED rows. The FLUX.2 Klein and Z-Image CAPPED rows were recorded while the -fit check still used the 1.5 GiB reserve, and the stricter budget changes none of their outcomes. FLUX.2-klein-base-4B still passes at 16 GB -(a predicted 8.54 GiB against 8.67), and at 24 GB it still keeps the compressed transformer loaded for the VAE decode -(12.69 GiB predicted against 14.00). FLUX.2-klein-base-9B at 24 GB and Z-Image-Turbo at 16 GB are still refused. -The Qwen-Image-2.1 rows were recorded under the current budget, which refuses both (a predicted 14.6 GiB while it -encodes the prompt). +rows and the watchdog ceiling on the CAPPED rows. The FLUX.2 Klein and Z-Image CAPPED rows were recorded when the fit +check used a 1.5 GiB reserve, and the 2 GiB reserve changes none of their outcomes. FLUX.2-klein-base-4B passes at +16 GB (a predicted 8.54 GiB against 8.67), and at 24 GB it keeps the compressed transformer loaded for the VAE decode +(12.69 GiB predicted against 14.00). FLUX.2-klein-base-9B at 24 GB and Z-Image-Turbo at 16 GB are refused. The +Qwen-Image-2.1 rows were recorded with the 2 GiB reserve, which refuses both (a predicted 14.6 GiB while it encodes +the prompt). "Peak (watched)" is the larger of the process footprint the OS reports and MLX's active plus cached memory, and a row's status is "target" when it stayed under the second column. "Peak MLX" is the larger of two @@ -646,9 +787,8 @@ mlx-dfloat generate --model flux2-klein-base-9b --steps 50 --guidance 4 --tier 2 ``` The last one runs with `--no-fit-check` because without it the fit check refuses the run: it predicts 16.3 GiB while -denoising against the 14.50 GiB budget of that run, and 14.00 GiB today -(`bench/results/refusals/flux2-klein-base-9b-1024-tier24.json`). It ends with -exit 70, like the Z-Image-Turbo run above. +denoising, against the 14.50 GiB budget recorded with that run and the 14.00 GiB budget `generate` applies +(`bench/results/refusals/flux2-klein-base-9b-1024-tier24.json`). It ends with exit 70, like the Z-Image-Turbo run above. The Qwen-Image-2.1 rows come from these commands, with the same prompt: @@ -696,7 +836,41 @@ table's copies are in `bench/results/tiers/aborts/`. Without `--no-fit-check` th exit 2 (`bench/results/refusals/ernie-image-turbo-1024-tier24.json`, `ernie-image-turbo-1024-tier16.json`, `ernie-image-1024-tier16.json` and `ernie-image-1024-tier24.json`). -Measured on an Apple M1 Max, 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0. +The Krea 2 rows come from these commands, with the same prompt: + +``` +mlx-dfloat generate --model krea-2 --tier 32 \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 --report bench/results/tiers/krea-2-1024.json +mlx-dfloat generate --model krea-2-raw --tier 32 \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 --report bench/results/tiers/krea-2-raw-1024.json +mlx-dfloat generate --model krea-2-raw --steps 52 --guidance 3.5 --tier 32 --wall-budget 7200 \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 --report bench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.json +mlx-dfloat generate --model krea-2 --tier 16 --no-fit-check \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 +mlx-dfloat generate --model krea-2-raw --tier 16 --no-fit-check \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 +mlx-dfloat generate --model krea-2 --tier 24 --no-fit-check \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 +mlx-dfloat generate --model krea-2-raw --tier 24 --no-fit-check \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 +``` + +The third command writes outside `bench/results/tiers/`, so it is not a table row: the table shows no step count, and +two Krea 2 Raw rows would look the same. In the last four the watchdog stops the run (exit 70) and writes its +`abort.json` next to the output image; the table's copies are in `bench/results/tiers/aborts/`. Without +`--no-fit-check` the last four are refused with exit 2 (`bench/results/refusals/krea-2-1024-tier16.json`, +`krea-2-raw-1024-tier16.json`, `krea-2-1024-tier24.json` and `krea-2-raw-1024-tier24.json`). + +Measured on an Apple M1 Max, 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0. The `git` fields in the committed result +files name development commits that were folded into a release; they tell the runs apart, but they cannot be fetched +from the public repository. | Mac | Working set − reserve | Model | DF11 size | Peak (watched) | Peak footprint | Peak MLX (sampled active + cache, or exact phase peak) | Label | Status | Limits | Result | @@ -711,6 +885,8 @@ Measured on an Apple M1 Max, 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0. | 24 GB | 14.50 GiB | FLUX.2-klein-base-4B | 4.90 GiB | 12.29 GiB | 12.29 GiB | 10.91 GiB | CAPPED | target | mlx-dfloat caps for the tier | `bench/results/tiers/flux2-klein-base-4b-1024-tier24.json` | | 32 GB | 22.96 GiB | FLUX.2-klein-base-4B | 4.90 GiB | 12.78 GiB | 12.78 GiB | 11.91 GiB | MEASURED | target | host caps | `bench/results/tiers/flux2-klein-base-4b-1024.json` | | 32 GB | 22.96 GiB | FLUX.2-klein-base-9B | 11.47 GiB | 17.60 GiB | 17.60 GiB | 16.58 GiB | MEASURED | target | host caps | `bench/results/tiers/flux2-klein-base-9b-1024.json` | +| 32 GB | 22.96 GiB | Krea 2 Turbo | 16.34 GiB | 21.17 GiB | 21.17 GiB | 20.64 GiB | MEASURED | target | host caps | `bench/results/tiers/krea-2-1024.json` | +| 32 GB | 22.96 GiB | Krea 2 Raw | 16.34 GiB | 21.15 GiB | 21.15 GiB | 20.64 GiB | MEASURED | target | host caps | `bench/results/tiers/krea-2-raw-1024.json` | | 32 GB | 22.96 GiB | FLUX.1-Krea-dev | 15.21 GiB | 20.04 GiB | 20.04 GiB | 19.19 GiB | MEASURED | target | host caps | `bench/results/tiers/krea-dev-1024.json` | | 32 GB | 22.96 GiB | Qwen-Image-2.1 | 9.06 GiB | 20.42 GiB | 20.42 GiB | 19.72 GiB | MEASURED | target | host caps | `bench/results/tiers/qwen-image-2.1-1024.json` | | 32 GB | 22.96 GiB | FLUX.1-schnell | 15.19 GiB | 19.78 GiB | 19.78 GiB | 19.03 GiB | MEASURED | target | host caps | `bench/results/tiers/schnell-1024.json` | @@ -720,6 +896,10 @@ Measured on an Apple M1 Max, 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0. | 24 GB | 14.50 GiB | ERNIE-Image | not recorded | at least 14.80 GiB (stopped after 20.6 s) | at least 14.80 GiB | at least 14.25 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/ernie-image-1024-tier24.json` | | 16 GB | 9.17 GiB | ERNIE-Image-Turbo | not recorded | at least 9.22 GiB (stopped after 18.6 s) | at least 9.22 GiB | at least 8.75 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/ernie-image-turbo-1024-tier16.json` | | 24 GB | 14.50 GiB | FLUX.2-klein-base-9B | not recorded | at least 14.64 GiB (stopped after 28.5 s) | at least 14.64 GiB | at least 14.26 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/flux2-klein-base-9b-1024-tier24.json` | +| 16 GB | 9.17 GiB | Krea 2 Turbo | not recorded | at least 9.50 GiB (stopped after 17.1 s) | at least 9.50 GiB | at least 9.39 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-1024-tier16.json` | +| 24 GB | 14.50 GiB | Krea 2 Turbo | not recorded | at least 15.05 GiB (stopped after 29.7 s) | at least 15.05 GiB | at least 14.95 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-1024-tier24.json` | +| 16 GB | 9.17 GiB | Krea 2 Raw | not recorded | at least 9.49 GiB (stopped after 17.7 s) | at least 9.49 GiB | at least 9.39 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-raw-1024-tier16.json` | +| 24 GB | 14.50 GiB | Krea 2 Raw | not recorded | at least 14.91 GiB (stopped after 32.0 s) | at least 14.91 GiB | at least 14.41 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-raw-1024-tier24.json` | | 16 GB | 9.17 GiB | Qwen-Image-2.1 | not recorded | at least 9.28 GiB (stopped after 4.4 s) | at least 8.90 GiB | at least 9.28 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/qwen-image-2.1-1024-tier16.json` | | 24 GB | 14.50 GiB | Qwen-Image-2.1 | not recorded | at least 14.59 GiB (stopped after 4.8 s) | at least 14.59 GiB | at least 14.31 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/qwen-image-2.1-1024-tier24.json` | | 16 GB | 9.17 GiB | Z-Image-Turbo | not recorded | at least 9.26 GiB (stopped after 4.7 s) | at least 9.26 GiB | at least 8.81 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/z-image-turbo-1024-tier16.json` | @@ -730,7 +910,8 @@ every condition in its own process. `df11` decodes each transformer block's weig the block runs and evaluates after every block. `control` runs the same graph with the same per-block evaluation, but hands every block the weights of one double and one single block, decoded once before timing started, so the gap between the two is the cost of decoding just in time. At the earlier 1.4 GB cache limit most of that gap was allocating a fresh -buffer for each block's decoded weights; at 2.5 GB the share was not measured. The depth-2 pair evaluates one block +buffer for each block's decoded weights (measured on 2026-09-27, not committed); at 2.5 GB the share was not +measured. The depth-2 pair evaluates one block behind instead, so the CPU can queue the next block, decode included, while the GPU runs the current one. The eval policy cost is what evaluating after every block adds by itself: `control` minus a control that evaluates once per step, with no decode involved. A negative value means the per-block control was the faster of the two. The q8 ratio is DF11 @@ -754,7 +935,7 @@ output. The control was validated once with `scripts/bench_control_validation.py`, at the earlier 1.4 GB cache limit and on a reduced-depth transformer (4 double and 8 single blocks instead of 19 and 38): it agreed within 0.13 % with a run -whose BF16 weights were all resident. +whose BF16 weights were all resident (`bench/results/validation/control-vs-bf16-schnell-1024.json`). FLUX.1-dev, 1024², per-block evaluation: +3.5 % (depth-2: +4.6 %); eval policy cost -0.07 s/step; DF11 (per-block) over the mflux q8 step (one eval per step, same cache limit): 1.26× @@ -838,3 +1019,15 @@ encoder copied from the DFloat11 repository, each with its licence. depends on the model and is slightly different for each one. [^flux]: The FLUX.1-dev transformer has about 12 billion parameters. In BF16 each takes 2 bytes, so about 24 GB for the weights alone, before activations and before the text encoders. +[^krea-defaults]: mflux's own Krea 2 command runs Turbo only; it refuses `krea-2-raw`. For Raw, 25 steps is what + mflux's table of step defaults gives a model it has no entry for, and guidance 1.0 is the default of mflux's + Python `generate_image`. +[^krea-tokens]: Tokens are counted after mflux strips Krea 2's chat template from the encoded prompt. The estimate's + cache term grows by about 0.84 MiB per token, and at 30 tokens the estimate is 0.38 GiB under the budget, so the + limit falls near 496 tokens (the constants are in `src/mlx_dfloat/mflux/krea2/memory.py`). The tokenizer itself + stops at 1,024 tokens, chat template included. A guided step is sized by the longer of its two prompts. +[^krea-groups]: The four blocks are `txtfusion.layerwise_blocks.0`, `.1` and `txtfusion.refiner_blocks.0`, `.1`; the + three layers are `tmlp`, `tproj` and `txtmlp`. `tproj` alone is 0.42 GiB decoded. +[^krea-nonblock]: Decoding per call without emptying the cache lowered that peak by only 0.14 GiB, to 22.42 GiB (a + third single-step run in the same file). In a guided step the second call otherwise starts on the buffers the + first call left in MLX's cache. diff --git a/bench/results/calibration/krea-2-1024.json b/bench/results/calibration/krea-2-1024.json new file mode 100644 index 0000000..c76ab56 --- /dev/null +++ b/bench/results/calibration/krea-2-1024.json @@ -0,0 +1,179 @@ +{ + "model": "krea-2", + "purpose": "the 1024² runs the Krea 2 memory constants were calibrated on (mlx_dfloat.mflux.krea2.memory); one set of constants serves both models, the runs are recorded per model. The constants come from the runs of the shipped mechanism only (mechanism 'shipped'); the VAE term also counts the resident run, whose VAE phase ran with the set dropped like every other. The superseded runs measured earlier per-call code and are kept as history, never as a constant's source. The git fields name development commits that were folded into the 0.2.0 release; they are kept to tell the runs apart and cannot be fetched from the public repository.", + "df11": "mingyi456/Krea-2-Turbo-DF11-ComfyUI@978da5fb7647bd222d33125993abd8fdc2840cfc", + "base": "krea/Krea-2-Turbo@98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7", + "calls_per_step": 1, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "size": 1024, + "seed": 42, + "text_tokens": 30, + "device": "Apple M1 Max, 32 GB", + "macos": "27.0.1", + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "footprint_peaks_label": "OS phys_footprint in bytes, the peak of each phase (the de-risk runs: 1024², the set dropped before the VAE decode; a multi-step run's denoise peak is over all its steps); for the MEASURED, card-recipe and identity runs the run's watched peak, reached in the denoise phase", + "sizes_label": "FamilySizes as mlx_dfloat.mflux.krea2.memory.sizes_for computes them in the run's non-block mode: nonblock is 0 when the seven groups are decoded per call, 1_325_400_064 when they stay resident; encoders is the language model's evaluated layers (layer 35 and the final norm never run), not the 8_044_936_192 B of every language-model tensor", + "overhead_from": "derisk-turbo-2", + "constants_from": { + "denoise": "measured-8-steps" + }, + "runs": [ + { + "label": "derisk-turbo-2", + "date": "2026-10-08", + "git": "c555f5031aecd06d9c2410c0ac5253b970254580", + "mechanism": "shipped: the non-block groups decoded per call, MLX's cache emptied and drained at each call's start and again before block 0's decode, after the non-block release", + "run": "the de-risk run of the shipped code: two steps at guidance 1.0 (one batch-1 call each; the second call starts after the first's activations), at the planner's limit for the probe's allowance", + "steps": 2, + "calls_per_step": 1, + "guidance": 1.0, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "encode": 8537103128, + "set_load": 18097157944, + "denoise": 22710482720, + "vae": 10223310128 + }, + "after_build": { + "footprint": 529171680, + "mlx_active": 4805792, + "mlx_cache": 4974 + }, + "mlx_encode_held": 7907524416, + "sizes": { + "compressed": 17540075298, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + }, + "denoise_footprint_after_each_step": [ + 22708303672, + 22710482720 + ] + }, + { + "label": "measured-8-steps", + "date": "2026-10-08", + "git": "f0f348fd2f333304a4cee47bb9cd64db9dd87d3c", + "mechanism": "shipped: the non-block groups decoded per call, MLX's cache emptied and drained at each call's start and again before block 0's decode, after the non-block release", + "run": "the MEASURED row (mlx-dfloat generate --tier 32, the final code): the default call, 8 steps at guidance 1.0 (one call per step), the set dropped before the VAE decode; the watched peak, reached in the denoise phase; bench/results/tiers/krea-2-1024.json", + "steps": 8, + "calls_per_step": 1, + "guidance": 1.0, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "denoise": 22736402568 + }, + "sizes": { + "compressed": 17540075298, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + } + }, + { + "label": "confirm-turbo-24711fb", + "date": "2026-10-08", + "git": "24711fbe86cc8c72e233735c26fb0f03ef11d2d6", + "mechanism": "shipped: the non-block groups decoded per call, MLX's cache emptied and drained at each call's start and again before block 0's decode, after the non-block release", + "run": "the confirmation run of the final code (the drain also waits for the GPU first): two steps at guidance 1.0 (one batch-1 call each), at the planner's limit for the probe's allowance", + "steps": 2, + "calls_per_step": 1, + "guidance": 1.0, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "encode": 8531319552, + "set_load": 18096125800, + "denoise": 22703486824, + "vae": 10222114120 + }, + "after_build": { + "footprint": 528598240, + "mlx_active": 4805792, + "mlx_cache": 4974 + }, + "mlx_encode_held": 7936899688, + "sizes": { + "compressed": 17540075298, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + }, + "denoise_footprint_after_each_step": [ + 22701307752, + 22703486824 + ] + }, + { + "label": "history-v1-turbo", + "date": "2026-10-08", + "git": "81e79f3abd955a7d3aa4a11015f4ec455cdc9b36", + "mechanism": "superseded (history): per-call decode released after block 0, no cache emptying", + "run": "the first per-call run (one step)", + "steps": 1, + "calls_per_step": 1, + "guidance": 1.0, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "encode": 8535087920, + "set_load": 18100025528, + "denoise": 24041142456, + "vae": 10129426472 + }, + "after_build": { + "footprint": 534185232, + "mlx_active": 4805792, + "mlx_cache": 4974 + }, + "mlx_encode_held": 7943016040, + "sizes": { + "compressed": 17540075298, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + } + }, + { + "label": "history-mm1-turbo", + "date": "2026-10-08", + "git": "bedba36cd0734655f48f7ac9b425ef3ae47235c0", + "mechanism": "superseded (history): per-call decode released before block 0 with a cache clear there, none at the call's start", + "run": "the second per-call run (one step: one call, which starts with an empty cache)", + "steps": 1, + "calls_per_step": 1, + "guidance": 1.0, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "encode": 8533121816, + "set_load": 18098091976, + "denoise": 22703093704, + "vae": 11059664992 + }, + "after_build": { + "footprint": 530384120, + "mlx_active": 4805792, + "mlx_cache": 4974 + }, + "mlx_encode_held": 7868153872, + "sizes": { + "compressed": 17540075298, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + } + } + ] +} diff --git a/bench/results/calibration/krea-2-raw-1024.json b/bench/results/calibration/krea-2-raw-1024.json new file mode 100644 index 0000000..dfe0cf9 --- /dev/null +++ b/bench/results/calibration/krea-2-raw-1024.json @@ -0,0 +1,252 @@ +{ + "model": "krea-2-raw", + "purpose": "the 1024² runs the Krea 2 memory constants were calibrated on (mlx_dfloat.mflux.krea2.memory); one set of constants serves both models, the runs are recorded per model. The constants come from the runs of the shipped mechanism only (mechanism 'shipped'); the VAE term also counts the resident run, whose VAE phase ran with the set dropped like every other. The superseded runs measured earlier per-call code and are kept as history, never as a constant's source. The git fields name development commits that were folded into the 0.2.0 release; they are kept to tell the runs apart and cannot be fetched from the public repository.", + "df11": "mingyi456/Krea-2-Raw-DF11-ComfyUI@8320616b25ac9340a830a7fb21f1b0237e160e66", + "base": "krea/Krea-2-Raw@6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "calls_per_step": 2, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "size": 1024, + "seed": 42, + "text_tokens": 30, + "device": "Apple M1 Max, 32 GB", + "macos": "27.0.1", + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "footprint_peaks_label": "OS phys_footprint in bytes, the peak of each phase (the de-risk runs: 1024², the set dropped before the VAE decode; a multi-step run's denoise peak is over all its steps); for the MEASURED, card-recipe and identity runs the run's watched peak, reached in the denoise phase", + "sizes_label": "FamilySizes as mlx_dfloat.mflux.krea2.memory.sizes_for computes them in the run's non-block mode: nonblock is 0 when the seven groups are decoded per call, 1_325_400_064 when they stay resident; encoders is the language model's evaluated layers (layer 35 and the final norm never run), not the 8_044_936_192 B of every language-model tensor", + "constants_from": { + "encode": "confirm-raw-24711fb", + "vae": "derisk-raw-1" + }, + "runs": [ + { + "label": "derisk-raw-2", + "date": "2026-10-08", + "git": "c555f5031aecd06d9c2410c0ac5253b970254580", + "mechanism": "shipped: the non-block groups decoded per call, MLX's cache emptied and drained at each call's start and again before block 0's decode, after the non-block release", + "run": "the de-risk run of the shipped code: one step at guidance 3.5 with mflux's \" \" negative (two batch-1 transformer calls), at the planner's limit for the probe's allowance", + "steps": 1, + "calls_per_step": 2, + "guidance": 3.5, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "encode": 8595823360, + "set_load": 18097109152, + "denoise": 22701701280, + "vae": 10163669128 + }, + "after_build": { + "footprint": 527484128, + "mlx_active": 4805792, + "mlx_cache": 4974 + }, + "mlx_encode_held": 7946299264, + "sizes": { + "compressed": 17540930560, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + }, + "denoise_footprint_after_each_step": [ + 22701701280 + ] + }, + { + "label": "measured-25-steps", + "date": "2026-10-08", + "git": "f0f348fd2f333304a4cee47bb9cd64db9dd87d3c", + "mechanism": "shipped: the non-block groups decoded per call, MLX's cache emptied and drained at each call's start and again before block 0's decode, after the non-block release", + "run": "the MEASURED row (mlx-dfloat generate --tier 32, the final code): the default call, 25 steps at guidance 1.0 (one call per step), the set dropped before the VAE decode; the watched peak, reached in the denoise phase; bench/results/tiers/krea-2-raw-1024.json", + "steps": 25, + "calls_per_step": 1, + "guidance": 1.0, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "denoise": 22713743304 + }, + "sizes": { + "compressed": 17540930560, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + } + }, + { + "label": "card-52-steps-cfg", + "date": "2026-10-08", + "git": "f0f348fd2f333304a4cee47bb9cd64db9dd87d3c", + "mechanism": "shipped: the non-block groups decoded per call, MLX's cache emptied and drained at each call's start and again before block 0's decode, after the non-block release", + "run": "the model card's recipe (mlx-dfloat generate --tier 32, the final code): 52 steps at guidance 3.5 with mflux's \" \" negative (two calls per step), the set dropped before the VAE decode; the watched peak, reached in the denoise phase; bench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.json", + "steps": 52, + "calls_per_step": 2, + "guidance": 3.5, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "denoise": 22714513664 + }, + "sizes": { + "compressed": 17540930560, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + } + }, + { + "label": "identity-4-steps-cfg", + "date": "2026-10-08", + "git": "e3a508e6d37a3946253801258a9200ff03a708a4", + "mechanism": "shipped: the non-block groups decoded per call, MLX's cache emptied and drained at each call's start and again before block 0's decode, after the non-block release", + "run": "the identity check's df11 side: 4 steps at guidance 3.5 with \" \" (two calls per step), the set dropped before the VAE decode; the watched peak, reached in the denoise phase; bench/results/identity/krea-2-raw-1024/df11-result.json", + "steps": 4, + "calls_per_step": 2, + "guidance": 3.5, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "denoise": 22711515080 + }, + "sizes": { + "compressed": 17540930560, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + } + }, + { + "label": "confirm-raw-24711fb", + "date": "2026-10-08", + "git": "24711fbe86cc8c72e233735c26fb0f03ef11d2d6", + "mechanism": "shipped: the non-block groups decoded per call, MLX's cache emptied and drained at each call's start and again before block 0's decode, after the non-block release", + "run": "the confirmation run of the final code (the drain also waits for the GPU first): one step at guidance 3.5 with mflux's \" \" negative (two batch-1 calls), at the planner's limit for the probe's allowance", + "steps": 1, + "calls_per_step": 2, + "guidance": 3.5, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "encode": 8598264624, + "set_load": 18099828728, + "denoise": 22704306168, + "vae": 10525591448 + }, + "after_build": { + "footprint": 528401632, + "mlx_active": 4805792, + "mlx_cache": 4974 + }, + "mlx_encode_held": 7953193856, + "sizes": { + "compressed": 17540930560, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + }, + "denoise_footprint_after_each_step": [ + 22704306168 + ] + }, + { + "label": "derisk-raw-1", + "date": "2026-10-08", + "git": "d1aa9df996cee99356ffa3f2f01fdb0ebd16469c", + "mechanism": "resident: the non-block groups decoded once at set load and kept with the set", + "run": "the de-risk run, rung 1: the same step with the seven non-block groups resident (its denoise peak was over the gate line, so per-call decode followed); a VAE sample (the set is dropped before the decode in every mechanism)", + "steps": 1, + "calls_per_step": 2, + "guidance": 3.5, + "nonblock": "resident", + "cache_limit": 4500000000, + "footprint_peaks": { + "encode": 8599804744, + "set_load": 19426654152, + "denoise": 24218219864, + "vae": 11869477608 + }, + "after_build": { + "footprint": 530826464, + "mlx_active": 4805792, + "mlx_cache": 4974 + }, + "mlx_encode_held": 7942988392, + "sizes": { + "compressed": 17540930560, + "extras": 4560024, + "nonblock": 1325400064, + "encoders": 7843074560, + "vae": 507591892 + } + }, + { + "label": "history-v1-raw", + "date": "2026-10-08", + "git": "59c961bfb5b58cdb95198ef1bf4a5fb548300262", + "mechanism": "superseded (history): per-call decode released after block 0, no cache emptying", + "run": "the first per-call run (the same step)", + "steps": 1, + "calls_per_step": 2, + "guidance": 3.5, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "encode": 8609733448, + "set_load": 18103810064, + "denoise": 24072009744, + "vae": 10396076024 + }, + "after_build": { + "footprint": 532530400, + "mlx_active": 4805792, + "mlx_cache": 4974 + }, + "mlx_encode_held": 7939351776, + "sizes": { + "compressed": 17540930560, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + } + }, + { + "label": "history-mm1-raw", + "date": "2026-10-08", + "git": "bedba36cd0734655f48f7ac9b425ef3ae47235c0", + "mechanism": "superseded (history): per-call decode released before block 0 with a cache clear there, none at the call's start", + "run": "the second per-call run (the same step); its second call started on the first call's cached activations and the denoise peak sat on MLX's GC line", + "steps": 1, + "calls_per_step": 2, + "guidance": 3.5, + "nonblock": "per-call", + "cache_limit": 4500000000, + "footprint_peaks": { + "encode": 8599624448, + "set_load": 18099238928, + "denoise": 24135923728, + "vae": 10151200592 + }, + "after_build": { + "footprint": 528352480, + "mlx_active": 4805792, + "mlx_cache": 4974 + }, + "mlx_encode_held": 7946292096, + "sizes": { + "compressed": 17540930560, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + } + } + ] +} diff --git a/bench/results/decode/flux1-schnell.json b/bench/results/decode/flux1-schnell.json new file mode 100644 index 0000000..38f6b8c --- /dev/null +++ b/bench/results/decode/flux1-schnell.json @@ -0,0 +1,211 @@ +{ + "key": { + "df11": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", + "variants": [ + "direct", + "staged", + "reference" + ], + "reps": 5, + "source_hash": "c898767a8b14579825aaf44b29494d94107f72d7b6286bceb187b1b8b1ad487e", + "mlx": "0.32.2" + }, + "groups": { + "single_transformer_blocks.0": { + "n_elements": 141557760, + "n_bytes": 46635275, + "n_launch": 11386, + "bytes_out": 283115520, + "max_elements_per_block": 13182, + "direct_blocks": 0, + "guard": { + "rate_bps": 20288985858.109768, + "source": "rate-from" + }, + "variants": { + "direct": { + "parity": true, + "direct_blocks": 11386, + "reps_s": [ + 0.03463154099881649, + 0.015058916993439198, + 0.015175417007412761, + 0.015008416958153248, + 0.015031082963105291 + ], + "median_s": 0.015058916993439198, + "spread": 1.303089993072713, + "gbps": 18.80052331275526 + }, + "staged": { + "parity": true, + "direct_blocks": 0, + "reps_s": [ + 0.0053494589519687, + 0.0053294579847715795, + 0.005460083018988371, + 0.005347124999389052, + 0.005424332979600877 + ], + "median_s": 0.0053494589519687, + "spread": 0.02441836368680219, + "gbps": 52.924141028469485 + }, + "reference": { + "reps_s": [ + 8.210994875000324, + 8.154502124991268, + 8.168234416982159, + 8.201749333005864, + 8.127500457980204 + ], + "median_s": 8.168234416982159, + "spread": 0.010221843884221876, + "gbps": 0.034660552764179856 + } + }, + "seconds": 57.7484415830113, + "footprint_peak": 2846313256, + "provenance": { + "git": "22a40354ced7d82168dd35461754f4a8cfdc2934", + "source_hash": "c898767a8b14579825aaf44b29494d94107f72d7b6286bceb187b1b8b1ad487e", + "mlx": "0.32.2", + "mflux": null, + "macos": "27.0", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 1400000000, + "cache_memory": 860, + "power": { + "ac": true, + "battery_pct": 100 + } + } + }, + "transformer_blocks.0": { + "n_elements": 339738624, + "n_bytes": 121650927, + "n_launch": 29700, + "bytes_out": 679477248, + "max_elements_per_block": 12832, + "direct_blocks": 0, + "guard": { + "rate_bps": 20288985858.109768, + "source": "rate-from" + }, + "variants": { + "direct": { + "parity": true, + "direct_blocks": 29700, + "reps_s": [ + 0.05665195797337219, + 0.028499290987383574, + 0.028091332991607487, + 0.028368040977511555, + 0.028444666997529566 + ], + "median_s": 0.028444666997529566, + "spread": 1.0040766159872854, + "gbps": 23.887685099601025 + }, + "staged": { + "parity": true, + "direct_blocks": 0, + "reps_s": [ + 0.01258870802121237, + 0.01261295803124085, + 0.012497207964770496, + 0.012474165996536613, + 0.012544416997116059 + ], + "median_s": 0.012544416997116059, + "spread": 0.01106404823246439, + "gbps": 54.16570958667993 + }, + "reference": { + "reps_s": [ + 20.507504915993195, + 20.559129791974556, + 20.408495916984975, + 20.397481957974378, + 20.34162062499672 + ], + "median_s": 20.408495916984975, + "spread": 0.010657775460895769, + "gbps": 0.033293842464623025 + } + }, + "seconds": 144.79180616600206, + "footprint_peak": 7189452368, + "provenance": { + "git": "22a40354ced7d82168dd35461754f4a8cfdc2934", + "source_hash": "c898767a8b14579825aaf44b29494d94107f72d7b6286bceb187b1b8b1ad487e", + "mlx": "0.32.2", + "mflux": null, + "macos": "27.0", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 1400000000, + "cache_memory": 860, + "power": { + "ac": true, + "battery_pct": 100 + } + } + } + }, + "df11": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 1400000000, + "provenance": { + "git": "22a40354ced7d82168dd35461754f4a8cfdc2934", + "source_hash": "c898767a8b14579825aaf44b29494d94107f72d7b6286bceb187b1b8b1ad487e", + "mlx": "0.32.2", + "mflux": null, + "macos": "27.0", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 1400000000, + "cache_memory": 860, + "power": { + "ac": true, + "battery_pct": 100 + } + }, + "gbps": 18.80052331275526, + "errors": [], + "exit_code": 0 +} diff --git a/bench/results/decode/qwen3-4b-layer0.json b/bench/results/decode/qwen3-4b-layer0.json new file mode 100644 index 0000000..a715c35 --- /dev/null +++ b/bench/results/decode/qwen3-4b-layer0.json @@ -0,0 +1,314 @@ +{ + "key": { + "df11": "~/.cache/huggingface/hub/models--DFloat11--Qwen3-4B-DF11/snapshots/618eddad93ca71dacf56cef22e69c5968357dd6e", + "variants": [ + "direct", + "staged", + "reference" + ], + "reps": 5, + "source_hash": "c898767a8b14579825aaf44b29494d94107f72d7b6286bceb187b1b8b1ad487e", + "mlx": "0.32.2" + }, + "groups": { + "model.layers.0": { + "n_elements": 100925440, + "n_bytes": 33159834, + "n_launch": 8096, + "bytes_out": 201850880, + "max_elements_per_block": 12848, + "direct_blocks": 0, + "guard": { + "rate_bps": 20101415844.237495, + "source": "calibration", + "ramp": [ + { + "k": 1, + "bytes": 20968, + "first_ms": { + "direct": 6.4494170364923775, + "staged": 0.7082920055836439 + }, + "timed_ms": { + "direct": 0.34075003350153565, + "staged": 0.311334035359323 + }, + "rate_bps": 61534843.54655391 + }, + { + "k": 2, + "bytes": 43840, + "first_ms": { + "direct": 0.44712499948218465, + "staged": 0.3017919952981174 + }, + "timed_ms": { + "direct": 0.3071249811910093, + "staged": 0.30570803210139275 + }, + "rate_bps": 142743191.48508054 + }, + { + "k": 4, + "bytes": 92192, + "first_ms": { + "direct": 0.30866602901369333, + "staged": 0.3082919865846634 + }, + "timed_ms": { + "direct": 0.31504203798249364, + "staged": 0.29437500052154064 + }, + "rate_bps": 292633962.7257076 + }, + { + "k": 8, + "bytes": 188888, + "first_ms": { + "direct": 0.6767080049030483, + "staged": 0.690166954882443 + }, + "timed_ms": { + "direct": 0.3540840116329491, + "staged": 0.31745899468660355 + }, + "rate_bps": 533455320.75536144 + }, + { + "k": 16, + "bytes": 379154, + "first_ms": { + "direct": 0.34775002859532833, + "staged": 0.3179580089636147 + }, + "timed_ms": { + "direct": 0.35320798633620143, + "staged": 0.3113750135526061 + }, + "rate_bps": 1073458173.8452024 + }, + { + "k": 32, + "bytes": 773132, + "first_ms": { + "direct": 0.36087498301640153, + "staged": 0.35183399450033903 + }, + "timed_ms": { + "direct": 0.388041022233665, + "staged": 0.3094999701716006 + }, + "rate_bps": 1992397596.3923898 + }, + { + "k": 64, + "bytes": 1559738, + "first_ms": { + "direct": 0.42495899833738804, + "staged": 0.3626249963417649 + }, + "timed_ms": { + "direct": 0.4549170262180269, + "staged": 0.33520901342853904 + }, + "rate_bps": 3428620847.5575247 + }, + { + "k": 128, + "bytes": 3133008, + "first_ms": { + "direct": 0.712334003765136, + "staged": 0.5285409861244261 + }, + "timed_ms": { + "direct": 0.7142079994082451, + "staged": 0.44604099821299314 + }, + "rate_bps": 4386688475.3403 + }, + { + "k": 256, + "bytes": 6338716, + "first_ms": { + "direct": 1.1096660164184868, + "staged": 0.7340830052271485 + }, + "timed_ms": { + "direct": 1.0890419944189489, + "staged": 0.7252499926835299 + }, + "rate_bps": 5820451399.013295 + }, + { + "k": 512, + "bytes": 12754474, + "first_ms": { + "direct": 1.8795840442180634, + "staged": 11.295999982394278 + }, + "timed_ms": { + "direct": 1.8562080222181976, + "staged": 0.6617920007556677 + }, + "rate_bps": 6871252492.895815 + }, + { + "k": 1024, + "bytes": 25449322, + "first_ms": { + "direct": 1.8582080374471843, + "staged": 1.0497079929336905 + }, + "timed_ms": { + "direct": 1.8409999902360141, + "staged": 0.8624580223113298 + }, + "rate_bps": 13823640486.134617 + }, + { + "k": 2048, + "bytes": 50606468, + "first_ms": { + "direct": 3.4113749861717224, + "staged": 1.190749986562878 + }, + "timed_ms": { + "direct": 2.890957985073328, + "staged": 1.1820420040749013 + }, + "rate_bps": 17505085947.735207 + }, + { + "k": 4096, + "bytes": 101904652, + "first_ms": { + "direct": 5.3323329775594175, + "staged": 2.1951659582555294 + }, + "timed_ms": { + "direct": 5.247833963949233, + "staged": 2.1727500134147704 + }, + "rate_bps": 19418421524.013332 + }, + { + "k": 8096, + "bytes": 201850880, + "first_ms": { + "direct": 10.130250011570752, + "staged": 4.100417019799352 + }, + "timed_ms": { + "direct": 10.041625006124377, + "staged": 3.9670830010436475 + }, + "rate_bps": 20101415844.237495 + } + ] + }, + "variants": { + "direct": { + "parity": true, + "direct_blocks": 8096, + "reps_s": [ + 0.010116082965396345, + 0.009926124999765307, + 0.00994879100471735, + 0.009874499985016882, + 0.010030917008407414 + ], + "median_s": 0.00994879100471735, + "spread": 0.02428264703368605, + "gbps": 20.28898585810977 + }, + "staged": { + "parity": true, + "direct_blocks": 0, + "reps_s": [ + 0.00400195800466463, + 0.004044540983159095, + 0.004029541974887252, + 0.004073375021107495, + 0.0040474169654771686 + ], + "median_s": 0.004044540983159095, + "spread": 0.017657632037908787, + "gbps": 49.90699336228238 + }, + "reference": { + "reps_s": [ + 5.880062541982625, + 5.92360183398705, + 5.89788637496531, + 5.8660767920082435, + 5.883949999988545 + ], + "median_s": 5.883949999988545, + "spread": 0.009776602788758955, + "gbps": 0.03430533570142387 + } + }, + "seconds": 41.87269316700986, + "footprint_peak": 1977632880, + "provenance": { + "git": "22a40354ced7d82168dd35461754f4a8cfdc2934", + "source_hash": "c898767a8b14579825aaf44b29494d94107f72d7b6286bceb187b1b8b1ad487e", + "mlx": "0.32.2", + "mflux": null, + "macos": "27.0", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 1400000000, + "cache_memory": 860, + "power": { + "ac": true, + "battery_pct": 100 + } + } + } + }, + "df11": "~/.cache/huggingface/hub/models--DFloat11--Qwen3-4B-DF11/snapshots/618eddad93ca71dacf56cef22e69c5968357dd6e", + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 1400000000, + "provenance": { + "git": "22a40354ced7d82168dd35461754f4a8cfdc2934", + "source_hash": "c898767a8b14579825aaf44b29494d94107f72d7b6286bceb187b1b8b1ad487e", + "mlx": "0.32.2", + "mflux": null, + "macos": "27.0", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 1400000000, + "cache_memory": 860, + "power": { + "ac": true, + "battery_pct": 100 + } + }, + "gbps": 20.28898585810977, + "errors": [], + "exit_code": 0 +} diff --git a/bench/results/identity/flux1-schnell-1024/bf16-result.json b/bench/results/identity/flux1-schnell-1024/bf16-result.json new file mode 100644 index 0000000..dbb29f1 --- /dev/null +++ b/bench/results/identity/flux1-schnell-1024/bf16-result.json @@ -0,0 +1,55 @@ +{ + "exit_code": 0, + "reads": 1976, + "output_shape": [ + 1, + 4096, + 64 + ], + "output_dtype": "mlx.core.float32", + "footprint_peak_bytes": 10350445792, + "mlx_peak_bytes_process": 9988562688, + "key": { + "model": "schnell", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 3.5, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", + "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", + "source": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "mlx": "0.32.2" + }, + "memory_caps_gb": [ + 20, + 22 + ], + "provenance": { + "git": "89ff46a6614e9fd4611cfea28c809c84243cca11", + "source_hash": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "mlx": "0.32.2", + "mflux": "0.20.0", + "macos": "27.0.1", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 0, + "cache_memory": 0, + "power": { + "ac": true, + "battery_pct": 94 + }, + "date": "2026-10-08" + } +} diff --git a/bench/results/identity/flux1-schnell-1024/compare.json b/bench/results/identity/flux1-schnell-1024/compare.json new file mode 100644 index 0000000..41701ec --- /dev/null +++ b/bench/results/identity/flux1-schnell-1024/compare.json @@ -0,0 +1,34 @@ +{ + "exit_code": 0, + "equal": true, + "degenerate": [], + "df11_key": { + "model": "schnell", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 3.5, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", + "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", + "source": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "mlx": "0.32.2" + }, + "bf16_key": { + "model": "schnell", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 3.5, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", + "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", + "source": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "mlx": "0.32.2" + }, + "pixels_identical": true, + "df11_png_sha256": "75b9af5d084de91d69e11a5a5e1a444aeba3a498cadc3a106ac6e2492db7e2c4", + "bf16_png_sha256": "e7c07257317efb656b88f18d65a45be0e914d44603888ca3312f0fc290e767c1" +} diff --git a/bench/results/identity/flux1-schnell-1024/df11-result.json b/bench/results/identity/flux1-schnell-1024/df11-result.json new file mode 100644 index 0000000..ed98c35 --- /dev/null +++ b/bench/results/identity/flux1-schnell-1024/df11-result.json @@ -0,0 +1,128 @@ +{ + "exit_code": 0, + "degenerate": [], + "report": { + "model": "schnell", + "df11": { + "root": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", + "repo_id": null, + "revision": "51a428b928197e0531cb93d6e438941e2d0b247e" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", + "repo_id": null, + "revision": "741f7c3ce8b383c54771c7003378a50191e9efe9" + }, + "eval_policy": "per-block", + "cache_limit_in_force": 2500000000, + "drop_set_before_vae": true, + "fit": { + "label": "predicted", + "phases": { + "encode": 11518505446, + "denoise": 21411230501, + "vae": 8808361940 + }, + "peak_phase": "denoise", + "peak_bytes": 21411230501, + "budget_bytes": 24653119488, + "fits": true + }, + "sizes": { + "compressed": 16195141095, + "extras": 113899648, + "encoders": 9770792936, + "vae": 167666902 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 363021464, + "mlx_peak": 10976952538, + "footprint_end": 9915833616, + "active_at_start": 115538184 + }, + "set_load": { + "footprint_start": 9915833616, + "mlx_peak": 16315890952, + "footprint_end": 16608785104, + "active_at_start": 119734024 + }, + "denoise": { + "footprint_start": 16608785104, + "mlx_peak": 18725842260, + "footprint_end": 19131576992, + "active_at_start": 16315890952 + }, + "vae": { + "footprint_start": 9832683784, + "mlx_peak": 10257014912, + "footprint_end": 583550464, + "active_at_start": 121831228 + } + }, + "decode_launches": 228, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 26.188013666018378, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + } + }, + "output_shape": [ + 1, + 4096, + 64 + ], + "output_dtype": "mlx.core.float32", + "footprint_peak_bytes": 21237592736, + "mlx_peak_bytes_max_over_phases": 18725842260, + "key": { + "model": "schnell", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 3.5, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", + "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", + "source": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "mlx": "0.32.2" + }, + "memory_caps_gb": [ + 20, + 22 + ], + "provenance": { + "git": "89ff46a6614e9fd4611cfea28c809c84243cca11", + "source_hash": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "mlx": "0.32.2", + "mflux": "0.20.0", + "macos": "27.0.1", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 32641751449, + "cache_memory": 1836212, + "power": { + "ac": true, + "battery_pct": 94 + }, + "date": "2026-10-08" + } +} diff --git a/bench/results/identity/krea-2-raw-1024/bf16-result.json b/bench/results/identity/krea-2-raw-1024/bf16-result.json new file mode 100644 index 0000000..07330f1 --- /dev/null +++ b/bench/results/identity/krea-2-raw-1024/bf16-result.json @@ -0,0 +1,59 @@ +{ + "exit_code": 0, + "reads": 1792, + "replaced_from_base": 206, + "calls_per_step": 2, + "output_shape": [ + 1, + 16, + 128, + 128 + ], + "output_dtype": "mlx.core.float32", + "footprint_peak_bytes": 10193290800, + "mlx_peak_bytes_process": 9576703062, + "key": { + "model": "krea-2-raw", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 3.5, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Raw-DF11-ComfyUI/snapshots/8320616b25ac9340a830a7fb21f1b0237e160e66", + "base": "~/.cache/huggingface/hub/models--krea--Krea-2-Raw/snapshots/6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "source": "bf1ad133debe1c39ccdb36eb83553035cbead8cbfac3b458ebdf6d466191086f", + "mlx": "0.32.2", + "negative_prompt": " " + }, + "memory_caps_gb": [ + 20, + 22 + ], + "provenance": { + "git": "e3a508e6d37a3946253801258a9200ff03a708a4", + "source_hash": "bf1ad133debe1c39ccdb36eb83553035cbead8cbfac3b458ebdf6d466191086f", + "mlx": "0.32.2", + "mflux": "0.20.0", + "macos": "27.0.1", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 0, + "cache_memory": 0, + "power": { + "ac": true, + "battery_pct": 94 + }, + "date": "2026-10-08" + } +} diff --git a/bench/results/identity/krea-2-raw-1024/compare.json b/bench/results/identity/krea-2-raw-1024/compare.json new file mode 100644 index 0000000..d23fdbf --- /dev/null +++ b/bench/results/identity/krea-2-raw-1024/compare.json @@ -0,0 +1,36 @@ +{ + "exit_code": 0, + "equal": true, + "degenerate": [], + "df11_key": { + "model": "krea-2-raw", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 3.5, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Raw-DF11-ComfyUI/snapshots/8320616b25ac9340a830a7fb21f1b0237e160e66", + "base": "~/.cache/huggingface/hub/models--krea--Krea-2-Raw/snapshots/6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "source": "bf1ad133debe1c39ccdb36eb83553035cbead8cbfac3b458ebdf6d466191086f", + "mlx": "0.32.2", + "negative_prompt": " " + }, + "bf16_key": { + "model": "krea-2-raw", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 3.5, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Raw-DF11-ComfyUI/snapshots/8320616b25ac9340a830a7fb21f1b0237e160e66", + "base": "~/.cache/huggingface/hub/models--krea--Krea-2-Raw/snapshots/6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "source": "bf1ad133debe1c39ccdb36eb83553035cbead8cbfac3b458ebdf6d466191086f", + "mlx": "0.32.2", + "negative_prompt": " " + }, + "pixels_identical": true, + "df11_png_sha256": "fb26f5d39a39671197ce254d069d9fd23e04fd6fe34073678e1ee4c4a3190cc7", + "bf16_png_sha256": "d0810d46bc2625744ac84943f46ff4635fec858fe134a7e9a5f879227fed33f9" +} diff --git a/bench/results/identity/krea-2-raw-1024/df11-result.json b/bench/results/identity/krea-2-raw-1024/df11-result.json new file mode 100644 index 0000000..7b20463 --- /dev/null +++ b/bench/results/identity/krea-2-raw-1024/df11-result.json @@ -0,0 +1,148 @@ +{ + "exit_code": 0, + "degenerate": [], + "report": { + "family": "krea2", + "model": "krea-2-raw", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Raw-DF11-ComfyUI/snapshots/8320616b25ac9340a830a7fb21f1b0237e160e66", + "repo_id": null, + "revision": "8320616b25ac9340a830a7fb21f1b0237e160e66", + "config_source": "header and spot checks match layout krea-2-raw-comfyui (mingyi456/Krea-2-Raw-DF11-ComfyUI@8320616b25ac9340a830a7fb21f1b0237e160e66)" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--krea--Krea-2-Raw/snapshots/6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "repo_id": null, + "revision": "6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603" + }, + "layout": "header and spot checks match layout krea-2-raw-comfyui (mingyi456/Krea-2-Raw-DF11-ComfyUI@8320616b25ac9340a830a7fb21f1b0237e160e66)", + "nonblock": "per-call", + "eval_policy": "per-block", + "eval_policy_status": "measured", + "cfg_calls_per_step": 2, + "text_tokens": 30, + "predict": "uncompiled", + "nonblock_groups": [ + "tmlp", + "tproj", + "txtfusion.layerwise_blocks.0", + "txtfusion.layerwise_blocks.1", + "txtfusion.refiner_blocks.0", + "txtfusion.refiner_blocks.1", + "txtmlp" + ], + "cache_limit_in_force": 4500000000, + "drop_set_before_vae": true, + "fit": { + "label": "predicted", + "phases": { + "encode": 8595823360, + "denoise": 24243378794, + "vae": 11869477608 + }, + "peak_phase": "denoise", + "peak_bytes": 24243378794, + "budget_bytes": 24653119488, + "fits": true + }, + "sizes": { + "compressed": 17540930560, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 541000952, + "mlx_peak": 7947813964, + "footprint_end": 7770725072, + "active_at_start": 4823348 + }, + "set_load": { + "footprint_start": 7770725072, + "mlx_peak": 17552853332, + "footprint_end": 18102204456, + "active_at_start": 11212068 + }, + "denoise": { + "footprint_start": 18102204456, + "mlx_peak": 21246251050, + "footprint_end": 22706911176, + "active_at_start": 17552853332 + }, + "vae": { + "footprint_start": 10819573640, + "mlx_peak": 9585686678, + "footprint_end": 4938287024, + "active_at_start": 14358666 + } + }, + "decode_launches": 224, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 26.930201666022185, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + } + }, + "output_shape": [ + 1, + 16, + 128, + 128 + ], + "output_dtype": "mlx.core.float32", + "footprint_peak_bytes": 22711515080, + "mlx_peak_bytes_max_over_phases": 21246251050, + "key": { + "model": "krea-2-raw", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 3.5, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Raw-DF11-ComfyUI/snapshots/8320616b25ac9340a830a7fb21f1b0237e160e66", + "base": "~/.cache/huggingface/hub/models--krea--Krea-2-Raw/snapshots/6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "source": "bf1ad133debe1c39ccdb36eb83553035cbead8cbfac3b458ebdf6d466191086f", + "mlx": "0.32.2", + "negative_prompt": " " + }, + "memory_caps_gb": [ + 20, + 22 + ], + "provenance": { + "git": "e3a508e6d37a3946253801258a9200ff03a708a4", + "source_hash": "bf1ad133debe1c39ccdb36eb83553035cbead8cbfac3b458ebdf6d466191086f", + "mlx": "0.32.2", + "mflux": "0.20.0", + "macos": "27.0.1", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 32641751449, + "cache_memory": 5113016, + "power": { + "ac": true, + "battery_pct": 94 + }, + "date": "2026-10-08" + } +} diff --git a/bench/results/parity/krea-2-raw/full-vs-original.json b/bench/results/parity/krea-2-raw/full-vs-original.json new file mode 100644 index 0000000..c0470cf --- /dev/null +++ b/bench/results/parity/krea-2-raw/full-vs-original.json @@ -0,0 +1,128 @@ +{ + "df11_repo": "mingyi456/Krea-2-Raw-DF11-ComfyUI", + "df11_revision": "8320616b25ac9340a830a7fb21f1b0237e160e66", + "bf16_repo": "krea/Krea-2-Raw", + "bf16_revision": "6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "mode": "parity, by HTTP range reads, 18 processes", + "config_source": [ + "layout krea-2-raw-comfyui" + ], + "control": { + "b00": "bf16", + "b01": "bf16", + "b02": "bf16", + "b03": "bf16", + "b04": "bf16", + "b05": "bf16", + "b06": "bf16", + "b07": "bf16", + "b08": "bf16", + "b09": "bf16", + "b10": "bf16", + "b11": "bf16", + "b12": "bf16", + "b13": "bf16", + "ext": "fp32 rounded to bf16 (nearest even)", + "f32": "fp32 rounded to bf16 (nearest even)", + "tf0": "bf16", + "tf1": "bf16" + }, + "exit_codes": { + "b00": 0, + "b01": 0, + "b02": 0, + "b03": 0, + "b04": 0, + "b05": 0, + "b06": 0, + "b07": 0, + "b08": 0, + "b09": 0, + "b10": 0, + "b11": 0, + "b12": 0, + "b13": 0, + "ext": 0, + "f32": 0, + "tf0": 0, + "tf1": 0 + }, + "groups": [ + "blocks.0", + "blocks.1", + "blocks.10", + "blocks.11", + "blocks.12", + "blocks.13", + "blocks.14", + "blocks.15", + "blocks.16", + "blocks.17", + "blocks.18", + "blocks.19", + "blocks.2", + "blocks.20", + "blocks.21", + "blocks.22", + "blocks.23", + "blocks.24", + "blocks.25", + "blocks.26", + "blocks.27", + "blocks.3", + "blocks.4", + "blocks.5", + "blocks.6", + "blocks.7", + "blocks.8", + "blocks.9", + "tmlp", + "tproj", + "txtfusion.layerwise_blocks.0", + "txtfusion.layerwise_blocks.1", + "txtfusion.refiner_blocks.0", + "txtfusion.refiner_blocks.1", + "txtmlp" + ], + "compared": 261, + "mismatched": 0, + "extras_compared": 169, + "extras_mismatched": 0, + "per_group_status": { + "blocks.0": "equal", + "blocks.1": "equal", + "blocks.2": "equal", + "blocks.3": "equal", + "blocks.4": "equal", + "blocks.5": "equal", + "blocks.6": "equal", + "blocks.7": "equal", + "blocks.8": "equal", + "blocks.9": "equal", + "blocks.10": "equal", + "blocks.11": "equal", + "blocks.12": "equal", + "blocks.13": "equal", + "blocks.14": "equal", + "blocks.15": "equal", + "blocks.16": "equal", + "blocks.17": "equal", + "blocks.18": "equal", + "blocks.19": "equal", + "blocks.20": "equal", + "blocks.21": "equal", + "blocks.22": "equal", + "blocks.23": "equal", + "blocks.24": "equal", + "blocks.25": "equal", + "blocks.26": "equal", + "blocks.27": "equal", + "tmlp": "equal", + "tproj": "equal", + "txtmlp": "equal", + "txtfusion.layerwise_blocks.0": "equal", + "txtfusion.layerwise_blocks.1": "equal", + "txtfusion.refiner_blocks.0": "equal", + "txtfusion.refiner_blocks.1": "equal" + } +} diff --git a/bench/results/parity/krea-2-raw/kernel-vs-reference.json b/bench/results/parity/krea-2-raw/kernel-vs-reference.json new file mode 100644 index 0000000..8e13c11 --- /dev/null +++ b/bench/results/parity/krea-2-raw/kernel-vs-reference.json @@ -0,0 +1,63 @@ +{ + "key": { + "df11_revision": "8320616b25ac9340a830a7fb21f1b0237e160e66", + "bf16_revision": "none", + "mode": "kernel-vs-reference", + "decoder": "metal", + "source": "44a590c5786c349fc49b55f38a9d7a5754f047bbf57868bb09d088d5f381f785", + "git": "d1aa9df996cee99356ffa3f2f01fdb0ebd16469c", + "mlx": "0.32.2" + }, + "exit_code": 0, + "groups": 35, + "partial": false, + "selected_groups": [ + "blocks.0", + "blocks.1", + "blocks.2", + "blocks.3", + "blocks.4", + "blocks.5", + "blocks.6", + "blocks.7", + "blocks.8", + "blocks.9", + "blocks.10", + "blocks.11", + "blocks.12", + "blocks.13", + "blocks.14", + "blocks.15", + "blocks.16", + "blocks.17", + "blocks.18", + "blocks.19", + "blocks.20", + "blocks.21", + "blocks.22", + "blocks.23", + "blocks.24", + "blocks.25", + "blocks.26", + "blocks.27", + "tmlp", + "tproj", + "txtfusion.layerwise_blocks.0", + "txtfusion.layerwise_blocks.1", + "txtfusion.refiner_blocks.0", + "txtfusion.refiner_blocks.1", + "txtmlp" + ], + "compared": 261, + "expected": 261, + "extras_compared": 0, + "mismatched": 0, + "missing_originals": [], + "uncovered_originals": [], + "ignored_originals": [], + "memory_caps_gb": [ + 20, + 22 + ], + "seconds": 791.46844495798 +} diff --git a/bench/results/parity/krea-2/full-vs-original.json b/bench/results/parity/krea-2/full-vs-original.json new file mode 100644 index 0000000..ce88327 --- /dev/null +++ b/bench/results/parity/krea-2/full-vs-original.json @@ -0,0 +1,128 @@ +{ + "df11_repo": "mingyi456/Krea-2-Turbo-DF11-ComfyUI", + "df11_revision": "978da5fb7647bd222d33125993abd8fdc2840cfc", + "bf16_repo": "krea/Krea-2-Turbo", + "bf16_revision": "98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7", + "mode": "parity, by HTTP range reads, 18 processes", + "config_source": [ + "layout krea-2-turbo-comfyui" + ], + "control": { + "b00": "bf16", + "b01": "bf16", + "b02": "bf16", + "b03": "bf16", + "b04": "bf16", + "b05": "bf16", + "b06": "bf16", + "b07": "bf16", + "b08": "bf16", + "b09": "bf16", + "b10": "bf16", + "b11": "bf16", + "b12": "bf16", + "b13": "bf16", + "ext": "fp32 rounded to bf16 (nearest even)", + "f32": "fp32 rounded to bf16 (nearest even)", + "tf0": "bf16", + "tf1": "bf16" + }, + "exit_codes": { + "b00": 0, + "b01": 0, + "b02": 0, + "b03": 0, + "b04": 0, + "b05": 0, + "b06": 0, + "b07": 0, + "b08": 0, + "b09": 0, + "b10": 0, + "b11": 0, + "b12": 0, + "b13": 0, + "ext": 0, + "f32": 0, + "tf0": 0, + "tf1": 0 + }, + "groups": [ + "blocks.0", + "blocks.1", + "blocks.10", + "blocks.11", + "blocks.12", + "blocks.13", + "blocks.14", + "blocks.15", + "blocks.16", + "blocks.17", + "blocks.18", + "blocks.19", + "blocks.2", + "blocks.20", + "blocks.21", + "blocks.22", + "blocks.23", + "blocks.24", + "blocks.25", + "blocks.26", + "blocks.27", + "blocks.3", + "blocks.4", + "blocks.5", + "blocks.6", + "blocks.7", + "blocks.8", + "blocks.9", + "tmlp", + "tproj", + "txtfusion.layerwise_blocks.0", + "txtfusion.layerwise_blocks.1", + "txtfusion.refiner_blocks.0", + "txtfusion.refiner_blocks.1", + "txtmlp" + ], + "compared": 261, + "mismatched": 0, + "extras_compared": 169, + "extras_mismatched": 0, + "per_group_status": { + "blocks.0": "equal", + "blocks.1": "equal", + "blocks.2": "equal", + "blocks.3": "equal", + "blocks.4": "equal", + "blocks.5": "equal", + "blocks.6": "equal", + "blocks.7": "equal", + "blocks.8": "equal", + "blocks.9": "equal", + "blocks.10": "equal", + "blocks.11": "equal", + "blocks.12": "equal", + "blocks.13": "equal", + "blocks.14": "equal", + "blocks.15": "equal", + "blocks.16": "equal", + "blocks.17": "equal", + "blocks.18": "equal", + "blocks.19": "equal", + "blocks.20": "equal", + "blocks.21": "equal", + "blocks.22": "equal", + "blocks.23": "equal", + "blocks.24": "equal", + "blocks.25": "equal", + "blocks.26": "equal", + "blocks.27": "equal", + "tmlp": "equal", + "tproj": "equal", + "txtmlp": "equal", + "txtfusion.layerwise_blocks.0": "equal", + "txtfusion.layerwise_blocks.1": "equal", + "txtfusion.refiner_blocks.0": "equal", + "txtfusion.refiner_blocks.1": "equal" + } +} diff --git a/bench/results/parity/krea-2/kernel-vs-reference.json b/bench/results/parity/krea-2/kernel-vs-reference.json new file mode 100644 index 0000000..0b1d078 --- /dev/null +++ b/bench/results/parity/krea-2/kernel-vs-reference.json @@ -0,0 +1,63 @@ +{ + "key": { + "df11_revision": "978da5fb7647bd222d33125993abd8fdc2840cfc", + "bf16_revision": "none", + "mode": "kernel-vs-reference", + "decoder": "metal", + "source": "0a9c8fd427322a5707a1699d760892523857e590414093fbf9d9194216182a1d", + "git": "0f82cca74efabcee0f87f3a47b8b74e75ebe9b24", + "mlx": "0.32.2" + }, + "exit_code": 0, + "groups": 35, + "partial": false, + "selected_groups": [ + "blocks.0", + "blocks.1", + "blocks.2", + "blocks.3", + "blocks.4", + "blocks.5", + "blocks.6", + "blocks.7", + "blocks.8", + "blocks.9", + "blocks.10", + "blocks.11", + "blocks.12", + "blocks.13", + "blocks.14", + "blocks.15", + "blocks.16", + "blocks.17", + "blocks.18", + "blocks.19", + "blocks.20", + "blocks.21", + "blocks.22", + "blocks.23", + "blocks.24", + "blocks.25", + "blocks.26", + "blocks.27", + "tmlp", + "tproj", + "txtfusion.layerwise_blocks.0", + "txtfusion.layerwise_blocks.1", + "txtfusion.refiner_blocks.0", + "txtfusion.refiner_blocks.1", + "txtmlp" + ], + "compared": 261, + "expected": 261, + "extras_compared": 0, + "mismatched": 0, + "missing_originals": [], + "uncovered_originals": [], + "ignored_originals": [], + "memory_caps_gb": [ + 20, + 22 + ], + "seconds": 790.681125416013 +} diff --git a/bench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.json b/bench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.json new file mode 100644 index 0000000..8992387 --- /dev/null +++ b/bench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.json @@ -0,0 +1,139 @@ +{ + "exit_code": 0, + "output": "krea-2-raw-card.png", + "memory_caps_gb": [ + 20, + 22 + ], + "model_kwargs": { + "model": "krea-2-raw", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Raw-DF11-ComfyUI/snapshots/8320616b25ac9340a830a7fb21f1b0237e160e66", + "base_path": "~/.cache/huggingface/hub/models--krea--Krea-2-Raw/snapshots/6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 32, + "label": "MEASURED", + "watchdog_ceiling_bytes": 30064771072, + "limits": { + "tier": { + "tier_gb": 32, + "ram_bytes": 34359738368, + "recommended_bytes": 26800603136, + "memory_limit_bytes": 32641751449, + "cache_limit_bytes": 25460572979, + "wired_limit_bytes": 0, + "reserve_bytes": 2147483648, + "budget_bytes": 26800603136, + "ceiling_bytes": 24653119488, + "fit_budget_bytes": 24653119488, + "budget_source": "device", + "is_host": true, + "label": "MEASURED" + }, + "effective": { + "memory_limit_bytes": 23622320128, + "cache_limit_bytes": 32641751449, + "wired_limit_bytes": 21474836480 + }, + "applied": "host-caps" + }, + "family": "krea2", + "model": "krea-2-raw", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Raw-DF11-ComfyUI/snapshots/8320616b25ac9340a830a7fb21f1b0237e160e66", + "repo_id": null, + "revision": "8320616b25ac9340a830a7fb21f1b0237e160e66", + "config_source": "header and spot checks match layout krea-2-raw-comfyui (mingyi456/Krea-2-Raw-DF11-ComfyUI@8320616b25ac9340a830a7fb21f1b0237e160e66)" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--krea--Krea-2-Raw/snapshots/6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "repo_id": null, + "revision": "6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603" + }, + "layout": "header and spot checks match layout krea-2-raw-comfyui (mingyi456/Krea-2-Raw-DF11-ComfyUI@8320616b25ac9340a830a7fb21f1b0237e160e66)", + "nonblock": "per-call", + "eval_policy": "per-block", + "eval_policy_status": "measured", + "cfg_calls_per_step": 2, + "text_tokens": 30, + "predict": "uncompiled", + "nonblock_groups": [ + "tmlp", + "tproj", + "txtfusion.layerwise_blocks.0", + "txtfusion.layerwise_blocks.1", + "txtfusion.refiner_blocks.0", + "txtfusion.refiner_blocks.1", + "txtmlp" + ], + "cache_limit_in_force": 4500000000, + "drop_set_before_vae": true, + "fit": { + "label": "predicted", + "phases": { + "encode": 8595823360, + "denoise": 24243378794, + "vae": 11869477608 + }, + "peak_phase": "denoise", + "peak_bytes": 24243378794, + "budget_bytes": 24653119488, + "fits": true + }, + "sizes": { + "compressed": 17540930560, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 537150688, + "mlx_peak": 7958027620, + "footprint_end": 7622171392, + "active_at_start": 4822320 + }, + "set_load": { + "footprint_start": 7622171392, + "mlx_peak": 17552853332, + "footprint_end": 18099501336, + "active_at_start": 11212068 + }, + "denoise": { + "footprint_start": 18099501336, + "mlx_peak": 21246267626, + "footprint_end": 22572742864, + "active_at_start": 17552853332 + }, + "vae": { + "footprint_start": 4930422776, + "mlx_peak": 9585686870, + "footprint_end": 4844816520, + "active_at_start": 14358858 + } + }, + "decode_launches": 2912, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 30.346461208013352, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "footprint_peak_bytes": 22714513664, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 22714513664, + "mlx_peak_bytes": 22161954042, + "elapsed_seconds": 3373.426 +} diff --git a/bench/results/refusals/ernie-image-turbo-1024-tier16.json b/bench/results/refusals/ernie-image-turbo-1024-tier16.json index 9a11d72..2079b61 100644 --- a/bench/results/refusals/ernie-image-turbo-1024-tier16.json +++ b/bench/results/refusals/ernie-image-turbo-1024-tier16.json @@ -40,10 +40,10 @@ }, "applied": "tier-caps" }, - "error": "DFloatResourceError: predicted peak 13.9 GiB in the denoise phase exceeds the budget 8.7 GiB (encode 6.8 GiB, denoise 13.9 GiB, vae 7.4 GiB); pass fit_check=False to run anyway", - "footprint_peak_bytes": 707167720, + "error": "DFloatResourceError: predicted peak 14.4 GiB in the denoise phase exceeds the budget 8.7 GiB (encode 6.8 GiB, denoise 14.4 GiB, vae 7.4 GiB); pass fit_check=False to run anyway", + "footprint_peak_bytes": 639436264, "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", - "watched_peak_bytes": 707167720, - "mlx_peak_bytes": 28533646, - "elapsed_seconds": 3.644 + "watched_peak_bytes": 639436264, + "mlx_peak_bytes": 28483214, + "elapsed_seconds": 4.53 } diff --git a/bench/results/refusals/krea-2-1024-tier16.json b/bench/results/refusals/krea-2-1024-tier16.json new file mode 100644 index 0000000..99430ca --- /dev/null +++ b/bench/results/refusals/krea-2-1024-tier16.json @@ -0,0 +1,49 @@ +{ + "exit_code": 2, + "output": "t.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "krea-2", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Turbo-DF11-ComfyUI/snapshots/978da5fb7647bd222d33125993abd8fdc2840cfc", + "base_path": "~/.cache/huggingface/hub/models--krea--Krea-2-Turbo/snapshots/98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true, + "budget_bytes": 9305762474 + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 16, + "label": "CAPPED", + "watchdog_ceiling_bytes": 9842633386, + "limits": { + "tier": { + "tier_gb": 16, + "ram_bytes": 17179869184, + "recommended_bytes": 11453246122, + "memory_limit_bytes": 10737418240, + "cache_limit_bytes": 10737418240, + "wired_limit_bytes": 8589934592, + "reserve_bytes": 1610612736, + "budget_bytes": 11453246122, + "ceiling_bytes": 9842633386, + "fit_budget_bytes": 9305762474, + "budget_source": "ratio", + "is_host": false, + "label": "CAPPED" + }, + "effective": { + "memory_limit_bytes": 10737418240, + "cache_limit_bytes": 10737418240, + "wired_limit_bytes": 8589934592 + }, + "applied": "tier-caps" + }, + "error": "DFloatResourceError: predicted peak 22.6 GiB in the denoise phase exceeds the budget 8.7 GiB (encode 8.0 GiB, denoise 22.6 GiB, vae 11.1 GiB); pass fit_check=False to run anyway", + "footprint_peak_bytes": 561186040, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 561186040, + "mlx_peak_bytes": 5230518, + "elapsed_seconds": 3.164 +} diff --git a/bench/results/refusals/krea-2-1024-tier24.json b/bench/results/refusals/krea-2-1024-tier24.json new file mode 100644 index 0000000..ec52d4b --- /dev/null +++ b/bench/results/refusals/krea-2-1024-tier24.json @@ -0,0 +1,49 @@ +{ + "exit_code": 2, + "output": "t.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "krea-2", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Turbo-DF11-ComfyUI/snapshots/978da5fb7647bd222d33125993abd8fdc2840cfc", + "base_path": "~/.cache/huggingface/hub/models--krea--Krea-2-Turbo/snapshots/98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true, + "budget_bytes": 15032385536 + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 24, + "label": "CAPPED", + "watchdog_ceiling_bytes": 15569256448, + "limits": { + "tier": { + "tier_gb": 24, + "ram_bytes": 25769803776, + "recommended_bytes": 17179869184, + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536, + "reserve_bytes": 1610612736, + "budget_bytes": 17179869184, + "ceiling_bytes": 15569256448, + "fit_budget_bytes": 15032385536, + "budget_source": "ratio", + "is_host": false, + "label": "CAPPED" + }, + "effective": { + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536 + }, + "applied": "tier-caps" + }, + "error": "DFloatResourceError: predicted peak 22.6 GiB in the denoise phase exceeds the budget 14.0 GiB (encode 8.0 GiB, denoise 22.6 GiB, vae 11.1 GiB); pass fit_check=False to run anyway", + "footprint_peak_bytes": 554517752, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 554517752, + "mlx_peak_bytes": 5229238, + "elapsed_seconds": 3.523 +} diff --git a/bench/results/refusals/krea-2-raw-1024-tier16.json b/bench/results/refusals/krea-2-raw-1024-tier16.json new file mode 100644 index 0000000..ce3b409 --- /dev/null +++ b/bench/results/refusals/krea-2-raw-1024-tier16.json @@ -0,0 +1,49 @@ +{ + "exit_code": 2, + "output": "r.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "krea-2-raw", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Raw-DF11-ComfyUI/snapshots/8320616b25ac9340a830a7fb21f1b0237e160e66", + "base_path": "~/.cache/huggingface/hub/models--krea--Krea-2-Raw/snapshots/6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true, + "budget_bytes": 9305762474 + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 16, + "label": "CAPPED", + "watchdog_ceiling_bytes": 9842633386, + "limits": { + "tier": { + "tier_gb": 16, + "ram_bytes": 17179869184, + "recommended_bytes": 11453246122, + "memory_limit_bytes": 10737418240, + "cache_limit_bytes": 10737418240, + "wired_limit_bytes": 8589934592, + "reserve_bytes": 1610612736, + "budget_bytes": 11453246122, + "ceiling_bytes": 9842633386, + "fit_budget_bytes": 9305762474, + "budget_source": "ratio", + "is_host": false, + "label": "CAPPED" + }, + "effective": { + "memory_limit_bytes": 10737418240, + "cache_limit_bytes": 10737418240, + "wired_limit_bytes": 8589934592 + }, + "applied": "tier-caps" + }, + "error": "DFloatResourceError: predicted peak 22.6 GiB in the denoise phase exceeds the budget 8.7 GiB (encode 8.0 GiB, denoise 22.6 GiB, vae 11.1 GiB); pass fit_check=False to run anyway", + "footprint_peak_bytes": 558204128, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 558204128, + "mlx_peak_bytes": 5180086, + "elapsed_seconds": 3.52 +} diff --git a/bench/results/refusals/krea-2-raw-1024-tier24.json b/bench/results/refusals/krea-2-raw-1024-tier24.json new file mode 100644 index 0000000..c562bc6 --- /dev/null +++ b/bench/results/refusals/krea-2-raw-1024-tier24.json @@ -0,0 +1,49 @@ +{ + "exit_code": 2, + "output": "r.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "krea-2-raw", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Raw-DF11-ComfyUI/snapshots/8320616b25ac9340a830a7fb21f1b0237e160e66", + "base_path": "~/.cache/huggingface/hub/models--krea--Krea-2-Raw/snapshots/6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true, + "budget_bytes": 15032385536 + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 24, + "label": "CAPPED", + "watchdog_ceiling_bytes": 15569256448, + "limits": { + "tier": { + "tier_gb": 24, + "ram_bytes": 25769803776, + "recommended_bytes": 17179869184, + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536, + "reserve_bytes": 1610612736, + "budget_bytes": 17179869184, + "ceiling_bytes": 15569256448, + "fit_budget_bytes": 15032385536, + "budget_source": "ratio", + "is_host": false, + "label": "CAPPED" + }, + "effective": { + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536 + }, + "applied": "tier-caps" + }, + "error": "DFloatResourceError: predicted peak 22.6 GiB in the denoise phase exceeds the budget 14.0 GiB (encode 8.0 GiB, denoise 22.6 GiB, vae 11.1 GiB); pass fit_check=False to run anyway", + "footprint_peak_bytes": 556565800, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 556565800, + "mlx_peak_bytes": 5229238, + "elapsed_seconds": 3.503 +} diff --git a/bench/results/repeats/ernie-image-turbo-1024-tier24/run-1.json b/bench/results/repeats/ernie-image-turbo-1024-tier24/run-1.json new file mode 100644 index 0000000..d2a0fcd --- /dev/null +++ b/bench/results/repeats/ernie-image-turbo-1024-tier24/run-1.json @@ -0,0 +1,131 @@ +{ + "exit_code": 0, + "output": "t.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "ernie-image-turbo", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-Turbo-DF11/snapshots/27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "base_path": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image-Turbo/snapshots/bc68c81e2a1730a394d5fc9fae70713dee940140", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true, + "budget_bytes": 15032385536 + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 24, + "label": "CAPPED", + "watchdog_ceiling_bytes": 15569256448, + "limits": { + "tier": { + "tier_gb": 24, + "ram_bytes": 25769803776, + "recommended_bytes": 17179869184, + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536, + "reserve_bytes": 1610612736, + "budget_bytes": 17179869184, + "ceiling_bytes": 15569256448, + "fit_budget_bytes": 15032385536, + "budget_source": "ratio", + "is_host": false, + "label": "CAPPED" + }, + "effective": { + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536 + }, + "applied": "tier-caps" + }, + "family": "ernie", + "model": "ernie-image-turbo", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-Turbo-DF11/snapshots/27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "repo_id": null, + "revision": "27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "config_source": "config.json" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image-Turbo/snapshots/bc68c81e2a1730a394d5fc9fae70713dee940140", + "repo_id": null, + "revision": "bc68c81e2a1730a394d5fc9fae70713dee940140" + }, + "eval_policy": "per-block", + "eval_policy_status": "measured", + "cfg_batch": 1, + "text_tokens": 26, + "predict": "uncompiled", + "nonblock_groups": [ + "adaLN_modulation.1", + "final_norm.linear", + "time_embedding" + ], + "cache_limit_in_force": 2600000000, + "drop_set_before_vae": true, + "fit": { + "label": "predicted", + "phases": { + "encode": 7271583948, + "denoise": 14974035520, + "vae": 7932725729 + }, + "peak_phase": "denoise", + "peak_bytes": 14974035520, + "budget_bytes": 15032385536, + "fits": true + }, + "sizes": { + "compressed": 10894810183, + "extras": 27961600, + "nonblock": 335544320, + "encoders": 6625216512, + "vae": 168120878 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 642598376, + "mlx_peak": 6696780044, + "footprint_end": 6919297120, + "active_at_start": 27961748 + }, + "set_load": { + "footprint_start": 6919297120, + "mlx_peak": 11259596372, + "footprint_end": 11991426792, + "active_at_start": 28125580 + }, + "denoise": { + "footprint_start": 11991426792, + "mlx_peak": 13399152656, + "footprint_end": 14541514352, + "active_at_start": 11026605964 + }, + "vae": { + "footprint_start": 5940154328, + "mlx_peak": 7428292746, + "footprint_end": 2237614144, + "active_at_start": 32365064 + } + }, + "decode_launches": 288, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 21.66897004103521, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "footprint_peak_bytes": 15069554288, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 15069554288, + "mlx_peak_bytes": 14273543200, + "elapsed_seconds": 169.161 +} diff --git a/bench/results/tiers/aborts/krea-2-1024-tier16.json b/bench/results/tiers/aborts/krea-2-1024-tier16.json new file mode 100644 index 0000000..866fd3c --- /dev/null +++ b/bench/results/tiers/aborts/krea-2-1024-tier16.json @@ -0,0 +1,30 @@ +{ + "reason": "memory", + "footprint": 10197902072, + "peak_footprint": 10197902072, + "ceiling": 9842633386, + "elapsed": 17.100785250018816, + "budget": 1800.0, + "rss": 9530130432, + "mlx_active": 10083739684, + "mlx_cache": 0, + "verdict_memory": 10197902072, + "verdict_counter": "footprint", + "peak_watched": 10197902072, + "peak_mlx": 10083739684, + "context": { + "model": "krea-2", + "height": 1024, + "width": 1024, + "seed": 42, + "steps": 8, + "tier_gb": 16, + "label": "CAPPED", + "limits": { + "memory": 10737418240, + "cache": 10737418240, + "wired": 8589934592 + }, + "phase": "set_load" + } +} diff --git a/bench/results/tiers/aborts/krea-2-1024-tier24.json b/bench/results/tiers/aborts/krea-2-1024-tier24.json new file mode 100644 index 0000000..231ff89 --- /dev/null +++ b/bench/results/tiers/aborts/krea-2-1024-tier24.json @@ -0,0 +1,30 @@ +{ + "reason": "memory", + "footprint": 16158880632, + "peak_footprint": 16158880632, + "ceiling": 15569256448, + "elapsed": 29.714266665978357, + "budget": 1800.0, + "rss": 1557839872, + "mlx_active": 16048234020, + "mlx_cache": 0, + "verdict_memory": 16158880632, + "verdict_counter": "footprint", + "peak_watched": 16158880632, + "peak_mlx": 16048234020, + "context": { + "model": "krea-2", + "height": 1024, + "width": 1024, + "seed": 42, + "steps": 8, + "tier_gb": 24, + "label": "CAPPED", + "limits": { + "memory": 17179869184, + "cache": 16320875724, + "wired": 15032385536 + }, + "phase": "set_load" + } +} diff --git a/bench/results/tiers/aborts/krea-2-raw-1024-tier16.json b/bench/results/tiers/aborts/krea-2-raw-1024-tier16.json new file mode 100644 index 0000000..d2105c5 --- /dev/null +++ b/bench/results/tiers/aborts/krea-2-raw-1024-tier16.json @@ -0,0 +1,30 @@ +{ + "reason": "memory", + "footprint": 10187629328, + "peak_footprint": 10187629328, + "ceiling": 9842633386, + "elapsed": 17.675457959005143, + "budget": 1800.0, + "rss": 10326802432, + "mlx_active": 10084067876, + "mlx_cache": 0, + "verdict_memory": 10187629328, + "verdict_counter": "footprint", + "peak_watched": 10187629328, + "peak_mlx": 10084067876, + "context": { + "model": "krea-2-raw", + "height": 1024, + "width": 1024, + "seed": 42, + "steps": 25, + "tier_gb": 16, + "label": "CAPPED", + "limits": { + "memory": 10737418240, + "cache": 10737418240, + "wired": 8589934592 + }, + "phase": "set_load" + } +} diff --git a/bench/results/tiers/aborts/krea-2-raw-1024-tier24.json b/bench/results/tiers/aborts/krea-2-raw-1024-tier24.json new file mode 100644 index 0000000..61eda83 --- /dev/null +++ b/bench/results/tiers/aborts/krea-2-raw-1024-tier24.json @@ -0,0 +1,30 @@ +{ + "reason": "memory", + "footprint": 16013292168, + "peak_footprint": 16013292168, + "ceiling": 15569256448, + "elapsed": 32.002064000000246, + "budget": 1800.0, + "rss": 2334359552, + "mlx_active": 15468651044, + "mlx_cache": 0, + "verdict_memory": 16013292168, + "verdict_counter": "footprint", + "peak_watched": 16013292168, + "peak_mlx": 15468651044, + "context": { + "model": "krea-2-raw", + "height": 1024, + "width": 1024, + "seed": 42, + "steps": 25, + "tier_gb": 24, + "label": "CAPPED", + "limits": { + "memory": 17179869184, + "cache": 16320875724, + "wired": 15032385536 + }, + "phase": "set_load" + } +} diff --git a/bench/results/tiers/krea-2-1024.json b/bench/results/tiers/krea-2-1024.json new file mode 100644 index 0000000..1392762 --- /dev/null +++ b/bench/results/tiers/krea-2-1024.json @@ -0,0 +1,139 @@ +{ + "exit_code": 0, + "output": "krea-2.png", + "memory_caps_gb": [ + 20, + 22 + ], + "model_kwargs": { + "model": "krea-2", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Turbo-DF11-ComfyUI/snapshots/978da5fb7647bd222d33125993abd8fdc2840cfc", + "base_path": "~/.cache/huggingface/hub/models--krea--Krea-2-Turbo/snapshots/98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 32, + "label": "MEASURED", + "watchdog_ceiling_bytes": 30064771072, + "limits": { + "tier": { + "tier_gb": 32, + "ram_bytes": 34359738368, + "recommended_bytes": 26800603136, + "memory_limit_bytes": 32641751449, + "cache_limit_bytes": 25460572979, + "wired_limit_bytes": 0, + "reserve_bytes": 2147483648, + "budget_bytes": 26800603136, + "ceiling_bytes": 24653119488, + "fit_budget_bytes": 24653119488, + "budget_source": "device", + "is_host": true, + "label": "MEASURED" + }, + "effective": { + "memory_limit_bytes": 23622320128, + "cache_limit_bytes": 32641751449, + "wired_limit_bytes": 21474836480 + }, + "applied": "host-caps" + }, + "family": "krea2", + "model": "krea-2", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Turbo-DF11-ComfyUI/snapshots/978da5fb7647bd222d33125993abd8fdc2840cfc", + "repo_id": null, + "revision": "978da5fb7647bd222d33125993abd8fdc2840cfc", + "config_source": "header and spot checks match layout krea-2-turbo-comfyui (mingyi456/Krea-2-Turbo-DF11-ComfyUI@978da5fb7647bd222d33125993abd8fdc2840cfc)" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--krea--Krea-2-Turbo/snapshots/98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7", + "repo_id": null, + "revision": "98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7" + }, + "layout": "header and spot checks match layout krea-2-turbo-comfyui (mingyi456/Krea-2-Turbo-DF11-ComfyUI@978da5fb7647bd222d33125993abd8fdc2840cfc)", + "nonblock": "per-call", + "eval_policy": "per-block", + "eval_policy_status": "measured", + "cfg_calls_per_step": 1, + "text_tokens": 30, + "predict": "uncompiled", + "nonblock_groups": [ + "tmlp", + "tproj", + "txtfusion.layerwise_blocks.0", + "txtfusion.layerwise_blocks.1", + "txtfusion.refiner_blocks.0", + "txtfusion.refiner_blocks.1", + "txtmlp" + ], + "cache_limit_in_force": 4500000000, + "drop_set_before_vae": true, + "fit": { + "label": "predicted", + "phases": { + "encode": 8595823360, + "denoise": 24242523532, + "vae": 11869477608 + }, + "peak_phase": "denoise", + "peak_bytes": 24242523532, + "budget_bytes": 24653119488, + "fits": true + }, + "sizes": { + "compressed": 17540075298, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 562808056, + "mlx_peak": 7955707996, + "footprint_end": 7916919552, + "active_at_start": 4823348 + }, + "set_load": { + "footprint_start": 7916919552, + "mlx_peak": 17549525844, + "footprint_end": 18120767600, + "active_at_start": 8754468 + }, + "denoise": { + "footprint_start": 18120767600, + "mlx_peak": 21242923578, + "footprint_end": 22736402568, + "active_at_start": 17549525844 + }, + "vae": { + "footprint_start": 10848966560, + "mlx_peak": 9583229094, + "footprint_end": 4962453400, + "active_at_start": 11901082 + } + }, + "decode_launches": 224, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 30.84187958302209, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "footprint_peak_bytes": 22736402568, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 22736402568, + "mlx_peak_bytes": 22158626378, + "elapsed_seconds": 305.676 +} diff --git a/bench/results/tiers/krea-2-raw-1024.json b/bench/results/tiers/krea-2-raw-1024.json new file mode 100644 index 0000000..b69cdbd --- /dev/null +++ b/bench/results/tiers/krea-2-raw-1024.json @@ -0,0 +1,139 @@ +{ + "exit_code": 0, + "output": "krea-2-raw.png", + "memory_caps_gb": [ + 20, + 22 + ], + "model_kwargs": { + "model": "krea-2-raw", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Raw-DF11-ComfyUI/snapshots/8320616b25ac9340a830a7fb21f1b0237e160e66", + "base_path": "~/.cache/huggingface/hub/models--krea--Krea-2-Raw/snapshots/6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 32, + "label": "MEASURED", + "watchdog_ceiling_bytes": 30064771072, + "limits": { + "tier": { + "tier_gb": 32, + "ram_bytes": 34359738368, + "recommended_bytes": 26800603136, + "memory_limit_bytes": 32641751449, + "cache_limit_bytes": 25460572979, + "wired_limit_bytes": 0, + "reserve_bytes": 2147483648, + "budget_bytes": 26800603136, + "ceiling_bytes": 24653119488, + "fit_budget_bytes": 24653119488, + "budget_source": "device", + "is_host": true, + "label": "MEASURED" + }, + "effective": { + "memory_limit_bytes": 23622320128, + "cache_limit_bytes": 32641751449, + "wired_limit_bytes": 21474836480 + }, + "applied": "host-caps" + }, + "family": "krea2", + "model": "krea-2-raw", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--Krea-2-Raw-DF11-ComfyUI/snapshots/8320616b25ac9340a830a7fb21f1b0237e160e66", + "repo_id": null, + "revision": "8320616b25ac9340a830a7fb21f1b0237e160e66", + "config_source": "header and spot checks match layout krea-2-raw-comfyui (mingyi456/Krea-2-Raw-DF11-ComfyUI@8320616b25ac9340a830a7fb21f1b0237e160e66)" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--krea--Krea-2-Raw/snapshots/6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + "repo_id": null, + "revision": "6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603" + }, + "layout": "header and spot checks match layout krea-2-raw-comfyui (mingyi456/Krea-2-Raw-DF11-ComfyUI@8320616b25ac9340a830a7fb21f1b0237e160e66)", + "nonblock": "per-call", + "eval_policy": "per-block", + "eval_policy_status": "measured", + "cfg_calls_per_step": 1, + "text_tokens": 30, + "predict": "uncompiled", + "nonblock_groups": [ + "tmlp", + "tproj", + "txtfusion.layerwise_blocks.0", + "txtfusion.layerwise_blocks.1", + "txtfusion.refiner_blocks.0", + "txtfusion.refiner_blocks.1", + "txtmlp" + ], + "cache_limit_in_force": 4500000000, + "drop_set_before_vae": true, + "fit": { + "label": "predicted", + "phases": { + "encode": 8595823360, + "denoise": 24243378794, + "vae": 11869477608 + }, + "peak_phase": "denoise", + "peak_bytes": 24243378794, + "budget_bytes": 24653119488, + "fits": true + }, + "sizes": { + "compressed": 17540930560, + "extras": 4560024, + "nonblock": 0, + "encoders": 7843074560, + "vae": 507591892 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 540116240, + "mlx_peak": 7958029668, + "footprint_end": 7977638680, + "active_at_start": 4822320 + }, + "set_load": { + "footprint_start": 7977638680, + "mlx_peak": 17550395732, + "footprint_end": 18098124720, + "active_at_start": 8754468 + }, + "denoise": { + "footprint_start": 18098124720, + "mlx_peak": 21243809918, + "footprint_end": 22713743304, + "active_at_start": 17550395732 + }, + "vae": { + "footprint_start": 8471728208, + "mlx_peak": 9583229162, + "footprint_end": 873498448, + "active_at_start": 11901150 + } + }, + "decode_launches": 700, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 29.83431070804363, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "footprint_peak_bytes": 22713743304, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 22713743304, + "mlx_peak_bytes": 22159496334, + "elapsed_seconds": 844.707 +} diff --git a/bench/results/validation/control-vs-bf16-schnell-1024.json b/bench/results/validation/control-vs-bf16-schnell-1024.json new file mode 100644 index 0000000..2f31274 --- /dev/null +++ b/bench/results/validation/control-vs-bf16-schnell-1024.json @@ -0,0 +1,141 @@ +{ + "rounds": { + "1": { + "T_bf16": 3.684958625002764, + "T_control": 3.693141083000228, + "T_df11": 3.992285207961686 + }, + "2": { + "T_bf16": 3.6915409169741906, + "T_control": 3.6984091250342317, + "T_df11": 4.014075957995374 + }, + "3": { + "T_bf16": 3.688040791952517, + "T_control": 3.6900886660441756, + "T_df11": 3.982614917040337 + } + }, + "complete_rounds": [ + 1, + 2, + 3 + ], + "control_vs_bf16_per_round": { + "1": 0.002215582295279224, + "2": 0.0018570709264020377, + "3": 0.000554966093498802 + }, + "pooled": { + "bf16": { + "median": 3.688304833951406, + "spread": 0.009334888941151141, + "n": 15, + "verify_median_s": 5.249981768429279e-06 + }, + "control": { + "median": 3.693141083000228, + "spread": 0.01274607223277016, + "n": 15, + "verify_median_s": 5.416979547590017e-06 + }, + "df11": { + "median": 3.994620500016026, + "spread": 0.013406884831935348, + "n": 15, + "verify_median_s": 0.0002883330453187227 + } + }, + "control_vs_bf16": 0.001309521878566514, + "spread_bound": 0.01274607223277016, + "within_spread": true, + "df11_vs_bf16": 0.08305052859105833, + "stopped": null, + "key": { + "model": "schnell", + "size": 1024, + "steps": 5, + "warmup": 2, + "seed": 42, + "df11": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", + "embeds": "schnell-embeds.safetensors", + "embeds_meta": { + "token_length": "256", + "synthetic": "false", + "seed": "42", + "mflux": "0.20.0", + "root": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "model": "schnell", + "mlx": "0.32.2" + }, + "source": "9b4803f1ff2e69eb3faa00ca975d9a935f25fa6c62615050d986e3f6a0688568", + "mlx": "0.32.2", + "double": 4, + "single": 8 + }, + "rounds_requested": 3, + "runs_complete": [ + [ + 1, + "bf16" + ], + [ + 1, + "control" + ], + [ + 1, + "df11" + ], + [ + 2, + "bf16" + ], + [ + 2, + "control" + ], + [ + 2, + "df11" + ], + [ + 3, + "bf16" + ], + [ + 3, + "control" + ], + [ + 3, + "df11" + ] + ], + "provenance": { + "git": "7c5b410296bc01295fba93fd4ecea187609d6456", + "source_hash": "9b4803f1ff2e69eb3faa00ca975d9a935f25fa6c62615050d986e3f6a0688568", + "mlx": "0.32.2", + "mflux": "0.20.0", + "macos": "27.0", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 32641751449, + "cache_memory": 0, + "power": { + "ac": true, + "battery_pct": 100 + } + } +} diff --git a/pyproject.toml b/pyproject.toml index 1e005f3..af77653 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "mlx-dfloat" -description = "Run DFloat11 losslessly compressed BF16 checkpoints on Apple Silicon with MLX and a Metal decode kernel; today that runs FLUX.1 on a 32 GB Mac with weights about 30% smaller" +description = "Run DFloat11 losslessly compressed BF16 checkpoints on Apple Silicon with MLX and a Metal decode kernel: image generation with FLUX.1, Z-Image, FLUX.2 Klein, Qwen-Image-2.1, ERNIE-Image and Krea 2 on a 32 GB Mac, with weights about 30% smaller" readme = "README.md" requires-python = ">=3.11" license = "Apache-2.0" @@ -167,8 +167,6 @@ omit = [ # tests/test_q8_rig.py, the real build by the slow test. "scripts/_q8_rig.py", # every test here is @pytest.mark.mflux; measured by the integration lane instead, see ci.yml. - # init.py stays OUT of this list: most of it has real offline (non-mflux) tests in - # tests/test_flux1_init.py, so it must keep gating the required 85% threshold too. "src/mlx_dfloat/mflux/flux1/model.py", # the same for Z-Image's model class: every test is @pytest.mark.mflux, gated in the integration lane. "src/mlx_dfloat/mflux/zimage/model.py", @@ -178,6 +176,22 @@ omit = [ "src/mlx_dfloat/mflux/qwen21/model.py", # the same for ERNIE-Image's model class: every test is @pytest.mark.mflux, gated in the integration lane. "src/mlx_dfloat/mflux/ernie/model.py", + # the same for Krea 2's model class: every test is @pytest.mark.mflux, gated in the integration lane. + "src/mlx_dfloat/mflux/krea2/model.py", + # Each family's base loaders (init.py) and, where its tests need mflux's own transformer class, transformer.py: + # without mflux their tests reach only 40-60 %, so the integration lane measures and gates each at 80 % on its + # own (ci.yml). names.py and memory.py run without mflux and stay in this gate, and so do flux1/transformer.py + # and zimage/transformer.py, which the seam fakes cover past 90 % offline. + "src/mlx_dfloat/mflux/flux1/init.py", + "src/mlx_dfloat/mflux/zimage/init.py", + "src/mlx_dfloat/mflux/flux2/init.py", + "src/mlx_dfloat/mflux/flux2/transformer.py", + "src/mlx_dfloat/mflux/qwen21/init.py", + "src/mlx_dfloat/mflux/qwen21/transformer.py", + "src/mlx_dfloat/mflux/ernie/init.py", + "src/mlx_dfloat/mflux/ernie/transformer.py", + "src/mlx_dfloat/mflux/krea2/init.py", + "src/mlx_dfloat/mflux/krea2/transformer.py", ] [tool.coverage.report] diff --git a/scripts/verify_image.py b/scripts/verify_image.py index 49799b7..98b02c0 100644 --- a/scripts/verify_image.py +++ b/scripts/verify_image.py @@ -1,7 +1,7 @@ """Image identity: DFloat11 latents against the same seam streaming the BF16 weights, for every BF16 family. -The ``df11`` side generates through ``DFloatFlux1``, ``DFloatZImage``, ``DFloatFlux2Klein``, ``DFloatQwenImage21`` or -``DFloatErnieImage`` +The ``df11`` side generates through ``DFloatFlux1``, ``DFloatZImage``, ``DFloatFlux2Klein``, ``DFloatQwenImage21``, +``DFloatErnieImage`` or ``DFloatKrea2`` (the prompt encoded by its own encoders; the final latents captured by an after-loop callback and saved with the embeddings and the image). The ``bf16`` side builds the same seamed transformer with its extras read from the base's BF16 transformer shards, drives it with a ``StreamingBF16Provider`` (each block's matrices read from @@ -22,6 +22,11 @@ sense as Klein's. ERNIE-Image is checked the same way (``--base`` holds its two-shard BF16 transformer for the bf16 side; both sides take mflux's step function uncompiled); its bf16 loop body is measurement glue too, so a drift fails as a false mismatch, not a false pass. ERNIE-Image-Turbo runs at guidance 1.0 only and takes no ``--negative-prompt``. +Krea 2 Raw's bf16 side reads the native single file ``raw.safetensors`` at the base's root: its BF16 block matrices +stream as the blocks run, and its FP32 extras and five non-block matrices are cast to BF16 as mflux casts them at load +(``astype``, nearest even); nothing of the DF11 file is read there but its group shapes. Both Krea sides run the step +function uncompiled; with guidance other than 1.0 each step is two batch-1 transformer calls (mflux's " " negative when +none is given). Krea 2 Turbo (``turbo.safetensors``) is checked at guidance 1.0 only. Usage (from the repository root of a synced checkout, ``--group bench``): uv run python -m scripts.verify_image --orchestrate --model schnell --df11 DIR --base DIR --out DIR \ [--prompt "..."] [--seed 42] [--steps 4] [--size 1024] [--guidance 3.5] [--negative-prompt "..."] [--eval-policy per-block] @@ -180,6 +185,28 @@ def ernie_base_problem(model: str, base: Path) -> str | None: return _sharded_transformer_problem(e, base) +# The native single-file transformer each Krea 2 base holds at its root. +KREA2_NATIVE_FILE = {"krea-2-raw": "raw.safetensors", "krea-2": "turbo.safetensors"} + + +def krea2_base_problem(model: str, base: Path) -> str | None: + """Why ``base`` cannot feed a Krea 2 bf16 side, or None (and None for every other family). + + The bf16 side reads the model's own native transformer file at the base's root (``raw.safetensors`` for Raw, + ``turbo.safetensors`` for Turbo); it must be a regular file once links resolve. The df11 side never reads it. + """ + e = entry(model) + if e.family != "krea2": + return None + native = KREA2_NATIVE_FILE[model] + if (base / native).is_file(): + return None + return ( + f"--base {base}: the bf16 side streams the transformer from {native} at the base's root (download " + f"{e.base_repo}'s {native} first)" + ) + + def parse_args(argv: list[str] | None = None) -> argparse.Namespace: """Parse the command line: one side (``df11``, ``bf16`` or ``compare``), or ``--orchestrate``.""" p = argparse.ArgumentParser( @@ -204,12 +231,14 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: default=None, help="guidance (default per model, as mflux: FLUX.1 3.5; Z-Image unset: the model's own rule; " "FLUX.2 Klein 1.0, base models with CFG at 4; Qwen-Image 2.1 1.0, CFG above 1 with --negative-prompt; " - "ERNIE-Image 4.0, CFG above 1; ERNIE-Image-Turbo 1.0 only)", + "ERNIE-Image 4.0, CFG above 1; ERNIE-Image-Turbo 1.0 only; Krea 2 Raw 1.0, CFG at any other value, as " + "two calls per step; Krea 2 Turbo 1.0 only in this check)", ) p.add_argument( "--negative-prompt", default=None, - help="the negative prompt (used by z-image, a base model, qwen-image-2.1 and ernie-image above guidance 1)", + help="the negative prompt (used by z-image, a base model, qwen-image-2.1 and ernie-image above guidance 1, " + "and by krea-2-raw at any guidance other than 1)", ) p.add_argument("--eval-policy", choices=("per-block", "depth2"), default="per-block") p.add_argument( @@ -232,6 +261,13 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: p.error( f"--guidance: {args.model} runs at guidance {e.fixed_guidance} only, as its mflux command does" ) + # Krea 2 Turbo is distilled: its recipe is guidance 1.0 (mflux's default), so the check covers that path only. + if args.model == "krea-2" and args.guidance is not None and args.guidance != 1.0: + p.error("--guidance: krea-2 runs at guidance 1.0 only in this check (its distilled recipe)") + if args.model == "krea-2" and args.negative_prompt is not None: + p.error( + "--negative-prompt: krea-2 runs at guidance 1.0 only here; there is no negative branch" + ) if entry(args.model).family == "flux2" and args.negative_prompt is not None: p.error( f"--negative-prompt: {args.model} takes none (mflux's FLUX.2 Klein encodes its own blank negative)" @@ -242,7 +278,11 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: # Only the runs that stream the BF16 transformer need it; the df11 side and compare never read it. if args.orchestrate or args.side == "bf16": base = Path(args.base).expanduser() - problem = qwen21_base_problem(args.model, base) or ernie_base_problem(args.model, base) + problem = ( + qwen21_base_problem(args.model, base) + or ernie_base_problem(args.model, base) + or krea2_base_problem(args.model, base) + ) if problem is not None: p.error(f"--model {args.model}: {problem}") if args.guidance is None: # resolved before run_key, so a stored FLUX.1 side keeps its 3.5 @@ -1215,12 +1255,172 @@ def run_bf16_ernie(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, An } +def run_df11_krea2(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, Any]: + """Generate through ``DFloatKrea2``, capture the final latents, and save the image, latents and embeddings. + + The embeddings file holds the prompt's embeddings (``embeds``) and, when classifier-free guidance ran (two + batch-1 transformer calls per step), the negative prompt's (``neg_embeds``), as the model encoded them. + + Raises: + VerifyImageError: The after-loop callback never ran (mflux's own loop never called it). + """ + from mflux.models.krea2.latent_creator.krea2_latent_creator import Krea2LatentCreator + + from mlx_dfloat.mflux.krea2.model import DFloatKrea2 + + out_dir = args.out / "df11" + out_dir.mkdir(parents=True, exist_ok=True) + model = DFloatKrea2( + args.model, df11_path=args.df11, base_path=args.base, eval_policy=args.eval_policy + ) + capture = _LatentCapture() + model.callbacks.register(capture) + image = model.generate_image( + args.seed, + args.prompt, + num_inference_steps=args.steps, + height=args.size, + width=args.size, + guidance=args.guidance, + negative_prompt=args.negative_prompt, + ) + captured = capture.latents + if captured is None: + raise VerifyImageError("the after-loop callback never ran; no latents were captured") + mx.eval(captured) + image.save(str(out_dir / "image.png"), export_json_metadata=False, overwrite=True) + mx.save_safetensors(str(out_dir / "latents.safetensors"), {"latents": captured}) + prompts = model.cfg_prompts( + args.prompt, negative_prompt=args.negative_prompt, guidance=args.guidance + ) + embeds = {"embeds": model._embeds[prompts[0]][0]} + if len(prompts) == 2: + embeds["neg_embeds"] = model._embeds[prompts[1]][0] + mx.save_safetensors(str(out_dir / "embeds.safetensors"), embeds) + noise = Krea2LatentCreator.create_noise(args.seed, args.size, args.size) + degenerate = nondegenerate(captured, noise) + report = model.report() + return { + "exit_code": EXIT_OK if not degenerate else EXIT_ERROR, + "degenerate": degenerate, + "report": report, + "output_shape": list(captured.shape), + "output_dtype": str(captured.dtype), + "footprint_peak_bytes": max(watchdog.peak_footprint, phys_footprint()), + "mlx_peak_bytes_max_over_phases": max_phase_peak(report["peaks"]), + } + + +def run_bf16_krea2(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, Any]: + """Stream the base's native file through the seam and mflux's Krea 2 loop body, then the VAE with the set absent. + + The transformer is built with the DF11 file's extras and placeholders, then every extra and every non-block + matrix is replaced by the base's (FP32 cast to BF16 as mflux casts it), so this side reads nothing of the DF11 + file but its group shapes; the 28 blocks stream from the base as they run. The step function is mflux's own + ``Krea2._predict``, taken uncompiled as on the df11 side, with the er_sde stepper at the run's seed. + + Raises: + VerifyImageError: The df11 side's saved embeddings are not present yet. + """ + import gc + + from mflux.models.common.config.config import Config + from mflux.models.common.config.model_config import ModelConfig + from mflux.models.common.vae.vae_util import VAEUtil + from mflux.models.krea2.latent_creator.krea2_latent_creator import Krea2LatentCreator + from mflux.models.krea2.model.krea2_sampler import Krea2Sampler + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + from mflux.utils.image_util import ImageUtil + + from mlx_dfloat._safetensors import read_header + from mlx_dfloat.mflux._compile import uncompiled + from mlx_dfloat.mflux.krea2 import init as kinit + from mlx_dfloat.mflux.krea2 import transformer as ktf + from mlx_dfloat.mflux.krea2.names import krea2_name_map + + out_dir = args.out / "bf16" + out_dir.mkdir(parents=True, exist_ok=True) + embeds_path = args.out / "df11" / "embeds.safetensors" + if not embeds_path.exists(): + raise VerifyImageError(f"{embeds_path}: run the df11 side first") + + df11_root = Path(args.df11).expanduser() + base_root = Path(args.base).expanduser() + native = (base_root / KREA2_NATIVE_FILE[args.model]).resolve() + names = krea2_name_map() + ckpt = open_checkpoint(df11_root) + build = ktf.build_transformer(ckpt, model=args.model) + replaced = ktf.install_base_weights(build.transformer, ckpt, native, names, build.nonblock) + index = {name: (native, info) for name, info in read_header(native).items()} + provider = StreamingBF16Provider( + index, {n: ckpt.groups[n].matrix_names for n in build.shapes}, names + ) + transformer = build.transformer + transformer.attach(provider, build.shapes, eval_policy=args.eval_policy, verify_in_call=True) + + saved = mx.load(str(embeds_path)) + embeds, neg_embeds = saved["embeds"], saved.get("neg_embeds") + scheduler = Krea2._resolve_scheduler( + None + ) # mflux's default for Krea 2, as the df11 side ran it + config = Config( + model_config=ModelConfig.from_name(model_name=args.model, base_model=None), + num_inference_steps=args.steps, + height=args.size, + width=args.size, + guidance=args.guidance, + scheduler=scheduler, + ) + sigmas = config.scheduler.sigmas + # mflux krea2.py generate_image: float32 noise, the embeddings evaluated before the loop. + latents = Krea2LatentCreator.create_noise(args.seed, args.size, args.size) + mx.eval(latents, embeds) + if neg_embeds is not None: + mx.eval(neg_embeds) + stepper = Krea2Sampler.make_stepper(scheduler, sigmas, args.seed) + predict = uncompiled(Krea2._predict, transformer, embeds, neg_embeds, args.guidance) + for t in config.time_steps: + # mflux krea2.py's loop body: predict (two batch-1 calls when CFG runs), denoised, the stepper, eval. + v = predict(latents=latents, timestep=sigmas[t].reshape(1)) + denoised = latents - sigmas[t] * v + latents = stepper.step(t, latents, v, denoised) + mx.eval(latents) + transformer.verify_step() + + mx.save_safetensors(str(out_dir / "latents.safetensors"), {"latents": latents}) + + provider_reads = provider.reads + # The VAE decodes with the set absent, as the model's does at 1024² on a 32 GB Mac. + transformer.detach() + del predict, transformer, build, provider, index + gc.collect() + mx.clear_cache() + mx.set_cache_limit(0) + vae = kinit.load_vae(base_root) + decoded = VAEUtil.decode(vae=vae, latent=latents, tiling_config=None) + ImageUtil.to_pil(decoded).save( + str(out_dir / "image.png") + ) # ImageUtil.to_image's own conversion + + return { + "exit_code": EXIT_OK, + "reads": provider_reads, + "replaced_from_base": len(replaced), + "calls_per_step": 1 if neg_embeds is None else 2, + "output_shape": list(latents.shape), + "output_dtype": str(latents.dtype), + "footprint_peak_bytes": max(watchdog.peak_footprint, phys_footprint()), + "mlx_peak_bytes_process": int(mx.get_peak_memory()), # never reset on this side + } + + SIDE_RUNNERS: dict[str, tuple[Callable[..., dict[str, Any]], Callable[..., dict[str, Any]]]] = { "flux1": (run_df11, run_bf16), "zimage": (run_df11_zimage, run_bf16_zimage), "flux2": (run_df11_flux2, run_bf16_flux2), "qwen21": (run_df11_qwen21, run_bf16_qwen21), "ernie": (run_df11_ernie, run_bf16_ernie), + "krea2": (run_df11_krea2, run_bf16_krea2), } diff --git a/scripts/verify_remote_group.py b/scripts/verify_remote_group.py index d883b5e..dca5515 100644 --- a/scripts/verify_remote_group.py +++ b/scripts/verify_remote_group.py @@ -12,7 +12,9 @@ A repository without a ``config.json`` (a single-file ComfyUI export) is read through a pinned layout: its header sha256 and spot checks of its stored bytes, fetched by range reads, must match one; a fused stored matrix is then cut into its row blocks and each block compared with the original of its own name. -Exit codes: 0 all sampled matrices equal, 1 a mismatch, 2 an error, 70/71 watchdog abort. +``--extras`` also compares every uncompressed tensor (the extras) with its original of the same name, BF16 as is or, +with ``--cast-fp32-to-bf16``, FP32 rounded to BF16; ``--groups none --extras`` compares the extras alone. +Exit codes: 0 all sampled matrices (and extras) equal, 1 a mismatch, 2 an error, 70/71 watchdog abort. """ import argparse @@ -382,7 +384,13 @@ def _group_arrays(src: RangeSource, file: str, header: dict, group: str) -> Grou def pick_groups(stats: dict[str, dict], selector: str) -> list[str]: - """Resolve a selector list to group names (natural order; KeyError for unknown names).""" + """Resolve a selector list to group names (natural order; KeyError for unknown names). + + ``none`` on its own selects no group (an extras-only run); mixed with other selectors it is + refused as an unknown name. + """ + if selector == "none": + return [] ordered = sorted(stats, key=natural_key) picked: list[str] = [] for token in selector.split(","): @@ -474,6 +482,98 @@ def verify_group( return record +def verify_extras( + df11: RangeSource, + bf16: RangeSource, + *, + dindex: dict[str, tuple[str, dict]], + bindex: dict[str, str], + cast_fp32: bool = False, + require_any: bool = False, +) -> dict: + """Compare every DF11 extra (each tensor that is not a group field) with its original, by name. + + The DF11 files are the checkpoint's own: the group index's files plus the shards a + ``*.safetensors.index.json`` lists, or, without an index, every safetensors file at the top + (an extras-only shard holds no group); a file outside an index is not the checkpoint's. Each + header is parsed once. A DF11 extra must be BF16, with the same shape as its original, which is + read as is when BF16, or rounded to BF16 (nearest even) when FP32 and ``cast_fp32`` is set. + Extras are compared in name order. ``require_any`` (an extras-only run) makes a checkpoint + without extras an error. + """ + record: dict = {"compared": 0, "mismatched": 0, "extras": []} + try: + headers: dict[str, dict[str, TensorInfo]] = dict(dindex.values()) + names = df11.files() + indexed: set[str] = set() + for name in names: + if name.endswith(".safetensors.index.json"): + indexed |= set(json.loads(df11.text(name))["weight_map"].values()) + own = indexed or {n for n in names if n.endswith(".safetensors")} + for name in sorted(own): + if name not in headers: + length = _header_length(df11, name) + headers[name] = parse_header( + df11.read(name, 8, length), + data_start=8 + length, + file_size=df11.size(name), + source=name, + ) + fields = {f"{group}.{field}" for group in dindex for field in GROUP_FIELDS} + extras = sorted( + ( + (tensor, file, info) + for file, header in headers.items() + for tensor, info in header.items() + if tensor not in fields + ), + key=lambda item: item[0], + ) + if require_any and not extras: + raise VerifyError( + "the checkpoint holds no extras: --groups none --extras compares nothing" + ) + header_cache: dict[str, tuple[dict, int]] = {} + for tensor, file, info in extras: + if info.dtype != "BF16": + raise VerifyError(f"extra {tensor} is {info.dtype}, not BF16") + shard = bindex.get(tensor) + if shard is None: + raise VerifyError(f"no original for extra {tensor}") + if shard not in header_cache: + header_cache[shard] = independent_header(bf16, shard) + bheader, base = header_cache[shard] + start, length, _n, dtype = original_range( + bheader.get(tensor), tensor, cast_fp32=cast_fp32 + ) + original_shape = tuple(bheader[tensor]["shape"]) + if tuple(info.shape) != original_shape: + raise VerifyError( + f"extra {tensor}: shape {tuple(info.shape)}, original has {original_shape}" + ) + got = np.frombuffer(df11.read(file, info.offset, info.nbytes), " dict[str, Any]: + """A private copy of the Krea 2 config, so the two layouts share no mutable state.""" + copy: dict[str, Any] = json.loads(json.dumps(_KREA2_CONFIG)) + return copy + + +KREA_2_RAW_COMFYUI = SynthesizedLayout( + key="krea-2-raw-comfyui", + label="Krea 2 Raw (ComfyUI single-file DF11)", + repo_id="mingyi456/Krea-2-Raw-DF11-ComfyUI", + revision="8320616b25ac9340a830a7fb21f1b0237e160e66", + file_name="krea2_raw_bf16-DF11.safetensors", + header_sha256="a0eb86513ccf0ba8baca0625147b62109707650b480046cd66043ea9a8700578", + file_sha256="f62cf31921ad0377f9d6e5bd2e3191a1ac9950d92f487b26fa7c052f503d9b97", + groups=35, + extras=169, + raw_config=_krea2_config(), + row_splits={}, + probes=( + ContentProbe( # blocks.0.attn.wq + tensor="blocks.0.sign_mantissa", + offset=0, + length=4096, + sha256="363d5ac87fa64889de9ef810e04d0b1e905ceb89fee078fa0e589c254298e35f", + ), + ContentProbe( # blocks.0.attn.wk + tensor="blocks.0.sign_mantissa", + offset=37_748_736, + length=4096, + sha256="ca6dd074e2d1e37891e94e867dfacc6ddb5c19f878f78a60f56e1e2f6f5e06bd", + ), + ContentProbe( # blocks.0.attn.wv + tensor="blocks.0.sign_mantissa", + offset=47_185_920, + length=4096, + sha256="6480a299d38e684bdb3a8d498752ed62afda615de694aab4b69a16a974ba518c", + ), + ContentProbe( # blocks.0.attn.gate + tensor="blocks.0.sign_mantissa", + offset=56_623_104, + length=4096, + sha256="3962f4172d9e4007c709729c5d262c1105748e472c17d714ef41ced9fd2aa40f", + ), + ContentProbe( # blocks.0.attn.wo + tensor="blocks.0.sign_mantissa", + offset=94_371_840, + length=4096, + sha256="585411a597165de162951bdd1d68cef4fb03a434caf433d33d3f60e7625d1ffe", + ), + ContentProbe( # blocks.0.mlp.gate + tensor="blocks.0.sign_mantissa", + offset=132_120_576, + length=4096, + sha256="2d41ddb871e747add117545678a7097e19e9901744434de7c54b885cfa049b10", + ), + ContentProbe( # blocks.0.mlp.up + tensor="blocks.0.sign_mantissa", + offset=232_783_872, + length=4096, + sha256="01bfce386ced98d20b1bceec8ccf51ec586036c730125c9c3ef4ea76764d19a0", + ), + ContentProbe( # blocks.0.mlp.down + tensor="blocks.0.sign_mantissa", + offset=333_447_168, + length=4096, + sha256="9489cb4060011bc2edca69148499606df538751f961e35dc7ba3c2d90954b3b5", + ), + ContentProbe( # blocks.27.attn.wq + tensor="blocks.27.sign_mantissa", + offset=0, + length=4096, + sha256="d011e6504676d98bde3696da028b03d8f6cbc7d4ad5ece7e67afda38c2cc22ad", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.attn.wq + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=0, + length=4096, + sha256="edcc578d390f9a375a43660de265524fba96f8cdfb7422c2260fe7ae3c81bf95", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.attn.wk + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=6_553_600, + length=4096, + sha256="c2f20aaf9f3e23b66e4d74e26a4239e75b2409fb4baffb875344084ae8cd1920", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.attn.wv + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=13_107_200, + length=4096, + sha256="cf78e88a47707eab0acab1561a31b3c7fc75174d7ee1c872bb1a5558cf2c3d6a", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.attn.gate + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=19_660_800, + length=4096, + sha256="92fedf37baba3c13617ec9bacb6017aafacf5e61996791edb92720ebd1223cb8", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.attn.wo + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=26_214_400, + length=4096, + sha256="c74085fa9516ac66243f038846454e9f4911b5ee54891e3d6e634b27cbdacb49", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.mlp.gate + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=32_768_000, + length=4096, + sha256="9b765f28f96af426c2c3ca7d48d08c11b38aab7b02c010dde4935a45fd421d55", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.mlp.up + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=50_462_720, + length=4096, + sha256="6c12c5521469efccaea92c86cf1232bf69cd7097033c23c77ea36d6899141a28", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.mlp.down + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=68_157_440, + length=4096, + sha256="5998851cd52398972b745d340839d779e5f45ba4189c08b2cd3d0493fa36e393", + ), + ContentProbe( # txtfusion.layerwise_blocks.1.attn.wq + tensor="txtfusion.layerwise_blocks.1.sign_mantissa", + offset=0, + length=4096, + sha256="4b0fd5b9e0ea3356bc8d71b3ef3008bd4cdee28e485c89c67188d533604885db", + ), + ContentProbe( # txtfusion.refiner_blocks.0.attn.wq + tensor="txtfusion.refiner_blocks.0.sign_mantissa", + offset=0, + length=4096, + sha256="7470ae9e5215bb3bd6f2dc75620cc6935e4d7a957ac05358838f6b851a01dff8", + ), + ContentProbe( # txtfusion.refiner_blocks.1.attn.wq + tensor="txtfusion.refiner_blocks.1.sign_mantissa", + offset=0, + length=4096, + sha256="a7dd2de120a3a713fe6eb782ad926c36d2718498c739ee3baac4cd5e26577214", + ), + ContentProbe( # tmlp.0 + tensor="tmlp.sign_mantissa", + offset=0, + length=4096, + sha256="356d0f90c117d207bdef0a26d602b4ca0a36a52f1e4d2588849e2dff3df4f6aa", + ), + ContentProbe( # tmlp.2 + tensor="tmlp.sign_mantissa", + offset=1_572_864, + length=4096, + sha256="305bc2e99d6161e4674a8f2a222c25e958f9b2b2d9efb36f5db799a3a0a3e4d3", + ), + ContentProbe( # tproj.1 + tensor="tproj.sign_mantissa", + offset=0, + length=4096, + sha256="9854423ef01b650b06113ec959298735bcf4cc98d79a5b628873ed2b20a39b9d", + ), + ContentProbe( # txtmlp.1 + tensor="txtmlp.sign_mantissa", + offset=0, + length=4096, + sha256="86beee24cdec56afaf7d367735d53153aed1a00956955a5aff8d1abc2b9d9182", + ), + ContentProbe( # txtmlp.3 + tensor="txtmlp.sign_mantissa", + offset=15_728_640, + length=4096, + sha256="ca3c3a6ee9e11db6c5b419f351187e46505403fdd230c3fdade6f3377527a48c", + ), + ), +) + +KREA_2_TURBO_COMFYUI = SynthesizedLayout( + key="krea-2-turbo-comfyui", + label="Krea 2 Turbo (ComfyUI single-file DF11)", + repo_id="mingyi456/Krea-2-Turbo-DF11-ComfyUI", + revision="978da5fb7647bd222d33125993abd8fdc2840cfc", + file_name="krea2_turbo_bf16-DF11.safetensors", + header_sha256="1de7b6f3f0dc4c04a05c3bb1e7972f0d3761902ea2ddd42916f14183aecbcafd", + file_sha256="ccfd85551ba0f37e7f6706d86d1c00276457d2965ac72fba6fef046653ad6498", + groups=35, + extras=169, + raw_config=_krea2_config(), + row_splits={}, + probes=( + ContentProbe( # blocks.0.attn.wq + tensor="blocks.0.sign_mantissa", + offset=0, + length=4096, + sha256="11f6150d52c2100d4e46f8a8962a8b1a6acf2c38f477cbe9e113643b095caef8", + ), + ContentProbe( # blocks.0.attn.wk + tensor="blocks.0.sign_mantissa", + offset=37_748_736, + length=4096, + sha256="7dc9c9a7f947dd621467761eba193174a35913e35c09ea273109ce1c7d1995e0", + ), + ContentProbe( # blocks.0.attn.wv + tensor="blocks.0.sign_mantissa", + offset=47_185_920, + length=4096, + sha256="c8fd1bc47699725e3a1cfabaddd2cdea7e563a01bf1b0ffa10617f8456ba6f0a", + ), + ContentProbe( # blocks.0.attn.gate + tensor="blocks.0.sign_mantissa", + offset=56_623_104, + length=4096, + sha256="35e55744d6156352bc36f3b5e2d716679d4d2f8c6dc7c6ee5795d4deec7432d0", + ), + ContentProbe( # blocks.0.attn.wo + tensor="blocks.0.sign_mantissa", + offset=94_371_840, + length=4096, + sha256="f905a0181ece0d857a95cbe554bcf98968b847ccebacb332b89755f6545d0dd3", + ), + ContentProbe( # blocks.0.mlp.gate + tensor="blocks.0.sign_mantissa", + offset=132_120_576, + length=4096, + sha256="171a5c675fb5bd9aff19ec75f77eedfef7dc17e6861568e2df3144dec58caaf3", + ), + ContentProbe( # blocks.0.mlp.up + tensor="blocks.0.sign_mantissa", + offset=232_783_872, + length=4096, + sha256="356076a9cd44ddcd624da9336dd56b0dc233ce3bea51e473380d739605292747", + ), + ContentProbe( # blocks.0.mlp.down + tensor="blocks.0.sign_mantissa", + offset=333_447_168, + length=4096, + sha256="7ccae2ff710cc0dbf818ca8fe768066e79b2a64881f8fb514dedcb2307eafc50", + ), + ContentProbe( # blocks.27.attn.wq + tensor="blocks.27.sign_mantissa", + offset=0, + length=4096, + sha256="9a865f96cbb3415c0d2fa6cbf04aac8e72cd914a4da7b30dfa346f48ae2f68b5", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.attn.wq + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=0, + length=4096, + sha256="8f836605424820e44b68ccd56399e183aa990ff0d3da3cec22250db6145cb3c5", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.attn.wk + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=6_553_600, + length=4096, + sha256="9a074d7d85a38cf157419326fe09f517982269097591597ffd5cf4b25438dc04", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.attn.wv + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=13_107_200, + length=4096, + sha256="2975a4ffd78330058c286a9c07d59db4f30fb30da0449cd19f770c96ca018029", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.attn.gate + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=19_660_800, + length=4096, + sha256="fe5976eb985efce2624bdfee10e3362074953a0197c4450dda7799314c905586", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.attn.wo + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=26_214_400, + length=4096, + sha256="6978c0ed44c394128072eb7671cc902112181f2f18a70a84e80d0b5a5aeed80a", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.mlp.gate + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=32_768_000, + length=4096, + sha256="6a8af83a863507b9de8683470115894dea0dd302fe507e4f7f976b3034204bb0", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.mlp.up + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=50_462_720, + length=4096, + sha256="3d7417959035994ad92f9036b3636b1b35e821f438468388dc4306a6221cee4e", + ), + ContentProbe( # txtfusion.layerwise_blocks.0.mlp.down + tensor="txtfusion.layerwise_blocks.0.sign_mantissa", + offset=68_157_440, + length=4096, + sha256="ff857b14597bfcf543d7938d71992f696177b52ee3a7574f63719a9a42dac94b", + ), + ContentProbe( # txtfusion.layerwise_blocks.1.attn.wq + tensor="txtfusion.layerwise_blocks.1.sign_mantissa", + offset=0, + length=4096, + sha256="06d9ba6294dbc0509ab38954d11052a3d0cdc1df21aff4f5fdc1a0d30269f02f", + ), + ContentProbe( # txtfusion.refiner_blocks.0.attn.wq + tensor="txtfusion.refiner_blocks.0.sign_mantissa", + offset=0, + length=4096, + sha256="200e8501fc64ada5d37bf10f2d2653f273367019075f3607f58b0cb54f7c2da2", + ), + ContentProbe( # txtfusion.refiner_blocks.1.attn.wq + tensor="txtfusion.refiner_blocks.1.sign_mantissa", + offset=0, + length=4096, + sha256="241ff415ea140b6838ea709c1317c9d5da6129231f23ca8dbdef946b7bf056c9", + ), + ContentProbe( # tmlp.0 + tensor="tmlp.sign_mantissa", + offset=0, + length=4096, + sha256="ee752a8f3b403449d3a1e95666369eba78aec410bf7879663215b58a54efac98", + ), + ContentProbe( # tmlp.2 + tensor="tmlp.sign_mantissa", + offset=1_572_864, + length=4096, + sha256="1b2aa81eca1c9b20b9f6b96018f53833acc5f9e9617bd37432bcd4920f6b6963", + ), + ContentProbe( # tproj.1 + tensor="tproj.sign_mantissa", + offset=0, + length=4096, + sha256="52dd92bbc434a9120d8038e8637bf5b350d69aba1d508aebaeb20da2d9b57dfe", + ), + ContentProbe( # txtmlp.1 + tensor="txtmlp.sign_mantissa", + offset=0, + length=4096, + sha256="72862c0bad66d0664fda6c8f8db121ceca7aafcac721268e6208d6a8fbf01f69", + ), + ContentProbe( # txtmlp.3 + tensor="txtmlp.sign_mantissa", + offset=15_728_640, + length=4096, + sha256="04724bf4d4567d5f72bb6e0e8e95a9e0560d9bd405ec19399f9b3aa84adfc2ba", + ), + ), +) + +KNOWN_LAYOUTS: tuple[SynthesizedLayout, ...] = ( + QWEN_IMAGE_21_COMFYUI, + KREA_2_RAW_COMFYUI, + KREA_2_TURBO_COMFYUI, +) def identify_layout( @@ -171,6 +559,8 @@ def identify_layout( __all__ = [ "KNOWN_LAYOUTS", + "KREA_2_RAW_COMFYUI", + "KREA_2_TURBO_COMFYUI", "QWEN_IMAGE_21_COMFYUI", "ContentProbe", "SynthesizedLayout", diff --git a/src/mlx_dfloat/integrate/blockseam.py b/src/mlx_dfloat/integrate/blockseam.py index 4cdc173..f71ac0f 100644 --- a/src/mlx_dfloat/integrate/blockseam.py +++ b/src/mlx_dfloat/integrate/blockseam.py @@ -46,6 +46,8 @@ class BlockSeam: def __call__(self, *args: Any, **kwargs: Any) -> Any: """Run the block with its weights assigned, evaluated per policy, placeholders restored. + The call's arguments reach the provider's optional ``before_block`` hook (``seam.BeforeBlock``). + Raises: DFloatIntegrationError: Nothing is attached. """ @@ -56,7 +58,13 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: f"{binding.name}: call attach(provider, shapes) before running a step" ) parent: Any = super() - return seam.run_block(state, binding.name, self, lambda: parent.__call__(*args, **kwargs)) + return seam.run_block( + state, + binding.name, + self, + lambda: parent.__call__(*args, **kwargs), + inputs=(args, kwargs), + ) @cache diff --git a/src/mlx_dfloat/integrate/resident.py b/src/mlx_dfloat/integrate/resident.py index 422fac0..ba9d95b 100644 --- a/src/mlx_dfloat/integrate/resident.py +++ b/src/mlx_dfloat/integrate/resident.py @@ -20,6 +20,10 @@ NonBlockShapes = dict[str, tuple[int, ...]] +# Called through this name so a test can count the evaluations. +_eval = mx.eval + + def _owner(module: Any, param: str) -> Any: if not param.endswith(".weight"): raise DFloatIntegrationError(f"{param}: a compressed matrix must land on a .weight") @@ -54,15 +58,21 @@ def decode_nonblock( name_map: NameMap, *, decode: Callable[[MxGroup], DecodeResult] | None = None, + eval_together: bool = False, ) -> dict[str, mx.array]: """Decode each non-block group (Metal by default), check its status, cut it into its parameters' shapes. + Each group is evaluated and its status checked before the next is queued; with ``eval_together`` every group is + queued first and evaluated in one ``mx.eval`` (a per-call decode pays one synchronisation, not one per group), then + the status words are checked in group order. + Raises: DFloatIntegrationError: A group is not resident, or its matrices do not fit their parameters. DFloatFormatError: A group's decode reported an error (the message names the group). """ run = decode if decode is not None else partial(decode_group, backend="metal") out: dict[str, mx.array] = {} + queued: list[tuple[str, mx.array, list[mx.array]]] = [] for group_name, names in matrix_names.items(): group = groups.get(group_name) if group is None: @@ -83,8 +93,15 @@ def decode_nonblock( ) out[param] = part.view(mx.bfloat16).reshape(shape) mine.append(out[param]) - mx.eval(result.status, *mine) + if eval_together: + queued.append((group_name, result.status, mine)) + continue + _eval(result.status, *mine) check_status(result.status, name=group_name) + if queued: + _eval(*(a for _name, status, mine in queued for a in (status, *mine))) + for group_name, status, _mine in queued: + check_status(status, name=group_name) return out diff --git a/src/mlx_dfloat/integrate/seam.py b/src/mlx_dfloat/integrate/seam.py index 133e967..7672ce2 100644 --- a/src/mlx_dfloat/integrate/seam.py +++ b/src/mlx_dfloat/integrate/seam.py @@ -1,7 +1,7 @@ """The block seam: assign one block's weights, run it, evaluate per policy, restore the placeholders.""" import time -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import Any, Protocol @@ -155,14 +155,40 @@ def _seam_eval(state: SeamState, out: Any) -> None: state.prev = out -def run_block(state: SeamState, block_name: str, block: Any, run: Callable[[], Any]) -> Any: +class BeforeBlock(Protocol): + """The optional provider hook ``before_block``, called with a block's inputs before its weights are requested. + + A provider that defines it can act on the block's inputs first (Krea 2's per-call decode evaluates block 0's + inputs and releases the non-block weights there); a provider without it is asked for the weights and nothing else. + """ + + def __call__(self, block_name: str, args: tuple[Any, ...], kwargs: Mapping[str, Any]) -> None: + """Act on ``block_name``'s call arguments; the weights are requested after it returns.""" + ... + + +def run_block( + state: SeamState, + block_name: str, + block: Any, + run: Callable[[], Any], + *, + inputs: tuple[tuple[Any, ...], Mapping[str, Any]] | None = None, +) -> Any: """Assign the block's weights, run it, evaluate per policy, restore the placeholders (also on a raise). + ``inputs`` are the block's call arguments, handed to the provider's ``before_block`` hook when it has one (see + ``BeforeBlock``); a seam that does not pass them hands the hook empty ones, new for each call. + Raises: DFloatIntegrationError: The provider's dict does not cover the block's matrices, or a weight's shape does not match the block. """ shapes = state.shapes[block_name] + before: BeforeBlock | None = getattr(state.provider, "before_block", None) + if before is not None: + args, kwargs = ((), {}) if inputs is None else inputs + before(block_name, args, kwargs) t_decode_start = time.perf_counter() weights = state.provider.weights_for(block_name, shapes) t_decode_end = time.perf_counter() diff --git a/src/mlx_dfloat/mflux/__init__.py b/src/mlx_dfloat/mflux/__init__.py index 780efe7..47c4a77 100644 --- a/src/mlx_dfloat/mflux/__init__.py +++ b/src/mlx_dfloat/mflux/__init__.py @@ -21,8 +21,8 @@ def require_mflux() -> None: def __getattr__(name: str) -> Any: """``DFloatModel`` and the model classes are imported on first use. - The classes are ``DFloatFlux1``, ``DFloatZImage``, ``DFloatFlux2Klein``, ``DFloatQwenImage21`` and - ``DFloatErnieImage``. So ``import mlx_dfloat.mflux`` needs neither mflux nor MLX. + The classes are ``DFloatFlux1``, ``DFloatZImage``, ``DFloatFlux2Klein``, ``DFloatQwenImage21``, + ``DFloatErnieImage`` and ``DFloatKrea2``. So ``import mlx_dfloat.mflux`` needs neither mflux nor MLX. """ if name == "DFloatModel": from mlx_dfloat.mflux.families import DFloatModel @@ -48,10 +48,14 @@ def __getattr__(name: str) -> Any: from mlx_dfloat.mflux.ernie.model import DFloatErnieImage return DFloatErnieImage + if name == "DFloatKrea2": + from mlx_dfloat.mflux.krea2.model import DFloatKrea2 + + return DFloatKrea2 raise AttributeError(f"module {__name__!r} has no attribute {name!r}") # The model classes stay out of __all__: a star-import must not need mflux; `mlx_dfloat.mflux.DFloatFlux1`, -# `mlx_dfloat.mflux.DFloatZImage`, `mlx_dfloat.mflux.DFloatFlux2Klein`, `mlx_dfloat.mflux.DFloatQwenImage21` and -# `mlx_dfloat.mflux.DFloatErnieImage` still work. +# `mlx_dfloat.mflux.DFloatZImage`, `mlx_dfloat.mflux.DFloatFlux2Klein`, `mlx_dfloat.mflux.DFloatQwenImage21`, +# `mlx_dfloat.mflux.DFloatErnieImage` and `mlx_dfloat.mflux.DFloatKrea2` still work. __all__ = ["DFloatModel", "require_mflux"] diff --git a/src/mlx_dfloat/mflux/ernie/memory.py b/src/mlx_dfloat/mflux/ernie/memory.py index b9d08bf..44f1673 100644 --- a/src/mlx_dfloat/mflux/ernie/memory.py +++ b/src/mlx_dfloat/mflux/ernie/memory.py @@ -32,7 +32,7 @@ TEXT_MAX_LENGTH = 2048 MAX_MEASURED_PIXELS = 1024 * 1024 # no run above 1024² on this path -# Measured on 2026-10-08 (git a771e58, M1 Max 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0): two one-step de-risk runs at +# Measured on 2026-10-08 (git a771e58, M1 Max 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0): two one-step calibration runs at # 1024², seed 42, the 26-token lighthouse prompt; one process each: build, encode, drop, set load, one step at the # planner's cache limit for the block-cache probe's allowance, VAE decode on the resident set. ERNIE-Image-Turbo # (mingyi456/ERNIE-Image-Turbo-DF11 @ 27f84b4 over baidu/ERNIE-Image-Turbo @ bc68c81) at guidance 1.0: batch 1, run @@ -48,7 +48,7 @@ # the two runs: Turbo 644_400_592 - 27_961_608 - 4_806 = 616_434_178 (the base's 548_194_818). One value for both # batches: the build is the same. OVERHEAD = 616_434_178 # 0.574 GiB -# VAE_TRANSIENT, by the overall-peak rule over every 1024² sample (the two de-risk runs so far): the highest measured +# VAE_TRANSIENT, by the overall-peak rule over every 1024² sample (the two calibration runs so far): the highest measured # VAE-phase peak on the resident set minus the VAE phase's other terms (compressed + extras + non-block + VAE file + # OVERHEAD), from the Turbo run: 19_163_080_232 - (10_894_810_183 + 27_961_600 + 335_544_320 + 168_120_878 + # 616_434_178). The base run's peak (19_127_199_104, with its own sizes) would give a term 35_821_552 B lower. A fit @@ -61,9 +61,9 @@ # CAPPED 24 GB run of ERNIE-Image-Turbo (label "capped-24-gb-2" in ernie-image-turbo-1024.json; git f4cd325), whose # watched peak was in the denoise phase: 15_446_156_600 - (10_894_810_183 + 27_961_600 + 335_544_320 + 436_207_616 + # 2_600_000_000 + 616_434_178) = 535_198_703. The first identical run ("capped-24-gb-1", git 5e2e7c2) peaked at -# 15_069_554_288 (term 158_596_391), 376_602_312 B (0.35 GiB) lower, and the de-risk run's 14_974_035_520 gave +# 15_069_554_288 (term 158_596_391), 376_602_312 B (0.35 GiB) lower, and the calibration run's 14_974_035_520 gave # 63_077_623: the denoise phase varies that much between identical runs. A fit check must not under-predict, so the -# term covers the highest sample. Batch 2 (base, limit 3_600_000_000), its de-risk +# term covers the highest sample. Batch 2 (base, limit 3_600_000_000), its calibration # run, the only denoise-phase sample: 16_361_874_888 - (10_894_750_607 + 27_961_600 + 335_544_320 + 436_207_616 + # 3_600_000_000 + 616_434_178) = 450_976_567. DENOISE_ACTIVATION: dict[int, int] = {1: 535_198_703, 2: 450_976_567} @@ -109,7 +109,7 @@ def _constants(batch: int) -> PhaseConstants: # (Turbo) and 332_466_044 (base) under the bound, and an encode that also loads the vision tower and projector # (840_167_424 B more) is over it by 507_701_380 at every prompt length. ENCODE_SLACK_BYTES = 256 * 1024**2 -# The prompt length ENCODE_ACTIVATION was measured at (the de-risk prompt's text tokens). +# The prompt length ENCODE_ACTIVATION was measured at (the calibration prompt's text tokens). ENCODE_ACTIVATION_TOKENS = 26 _LAYER = re.compile(r"^language_model\.model\.layers\.(\d+)\.") diff --git a/src/mlx_dfloat/mflux/families.py b/src/mlx_dfloat/mflux/families.py index 155ae24..73f36c4 100644 --- a/src/mlx_dfloat/mflux/families.py +++ b/src/mlx_dfloat/mflux/families.py @@ -37,6 +37,11 @@ class ModelEntry: # The commit the default base repository is pinned to (None: the default branch); a user's own base is taken as # given. base_revision: str | None = None + # Classifier-free guidance runs for any guidance other than 1.0, below 1.0 too (mflux's rule for Krea 2), rather + # than only above 1.0. + cfg_unless_guidance_one: bool = False + # The schedulers the model's pipeline accepts (the CLI refuses any other before the build); None: not checked here. + schedulers: tuple[str, ...] | None = None @dataclass(frozen=True, slots=True, kw_only=True) @@ -109,6 +114,18 @@ def _ernie_names() -> NameMap: return ernie_name_map() +def _krea2_class() -> type: + from mlx_dfloat.mflux.krea2.model import DFloatKrea2 + + return DFloatKrea2 + + +def _krea2_names() -> NameMap: + from mlx_dfloat.mflux.krea2.names import krea2_name_map + + return krea2_name_map() + + FAMILIES: dict[str, FamilySpec] = { "flux1": FamilySpec( name="flux1", @@ -143,6 +160,12 @@ def _ernie_names() -> NameMap: load_model_class=_ernie_class, load_name_map=_ernie_names, ), + "krea2": FamilySpec( + name="krea2", + refused_flags={}, + load_model_class=_krea2_class, + load_name_map=_krea2_names, + ), } @@ -295,7 +318,7 @@ def _klein(name: str, label: str, steps: int, revision: str, *, base: bool) -> M cfg_two_calls=False, uses_negative_prompt=False, fixed_guidance=1.0, - # The snapshot the parity, de-risk and MEASURED runs used (text encoder, VAE, tokenizer). + # The snapshot the parity, calibration and MEASURED runs used (text encoder, VAE, tokenizer). base_revision="bc68c81e2a1730a394d5fc9fae70713dee940140", ), # mflux 0.20.0: 50 steps (cli/defaults/defaults.py:35); guidance 4.0 (models/ernie_image/cli/ @@ -313,9 +336,49 @@ def _klein(name: str, label: str, steps: int, revision: str, *, base: bool) -> M default_scheduler="linear", cfg_two_calls=False, # CFG runs as one batch-2 call (ernie_image.py:239-250) uses_negative_prompt=True, - # The snapshot the parity, identity, de-risk and MEASURED runs used. + # The snapshot the parity, identity, calibration and MEASURED runs used. base_revision="5346b31d68c9c23758ba56ef8be5e9dc174c7f99", ), + # mflux 0.20.0: 8 steps (cli/defaults/defaults.py:47), guidance 1.0 (models/krea2/cli/krea2_generate.py:17; + # krea2.py:58), scheduler er_sde (krea2.py:206-212: None and "linear" resolve to er_sde). Classifier-free guidance + # runs for any guidance other than 1.0 (prompt_encoder.py:33), as two transformer calls (krea2.py:195-200), with + # mflux's " " when the negative prompt is missing or blank. + ModelEntry( + name="krea-2", + family="krea2", + label="Krea 2 Turbo", + df11_repo="mingyi456/Krea-2-Turbo-DF11-ComfyUI", + df11_revision="978da5fb7647bd222d33125993abd8fdc2840cfc", + base_repo="krea/Krea-2-Turbo", + default_steps=8, + default_guidance=1.0, + default_scheduler="er_sde", + cfg_two_calls=True, + uses_negative_prompt=True, + cfg_unless_guidance_one=True, + schedulers=("er_sde", "euler", "linear"), + base_revision="98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7", + ), + # mflux's Krea command refuses this name (krea2_generate.py:46-49, config_resolution.py:85-91), so the step + # table's fallback is mflux's only step default for it: 25 (cli/defaults/defaults.py:19,111-118, no "krea-2-raw" + # entry); guidance 1.0 is mflux's generate_image default (krea2.py:58), so no CFG by default. The model card's + # 52 steps at 3.5 is a recipe the README names. + ModelEntry( + name="krea-2-raw", + family="krea2", + label="Krea 2 Raw", + df11_repo="mingyi456/Krea-2-Raw-DF11-ComfyUI", + df11_revision="8320616b25ac9340a830a7fb21f1b0237e160e66", + base_repo="krea/Krea-2-Raw", + default_steps=25, + default_guidance=1.0, + default_scheduler="er_sde", + cfg_two_calls=True, + uses_negative_prompt=True, + cfg_unless_guidance_one=True, + schedulers=("er_sde", "euler", "linear"), + base_revision="6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + ), ) } @@ -344,7 +407,8 @@ def DFloatModel(name: str, /, **kwargs: Any) -> Any: # noqa: N802 Imports mflux. ``DFloatModel("z-image-turbo")`` returns a ``DFloatZImage``; ``DFloatModel("schnell")`` a ``DFloatFlux1``; ``DFloatModel("flux2-klein-4b")`` a ``DFloatFlux2Klein``; ``DFloatModel("qwen-image-2.1")`` a - ``DFloatQwenImage21``; ``DFloatModel("ernie-image")`` a ``DFloatErnieImage``. + ``DFloatQwenImage21``; ``DFloatModel("ernie-image")`` a ``DFloatErnieImage``; ``DFloatModel("krea-2")`` a + ``DFloatKrea2``. Args: name: A registered model name, such as ``"schnell"`` or ``"z-image-turbo"``. diff --git a/src/mlx_dfloat/mflux/generate.py b/src/mlx_dfloat/mflux/generate.py index 7d6550d..88ed090 100644 --- a/src/mlx_dfloat/mflux/generate.py +++ b/src/mlx_dfloat/mflux/generate.py @@ -21,7 +21,7 @@ from mlx_dfloat.bench import capped from mlx_dfloat.bench.capped import TierLimits, host_tier_gb, limits_record, tier_limits from mlx_dfloat.errors import DFloatError -from mlx_dfloat.mflux.families import FAMILIES, MODELS, entry, family_of +from mlx_dfloat.mflux.families import FAMILIES, MODELS, ModelEntry, entry, family_of EXIT_OK, EXIT_ERROR = 0, 2 FOOTPRINT_PEAK_LABEL = "OS phys_footprint, sampled every 0.05 s by the watchdog" @@ -119,7 +119,8 @@ def _add_arguments(p: argparse.ArgumentParser) -> None: type=int, default=None, help="denoise steps (default per model, as mflux: schnell 4, dev 25, z-image 50, z-image-turbo 9, " - "FLUX.2 Klein base 50, distilled 4, qwen-image-2.1 40, ernie-image 50, ernie-image-turbo 8)", + "FLUX.2 Klein base 50, distilled 4, qwen-image-2.1 40, ernie-image 50, ernie-image-turbo 8, " + "krea-2-raw 25, krea-2 8)", ) p.add_argument("--height", type=int, default=1024) p.add_argument("--width", type=int, default=1024) @@ -131,19 +132,22 @@ def _add_arguments(p: argparse.ArgumentParser) -> None: "z-image 0, its model card suggests about 4; FLUX.2 Klein 1.0, where the base models take another " "value such as 4 and the distilled ones refuse it; Qwen-Image 2.1 1.0, where classifier-free guidance " "needs a value above 1 and --negative-prompt; ERNIE-Image 4.0 (classifier-free guidance above 1, with or " - "without --negative-prompt); ERNIE-Image-Turbo 1.0 only)", + "without --negative-prompt); ERNIE-Image-Turbo 1.0 only; Krea 2 Raw and Turbo 1.0 (Krea 2 Raw's model card " + "uses --steps 52 --guidance 3.5); Krea 2 runs classifier-free guidance, two transformer calls per step, for " + "any value other than 1.0)", ) p.add_argument( "--scheduler", default=None, help="scheduler (default per model: FLUX.1, Qwen-Image 2.1 and ERNIE-Image linear; FLUX.2 Klein always runs " - "flow_match_euler_discrete and refuses the flag)", + "flow_match_euler_discrete and refuses the flag; Krea 2: er_sde (default), euler, or linear, which mflux " + "runs as er_sde)", ) p.add_argument( "--negative-prompt", default=None, - help="used by z-image, qwen-image-2.1 and ernie-image with --guidance above 1; accepted and ignored for the " - "other models, as mflux does", + help="used by z-image, qwen-image-2.1 and ernie-image with --guidance above 1, and by krea-2 and krea-2-raw " + "with any --guidance other than 1; accepted and ignored for the other models, as mflux does", ) p.add_argument( "--output", @@ -284,20 +288,58 @@ def ceiling_for( return (default_ceiling_bytes if limits.is_host else limits.ceiling_bytes), limits, limits.label +# The families whose tokenizer gives an empty prompt no tokens at all, by the name the refusal uses. +_NO_TOKEN_FAMILIES = {"ernie": "ERNIE-Image", "krea2": "Krea 2"} + + def empty_prompt_refusal(args: argparse.Namespace) -> str | None: """Why an empty or blank ``--prompt`` is refused for this model, or None. - ERNIE-Image refuses an empty or blank prompt as a user error (its tokenizer gives an empty prompt no tokens at all, - not even a start token); the other families' mflux pipelines encode an empty prompt as it is. + ERNIE-Image and Krea 2 refuse an empty or blank prompt as a user error (their tokenizers give an empty prompt no + tokens at all, not even a start token); the other families' mflux pipelines encode an empty prompt as it is. """ - if entry(args.model).family != "ernie" or args.prompt.strip(): + name = _NO_TOKEN_FAMILIES.get(entry(args.model).family) + if name is None or args.prompt.strip(): return None return ( f"--prompt {args.prompt!r}: an empty or blank prompt is refused as a user error (an empty prompt gives " - "ERNIE-Image's tokenizer no tokens at all)" + f"{name}'s tokenizer no tokens at all)" ) +def runs_cfg(e: ModelEntry, guidance: float | None) -> bool: + """Whether a call at ``guidance`` runs classifier-free guidance under the model's rule. + + Krea 2 runs it for any guidance other than 1.0; the other families only above 1.0 (Qwen-Image 2.1 also needs a + negative prompt, which the caller checks). None, the model's own rule, runs none. + """ + if guidance is None: + return False + return guidance != 1.0 if e.cfg_unless_guidance_one else guidance > 1.0 + + +def scheduler_refusal(args: argparse.Namespace) -> str | None: + """Why ``--scheduler`` is refused for this model, or None (a model without an allow-list takes any name).""" + e = entry(args.model) + if e.schedulers is None or args.scheduler is None or args.scheduler in e.schedulers: + return None + return f"--scheduler: {e.label} runs {', '.join(e.schedulers)}" + + +def cfg_cost_note(args: argparse.Namespace) -> str | None: + """An info line when a Krea 2 call runs classifier-free guidance (two transformer calls per step), else None. + + Krea 2 runs it for any guidance other than 1.0 without asking for a negative prompt, so a guidance such as 0.5 + doubles the step cost on mflux's ``" "`` negative; the line says so. No other family gets it. + """ + e = entry(args.model) + guidance = _guidance(args) + if not e.cfg_unless_guidance_one or not runs_cfg(e, guidance): + return None + note = f"guidance {guidance} on {e.label} runs classifier-free guidance: two transformer calls per step" + return note if args.negative_prompt else note + ' (negative prompt " ")' + + def _steps(args: argparse.Namespace) -> int: """The denoise steps this call runs: ``--steps``, else the model's mflux default.""" return int(args.steps) if args.steps is not None else entry(args.model).default_steps @@ -401,7 +443,12 @@ def run( between them). """ started = clock() - refused = refused_option(args) or fixed_guidance_refusal(args) or empty_prompt_refusal(args) + refused = ( + refused_option(args) + or fixed_guidance_refusal(args) + or empty_prompt_refusal(args) + or scheduler_refusal(args) + ) if refused is not None: print(f"error: {refused}", file=sys.stderr) return EXIT_ERROR @@ -413,13 +460,23 @@ def run( else "this model has no negative branch" ) print(f"warning: --negative-prompt is ignored: {why}", file=sys.stderr) - elif args.negative_prompt and (_guidance(args) or 0.0) <= 1.0: + elif args.negative_prompt and not runs_cfg(entry(args.model), _guidance(args)): default = entry(args.model).default_guidance or 0 # None (Z-Image): mflux's own rule, 0 - print( - "warning: --negative-prompt has no effect: classifier-free guidance runs only above " - f"guidance 1.0 (the default is {default}). Pass --guidance above 1.0 to enable it.", - file=sys.stderr, - ) + if entry(args.model).cfg_unless_guidance_one: + print( + "warning: --negative-prompt has no effect: classifier-free guidance runs only for a guidance " + f"other than 1.0 (the default is {default}). Pass another --guidance to enable it.", + file=sys.stderr, + ) + else: + print( + "warning: --negative-prompt has no effect: classifier-free guidance runs only above " + f"guidance 1.0 (the default is {default}). Pass --guidance above 1.0 to enable it.", + file=sys.stderr, + ) + note = cfg_cost_note(args) + if note is not None: + print(f"info: {note}", file=sys.stderr) if ( entry(args.model).cfg_needs_negative and (_guidance(args) or 0.0) > 1.0 diff --git a/src/mlx_dfloat/mflux/krea2/__init__.py b/src/mlx_dfloat/mflux/krea2/__init__.py new file mode 100644 index 0000000..7f68195 --- /dev/null +++ b/src/mlx_dfloat/mflux/krea2/__init__.py @@ -0,0 +1 @@ +"""Krea 2 Raw and Krea 2 Turbo on mflux from a DFloat11 transformer.""" diff --git a/src/mlx_dfloat/mflux/krea2/init.py b/src/mlx_dfloat/mflux/krea2/init.py new file mode 100644 index 0000000..46d5df7 --- /dev/null +++ b/src/mlx_dfloat/mflux/krea2/init.py @@ -0,0 +1,194 @@ +"""The Krea 2 base repository without its transformer: text encoder, VAE and tokenizer, resolved and loaded. + +mflux's ``Krea2Initializer.init`` wants every component under one path, so it would download and size the transformer +(for Raw, mflux's own download patterns fetch the 26 GB diffusers ``transformer/`` shards); this module loads the other +components through mflux's own loader and applier with weight definitions that name only them. The download patterns +never name ``transformer/``, the root ``raw.safetensors`` / ``turbo.safetensors`` or the sample ``images/``. Krea 2 +Raw and Krea 2 Turbo share their text encoder, VAE and tokenizer byte for byte. +""" + +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from mlx.utils import tree_flatten + +from mlx_dfloat.errors import DFloatAccessError, DFloatFormatError +from mlx_dfloat.mflux import _hub, require_mflux +from mlx_dfloat.mflux._hub import ResolvedRepo as ResolvedRepo + +# The ComfyUI DF11 repository holds one weight file and no config. +DF11_PATTERNS: tuple[str, ...] = ("*.safetensors",) +ENCODER_PATTERNS: tuple[str, ...] = ("text_encoder/*.safetensors", "text_encoder/*.json") +VAE_PATTERNS: tuple[str, ...] = ("vae/*.safetensors", "vae/*.json") +# The tokenizer's files, named one by one (never mflux's "tokenizer/**", which would fetch anything a later revision +# drops in there). Both pinned base snapshots hold the first three (2026-10-08); the other two are files transformers' +# AutoTokenizer reads when present, named so a later revision that adds them still loads. +TOKENIZER_PATTERNS: tuple[str, ...] = ( + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + "tokenizer/chat_template.jinja", + "tokenizer/special_tokens_map.json", + "tokenizer/added_tokens.json", +) +BASE_PATTERNS: tuple[str, ...] = ENCODER_PATTERNS + VAE_PATTERNS + TOKENIZER_PATTERNS + +_WHAT = "Krea 2 base" +# Parameters no checkpoint holds: built at construction, so a base never replaces them. mflux 0.20.0: the text +# encoder's ``rotary_emb.inv_freq`` is computed in ``Qwen3TextRotaryEmbedding.__init__`` (the base file holds only +# ``language_model.*`` and ``visual.*`` tensors). The VAE has none: every one of its parameters is a target of mflux's +# VAE mapping (Qwen-Image 1's). +COMPUTED: dict[str, frozenset[str]] = { + "text_encoder": frozenset({"rotary_emb.inv_freq"}), + "vae": frozenset(), +} +_SHOWN = 3 # missing names quoted in the error + + +def _licence_error(repo_id: str) -> DFloatAccessError: + return DFloatAccessError( + f"{repo_id}: this repository is gated or this account may not read it; accept the licence at " + f"https://huggingface.co/{repo_id} and run `hf auth login`" + ) + + +def resolve(spec: str, *, patterns: tuple[str, ...], revision: str | None = None) -> ResolvedRepo: + """``_hub.resolve``, with the licence page named when the Hub refuses access (both Krea 2 bases are gated). + + Raises: + DFloatFormatError: Neither a directory nor a Hub id, or the Hub has no such repository. + DFloatAccessError: The repository is gated or this account may not read it. + """ + try: + return _hub.resolve(spec, patterns=patterns, revision=revision) + except DFloatAccessError as exc: + raise _licence_error(spec) from exc + + +@dataclass(slots=True, kw_only=True) +class Krea2Components: + """The loaded base components: mflux modules (weights still lazy) and the tokenizers.""" + + vae: Any + text_encoder: Any + tokenizers: dict[str, Any] + + +def _definition(names: tuple[str, ...], patterns: tuple[str, ...]) -> type: + """An mflux weight definition holding only the named Krea 2 components.""" + require_mflux() + from mflux.models.krea2.weights.krea2_weight_definition import Krea2WeightDefinition + + return _hub.subset_definition(Krea2WeightDefinition, names, patterns) + + +def encoder_definition() -> type: + """The text encoder component of mflux's Krea 2 definition (its prefix-stripping key transform kept).""" + return _definition(("text_encoder",), ENCODER_PATTERNS) + + +def vae_definition() -> type: + """The VAE component of mflux's Krea 2 definition, with its download patterns.""" + return _definition(("vae",), VAE_PATTERNS) + + +def _shapes(module: Any) -> dict[str, tuple[int, ...]]: + """Every parameter's shape by path (read without evaluating anything).""" + flat: dict[str, Any] = dict(tree_flatten(module.parameters())) + return {name: tuple(p.shape) for name, p in flat.items()} + + +def _require_replaced(root: Path, name: str, before: dict[str, Any], module: Any) -> None: + """Require every parameter of ``module`` to have been replaced by the base since ``before`` was taken. + + mflux's applier updates with ``strict=False``: a tensor the base lacks (a lost shard, a renamed key) leaves that + parameter at its random float32 init without an error. ``before`` holds each parameter as built; one still the + same object after the load was not in the base (``COMPUTED`` names the exceptions). + + Raises: + DFloatFormatError: A parameter was not replaced by the base. + """ + after: dict[str, Any] = dict(tree_flatten(module.parameters())) + missing = [n for n, init in before.items() if n not in COMPUTED[name] and after.get(n) is init] + if missing: + more = len(missing) - _SHOWN + shown = ", ".join(missing[:_SHOWN]) + (f" and {more} more" if more > 0 else "") + raise DFloatFormatError( + f"{root}: {name} {shown} not in the base: a missing shard or a renamed tensor?" + ) + + +def load_text_encoder(root: Path, *, dims: Mapping[str, Any] | None = None) -> Any: + """A fresh ``Krea2TextEncoder`` with the base's weights (lazy). + + ``dims`` are the encoder's constructor arguments (none: mflux's defaults, the published model). Only the language + model's tensors are loaded; the vision tower in the same file never is. mflux's applier assigns weights without a + shape check and ignores a tensor the base lacks, so after the load every parameter's shape is compared with the + freshly built module's (another model's base is refused here, before any encode) and every parameter must have + been replaced by the base. + + Raises: + DFloatFormatError: The directory lacks the text encoder, a loaded parameter's shape differs from the model's + (another model's base), or a parameter is not in the base. + """ + require_mflux() + from mflux.models.krea2.model.krea2_text_encoder.text_encoder import Krea2TextEncoder + + encoder = Krea2TextEncoder(**dict(dims or {})) + want = _shapes(encoder) + before: dict[str, Any] = dict(tree_flatten(encoder.parameters())) + _hub.load_into(root, encoder_definition(), {"text_encoder": encoder}, what=_WHAT) + for name, got in _shapes(encoder).items(): + if want.get(name) != got: + raise DFloatFormatError( + f"{root}: {name} has shape {got}, this model needs {want.get(name)}: another model's " + "base repository?" + ) + _require_replaced(root, "text_encoder", before, encoder) + return encoder + + +def load_vae(root: Path) -> Any: + """A fresh ``QwenVAE`` (Qwen-Image 1's VAE, which Krea 2 uses) with the base's weights (lazy). + + Raises: + DFloatFormatError: The directory lacks the VAE, or a parameter is not in the base. + """ + require_mflux() + from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE + + vae = QwenVAE() + before: dict[str, Any] = dict(tree_flatten(vae.parameters())) + _hub.load_into(root, vae_definition(), {"vae": vae}, what=_WHAT) + _require_replaced(root, "vae", before, vae) + return vae + + +def load_tokenizers(root: Path) -> dict[str, Any]: + """The Krea 2 tokenizer (one, keyed ``"qwen3vl"``, read from ``tokenizer/``). + + Raises: + DFloatFormatError: The tokenizer files are missing or incomplete. + """ + require_mflux() + from mflux.models.common.tokenizer.tokenizer_loader import TokenizerLoader + from mflux.models.krea2.weights.krea2_weight_definition import Krea2WeightDefinition + + try: + return dict( + TokenizerLoader.load_all( + definitions=Krea2WeightDefinition.get_tokenizers(), model_path=str(root) + ) + ) + except (OSError, RuntimeError) as exc: # OSError: transformers' error for a partial tokenizer/ + raise DFloatFormatError(f"{root}: no usable Krea 2 tokenizer: {exc}") from exc + + +def load_base(root: Path) -> Krea2Components: + """Text encoder, VAE and tokenizer of a base repository (weights lazy, every parameter from the base).""" + return Krea2Components( + vae=load_vae(root), + text_encoder=load_text_encoder(root), + tokenizers=load_tokenizers(root), + ) diff --git a/src/mlx_dfloat/mflux/krea2/memory.py b/src/mlx_dfloat/mflux/krea2/memory.py new file mode 100644 index 0000000..13b8861 --- /dev/null +++ b/src/mlx_dfloat/mflux/krea2/memory.py @@ -0,0 +1,196 @@ +"""Krea 2 memory rules: the constants, the text-token count and the text encoder's resident bytes. + +The phase arithmetic itself is the shared one (``mlx_dfloat.mflux._phases``). Krea 2 differs in three inputs. Its text +encoder's file also holds a vision tower that mflux never loads, and the prompt embedding taps hidden states up to the +output of layer 34, so the encoder's last layer and final norm never run (see ``encoder_bytes_used``). Its prompt +length varies per call (up to the tokenizer's 1024 tokens), and mflux strips the chat-template prefix before the +transformer sees the text, so the text tokens are counted per prompt after that strip (``prompt_tokens``). And +classifier-free guidance runs as two batch-1 transformer calls, so one set of constants serves both Krea 2 models and +a CFG step is sized by its longer prompt (``text_tokens_for``). + +The seven non-block groups (the four text-fusion blocks, ``tmlp``, ``tproj``, ``txtmlp``) are decoded at the start of +each transformer call and released before block 0's decode, so the model plans with ``sizes_for(..., +nonblock_per_call=True)`` (no resident non-block bytes) and the denoise term carries their call-start decode. Each +call starts with MLX's cache emptied and drained, and so does block 0 (``mlx_dfloat.mflux.krea2.transformer``). At +1024² that keeps a guided Krea 2 Raw step's denoise peak at 21.14 GiB, and two Krea 2 Turbo steps at 21.15 GiB: a +step's later calls follow the same path as its first. The peak is block 0's own evaluation, with the cache far below +the planner's cache limit, so the cache term over-counts it and the measured activation term is 0 (the +size-independent floor carries it). + +At 1024² on a 32 GB Mac the set is always dropped before the VAE decode. The transformer's residual stream runs float32 +(the latents are float32), and so does the VAE decode. +""" + +import re +from collections.abc import Sequence +from dataclasses import replace +from pathlib import Path +from typing import Any + +from mlx_dfloat._safetensors import read_header +from mlx_dfloat.format import DF11Checkpoint +from mlx_dfloat.mflux import require_mflux +from mlx_dfloat.mflux._phases import FamilySizes, PhaseConstants +from mlx_dfloat.mflux._phases import sizes_for as family_sizes_for +from mlx_dfloat.mflux.krea2.names import NONBLOCK_GROUPS + +# The text-encoder prefixes mflux keeps (its weight definition strips either and drops every other tensor). +ENCODER_PREFIXES: tuple[str, ...] = ("model.language_model.", "language_model.") +# The tokenizer's max_length (mflux's weight definition; tokenization truncates at it, the template included). +TEXT_MAX_LENGTH = 1024 +MAX_MEASURED_PIXELS = 1024 * 1024 # no run above 1024² on this path + +# Measured on 2026-10-08 (M1 Max 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0) at 1024², seed 42, the 30-token +# lighthouse prompt; one process per run: build, encode, drop the encoder, set load, the denoise steps at the planner's +# cache limit for the block-cache probe's allowance (4_500_000_000 B), the set dropped, the VAE decode. The constants +# rest on the runs of the shipped mechanism (per-call non-block decode, MLX's cache emptied and drained at each call's +# start and again before block 0's decode): Krea 2 Raw (mingyi456/Krea-2-Raw-DF11-ComfyUI @ 8320616 over +# krea/Krea-2-Raw @ 6b0ece7), one step at guidance 3.5 with mflux's " " negative (two batch-1 calls; "derisk-raw-2" and +# "confirm-raw-24711fb"), and Krea 2 Turbo (mingyi456/Krea-2-Turbo-DF11-ComfyUI @ 978da5f over krea/Krea-2-Turbo @ +# 98e0fe1), two steps at guidance 1.0 ("derisk-turbo-2" and "confirm-turbo-24711fb"). The VAE term also counts the Raw +# run with the non-block groups resident ("derisk-raw-1"): its VAE phase ran with the set dropped, as in every +# mechanism. Runs of earlier per-call code are recorded as history and source nothing. The git fields in the records +# name development commits folded into the 0.2.0 release. Records: bench/results/calibration/krea-2-raw-1024.json and +# krea-2-1024.json (constants_from names the run of each term). Bytes throughout; re-measure before changing any. +# +# REFERENCE_TOKENS: the calibration point, 1024² (4096 image tokens) + the runs' text tokens, 30 (the prompt's; the +# negative " " is 6, and a step is sized by its longer call). +REFERENCE_TOKENS = 4096 + 30 +# OVERHEAD: the footprint MLX's counters do not see after the build (footprint - MLX active - MLX cache), the largest +# of the four shipped runs: derisk-turbo-2, 529_171_680 - 4_805_792 - 4_974 (the others 522_673_362 to 523_787_474). +OVERHEAD = 524_360_914 # 0.488 GiB +# VAE_TRANSIENT, by the overall-peak rule over every 1024² VAE-phase sample of the shipped runs and the resident run +# (all decoded with the set dropped): the highest VAE footprint peak minus the dropped-set phase's other terms (extras +# + VAE file + OVERHEAD), from derisk-raw-1: 11_869_477_608 - (4_560_024 + 507_591_892 + 524_360_914). The shipped +# runs give 9_127_156_298 to 9_489_078_618; a fit check must not under-predict, so the term covers the highest sample +# (the superseded runs' VAE peaks are all below it too). +VAE_TRANSIENT = 10_832_964_778 # 10.089 GiB +# DENOISE_ACTIVATION, by the overall-peak rule over every shipped 1024² denoise sample (the four one-step runs, the two +# MEASURED rows, the 52-step card recipe and the identity check's df11 side): per run, the denoise footprint peak minus +# the phase's other terms with that run's own sizes at the planner's limit (compressed + extras + non-block 0 + one +# block 868_220_928 + cache limit 4_500_000_000 + OVERHEAD). Every residual is negative; the highest is the Turbo +# MEASURED row's, 22_736_402_568 - 23_437_217_164 = -700_814_596. With the cache emptied at each call's start and before +# block 0, the peak is block 0's and the cache never holds the cache-limit term's full 4.5e9 B there. The term is +# floored at 0, so the activation is the floor below at every size, and the 1024² estimate sits 1.51e9 to 1.54e9 B +# over every shipped peak. +DENOISE_ACTIVATION = 0 +# DENOISE_ACTIVATION_FLOOR: a size-independent lower bound on the denoise term. The residual stream is float32 and the +# block weights BF16, so each block's MLP makes MLX copy mlp.gate and mlp.up to float32 whatever the image size: 2 x +# 24576 x 4096 x 4, counted once, not per token (ERNIE-Image's rule for its batch-2 copies). +DENOISE_ACTIVATION_FLOOR = 805_306_368 # 0.75 GiB +# ALLOWANCE: the cache room for the activation buffers a block frees, at ALLOWANCE_TOKENS. The block-cache probe (one +# synthetic block at 1024² + 30 tokens, its decoded buffer allocated then released per iteration) saw the buffer reused +# at a 4.5e9 B limit; the allowance is that limit minus one decoded block, 4_500_000_000 - 868_220_928. The shipped +# runs' denoise peaks at this limit sat 1.40e9 to 1.42e9 B under the gate line (the fit budget minus 0.5 GiB). +ALLOWANCE = 3_631_779_072 # 3.382 GiB +ALLOWANCE_TOKENS = 4096 + 30 +# ENCODE_ACTIVATION: what a prompt encode held beyond the language model's evaluated tensors by footprint (encode +# footprint peak - encoder_bytes_used - OVERHEAD), the largest of the four shipped runs: confirm-raw-24711fb, +# 8_598_264_624 - 7_843_074_560 - 524_360_914 (the others 163_884_078 to 228_387_886). Taken by footprint so the encode +# phase never under-predicts a calibration run (by MLX's measure the runs held 64_449_856 to 110_119_296 beyond the +# language model). mflux encodes a guided call's two prompts one after the other, so the term is per prompt, at 30 +# tokens, unscaled. +ENCODE_ACTIVATION = 230_829_150 +CONSTANTS: PhaseConstants = PhaseConstants( + overhead_bytes=OVERHEAD, + vae_transient_bytes=VAE_TRANSIENT, + denoise_activation_at_reference=DENOISE_ACTIVATION, + reference_tokens=REFERENCE_TOKENS, + allowance_at_reference=ALLOWANCE, + allowance_reference_tokens=ALLOWANCE_TOKENS, + max_measured_pixels=MAX_MEASURED_PIXELS, + encode_activation_bytes=ENCODE_ACTIVATION, + denoise_activation_floor_bytes=DENOISE_ACTIVATION_FLOOR, +) +# What a prompt encode may hold by MLX's measure beyond the estimate's encode phase before the model warns (the bound +# is the evaluated language model + ENCODE_ACTIVATION + this slack, 8_342_339_166 B). 256 MiB: every recorded encode +# held 389_145_310 to 474_185_294 B under the bound, and an encode that also loads the vision tower (830_695_424 B more) +# is over it by 441_550_114 at every prompt length. +ENCODE_SLACK_BYTES = 256 * 1024**2 +# The prompt length ENCODE_ACTIVATION was measured at (the calibration prompt's text tokens). +ENCODE_ACTIVATION_TOKENS = 30 +_LAYER = re.compile(r"^(?:model\.)?language_model\.layers\.(\d+)\.") + + +def encode_peak_bound(encode_phase_bytes: int, constants: PhaseConstants = CONSTANTS) -> int: + """What a prompt encode may hold by MLX's measure before the model warns. + + That is the estimate's encode phase without its overhead (the evaluated language model and the encode + activations) plus ``ENCODE_SLACK_BYTES``. + """ + return encode_phase_bytes - constants.overhead_bytes + ENCODE_SLACK_BYTES + + +def encode_peak_warning(held: int, bound: int, *, text_tokens: int) -> str | None: + """The warning for an encode that held ``held`` bytes against ``bound``, or None at or under it. + + Two causes fit an excess: mflux loading the vision tower stored in the encoder's file, or a prompt longer than the + one the encode term was measured on (the term is not scaled with the prompt length). The message names the call's + longest prompt in tokens so the reader can tell which. + """ + if held <= bound: + return None + return ( + f"the prompt encode held {held / 1024**3:.2f} GiB for a {text_tokens}-token prompt, over the " + f"{bound / 1024**3:.2f} GiB the fit estimate allows: mflux may now load the vision tower stored in the text " + f"encoder's file, or this prompt holds more than the encode term, which was measured on one " + f"{ENCODE_ACTIVATION_TOKENS}-token prompt and is not scaled with the prompt length; the estimate " + "under-predicts the encode phase either way" + ) + + +def encoder_bytes_used(base_root: Path) -> int: + """The text encoder's bytes a Krea 2 prompt encode makes resident, read from the safetensors headers. + + Only the language model's tensors are loaded; the vision tower in the same file never is. The prompt embedding + stacks hidden states up to the output of the second-to-last layer, so under lazy evaluation the last layer (and + the final norm after it) never run: of ``n`` numbered layers the first ``n - 1`` count (35 of the published 36). + Every other language-model tensor (the embeddings, the final norm) counts in full. + + Raises: + DFloatFormatError: An encoder file's header cannot be read. + """ + tensors = [ + (name, info.nbytes) + for path in sorted((base_root / "text_encoder").glob("*.safetensors")) + for name, info in read_header(path).items() + if name.startswith(ENCODER_PREFIXES) + ] + indices = [int(m.group(1)) for name, _n in tensors if (m := _LAYER.match(name)) is not None] + last = max(indices, default=-1) + total = 0 + for name, nbytes in tensors: + layer = _LAYER.match(name) + if layer is None or int(layer.group(1)) < last: + total += nbytes + return total + + +def prompt_tokens(tokenizer: Any, prompt: str) -> int: + """The text tokens mflux feeds the transformer for ``prompt``: its ids after the chat-template prefix. + + The strip is the one ``Krea2TextEncoder.get_prompt_embeds`` applies (imports mflux). + """ + require_mflux() + from mflux.models.krea2.model.krea2_text_encoder.text_encoder import Krea2TextEncoder + + ids = tokenizer.tokenize(prompt).input_ids + return int(ids.shape[1]) - int(Krea2TextEncoder._template_end(ids)) + + +def text_tokens_for(tokenizer: Any, prompts: Sequence[str]) -> int: + """The text tokens a step is sized by: the longest of its calls' prompts (a CFG step runs one call per prompt).""" + return max(prompt_tokens(tokenizer, p) for p in prompts) + + +def sizes_for( + ckpt: DF11Checkpoint, base_root: Path, *, nonblock_per_call: bool = False +) -> FamilySizes: + """The family sizes of a Krea 2 checkpoint and base, the encoder counted by its evaluated language model. + + With ``nonblock_per_call`` the non-block groups are decoded per transformer call rather than kept decoded with the + set, so ``nonblock`` is 0; the call-start transient belongs to the denoise activation term. + """ + sizes = family_sizes_for(ckpt, base_root, nonblock_groups=NONBLOCK_GROUPS) + sizes = replace(sizes, encoders=encoder_bytes_used(base_root)) + return replace(sizes, nonblock=0) if nonblock_per_call else sizes diff --git a/src/mlx_dfloat/mflux/krea2/model.py b/src/mlx_dfloat/mflux/krea2/model.py new file mode 100644 index 0000000..e9253c0 --- /dev/null +++ b/src/mlx_dfloat/mflux/krea2/model.py @@ -0,0 +1,701 @@ +"""``DFloatKrea2``: mflux's Krea 2 pipeline over a DFloat11 transformer, decoded one block at a time. + +The class subclasses mflux's ``Krea2`` so its loop, scheduler, sampler, callbacks, VAE decode and image metadata run +unchanged; only construction and the prelude of ``generate_image`` differ. The text encoder and the compressed +transformer are never resident together (see ``mlx_dfloat.mflux.lifecycle``): a call's prompts are encoded first into +a per-prompt cache, which the loop reads without touching the encoder. +""" + +import logging +import sys +import traceback +import warnings +from collections.abc import Callable +from dataclasses import asdict +from importlib import metadata +from pathlib import Path +from typing import Any + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_flatten + +from mlx_dfloat import _metal_decode +from mlx_dfloat.decode import DecodeResult +from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError, DFloatUnsupportedError +from mlx_dfloat.format import DF11Checkpoint, MxGroup, open_checkpoint +from mlx_dfloat.integrate.coverage import load_resident_set +from mlx_dfloat.integrate.memory import CallPlan, budget_bytes, largest_decoded_bytes +from mlx_dfloat.integrate.names import NameMap +from mlx_dfloat.integrate.providers import DF11Provider +from mlx_dfloat.integrate.resident import clear_nonblock, decode_nonblock, install_nonblock +from mlx_dfloat.mflux import _pipeline, require_mflux +from mlx_dfloat.mflux import families as registry +from mlx_dfloat.mflux._compile import uncompiled +from mlx_dfloat.mflux._phases import FamilySizes +from mlx_dfloat.mflux.krea2 import init as kinit +from mlx_dfloat.mflux.krea2 import memory as kmem +from mlx_dfloat.mflux.krea2.init import Krea2Components, ResolvedRepo +from mlx_dfloat.mflux.krea2.names import NONBLOCK_GROUPS, check_variant, krea2_name_map +from mlx_dfloat.mflux.krea2.transformer import ( + Krea2Build, + PerCallNonBlock, + build_transformer, + seam_transformer_class, +) +from mlx_dfloat.mflux.lifecycle import Lifecycle + +require_mflux() +from mflux.models.krea2.variants.txt2img.krea2 import Krea2 # noqa: E402 + +log = logging.getLogger("mlx_dfloat.mflux.krea2") + +# (DF11 repository, base repository), derived from the registry. +MODELS: dict[str, tuple[str, str]] = { + n: (e.df11_repo, e.base_repo) for n, e in registry.MODELS.items() if e.family == "krea2" +} +POLICIES: tuple[str, ...] = ("per-block", "depth2") +# The policies a Krea 2 run has measured; depth2 is accepted, and the report labels it unmeasured. +MEASURED_POLICIES: frozenset[str] = frozenset({"per-block"}) +# What else may stay active across calls besides the assembly baseline, the cached embeddings and the VAE. It is +# larger than one transformer call's decoded non-block weights (1_325_400_064 B on the published model), so the +# retained-memory check does not catch one call's non-block decode left referenced after a drop; the per-call class's +# own cleanup and its tests guard that. +RETAINED_SLACK_BYTES = 2 * 1024**3 + + +def _clear_frames(exc: BaseException, *, outer: BaseException | None = None) -> None: + """Release the locals of the frames ``exc`` and every exception in its chain (causes and contexts) passed through. + + ``outer`` is the exception the caller was already handling when our call began (``sys.exc_info()`` at entry): the + walk stops there, so the caller's own frames keep their locals. The walk keeps a seen-set, so a chain that loops + back on itself ends. + """ + seen: set[int] = set() if outer is None else {id(outer)} + pending: list[BaseException] = [exc] + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + if current.__traceback__ is not None: + traceback.clear_frames(current.__traceback__) + pending.extend(e for e in (current.__cause__, current.__context__) if e is not None) + + +class _EncoderNotResident: + """Takes the text encoder's place while it is dropped: mflux calls the encoder only for a prompt not cached.""" + + def _refuse(self) -> Any: + raise DFloatIntegrationError( + "the text encoder is not resident: a prompt was not encoded before the call " + "(generate_image encodes the prompts it needs first; call encode() for any other use)" + ) + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Refuse: every prompt a call needs is encoded before mflux's loop starts. + + Raises: + DFloatIntegrationError: Always. + """ + del args, kwargs + return self._refuse() + + def get_prompt_embeds(self, *args: Any, **kwargs: Any) -> Any: + """Refuse, as ``__call__`` does (mflux's prompt encoder calls this method). + + Raises: + DFloatIntegrationError: Always. + """ + del args, kwargs + return self._refuse() + + +class DFloatKrea2(Krea2): # type: ignore[misc] # mflux ships no type information + """Krea 2 Raw and Krea 2 Turbo generation from a DFloat11 transformer through mflux. + + The transformer stays compressed and each of its 28 blocks is decoded on the GPU as it runs; the seven non-block + groups (the four text-fusion blocks, the timestep MLP and projection, the text MLP) are decoded at the start of + each transformer call (twice per step when classifier-free guidance runs) and released before the first block's + decode, once its inputs are evaluated, so they are not resident with the set between calls. Each call starts with + MLX's buffer cache emptied and the OS footprint drained, so a guided step's second call does not build on the + first call's cached activations (MLX would otherwise trim its cache only at the memory limit); at 1024² this + keeps a guided Krea 2 Raw step's denoise peak at 21.14 GiB. The step function runs uncompiled on every chip (the + per-block evaluation cannot run inside a compiled function). The decode is lossless. + + The text encoder and the compressed set are never resident together: a new prompt after a generation drops the + set, reloads the encoder, encodes and reloads the set. + """ + + def __init__( + self, + model: str = "krea-2", + *, + df11_path: str | None = None, + base_path: str | None = None, + eval_policy: str = "per-block", + cache_limit: int | None = None, + fit_check: bool = True, + budget_bytes: int | None = None, + quantize: int | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + bake_lora: bool = True, + ) -> None: + """Resolve the checkpoint and the base repository and build the model. + + The checkpoint header, the base's tokenizer and the transformer's extras are read; the encoder, VAE and + compressed weights stay lazy until used. The default checkpoint and base are pinned to the revisions this + package was checked against; a user's own paths are taken as given. ``budget_bytes`` replaces the device's + budget in every call's fit estimate and VAE decision. + + Raises: + DFloatUnsupportedError: ``quantize``, ``lora_paths``, ``lora_scales`` or ``bake_lora=False`` (accepted only + to refuse), an eval policy other than ``"per-block"`` / ``"depth2"``, or an unknown ``model``. + DFloatBackendError: The GPU decoder fails its check on the two packaged self-check groups (raised before + anything is resolved or loaded), or Metal is unavailable. + DFloatFormatError: The checkpoint is not a Krea 2 DF11 checkpoint (the other Krea 2 model's included), its + block count differs from the model's, or the base lacks a component or is another model's. + DFloatAccessError: A gated repository this account may not read (both bases are gated). + DFloatDependencyError: The ``mlx-dfloat[mflux]`` extra is not installed. + DFloatIntegrationError: mflux's own weight definition or mapping does not shape the way this adapter + expects. + """ + _pipeline.refuse_construction_args( + quantize=quantize, lora_paths=lora_paths, lora_scales=lora_scales, bake_lora=bake_lora + ) + _pipeline.check_eval_policy(eval_policy, POLICIES) + if model not in MODELS: + raise DFloatUnsupportedError(f"model {model!r}: this path runs {sorted(MODELS)}") + # The GPU decoder proves itself on the two packaged self-check groups before anything is resolved, downloaded or + # loaded, so a broken decoder refuses here and not minutes later at the first block. + _metal_decode.ensure_canary(force_direct=False) + from mflux.models.common.config.model_config import ModelConfig + + model_config = ModelConfig.from_name(model_name=model, base_model=None) + entry = registry.MODELS[model] + df11_default, base_default = MODELS[model] + pin = None if df11_path else entry.df11_revision + df11 = kinit.resolve(df11_path or df11_default, patterns=kinit.DF11_PATTERNS, revision=pin) + ckpt = open_checkpoint(df11.root) + # The other Krea 2 model's published file is refused here, before the base (about 16 GB) is resolved. + check_variant(model, ckpt) + base_pin = None if base_path else entry.base_revision + base = kinit.resolve( + base_path or base_default, patterns=kinit.BASE_PATTERNS, revision=base_pin + ) + # mflux builds the published transformer with its constructor defaults (Krea2Initializer); the variant check + # refuses the other Krea 2 model's published file. The non-block groups are decoded per call: with them + # resident the denoise step of Krea 2 Raw at 1024² measured over the 32 GB fit line. + build = build_transformer( + ckpt, model=model, transformer_class=seam_transformer_class(per_call=True) + ) + components = kinit.load_base(base.root) + self._assemble( + model=model, + model_config=model_config, + ckpt=ckpt, + df11=df11, + base=base, + components=components, + build=build, + sizes=kmem.sizes_for(ckpt, base.root, nonblock_per_call=True), + eval_policy=eval_policy, + cache_limit=cache_limit, + fit_check=fit_check, + budget_bytes=budget_bytes, + ) + + @classmethod + def _from_parts(cls, **parts: Any) -> "DFloatKrea2": + """Assemble a model from already loaded parts (what ``__init__`` ends with).""" + self = cls.__new__(cls) + self._assemble(**parts) + return self + + def _assemble( + self, + *, + model: str = "custom", + model_config: Any, + ckpt: DF11Checkpoint, + df11: ResolvedRepo, + base: ResolvedRepo, + components: Krea2Components, + build: Krea2Build, + sizes: FamilySizes, + name_map: NameMap | None = None, + eval_policy: str = "per-block", + cache_limit: int | None = None, + fit_check: bool = True, + budget_bytes: int | None = None, + decode: Callable[[MxGroup], DecodeResult] | None = None, + ) -> None: + from mflux.models.krea2.krea2_initializer import Krea2Initializer + + _pipeline.check_eval_policy(eval_policy, POLICIES) + nn.Module.__init__(self) # type: ignore[attr-defined] # not Krea2.__init__: mflux's initializer loads the whole base + # prompt_cache, model_config, callbacks, tiling_config = None (krea2_initializer.py:47-54) + Krea2Initializer._init_config(self, model_config) + self.vae = components.vae + self.text_encoder = components.text_encoder + self.tokenizers = components.tokenizers + self.transformer = build.transformer + self.bits = None + self.lora_paths: list[str] = [] + self.lora_scales: list[float] = [] + self._model = model + self._constants = kmem.CONSTANTS + self._ckpt, self._df11, self._base, self._sizes = ckpt, df11, base, sizes + self._shapes, self._nonblock = build.shapes, build.nonblock + # The build decides: a per-call transformer decodes the non-block groups at each call, any other keeps them + # decoded while the set is resident. + self._per_call = isinstance(build.transformer, PerCallNonBlock) + if any(g in ckpt.groups for g in NONBLOCK_GROUPS) and self._per_call != ( + sizes.nonblock == 0 + ): + mode = "per-call" if self._per_call else "resident" + raise DFloatIntegrationError( + f"the build runs the {mode} non-block decode but the sizes say nonblock={sizes.nonblock}: a " + "per-call build plans with no resident non-block bytes (sizes_for(..., nonblock_per_call=True)), " + "a resident one with them" + ) + self._decode = decode + self._names = krea2_name_map() if name_map is None else name_map + self._cfg_calls = 1 + self._text_tokens: int | None = None + self._predict_mode: str | None = None + self._largest = largest_decoded_bytes(ckpt, self._names, skip=frozenset(NONBLOCK_GROUPS)) + self._policy, self._cache_limit_override, self._fit_check = ( + eval_policy, + cache_limit, + fit_check, + ) + self._budget_override = budget_bytes + self._provider: DF11Provider | None = None + self._plan: CallPlan | None = None + self._tracker = _pipeline.PhaseTracker() + self._launches = 0 + self._baseline_active = int(mx.get_active_memory()) + # The embeddings, one entry per prompt as encoded; mflux's own prompt_cache (keyed by the call's prompt, + # negative prompt and guidance) stays empty, and the _encode_prompts override answers from this one. + self._embeds: dict[str, tuple[mx.array, ...]] = {} + self._lifecycle = Lifecycle( + load_encoders=self._load_encoders, + unload_encoders=self._unload_encoders, + encode=self._encode, + load_set=self._load_set, + unload_set=self._unload_set, + prompt_cache=self._embeds, + retained_bound=self._retained_bound, + encoders_loaded=True, + ) + self.callbacks.register( + _pipeline.VaePoolGuard( + plan=lambda: self._plan, + end_denoise=lambda: self._tracker.end("denoise"), + drop_set=lambda: self._lifecycle.drop_set(), + begin_vae=lambda: self._tracker.begin("vae"), + ) + ) + + # --- lifecycle callbacks ------------------------------------------------------------------------ + + def _retained_bound(self) -> int: + """What may stay active after a drop: the assembly baseline, the cached embeddings, the VAE, slack.""" + embeds = sum(int(a.nbytes) for arrays in self._embeds.values() for a in arrays) + flat: list[tuple[str, mx.array]] = list(tree_flatten(self.vae.parameters())) # type: ignore[arg-type] + vae = sum(int(v.nbytes) for _name, v in flat) + return self._baseline_active + embeds + vae + RETAINED_SLACK_BYTES + + def _load_encoders(self) -> None: + self.text_encoder = kinit.load_text_encoder(self._base.root) + + def _unload_encoders(self) -> None: + # Not None: mflux would call it on a cache miss and fail with a TypeError far from the cause. + self.text_encoder = _EncoderNotResident() + + def _encode(self, prompt: str) -> tuple[mx.array, ...]: + """The prompt's embeddings, as mflux's ``Krea2PromptEncoder.encode_prompt`` computes them.""" + from mflux.models.krea2.model.krea2_text_encoder.prompt_encoder import Krea2PromptEncoder + + return ( + Krea2PromptEncoder.encode_prompt(prompt, self.tokenizers["qwen3vl"], self.text_encoder), + ) + + def _load_set(self) -> None: + outer = sys.exc_info()[1] + try: + self._install_set() + except BaseException as exc: + # The traceback's frames hold the compressed set (_install_set's ``resident``) and any decoded non-block + # weight for as long as the caller keeps the exception; the retry's memory check would then refuse. + _clear_frames(exc, outer=outer) + self._unload_set() # nothing half-installed: the lifecycle believes no set is resident + raise + + def _install_set(self) -> None: + resident = load_resident_set(self._ckpt) + nonblock = {g: self._ckpt.groups[g].matrix_names for g in NONBLOCK_GROUPS if g in resident} + if self._per_call: + # Bound compressed; each transformer call decodes them and releases them after its first block. + self.transformer.bind_nonblock( + {g: resident[g] for g in nonblock}, self._nonblock, self._names, decode=self._decode + ) + else: + install_nonblock( + self.transformer, + decode_nonblock( + {g: resident[g] for g in nonblock}, + nonblock, + self._nonblock, + self._names, + decode=self._decode, + ), + ) + blocks = {n: g for n, g in resident.items() if n not in NONBLOCK_GROUPS} + self._provider = DF11Provider( + blocks, + {n: self._ckpt.groups[n].matrix_names for n in blocks}, + self._names, + decode=self._decode, + ) + # Classifier-free guidance calls the transformer twice per step: verify inside each call. + self.transformer.attach( + self._provider, self._shapes, eval_policy=self._policy, verify_in_call=True + ) + + def _unload_set(self) -> None: + if self._provider is not None: + self._launches += self._provider.launches + self.transformer.detach() + clear_nonblock(self.transformer, self._nonblock) + self._provider = None + + # --- phases ----------------------------------------------------------------------------------- + + @property + def open_phase(self) -> str | None: + """The phase running now (``encode``, ``set_load``, ``denoise`` or ``vae``); None outside them.""" + return self._tracker.open_phase + + def _budget(self) -> int: + """The fit budget: the constructor's ``budget_bytes``, else the device's (read at each call).""" + return self._budget_override if self._budget_override is not None else budget_bytes() + + # --- the public surface ----------------------------------------------------------------------- + + @_pipeline.with_call_caps + def encode(self, *prompts: str) -> None: + """Encode prompts now, so several generations pay the encoder reload once (drops a resident set first). + + Runs under the command's memory caps unless a wired limit is already set. + + Raises: + DFloatUnsupportedError: An empty prompt (Krea 2's tokenizer gives it no tokens at all). A blank one such + as ``" "``, mflux's own negative prompt, is encoded. + """ + if any(p == "" for p in prompts): + raise DFloatUnsupportedError( + "an empty prompt is refused: Krea 2's tokenizer gives it no tokens at all" + ) + self._lifecycle.ensure_embeddings(*dict.fromkeys(prompts)) + + @staticmethod + def cfg_prompts( + prompt: str, *, negative_prompt: str | None, guidance: float + ) -> tuple[str, ...]: + """The prompts a call encodes: the prompt, then the negative one when classifier-free guidance runs. + + mflux 0.20's rule for Krea 2: the negative branch runs for any guidance other than 1.0 (below 1.0 too); a + blank or missing negative prompt is encoded as ``" "``. + """ + if guidance == 1.0: + return (prompt,) + negative = negative_prompt if negative_prompt and negative_prompt.strip() else " " + return (prompt, negative) + + def _encode_prompts( + self, *, prompt: str, negative_prompt: str | None, guidance: float + ) -> tuple[mx.array, mx.array | None]: + """Mflux's hook, answered from the per-prompt cache (the encoder is not resident during the loop). + + Raises: + DFloatIntegrationError: A prompt of the call was not encoded first. + """ + prompts = self.cfg_prompts(prompt, negative_prompt=negative_prompt, guidance=guidance) + missing = [p for p in prompts if p not in self._embeds] + if missing: + raise DFloatIntegrationError(f"prompts {missing} were not encoded before the call") + embeds = self._embeds[prompts[0]][0] + return embeds, (self._embeds[prompts[1]][0] if len(prompts) == 2 else None) + + def _predict( + self, transformer: Any, embeds: mx.array, neg_embeds: mx.array | None, guidance: float + ) -> Any: + """Mflux's predict closure, uncompiled on every chip (the seam evaluates at block boundaries). + + Raises: + DFloatIntegrationError: mflux's factory still compiled it. + """ + fn = uncompiled(Krea2._predict, transformer, embeds, neg_embeds, guidance) + self._predict_mode = "uncompiled" + return fn + + def _entry(self) -> Any: + """The registry entry of the model this instance runs, or None for a model assembled under another name.""" + return registry.MODELS.get(self._model) + + def plan_call(self, *, height: int, width: int, text_tokens: int) -> CallPlan: + """The cache limit, the fit estimate and the VAE strategy of a call (sizes rounded down to multiples of 16). + + ``text_tokens`` is the call's longer prompt in tokens (``memory.text_tokens_for``): a guided step runs one + transformer call per prompt. The set stays resident through the VAE decode when the estimate allows it; + otherwise the call drops it before decoding and the next call reloads it. + + Raises: + DFloatResourceError: The size is above 1024² or the predicted peak exceeds the budget even with the set + dropped, and ``fit_check`` is on. + """ + plan = _pipeline.plan_call_for( + constants=self._constants, + sizes=self._sizes, + largest=self._largest, + policy=self._policy, + cache_limit_override=self._cache_limit_override, + fit_check=self._fit_check, + budget=self._budget(), + height=height, + width=width, + text_tokens=text_tokens, + log=log, + ) + self._plan = plan + return plan + + @staticmethod + def _refuse_blank(prompt: str) -> None: + """Refuse an empty or blank prompt as a user error. + + Raises: + DFloatUnsupportedError: ``prompt`` is empty or only whitespace. + """ + if not prompt.strip(): + raise DFloatUnsupportedError( + f"prompt {prompt!r}: an empty or blank prompt is refused as a user error (an empty prompt gives " + "Krea 2's tokenizer no tokens at all)" + ) + + @staticmethod + def _resolve_scheduler(scheduler: str | None) -> str: + """Mflux's scheduler resolution (None and ``"linear"`` become ``"er_sde"``), with an unknown name refused here. + + Raises: + DFloatUnsupportedError: A scheduler Krea 2's pipeline does not run. + """ + try: + return str(Krea2._resolve_scheduler(scheduler)) + except ValueError: + raise DFloatUnsupportedError( + f"scheduler {scheduler!r}: Krea 2 runs er_sde, euler, linear (linear is er_sde)" + ) from None + + @_pipeline.with_call_caps + def generate_image( + self, + seed: int, + prompt: str, + num_inference_steps: int | None = None, + height: int = 1024, + width: int = 1024, + guidance: float | None = None, + negative_prompt: str | None = None, + image_path: Path | str | None = None, + image_strength: float | None = None, + scheduler: str | None = None, + pid_decode: bool = False, + pid_degrade_sigma: float = 0.0, + ) -> Any: + """Mflux's ``generate_image`` behind a prelude: refusals, cache limit, fit check, prompt encoding, set load. + + ``num_inference_steps``, ``guidance`` and ``scheduler`` left as None take this model's defaults, mflux's per + variant: Krea 2 Raw 25 steps at guidance 1.0 and Krea 2 Turbo 8 steps at guidance 1.0, so one transformer call + per step; ``"er_sde"`` for both. Krea 2 Raw's model card runs 52 steps at guidance 3.5. Classifier-free + guidance runs for any guidance other than 1.0, below 1.0 too, as two transformer calls per step, with the + negative prompt or, when it is missing or blank, mflux's ``" "``. Each call runs under the command's memory + caps unless a wired limit is already set. + + Raises: + DFloatUnsupportedError: An empty or blank ``prompt`` (a user error), a scheduler other than ``er_sde``, + ``euler`` or ``linear``, ``image_path`` / ``image_strength`` (img2img) or ``pid_decode``; all before + anything loads. + DFloatResourceError: The fit check refuses the call (the estimate, or a size above 1024²), or memory stayed + active after a drop. + DFloatFormatError: A block's decode reported an error; the set is dropped for a clean retry. + DFloatIntegrationError: mflux compiled the step function despite the bypass. + """ + outer = sys.exc_info()[ + 1 + ] # the exception our caller is handling, if any: its frames are not ours to clear + self._refuse_blank(prompt) + if image_path is not None or image_strength is not None: + _pipeline.refuse("image_path/image_strength", "img2img") + if pid_decode: + _pipeline.refuse( + "pid_decode", + "mflux's alternative image decoder (PiD) loads its own caption encoder next to the compressed set", + ) + entry = self._entry() + if guidance is None: + guidance = ( + 1.0 if entry is None or entry.default_guidance is None else entry.default_guidance + ) + if num_inference_steps is None: + num_inference_steps = 8 if entry is None else entry.default_steps + if scheduler is None: + scheduler = ( + "er_sde" + if entry is None or entry.default_scheduler is None + else entry.default_scheduler + ) + # Before anything loads; mflux resolves the name again inside its own generate_image, the same way. + self._resolve_scheduler(scheduler) + prompts = self.cfg_prompts(prompt, negative_prompt=negative_prompt, guidance=guidance) + text_tokens = kmem.text_tokens_for(self.tokenizers["qwen3vl"], prompts) + plan = self.plan_call(height=height, width=width, text_tokens=text_tokens) + self._cfg_calls, self._text_tokens = len(prompts), text_tokens + self._tracker.begin("encode") + try: + self.encode(*prompts) + finally: # a failed encode closes its phase: the watchdog's context must not name it afterwards + self._tracker.end("encode") + self._check_encode_peak(plan) + self._tracker.begin("set_load") + try: + self._lifecycle.ensure_set() + finally: + self._tracker.end("set_load") + self._tracker.begin("denoise") # before set_cache_limit: a raise here must change nothing + previous = mx.set_cache_limit(plan.cache_limit) + try: + return super().generate_image( + seed=seed, + prompt=prompt, + num_inference_steps=num_inference_steps, + height=height, + width=width, + guidance=guidance, + negative_prompt=negative_prompt, + image_path=None, + image_strength=None, + scheduler=scheduler, + pid_decode=False, + pid_degrade_sigma=pid_degrade_sigma, + ) + except BaseException as exc: + # The traceback keeps the seam frames, and with them decoded weights or the resident set, alive for as + # long as the caller holds the exception (and past drop_set's own gc.collect()). Clear them for any + # exception and every one in its chain raised in this call (not the one the caller was handling). + _clear_frames(exc, outer=outer) + if isinstance(exc, DFloatFormatError): + self._lifecycle.drop_set() # a corrupt block: the retry starts from a clean load + raise + finally: + mx.set_cache_limit(previous) + self._tracker.end("denoise") + self._tracker.end("vae") + + def _check_encode_peak(self, plan: CallPlan) -> None: + """Warn when the encode held more than the estimate's encode phase (evaluated language model + term) + slack. + + The estimate counts only the encoder's evaluated language-model tensors (``memory.encoder_bytes_used``); an + mflux change that also loads the vision tower would make it under-predict the encode phase, and so may a prompt + much longer than the one the encode term was measured on (``memory.encode_peak_warning``). The measure is MLX's + peak minus what was active when the phase began, so an encode that first dropped a resident set reads low (no + warning), never high. + """ + record = self._tracker.peaks.get("encode") + if record is None: + return + message = kmem.encode_peak_warning( + record["mlx_peak"] - record["active_at_start"], + kmem.encode_peak_bound(plan.estimate.phases["encode"], self._constants), + text_tokens=self._text_tokens or 0, + ) + if message is not None: + warnings.warn(message, stacklevel=3) + + def report(self) -> dict[str, Any]: + """What this model is and what its calls cost: repositories, policy, limits, estimate, peaks, counts, versions. + + Peaks are sampled at phase boundaries (labelled ``"sampled"``); the fit estimate is a prediction. ``predict`` + is how the last call's step function ran (``"uncompiled"``; None before any call). ``cfg_calls_per_step`` and + ``text_tokens`` are the last call's (2 calls when classifier-free guidance ran; the longer prompt's tokens). + ``sizes["encoders"]`` is the text encoder's evaluated language-model bytes, not its file size. ``layout`` says + where the checkpoint's config came from (the pinned layout a published file matched). ``nonblock`` is + ``"per-call"`` (the seven non-block groups decoded at each transformer call; ``sizes["nonblock"]`` is 0) or + ``"resident"`` (decoded once when the set loads). + ``eval_policy_status`` is ``"unmeasured"`` for depth2, which no Krea 2 run has measured. + """ + launches = self._launches + (self._provider.launches if self._provider is not None else 0) + fit = None if self._plan is None else self._plan.estimate + try: + mflux_version: str | None = metadata.version("mflux") + except metadata.PackageNotFoundError: + mflux_version = None + return { + "family": "krea2", + "model": self._model, + "df11": { + "root": str(self._df11.root), + "repo_id": self._df11.repo_id, + "revision": self._df11.revision, + "config_source": self._ckpt.config_source, + }, + "base": { + "root": str(self._base.root), + "repo_id": self._base.repo_id, + "revision": self._base.revision, + }, + "layout": self._ckpt.config_source, + "nonblock": "per-call" if self._per_call else "resident", + "eval_policy": self._policy, + "eval_policy_status": "measured" if self._policy in MEASURED_POLICIES else "unmeasured", + "cfg_calls_per_step": self._cfg_calls, + "text_tokens": self._text_tokens, + "predict": self._predict_mode, + "nonblock_groups": sorted(g for g in NONBLOCK_GROUPS if g in self._ckpt.groups), + "cache_limit_in_force": None if self._plan is None else self._plan.cache_limit, + "drop_set_before_vae": None if self._plan is None else self._plan.drop_set_before_vae, + "fit": None + if fit is None + else { + "label": "predicted", + "phases": dict(fit.phases), + "peak_phase": fit.peak_phase, + "peak_bytes": fit.peak_bytes, + "budget_bytes": fit.budget_bytes, + "fits": fit.fits, + }, + "sizes": asdict(self._sizes), + "peaks": {"label": "sampled at phase boundaries", **self._tracker.peaks}, + "decode_launches": launches, + "lifecycle": self._lifecycle.counters.as_dict(), + "versions": {"mlx": mx.__version__, "mflux": mflux_version}, # type: ignore[attr-defined] + } + + def save_model(self, base_path: str) -> None: + """Refused: the DFloat11 checkpoint is the saved form; there is nothing of mflux's to write. + + Raises: + DFloatUnsupportedError: Always. + """ + del base_path + _pipeline.refuse("save_model", "the DFloat11 repository is the checkpoint") + + def freeze(self, **kwargs: Any) -> None: + """Freeze the components that are loaded (a dropped encoder is a stand-in, not a module).""" + del kwargs + for module in (self.vae, self.transformer, self.text_encoder): + if not isinstance(module, _EncoderNotResident): + module.freeze() diff --git a/src/mlx_dfloat/mflux/krea2/names.py b/src/mlx_dfloat/mflux/krea2/names.py new file mode 100644 index 0000000..a0773e1 --- /dev/null +++ b/src/mlx_dfloat/mflux/krea2/names.py @@ -0,0 +1,129 @@ +"""Krea 2 naming: the block kind, the decoded sub-paths, the map derived from mflux, the group and variant checks. + +The block kind and its run order: mflux 0.20.0 ``Krea2Transformer`` runs one list, ``blocks``, inline. The decoded +sub-paths and the seven non-block groups (the four text-fusion blocks, ``tmlp``, ``tproj``, ``txtmlp``) are the +``pattern_dict`` of the two ComfyUI single-file DF11 exports (Krea 2 Raw and Krea 2 Turbo, identical); their stored +order was checked against the originals' bytes. mflux's ``Krea2WeightMapping`` templates its block targets with +``{layer}`` and lists the text-fusion blocks by concrete index as identities, so the shared derivation reads it as it +is. The text-fusion blocks are not a seamed kind: they run once per call before the blocks, like ``tmlp``, ``tproj`` +and ``txtmlp``. The model decodes all seven at the start of each transformer call +(``mlx_dfloat.mflux.krea2.transformer.PerCallNonBlock``); the resident build keeps them decoded with the set. +""" + +import re +from typing import Any + +from mlx_dfloat.errors import DFloatFormatError +from mlx_dfloat.integrate.names import StaticNameMap +from mlx_dfloat.mflux import require_mflux +from mlx_dfloat.mflux._mapping import name_map_from_mapping + +KIND = "blocks" +RUN_ORDER: tuple[str, ...] = (KIND,) +BLOCK_SUBS: tuple[str, ...] = ( + "attn.wq", + "attn.wk", + "attn.wv", + "attn.gate", + "attn.wo", + "mlp.gate", + "mlp.up", + "mlp.down", +) +TEXT_SUBS: tuple[str, ...] = ( + BLOCK_SUBS # a text-fusion block stores the same eight matrices, in the same order +) +MATRIX_SUBS: dict[str, tuple[str, ...]] = {KIND: BLOCK_SUBS} +TEXT_FUSION_GROUPS: tuple[str, ...] = ( + "txtfusion.layerwise_blocks.0", + "txtfusion.layerwise_blocks.1", + "txtfusion.refiner_blocks.0", + "txtfusion.refiner_blocks.1", +) +NONBLOCK_GROUPS: dict[str, tuple[str, ...]] = { + "tmlp": ("tmlp.0.weight", "tmlp.2.weight"), + "tproj": ("tproj.1.weight",), + "txtmlp": ("txtmlp.1.weight", "txtmlp.3.weight"), + **{g: tuple(f"{g}.{s}.weight" for s in TEXT_SUBS) for g in TEXT_FUSION_GROUPS}, +} +DEFAULT_LAYERS = 28 # mflux 0.20.0 ``Krea2Transformer.__init__(layers=28)`` +LAYOUT_FOR_MODEL: dict[str, str] = { + "krea-2-raw": "krea-2-raw-comfyui", + "krea-2": "krea-2-turbo-comfyui", +} +_LABELS = {"krea-2-raw": "Krea 2 Raw", "krea-2": "Krea 2 Turbo"} +_SOURCE = re.compile(r"header and spot checks match layout (\S+) ") + + +def krea2_name_map() -> StaticNameMap: + """The map read from mflux's ``Krea2WeightMapping`` (imports mflux). + + Raises: + DFloatDependencyError: The optional ``mflux`` extra is not installed. + DFloatIntegrationError: mflux's mapping no longer shapes the way this adapter expects. + """ + require_mflux() + from mflux.models.krea2.weights.krea2_weight_mapping import Krea2WeightMapping + + return name_map_from_mapping( + Krea2WeightMapping.get_transformer_mapping(), matrix_subs=MATRIX_SUBS + ) + + +def check_krea2_groups(ckpt: Any) -> dict[str, int]: + """The block count: block groups contiguous from 0, each holding the eight matrices, plus non-block groups. + + Each non-block group is optional: a checkpoint that stores it as BF16 extras loads it as plain weights. mflux + builds the block list from its own count, so a hole would leave a block on placeholders. + + Raises: + DFloatFormatError: A group of another name, a hole in the block indices, no block groups, a block whose + matrices differ from the eight, or a non-block group with other matrices. + """ + indices: set[int] = set() + for name, group in ckpt.groups.items(): + if name in NONBLOCK_GROUPS: + if tuple(group.matrix_names) != NONBLOCK_GROUPS[name]: + raise DFloatFormatError( + f"{name}: holds {list(group.matrix_names)}, expected {list(NONBLOCK_GROUPS[name])}" + ) + continue + kind, _dot, idx = name.partition(".") + # A zero-padded index ("01") is refused: the seam asks for the group by its canonical name. + if kind != KIND or not (idx.isascii() and idx.isdigit() and idx == str(int(idx))): + raise DFloatFormatError( + f"{name}: not a Krea 2 group (blocks: {KIND}.; non-block groups accepted: " + f"{', '.join(NONBLOCK_GROUPS)})" + ) + if tuple(group.matrix_names) != tuple(f"{name}.{sub}.weight" for sub in BLOCK_SUBS): + raise DFloatFormatError( + f"{name}: holds {len(group.matrix_names)} matrices {list(group.matrix_names)}; a Krea 2 block " + f"holds {', '.join(BLOCK_SUBS)}" + ) + indices.add(int(idx)) + if not indices: + raise DFloatFormatError(f"no {KIND} groups in the checkpoint") + if indices != set(range(len(indices))): + raise DFloatFormatError(f"{KIND} groups are not contiguous: {sorted(indices)}") + return {KIND: len(indices)} + + +def check_variant(model: str, ckpt: Any) -> None: + """Refuse the other Krea 2 model's published checkpoint (Raw and Turbo share every name and shape). + + The checkpoint's ``config_source`` names the pinned layout it opened through; a source that names no layout (a + ``config.json``) or a layout of neither Krea model passes. + + Raises: + DFloatFormatError: The checkpoint is the other Krea 2 model's published file. + """ + found = _SOURCE.match(ckpt.config_source) + if found is None: + return + key = found.group(1) + for other, layout in LAYOUT_FOR_MODEL.items(): + if other != model and layout == key: + raise DFloatFormatError( + f"{ckpt.root.name}: this is the {_LABELS[other]} DF11 checkpoint (layout {key}); run it with " + f"--model {other}" + ) diff --git a/src/mlx_dfloat/mflux/krea2/transformer.py b/src/mlx_dfloat/mflux/krea2/transformer.py new file mode 100644 index 0000000..3e3b9da --- /dev/null +++ b/src/mlx_dfloat/mflux/krea2/transformer.py @@ -0,0 +1,480 @@ +"""mflux's Krea 2 transformer with the class-swap seam. + +mflux 0.20.0 ``Krea2Transformer.__call__`` runs its one block list, ``blocks``, inline, after the text fusion, the +timestep MLP and the text MLP have run once. The pipeline's ``_predict`` factory compiles the step function off base +and Pro M1/M2 chips (an M1 Max or Ultra, and every M3 or later, compile); the model runs it through +``mlx_dfloat.mflux._compile.uncompiled`` on every chip, because the block seam evaluates at block boundaries, which +cannot happen inside ``mx.compile``. A classifier-free guidance step is two transformer calls, so each block decodes +twice per step. + +The build installs placeholders on the eight block matrices and on the seven non-block groups' matrices (the four +text-fusion blocks, ``tmlp``, ``tproj``, ``txtmlp``) and loads every other parameter from the checkpoint's extras (or a +BF16 base's). The non-block groups run once per call, before the blocks. They are either decoded once when the set +loads and kept installed (the resident class), or decoded at the start of each transformer call +(``seam_transformer_class(per_call=True)``, ``PerCallNonBlock``) at the cost of decoding them twice per guided step. + +The per-call path empties MLX's buffer cache at two points, without changing any arithmetic: at the start of each +call (``_release_and_drain``), and before block 0's decode, once block 0's inputs are evaluated and the non-block +weights are back on their placeholders. Otherwise a guided step's second call would decode the non-block groups and +build its pre-block graph on top of the first call's freed activations, which MLX keeps cached until active plus +cached memory reaches the memory limit. Both points wait for the OS footprint to fall, because a released Metal +buffer leaves it only after the next GPU submission. At 1024² this keeps a guided Krea 2 Raw step's denoise peak at +21.14 GiB. +""" + +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from functools import cache +from pathlib import Path +from typing import Any + +import mlx.core as mx +import numpy as np +from mlx.utils import tree_flatten + +from mlx_dfloat._safetensors import TensorInfo, read_array, read_header +from mlx_dfloat._watchdog import phys_footprint +from mlx_dfloat.decode import DecodeResult +from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError +from mlx_dfloat.format import DF11Checkpoint, MxGroup +from mlx_dfloat.integrate.blockseam import StepSeam, seam_blocks +from mlx_dfloat.integrate.coverage import check_extras_cover, extras_plan, load_extras +from mlx_dfloat.integrate.names import BlockShapes, NameMap, Shapes +from mlx_dfloat.integrate.placeholders import PLACEHOLDER, install_placeholders +from mlx_dfloat.integrate.providers import WeightProvider, read_bf16 +from mlx_dfloat.integrate.resident import ( + NonBlockShapes, + clear_nonblock, + decode_nonblock, + install_nonblock, + install_nonblock_placeholders, +) +from mlx_dfloat.mflux import require_mflux +from mlx_dfloat.mflux.flux2.transformer import base_transformer_files_index +from mlx_dfloat.mflux.krea2.names import ( + DEFAULT_LAYERS, + KIND, + NONBLOCK_GROUPS, + RUN_ORDER, + check_krea2_groups, + check_variant, + krea2_name_map, +) + +MAX_BUILD_ACTIVE_BYTES = 2 * 1024**3 + +# mflux 0.20.0's Krea2Transformer builds no parameter of its own: the RoPE embedder holds no arrays, the timestep +# frequencies are computed per call, and the modulation tables and norm scales come from the checkpoint. Kept as a +# named constant so the build reads like the other families' and a future mflux table has one place to go. +COMPUTED_PARAMS: frozenset[str] = frozenset() + + +# The policies under which a call releases the non-block weights before block 0's decode. "none" promises no +# evaluation inside the call, so under it they stay installed until the call's own cleanup. +RELEASE_POLICIES: frozenset[str] = frozenset({"per-block", "depth2"}) +_BINDING = "_dfloat_nonblock" + +# The wrapper calls MLX through these names so a test can record the order of the evaluation and the drain. +_eval = mx.eval +_clear_cache = mx.clear_cache + +# The drain: MLX's cache, the OS footprint and the clock are read through these names so a test can script them. +_synchronize = mx.synchronize +_cache_memory = mx.get_cache_memory +_footprint = phys_footprint +_clock = time.perf_counter +_sleep = time.sleep +# Wait for the footprint to fall by this fraction of the released bytes, at most this long, polling at this period. +DRAIN_FRACTION = 0.9 +DRAIN_TIMEOUT_S = 0.25 +DRAIN_POLL_S = 0.0005 +# Below this many released bytes the drain clears and kicks but does not wait: a small release cannot stack into a +# peak, and a few MB may not show in the footprint at all (the wait would run to its cap). +DRAIN_MIN_BYTES = 64 * 1024**2 + + +def _kick() -> None: + """One tiny GPU evaluation: the driver frees released buffers after the next command submission.""" + mx.eval(mx.zeros((1,)) + 0) + + +def _release_and_drain() -> bool: + """Empty MLX's buffer cache and wait until the OS footprint has dropped; touches no arithmetic. + + A buffer MLX releases leaves the process footprint only after the next GPU submission (about 20 ms with the GPU + busy, 80-700 ms idle), so an allocation made in that window stacks on it. The drain waits for the GPU first, so + every buffer whose free was still pending (MLX frees a buffer the GPU used from the command's completion handler) + has reached the cache. Then it reads the cached bytes and the footprint, clears the cache, submits one tiny + evaluation, and polls the footprint until it has fallen by ``DRAIN_FRACTION`` of the released bytes or + ``DRAIN_TIMEOUT_S`` has passed. The tiny evaluation's own buffers reach the cache once its command completes, so + the drain ends by waiting for the GPU again and clearing the cache once more. Returns whether the drop was seen (a + release under ``DRAIN_MIN_BYTES`` counts as seen without a wait). + """ + _synchronize() + released = int(_cache_memory()) + waits = released >= DRAIN_MIN_BYTES + before = _footprint() if waits else 0 + _clear_cache() + _kick() + dropped = True + if waits: + deadline = _clock() + DRAIN_TIMEOUT_S + while before - _footprint() < DRAIN_FRACTION * released: + if _clock() >= deadline: + dropped = False + break + _sleep(DRAIN_POLL_S) + _synchronize() + _clear_cache() + return dropped + + +class _NonBlockBinding: + """The bound non-block groups and how to decode them (a plain object: a dict attribute would join the module tree).""" + + __slots__ = ("decode", "groups", "name_map", "shapes") + + def __init__( + self, + groups: Mapping[str, MxGroup], + shapes: NonBlockShapes, + name_map: NameMap, + decode: Callable[[MxGroup], DecodeResult] | None, + ) -> None: + self.groups = dict(groups) + self.shapes = dict(shapes) + self.name_map = name_map + self.decode = decode + + +class ReleaseAtFirstBlock: + """Wraps the seam's provider for one call and releases the non-block weights before block 0's decode. + + Through the seam's ``before_block`` hook, at block 0: evaluate the block's inputs (the whole pre-block graph, so + nothing lazy still references a decoded non-block array), put the non-block weights back on their placeholders, + empty MLX's buffer cache and wait for the footprint to fall (``_release_and_drain``: a released buffer otherwise + stays in the cache, and so in the footprint), and only then let block 0 request its weights. Everything else is + the wrapped provider's. Under a policy outside ``RELEASE_POLICIES`` (``"none"``) nothing is released mid-call; the + call's own cleanup does it. + """ + + def __init__( + self, inner: WeightProvider, transformer: Any, shapes: NonBlockShapes, *, policy: str + ) -> None: + """Wrap ``inner`` for one call of ``transformer`` under ``policy``.""" + self._inner = inner + self._transformer = transformer + self._shapes = shapes + self._releases = policy in RELEASE_POLICIES + self._released = False + self.launching = inner.launching + self.policies = inner.policies + + @property + def launches(self) -> int: + """The wrapped provider's launch count.""" + return self._inner.launches + + @launches.setter + def launches(self, value: int) -> None: + self._inner.launches = value + + @property + def pending(self) -> Any: + """The wrapped provider's unchecked status words (None when it keeps none), for the seam's ``begin_step``.""" + return getattr(self._inner, "pending", None) + + def before_block( + self, block_name: str, args: tuple[Any, ...], kwargs: Mapping[str, Any] + ) -> None: + """At the call's first block: evaluate its inputs, release the non-block weights, clear the cache.""" + if self._released or not self._releases: + return + flat: dict[str, Any] = dict(tree_flatten((args, dict(kwargs)))) + _eval(*(leaf for leaf in flat.values() if isinstance(leaf, mx.array))) + clear_nonblock(self._transformer, self._shapes) + _release_and_drain() + self._released = True + + def weights_for(self, block_name: str, shapes: BlockShapes) -> dict[str, mx.array]: + """The wrapped provider's weights.""" + return self._inner.weights_for(block_name, shapes) + + def verify(self) -> None: + """The wrapped provider's deferred checks.""" + self._inner.verify() + + def reset(self) -> None: + """The wrapped provider's reset.""" + self._inner.reset() + + +class PerCallNonBlock(StepSeam): + """``StepSeam`` whose transformer decodes the bound non-block groups at the start of each call. + + Each call first empties MLX's buffer cache and waits for the footprint to fall (``_release_and_drain``), so it + never starts on the previous call's cached activations. The decoded weights are installed for the call's pre-block + graph and put back on their placeholders before block 0 requests its own weights, once that graph is evaluated + (``ReleaseAtFirstBlock``), and again when the call ends however it ends. The call evaluates its output before it returns, so no lazy tail keeps this call's arrays alive + into the next call of a guided step. + """ + + def bind_nonblock( + self, + groups: Mapping[str, MxGroup], + shapes: NonBlockShapes, + name_map: NameMap, + *, + decode: Callable[[MxGroup], DecodeResult] | None = None, + ) -> None: + """Keep the compressed non-block groups to decode per call (Metal by default); ``detach`` forgets them.""" + vars(self)[_BINDING] = _NonBlockBinding(groups, shapes, name_map, decode) + + def nonblock_bound(self) -> bool: + """Whether non-block groups are bound.""" + return vars(self).get(_BINDING) is not None + + def detach(self) -> None: + """Forget the provider, the shapes and the bound non-block groups (the set is dropping).""" + super().detach() + vars(self).pop(_BINDING, None) + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Decode and install the non-block weights, run the seam step with the release wrapper, evaluate, clear them. + + Raises: + DFloatIntegrationError: No non-block groups are bound, or nothing is attached. + DFloatFormatError: A non-block group's decode reported an error. + """ + binding: _NonBlockBinding | None = vars(self).get(_BINDING) + if binding is None: + raise DFloatIntegrationError( + "call bind_nonblock(groups, shapes, name_map) before running a step" + ) + state = self._state() + inner = state.provider + # The previous call's freed activations leave MLX's cache before this call decodes and builds its pre-block + # graph, so every call follows a first call's trajectory instead of stacking on the last one's cache. + _release_and_drain() + try: + matrices = {g: NONBLOCK_GROUPS[g] for g in binding.groups} + weights = decode_nonblock( + binding.groups, + matrices, + binding.shapes, + binding.name_map, + decode=binding.decode, + eval_together=True, + ) + install_nonblock(self, weights) + del weights + state.provider = ReleaseAtFirstBlock(inner, self, binding.shapes, policy=state.policy) + out = super().__call__(*args, **kwargs) + _eval(out) + return out + finally: + state.provider = inner + clear_nonblock(self, binding.shapes) + + +@cache +def _seam_class(per_call: bool) -> type: + from mflux.models.krea2.model.krea2_transformer.transformer import Krea2Transformer + + if per_call: + return type("SeamKrea2TransformerPerCall", (PerCallNonBlock, Krea2Transformer), {}) + return type("SeamKrea2Transformer", (StepSeam, Krea2Transformer), {}) + + +def seam_transformer_class(*, per_call: bool = False) -> type: + """``StepSeam`` composed in front of mflux's ``Krea2Transformer`` (imports mflux); one class per mode. + + ``per_call`` puts ``PerCallNonBlock`` in front instead: the non-block groups are bound with ``bind_nonblock`` and + decoded at each call. Importing the class also registers the ``er_sde`` and ``euler`` schedulers (mflux registers + them in its Krea 2 package). + + Raises: + DFloatDependencyError: The optional ``mflux`` extra is not installed. + """ + require_mflux() + return _seam_class(per_call) + + +def block_lists(transformer: Any) -> list[tuple[str, Sequence[Any]]]: + """The one block list the forward pass runs, ``blocks``.""" + return [(kind, getattr(transformer, kind)) for kind in RUN_ORDER] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Krea2Build: + """A built transformer, its block shapes, the non-block matrices' shapes and the block count built.""" + + transformer: Any + shapes: Shapes + nonblock: NonBlockShapes + counts: dict[str, int] + + +def build_transformer( + ckpt: DF11Checkpoint, + *, + model: str | None = None, + transformer_kwargs: Mapping[str, Any] | None = None, + name_map: NameMap | None = None, + n_layers: int | None = None, + extras: Mapping[str, tuple[Path, TensorInfo]] | None = None, + nonblock_from_extras: bool = False, + transformer_class: Callable[..., Any] | None = None, +) -> Krea2Build: + """Construct the seamed transformer: block and non-block matrices as placeholders, every other parameter loaded. + + ``model`` (``"krea-2-raw"`` or ``"krea-2"``), when given, refuses the other model's published checkpoint before + anything is built. ``transformer_kwargs`` are the transformer's constructor arguments (none for the published + model); the checkpoint's block count must equal the depth they build (``layers``; mflux's default 28 when they + leave it out). ``n_layers`` then builds fewer blocks (a partial build for tests). ``extras`` replaces the + checkpoint's own extras (a BF16 base's tensors, via ``base_extras``); with ``nonblock_from_extras`` the non-block + matrices load from them as plain weights instead of placeholders (the BF16 side of an identity check). A + non-block group the checkpoint stores as extras loads as plain weights either way. The active memory the build + added is asserted. + + Raises: + DFloatFormatError: The other Krea 2 model's checkpoint; the groups are not Krea 2's; the checkpoint's block + count differs from the model's depth; an extra is not BF16 or has the wrong shape. + DFloatIntegrationError: ``n_layers`` out of range, an uncovered parameter or an extra without one, or too + much active memory added. + DFloatDependencyError: The ``mflux`` extra is missing (default class and map only). + """ + if model is not None: + check_variant(model, ckpt) + counts = check_krea2_groups(ckpt) + kwargs = dict(transformer_kwargs or {}) + want = int(kwargs.get("layers", DEFAULT_LAYERS)) + if counts[KIND] != want: + raise DFloatFormatError(f"checkpoint has {counts[KIND]} blocks; this model builds {want}") + n = counts[KIND] if n_layers is None else n_layers + if not 0 <= n <= counts[KIND]: + raise DFloatIntegrationError(f"asked for {n} blocks; the checkpoint has {counts[KIND]}") + names = krea2_name_map() if name_map is None else name_map + make = seam_transformer_class() if transformer_class is None else transformer_class + before = int(mx.get_active_memory()) + tf = make(**{**kwargs, "layers": n}) + lists = block_lists(tf) + shapes = install_placeholders(lists, names) + seam_blocks(lists, tf.seam_cell) + nonblock_matrices = ( + [] + if nonblock_from_extras + else [m for g in NONBLOCK_GROUPS if g in ckpt.groups for m in ckpt.groups[g].matrix_names] + ) + nonblock = install_nonblock_placeholders(tf, nonblock_matrices, names) + matrix_paths = {f"{b}.{a}.weight" for b, per in shapes.items() for a in per} | set(nonblock) + built = {KIND: n} + plan = extras_plan(ckpt, names, counts=built, extras=extras) + params = set(dict(tree_flatten(tf.parameters()))) - COMPUTED_PARAMS + check_extras_cover(params, (p for p, _f, _i in plan), matrix_paths) + load_extras(tf, plan, names) + mx.eval(PLACEHOLDER) + added = int(mx.get_active_memory()) - before + if added >= MAX_BUILD_ACTIVE_BYTES: + raise DFloatIntegrationError(f"extras added {added / 1024**3:.2f} GiB of active memory") + return Krea2Build(transformer=tf, shapes=shapes, nonblock=nonblock, counts=built) + + +def base_extras( + index: Mapping[str, tuple[Path, TensorInfo]], ckpt: DF11Checkpoint +) -> dict[str, tuple[Path, TensorInfo]]: + """The BF16 base's tensors minus the block groups' matrices (the non-block matrices stay, as plain weights).""" + block = { + m for name, g in ckpt.groups.items() if name not in NONBLOCK_GROUPS for m in g.matrix_names + } + return {name: entry for name, entry in index.items() if name not in block} + + +def base_bf16_weights(path: Path, names: Iterable[str]) -> dict[str, mx.array]: + """The named tensors of a base's native single file (``raw.safetensors``) as BF16, as mflux loads them. + + A BF16 tensor is read as it is (a bit view); an FP32 one is cast with ``astype(mx.bfloat16)``, mflux's own load + cast, which rounds to nearest even. The published Krea 2 files hold the transformer's matrices in BF16 and its + extras and five non-block matrices in FP32. + + Raises: + DFloatFormatError: A name the file does not hold, a tensor of another dtype, or an unreadable header. + """ + header = read_header(path) + out: dict[str, mx.array] = {} + for name in names: + info = header.get(name) + if info is None: + raise DFloatFormatError(f"{path.name} has no tensor {name!r}") + if info.dtype == "BF16": + out[name] = read_bf16(path, info) + elif info.dtype == "F32": + cast = mx.array(np.ascontiguousarray(read_array(path, info))).astype(mx.bfloat16) + _eval(cast) # one FP32 source alive at a time, not all of them until the caller's eval + out[name] = cast + else: + raise DFloatFormatError(f"{name} is {info.dtype}; the base must hold BF16 or F32") + return out + + +def install_base_weights( + transformer: Any, + ckpt: DF11Checkpoint, + path: Path, + name_map: NameMap, + nonblock: NonBlockShapes, +) -> list[str]: + """Replace every extra the transformer was built with, and every non-block matrix, with the base file's. + + The reference side of an image identity check then reads nothing of the DF11 file but its group shapes: the + extras the build loaded from the checkpoint and the non-block placeholders (``nonblock``, the build's shapes) all + take the base's tensors (``base_bf16_weights``); the block matrices stay placeholders for a streaming provider. + Returns the checkpoint names replaced. + + Raises: + DFloatFormatError: A base tensor missing, of another dtype, or of another shape than the model's parameter. + """ + params: dict[str, Any] = dict(tree_flatten(transformer.parameters())) + extras = [n for n in ckpt.extras if name_map.param_name(n) in params] + matrices = [ + m + for g in NONBLOCK_GROUPS + if g in ckpt.groups + for m in ckpt.groups[g].matrix_names + if name_map.param_name(m) in nonblock + ] + weights = base_bf16_weights(path, [*extras, *matrices]) + plain: list[tuple[str, mx.array]] = [] + installed: dict[str, mx.array] = {} + for name, weight in weights.items(): + param = name_map.param_name(name) + want = nonblock[param] if param in nonblock else tuple(params[param].shape) + if tuple(weight.shape) != tuple(want): + raise DFloatFormatError( + f"{name} has shape {tuple(weight.shape)}; the model needs {tuple(want)}" + ) + if param in nonblock: + installed[param] = weight + else: + plain.append((param, weight)) + transformer.load_weights(plain, strict=False) + install_nonblock(transformer, installed) + return [*extras, *matrices] + + +__all__ = [ + "COMPUTED_PARAMS", + "MAX_BUILD_ACTIVE_BYTES", + "RELEASE_POLICIES", + "Krea2Build", + "PerCallNonBlock", + "ReleaseAtFirstBlock", + "base_bf16_weights", + "base_extras", + "base_transformer_files_index", + "block_lists", + "build_transformer", + "install_base_weights", + "seam_transformer_class", +] diff --git a/src/mlx_dfloat/mflux/qwen21/memory.py b/src/mlx_dfloat/mflux/qwen21/memory.py index 7604007..9e9051d 100644 --- a/src/mlx_dfloat/mflux/qwen21/memory.py +++ b/src/mlx_dfloat/mflux/qwen21/memory.py @@ -28,7 +28,7 @@ TEXT_MAX_LENGTH = 2048 MAX_MEASURED_PIXELS = 1024 * 1024 # no run above 1024² on this path -# Measured on 2026-10-08 (the first one-step de-risk run, at the old cache limit): Qwen-Image-2.1 DF11 +# Measured on 2026-10-08 (the first one-step calibration run, at the old cache limit): Qwen-Image-2.1 DF11 # (mingyi456/Qwen-Image-2.1-DF11-ComfyUI @ 1b22a3a) over the Qwen/Qwen-Image-2.1 base (@ d26bb61), 1024², seed 42, # guidance 4 with the negative prompt " " (two transformer calls per step); one process: build, encode, drop, set load, # one step, VAE decode on the resident set; M1 Max 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0. Bytes throughout. @@ -43,9 +43,9 @@ # VAE_TRANSIENT, by the overall-peak rule over every 1024² sample: the highest measured peak of the VAE decode on the # resident set minus the VAE phase's other terms (compressed + extras + decoded non-block + VAE file + OVERHEAD): # 22_788_580_400 - (9_587_103_379 + 137_388_032 + 134_217_728 + 1_350_989_512 + 459_445_098). Eleven samples on -# 2026-10-08 (the de-risk run above, four de-risk reruns of one step at two cache limits, three 40-step MEASURED +# 2026-10-08 (the calibration run above, four calibration reruns of one step at two cache limits, three 40-step MEASURED # runs, three 4-step identity runs with CFG; the highest is the second one-step rerun at the old cache limit) -# peaked at 21.76 to 22.79 GB in the VAE phase, with MLX's own VAE peak the same in every de-risk run: about 1 GiB of +# peaked at 21.76 to 22.79 GB in the VAE phase, with MLX's own VAE peak the same in every calibration run: about 1 GiB of # run-to-run spread outside MLX's counters, as Z-Image's VAE phase shows too. A fit check must not under-predict, so # the term covers the highest sample; the first run alone (22_139_643_024) would put the VAE phase 0.60 GiB under it. # About 1.32x the FLUX.1 VAE's 7.82 GiB. @@ -53,15 +53,15 @@ # DENOISE_ACTIVATION: denoise footprint peak minus the denoise phase's other terms at the planner's cache limit # (compressed + extras + non-block + in-flight + cache limit + OVERHEAD), at REFERENCE_TOKENS: # 13_939_106_072 - (9_587_103_379 + 137_388_032 + 134_217_728 + 436_207_616 + 2_731_761_634 + 459_445_098). Measured -# 2026-10-08 (the first one-step de-risk run at the new cache limit, in the cache-limit A/B: the same model, prompt, +# 2026-10-08 (the first one-step calibration run at the new cache limit, in the cache-limit A/B: the same model, prompt, # size, seed and guidance as above, one step at the 2_731_761_634 limit); the larger of the two runs at that limit # (the other's denoise peak, 13_506_126_128, gives 20_002_641). DENOISE_ACTIVATION = 452_982_585 # 0.422 GiB -# ALLOWANCE: the cache room for the activation buffers a block frees, at ALLOWANCE_TOKENS (1024² + the de-risk prompt's +# ALLOWANCE: the cache room for the activation buffers a block frees, at ALLOWANCE_TOKENS (1024² + the calibration prompt's # 33 tokens). The block-cache probe (one synthetic block at 1024² + 33 tokens, the decoded buffer allocated then # released per iteration) saw the buffer reused at a 2605 MiB limit and released at 1773 and 2189 MiB: the limit is set # at 2_731_761_634 B, so the allowance is that minus one decoded block, 2_731_761_634 - 436_207_616. Set at the probe's -# reuse point: an A/B of real de-risk steps at this limit and the old 1_859_346_402 (two runs each, 2026-10-08) showed +# reuse point: an A/B of real calibration steps at this limit and the old 1_859_346_402 (two runs each, 2026-10-08) showed # no measurable difference in step time or denoise footprint. ALLOWANCE = 2_295_554_018 # 2.138 GiB ALLOWANCE_TOKENS = 4096 + 33 @@ -87,7 +87,7 @@ # run held exactly ENCODE_ACTIVATION, the whole slack under the bound; an encode that also loads the vision tower and # lm_head (~2_397_528_480 more) is over it by 2_129_093_024 at every prompt length (the bound does not grow with it). ENCODE_SLACK_BYTES = 256 * 1024**2 -# The prompt length ENCODE_ACTIVATION was measured at (the de-risk prompt's text tokens). +# The prompt length ENCODE_ACTIVATION was measured at (the calibration prompt's text tokens). ENCODE_ACTIVATION_TOKENS = 33 diff --git a/src/mlx_dfloat/mflux/qwen21/model.py b/src/mlx_dfloat/mflux/qwen21/model.py index 548c81e..db6a68b 100644 --- a/src/mlx_dfloat/mflux/qwen21/model.py +++ b/src/mlx_dfloat/mflux/qwen21/model.py @@ -50,7 +50,7 @@ n: (e.df11_repo, e.base_repo) for n, e in registry.MODELS.items() if e.family == "qwen21" } POLICIES: tuple[str, ...] = ("per-block", "depth2") -# The policies a Qwen-Image 2.1 run has measured (the de-risk, the MEASURED row, the identity check); depth2 is +# The policies a Qwen-Image 2.1 run has measured (the one-step calibration run, the MEASURED row, the identity check); depth2 is # accepted, and the report labels it unmeasured. MEASURED_POLICIES: frozenset[str] = frozenset({"per-block"}) RETAINED_SLACK_BYTES = 2 * 1024**3 diff --git a/tests/_krea2_tiny.py b/tests/_krea2_tiny.py new file mode 100644 index 0000000..9f51dcd --- /dev/null +++ b/tests/_krea2_tiny.py @@ -0,0 +1,357 @@ +"""Tiny Krea 2 transformer helpers for the mflux lane. Every mflux import happens inside a function.""" + +from types import SimpleNamespace + +import mlx.core as mx +import mlx.nn as nn +import numpy as np +from mlx.utils import tree_flatten +from tests._df11_fixtures import pin_layout, random_bf16, write_checkpoint + +__all__ = [ + "CHECKPOINT_NAME", + "TEMPLATE_IDS", + "TINY", + "TINY_SEED", + "StubTextEncoder", + "StubTokenizer", + "StubVAE", + "fake_model", + "stub_components", + "tiny_groups", + "tiny_inputs", + "tiny_predict", + "tiny_seamed", + "write_tiny_checkpoint", +] + +# Head dim 16 = 32 / 2 heads = the sum of the RoPE axes mflux derives from it ([16 - 12, 6, 6] = [4, 6, 6], +# transformer.py:39-40); one KV head (wk / wv 16 x 32, a GQA repeat of 2). SwiGLU width int(2 * 32 / 3) * 4 = 84, +# rounded up to 128 (feed_forward.py:7-9). 16 channels and patch 2: Krea2LatentCreator always makes 16-channel +# latents. txtlayers 12 at txtdim 32: the context the transformer unpacks is 12 x 32 = 384 wide (transformer.py:93-101). +TINY = { + "features": 32, + "tdim": 32, + "txtdim": 32, + "heads": 2, + "kvheads": 1, + "multiplier": 4, + "layers": 2, + "patch": 2, + "channels": 16, + "txtlayers": 12, + "txtheads": 2, + "txtkvheads": 2, +} + + +def tiny_seamed(n_layers=2): + """A seamed mflux ``Krea2Transformer`` at ``TINY`` size with placeholders installed; returns (tf, shapes).""" + from mlx_dfloat.integrate.blockseam import seam_blocks + from mlx_dfloat.integrate.placeholders import install_placeholders + from mlx_dfloat.mflux.krea2.names import krea2_name_map + from mlx_dfloat.mflux.krea2.transformer import block_lists, seam_transformer_class + + tf = seam_transformer_class()(**{**TINY, "layers": n_layers}) + shapes = install_placeholders(block_lists(tf), krea2_name_map()) + seam_blocks(block_lists(tf), tf.seam_cell) + return tf, shapes + + +def tiny_inputs(seed=0): + """An 8 x 8 latent of 16 channels in float32 (as ``Krea2LatentCreator.create_noise``), a sigma, and two prompts. + + The prompts are 6 and 3 tokens at the 384-wide context (12 taps x 32), bf16 as the text encoder returns them; + ``neg_embeds`` is the negative a CFG step encodes. 8 x 8 latents at patch 2 are 16 image tokens. + """ + keys = mx.random.split(mx.random.key(seed), 3) + return { + "latents": mx.random.normal((1, 16, 8, 8), key=keys[0]), + "timestep": mx.array([0.5]), + "embeds": mx.random.normal((1, 6, 384), key=keys[1]).astype(mx.bfloat16), + "neg_embeds": mx.random.normal((1, 3, 384), key=keys[2]).astype(mx.bfloat16), + } + + +def tiny_predict(tf, inputs, *, guidance, compiled=False): + """One call of mflux's ``Krea2._predict`` on ``tf`` with the chip check answering "not M1/M2". + + ``inputs["neg_embeds"]`` None runs one transformer call, as mflux at guidance 1.0. Unless ``compiled``, the factory + runs through ``_compile.uncompiled`` (the adapter's bypass); with ``compiled`` it is the stock factory, which + compiles off base and Pro M1/M2 (krea2.py:202-204). + """ + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + from mflux.utils.apple_silicon import AppleSiliconUtil + + from mlx_dfloat.mflux._compile import uncompiled + + original = AppleSiliconUtil.__dict__["is_m1_or_m2"] + setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: False)) # noqa: B010 + try: + args = (tf, inputs["embeds"], inputs["neg_embeds"], guidance) + fn = Krea2._predict(*args) if compiled else uncompiled(Krea2._predict, *args) + finally: + setattr(AppleSiliconUtil, "is_m1_or_m2", original) # noqa: B010 + return fn(latents=inputs["latents"], timestep=inputs["timestep"]) + + +def tiny_groups(shapes, rng): + """One real DF11 group per block of ``shapes``; Krea's sub-paths are its attribute paths, so the names match. + + Returns groups, names, sources (``tests/_family_fakes.py:compress_blocks``). + """ + from tests._family_fakes import compress_blocks + + return compress_blocks(shapes, rng) + + +# mflux parameter -> checkpoint name: the inverse of Krea2WeightMapping's eleven non-block renames +# (krea2_weight_mapping.py:59-66, 75-81). Every other parameter is stored under its own name. +CHECKPOINT_NAME = { + "tmlp.linear_in.weight": "tmlp.0.weight", + "tmlp.linear_in.bias": "tmlp.0.bias", + "tmlp.linear_out.weight": "tmlp.2.weight", + "tmlp.linear_out.bias": "tmlp.2.bias", + "tproj.linear.weight": "tproj.1.weight", + "tproj.linear.bias": "tproj.1.bias", + "txtmlp.norm.scale": "txtmlp.0.scale", + "txtmlp.linear_in.weight": "txtmlp.1.weight", + "txtmlp.linear_in.bias": "txtmlp.1.bias", + "txtmlp.linear_out.weight": "txtmlp.3.weight", + "txtmlp.linear_out.bias": "txtmlp.3.bias", +} + + +def write_tiny_checkpoint( + root, rng, *, n_layers=2, nonblock_as_groups=True, random_extras=None, key="tiny-krea" +): + """A config-less single-file DF11 checkpoint for mflux's ``Krea2Transformer`` at ``TINY`` size, as published. + + The published layout's ``pattern_dict``; one group per block in ``BLOCK_SUBS`` order; the seven non-block groups + (``tmlp``, ``tproj``, ``txtmlp``, the four text-fusion blocks) or, with ``nonblock_as_groups=False``, their + matrices as BF16 extras; every other parameter a BF16 extra under its checkpoint name (``CHECKPOINT_NAME``), + filled with a distinct constant ``0x3F80 + i``. ``random_extras`` (a NumPy generator) fills the extras with random + BF16 instead: constant matrices blind an RMSNorm-fed forward pass to its inputs. ``key`` names the pinned layout + (a published layout's key exercises the variant check). + + Returns ``(matrices, constants, layout)``: ``{group: {sub-path or mflux parameter: uint16 matrix}}``, ``{mflux + parameter: constant}`` (with ``random_extras``, the uint16 arrays written) and the ``SynthesizedLayout`` pinned + to the written file, for ``open_checkpoint(root, layouts=(layout,))``. + """ + from mflux.models.krea2.model.krea2_transformer.transformer import Krea2Transformer + + from mlx_dfloat._layouts import KREA_2_RAW_COMFYUI + from mlx_dfloat.integrate.placeholders import install_placeholders + from mlx_dfloat.mflux.krea2.names import BLOCK_SUBS, NONBLOCK_GROUPS, krea2_name_map + from mlx_dfloat.mflux.krea2.transformer import COMPUTED_PARAMS, block_lists + + names = krea2_name_map() + probe = Krea2Transformer(**{**TINY, "layers": n_layers}) + probe_params = dict(tree_flatten(probe.parameters())) + shapes = install_placeholders(block_lists(probe), names) + matrices, groups = {}, {} + for block, per in shapes.items(): + matrices[block] = {attr: random_bf16(rng, shape) for attr, shape in per.items()} + groups[block] = [matrices[block][sub] for sub in BLOCK_SUBS] + taken = {f"{b}.{a}.weight" for b, per in shapes.items() for a in per} | set(COMPUTED_PARAMS) + patterns = dict(KREA_2_RAW_COMFYUI.raw_config["pattern_dict"]) + if nonblock_as_groups: + for group, stored in NONBLOCK_GROUPS.items(): + params = [names.param_name(m) for m in stored] + matrices[group] = {p: random_bf16(rng, tuple(probe_params[p].shape)) for p in params} + groups[group] = [matrices[group][p] for p in params] + taken |= set(params) + else: + patterns = {k: v for k, v in patterns.items() if k == r"blocks\.\d+"} + constants, extras = {}, {} + for i, (param, array) in enumerate(sorted(probe_params.items())): + if param in taken: + continue + name = CHECKPOINT_NAME.get(param, param) + if random_extras is None: + constants[param] = 0x3F80 + i + extras[name] = np.full(array.shape, 0x3F80 + i, dtype=np.uint16) + else: + extras[name] = constants[param] = random_bf16(random_extras, tuple(array.shape)) + write_checkpoint( + root, groups=groups, patterns=patterns, extras=extras, single_file=True, write_config=False + ) + layout = pin_layout( + root / "model.safetensors", + pattern_dict=patterns, + row_splits={}, + groups=len(groups), + extras=len(extras), + key=key, + ) + return matrices, constants, layout + + +TINY_SEED = 5 # ``fake_model`` writes its checkpoint with this seed: a test regenerates the source matrices from it +# The stub's chat-template prefix: <|im_start|>, a system token, <|im_start|>, "user", newline. mflux's +# Krea2TextEncoder._template_end finds the second <|im_start|> (151644) followed by "user" (872) and a newline (198) +# and strips everything before the prompt (text_encoder.py:117-127): 5 ids. +TEMPLATE_IDS = (151644, 9, 151644, 872, 198) + + +class StubTokenizer: + """mflux's tokenizer protocol as ``Krea2PromptEncoder.encode_prompt`` uses it (prompt_encoder.py:9-15). + + ``tokenize(prompt)`` gives the template prefix ``TEMPLATE_IDS`` and then one id per character, ``1 + ord(c) % 7`` + (``lengths`` overrides a prompt's id count; the ids then cycle through its characters), and an all-ones mask. The + empty prompt gives no ids at all, as the real Qwen tokenizer does. + """ + + def __init__(self, lengths=None): + self.lengths = dict(lengths or {}) + + def tokenize(self, prompt, **_kwargs): + if not prompt: + ids = [] + else: + n = self.lengths.get(prompt, len(prompt)) + ids = [*TEMPLATE_IDS, *(1 + ord(prompt[i % len(prompt)]) % 7 for i in range(n))] + return SimpleNamespace( + input_ids=mx.array([ids], dtype=mx.int32).reshape(1, len(ids)), + attention_mask=mx.ones((1, len(ids)), dtype=mx.int32), + ) + + +class StubTextEncoder(nn.Module): + """Stands in for the Qwen3-VL stack: ``get_prompt_embeds(input_ids, attention_mask)`` -> (1, L', 384) bf16 rows. + + 384 = TINY's 12 taps x 32 (the context the transformer unpacks). Each row's direction depends on its id (a sine of + the id at 384 frequencies; a row that is one vector scaled by the id would look the same to every RMSNorm), and the + template prefix is stripped as mflux strips it (``Krea2TextEncoder._template_end``). + """ + + def __init__(self): + super().__init__() + self.freq = mx.arange(1, 385, dtype=mx.float32) / 7 + self.calls = 0 + + def get_prompt_embeds(self, input_ids, attention_mask=None): + from mflux.models.krea2.model.krea2_text_encoder.text_encoder import Krea2TextEncoder + + self.calls += 1 + rows = mx.sin(input_ids[..., None].astype(mx.float32) * self.freq).astype(mx.bfloat16) + return rows[:, Krea2TextEncoder._template_end(input_ids) :, :] + + +class StubVAE(nn.Module): + """One weight (so the retained bound can size it) and the decode ``VAEUtil.decode`` calls (vae_util.py:66): a + zero image of 8 pixels per latent cell (Krea 2's latents are (1, 16, h / 8, w / 8), unpacked).""" + + def __init__(self): + super().__init__() + self.weight = mx.zeros((4,), dtype=mx.bfloat16) + self.seen = [] + + def decode(self, latents): + self.seen.append(latents) + return mx.zeros((1, 3, 8 * latents.shape[-2], 8 * latents.shape[-1])) + + +def stub_components(tokenizer=None): + """Base components as ``Krea2Components`` with the stub encoder, VAE and tokenizer.""" + from mlx_dfloat.mflux.krea2.init import Krea2Components + + return Krea2Components( + vae=StubVAE(), + text_encoder=StubTextEncoder(), + tokenizers={"qwen3vl": tokenizer or StubTokenizer()}, + ) + + +def fake_model( + tmp_path, monkeypatch, *, model="krea-2", tokenizer=None, nonblock="per-call", **overrides +): + """A ``DFloatKrea2`` over the tiny real transformer (2 blocks, config-less checkpoint) and stub base parts. + + No weights, no network. As the other families' fakes: the budget is fixed at 23 GiB, the encoder reload builds a + fresh stub, every decode runs on the NumPy reference. ``nonblock`` is ``"per-call"`` (the model's own build: the + non-block groups decoded at each transformer call) or ``"resident"`` (decoded once at set load). + """ + from functools import partial + + from mflux.models.common.config.model_config import ModelConfig + + from mlx_dfloat.decode import decode_group + from mlx_dfloat.format import open_checkpoint + from mlx_dfloat.mflux._hub import ResolvedRepo + from mlx_dfloat.mflux._phases import FamilySizes + from mlx_dfloat.mflux.krea2 import init as kinit + from mlx_dfloat.mflux.krea2 import model as model_module + from mlx_dfloat.mflux.krea2.transformer import build_transformer, seam_transformer_class + + _m, _c, layout = write_tiny_checkpoint(tmp_path / "df11", np.random.default_rng(TINY_SEED)) + ckpt = open_checkpoint(tmp_path / "df11", layouts=(layout,)) + build = build_transformer( + ckpt, + transformer_kwargs=TINY, + transformer_class=seam_transformer_class(per_call=nonblock == "per-call"), + ) + (tmp_path / "base").mkdir(exist_ok=True) + monkeypatch.setattr(kinit, "load_text_encoder", lambda root, **kwargs: StubTextEncoder()) + # A fixed budget: the CI runner's device reports a 7 GiB working set; the tests are about the arithmetic. + monkeypatch.setattr(model_module, "budget_bytes", lambda: 23 * 1024**3) + parts = { + "model": model, + "model_config": ModelConfig.from_name(model_name=model, base_model=None), + "ckpt": ckpt, + "df11": ResolvedRepo(root=tmp_path / "df11", repo_id="mingyi456/tiny", revision="a" * 40), + "base": ResolvedRepo(root=tmp_path / "base", repo_id="krea/tiny", revision="b" * 40), + "components": stub_components(tokenizer), + "build": build, + # A per-call model keeps no decoded non-block bytes with the set. + "sizes": FamilySizes( + compressed=1_000, + extras=100, + nonblock=0 if nonblock == "per-call" else 10, + encoders=1_000, + vae=10, + ), + "decode": partial(decode_group, backend="reference"), + } + parts.update(overrides) + return model_module.DFloatKrea2._from_parts(**parts) + + +def stock_krea2(tokenizer, text_encoder, vae, matrices, extras): + """Stock mflux ``Krea2`` over a plain tiny ``Krea2Transformer`` holding the BF16 source of a tiny checkpoint. + + ``write_tiny_checkpoint`` returns the block matrices by sub-path and the non-block matrices and the extras by + mflux's parameter name already, so no rename is needed here. + """ + from mflux.models.common.config.model_config import ModelConfig + from mflux.models.krea2.krea2_initializer import Krea2Initializer + from mflux.models.krea2.model.krea2_transformer.transformer import Krea2Transformer + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + from mlx.utils import tree_flatten, tree_unflatten + + from mlx_dfloat.mflux.krea2.names import NONBLOCK_GROUPS + + transformer = Krea2Transformer(**TINY) + shapes = {name: p.shape for name, p in tree_flatten(transformer.parameters())} + weights = [ + ( + name if group in NONBLOCK_GROUPS else f"{group}.{name}.weight", + mx.array(m).view(mx.bfloat16), + ) + for group, per in matrices.items() + for name, m in per.items() + ] + weights += [(name, mx.array(bits).view(mx.bfloat16)) for name, bits in extras.items()] + assert sorted(n for n, _ in weights) == sorted( + shapes + ) # every parameter from the source, none twice + assert all(tuple(shapes[n]) == w.shape for n, w in weights) + transformer.update(tree_unflatten(weights)) + stock = Krea2.__new__(Krea2) + nn.Module.__init__(stock) + Krea2Initializer._init_config(stock, ModelConfig.krea2_raw()) + stock.vae, stock.text_encoder, stock.transformer = vae, text_encoder, transformer + stock.tokenizers = {"qwen3vl": tokenizer} + stock.bits = None + stock.lora_paths, stock.lora_scales = None, None + return stock diff --git a/tests/_qwen21_tiny.py b/tests/_qwen21_tiny.py index 638d364..2ed7cbd 100644 --- a/tests/_qwen21_tiny.py +++ b/tests/_qwen21_tiny.py @@ -94,7 +94,7 @@ def tiny_groups(shapes, rng): return compress_blocks(shapes, rng) -# The ComfyUI export's stored order (S0): gate_layer and proj stacked by rows as img_mlp.gate_up, gate first. +# The ComfyUI export's stored order (confirmed by a 2026-10-07 byte comparison): gate_layer and proj stacked by rows as img_mlp.gate_up, gate first. STORED_SUBS = ( "attn.to_q", "attn.to_k", diff --git a/tests/fixtures/krea-2-raw-df11-header.bin b/tests/fixtures/krea-2-raw-df11-header.bin new file mode 100644 index 0000000..1b4b9a0 Binary files /dev/null and b/tests/fixtures/krea-2-raw-df11-header.bin differ diff --git a/tests/fixtures/krea-2-turbo-df11-header.bin b/tests/fixtures/krea-2-turbo-df11-header.bin new file mode 100644 index 0000000..952bd40 Binary files /dev/null and b/tests/fixtures/krea-2-turbo-df11-header.bin differ diff --git a/tests/test_bench_capped.py b/tests/test_bench_capped.py index bb19caf..d126ff8 100644 --- a/tests/test_bench_capped.py +++ b/tests/test_bench_capped.py @@ -137,7 +137,7 @@ def test_device_data_overrides_the_ratio_and_is_labelled_device(): def test_a_tier_above_the_host_is_refused(): - # Review Focus 2: a 48 GB "cap" on a 32 GB Mac would set limits above RAM and look like a pass. + # A 48 GB "cap" on a 32 GB Mac would set limits above RAM and look like a pass. with pytest.raises(DFloatUnsupportedError, match="48"): tier_limits(48, host_ram_bytes=HOST_RAM, host_recommended_bytes=HOST_REC) diff --git a/tests/test_bench_flux1.py b/tests/test_bench_flux1.py index 1a40793..590514c 100644 --- a/tests/test_bench_flux1.py +++ b/tests/test_bench_flux1.py @@ -175,7 +175,7 @@ def test_plan_children_skips_complete_runs_in_interleaved_order(tmp_path): def test_plan_children_refuses_a_dir_whose_child_key_differs(tmp_path): - # Review Focus 4. Bug caught: plan_children skipping the resume-key check, so a run with another + # Bug caught: plan_children skipping the resume-key check, so a run with another # mlx version or source hash resumes into (and pools with) results it did not produce. s = scenario_from_mapping(SPEC) _child_json(tmp_path / "round2-q8.json", key={"model": "dev"}, round_no=2, mode="q8") diff --git a/tests/test_bench_results.py b/tests/test_bench_results.py index 32cbd16..d92d5c8 100644 --- a/tests/test_bench_results.py +++ b/tests/test_bench_results.py @@ -87,7 +87,7 @@ def test_a_pair_without_a_shared_round_yields_no_overhead(): def test_two_scenario_hashes_are_refused(): - # Review Focus 4: a stale round1-df11.json from another recipe must not be averaged in. + # A stale round1-df11.json from another recipe must not be averaged in. with pytest.raises(DFloatFormatError, match="scenario"): summarise([_r("df11", 1, [1.0]), _r("control", 1, [1.0], h="b" * 64)]) diff --git a/tests/test_bench_scenario.py b/tests/test_bench_scenario.py index 942ee5d..af3dfbc 100644 --- a/tests/test_bench_scenario.py +++ b/tests/test_bench_scenario.py @@ -48,7 +48,7 @@ def test_a_complete_mapping_builds_a_frozen_scenario_with_tuples(): @pytest.mark.parametrize( ("field", "value", "fragment"), [ - ("warmup_steps", 2, "warmup_steps"), # unknown key (Review Focus 1) + ("warmup_steps", 2, "warmup_steps"), # unknown key ("model", "flux2", "model"), ("model", "krea-dev", "model"), ("df11_revision", "main", "df11_revision"), # not a 40-hex SHA diff --git a/tests/test_bench_table.py b/tests/test_bench_table.py index ea1bc31..0131ab8 100644 --- a/tests/test_bench_table.py +++ b/tests/test_bench_table.py @@ -127,7 +127,7 @@ def test_splice_replaces_only_the_marked_block(): ], ) def test_splice_refuses_missing_or_duplicated_markers(doc): - # Review Focus 5: a README without exactly one marker pair must not be silently rewritten. + # A README without exactly one marker pair must not be silently rewritten. with pytest.raises(DFloatFormatError, match="bench:x"): splice(doc, "x", "new\n") diff --git a/tests/test_ci_config.py b/tests/test_ci_config.py index 6a7869f..d4b8bb2 100644 --- a/tests/test_ci_config.py +++ b/tests/test_ci_config.py @@ -87,8 +87,8 @@ def test_the_mflux_lane_installs_the_locked_extra_and_proves_it_ran_tests(): # Every `uv run` names the extra and the group, so no step depends on uv's sync mode keeping # what `uv sync --extra mflux` installed. uv_runs = [run for run in runs if "uv run" in run] - # sanity + lane + five per-model coverage gates + the collection check - assert len(uv_runs) == 8 + # sanity + lane + six per-model and ten per-module coverage gates + the collection check + assert len(uv_runs) == 19 assert all("uv run --extra mflux --group dev " in run for run in uv_runs) pytest_runs = [ shlex.split(run) @@ -99,13 +99,13 @@ def test_the_mflux_lane_installs_the_locked_extra_and_proves_it_ran_tests(): assert len(tests) == 1 assert {"-m", "mflux", "--run-network", "-rs"} <= set(tests[0]) # model.py has no offline test at all (every test is @pytest.mark.mflux), so only this lane - # measures it (the required job's gate omits it; see pyproject.toml). init.py is gated in the - # required job (it has real offline tests); this lane only reports its number. + # measures it (the required job's gate omits it; see pyproject.toml). The same holds for the + # init.py modules, which drive mflux's loaders; each is gated on its own after this step. assert "--cov=mlx_dfloat.mflux.flux1.model" in tests[0] assert "--cov=mlx_dfloat.mflux.flux1.init" in tests[0] assert "--cov-config=.coveragerc-integration" in tests[0] - # No combined threshold: on a tokenless runner init.py's live tests skip and would drag a - # combined number under 80 %; the per-file step after this one is the lane's gate. + # No combined threshold: one module's lines would dilute another's; the per-file steps after + # this one are the lane's gates. assert not any(arg.startswith("--cov-fail-under") for arg in tests[0]) guard = [run for run in runs if "--co" in run and "grep -c" in run] assert len(guard) == 1 @@ -200,7 +200,7 @@ def test_the_mflux_lane_measures_and_gates_the_ernie_model_module_on_its_own(): # Bug caught: ernie/model.py (every test @pytest.mark.mflux) omitted from the required job's gate but measured by # no job either, so its coverage could fall to zero with every check green; or its gate diluted by the other model # modules' lines, or run before the lane's pytest wrote the data; or ERNIE's init.py and transformer.py not - # reported in the lane beside its model (measured for consistency, not gated). + # measured in the lane beside its model (each is gated on its own, see the test below). import tomllib runs = [step.get("run", "") for step in _integration_job()["steps"]] @@ -221,6 +221,31 @@ def test_the_mflux_lane_measures_and_gates_the_ernie_model_module_on_its_own(): assert "src/mlx_dfloat/mflux/ernie/model.py" in pyproject["tool"]["coverage"]["run"]["omit"] +def test_the_mflux_lane_measures_and_gates_the_krea2_model_module_on_its_own(): + # Bug caught: krea2/model.py (every test @pytest.mark.mflux) omitted from the required job's gate but measured by + # no job either, so its coverage could fall to zero with every check green; or its gate diluted by the other model + # modules' lines, or run before the lane's pytest wrote the data; or Krea 2's init.py and transformer.py not + # measured in the lane beside its model (each is gated on its own, see the test below). + import tomllib + + runs = [step.get("run", "") for step in _integration_job()["steps"]] + lane = next( + i + for i, run in enumerate(runs) + if run.startswith("uv run --extra mflux --group dev pytest") + and "--co" not in shlex.split(run) + ) + args = shlex.split(runs[lane]) + assert "--cov=mlx_dfloat.mflux.krea2.model" in args + assert "--cov=mlx_dfloat.mflux.krea2.init" in args + assert "--cov=mlx_dfloat.mflux.krea2.transformer" in args + report = "uv run --extra mflux --group dev coverage report --include='*/mflux/krea2/model.py' --fail-under=80" + assert report in runs + assert lane < runs.index(report) < lane + 7 # the sixth per-model gate after the lane + pyproject = tomllib.loads((CI.parents[2] / "pyproject.toml").read_text()) + assert "src/mlx_dfloat/mflux/krea2/model.py" in pyproject["tool"]["coverage"]["run"]["omit"] + + def _config() -> dict: return yaml.safe_load(CI.read_text()) @@ -263,3 +288,57 @@ def test_the_workflow_token_is_read_only(): # repository default (read/write on older defaults), in a workflow whose mflux lane now carries a # Hub secret; or a write scope added without a step that needs it. assert _config()["permissions"] == {"contents": "read"} + + +# The family modules that drive mflux (its loaders, its transformer classes): their mflux-free tests leave them at +# 40-60 % in the required job, so that job's gate omits them and the mflux lane gates each one at 80 %, as it does +# model.py. flux1/transformer.py and zimage/transformer.py stay in the required gate: their offline tests (the seam +# fakes) cover them past 90 % without mflux, while the lane alone reaches only 57 % of flux1/transformer.py. +MFLUX_ONLY = [ + "flux1/init", + "zimage/init", + "flux2/init", + "flux2/transformer", + "qwen21/init", + "qwen21/transformer", + "ernie/init", + "ernie/transformer", + "krea2/init", + "krea2/transformer", +] + + +def test_the_mflux_only_modules_are_omitted_from_the_required_gate_and_gated_in_the_lane(): + # Bug caught: a module omitted from the required job's 85 % gate but measured or gated by no job (its coverage + # could fall to zero with every check green), a gate diluted by another module's lines, a gate run before the + # lane's pytest wrote the data, or a module that runs without mflux (names.py, memory.py, the two seam-covered + # transformer.py files) dropped from the required gate. + import tomllib + + omit = tomllib.loads((CI.parents[2] / "pyproject.toml").read_text())["tool"]["coverage"]["run"][ + "omit" + ] + runs = [step.get("run", "") for step in _integration_job()["steps"]] + lane = next( + i + for i, run in enumerate(runs) + if run.startswith("uv run --extra mflux --group dev pytest") + and "--co" not in shlex.split(run) + ) + args = shlex.split(runs[lane]) + collect = next(i for i, run in enumerate(runs) if "--co" in run and "grep -c" in run) + for module in MFLUX_ONLY: + family, name = module.split("/") + assert f"src/mlx_dfloat/mflux/{module}.py" in omit, module + assert f"--cov=mlx_dfloat.mflux.{family}.{name}" in args, module + report = f"uv run --extra mflux --group dev coverage report --include='*/mflux/{module}.py' --fail-under=80" + assert report in runs, module + assert lane < runs.index(report) < collect, module + for kept in ( + "flux1/transformer", + "zimage/transformer", + "krea2/names", + "krea2/memory", + "ernie/memory", + ): + assert f"src/mlx_dfloat/mflux/{kept}.py" not in omit, kept diff --git a/tests/test_cli.py b/tests/test_cli.py index 5ee3f44..cb398ef 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -160,7 +160,7 @@ def test_generate_builds_the_model_from_the_flags_writes_the_image_and_the_repor # Bug caught: a flag not forwarded (e.g. --cache-limit dropped), the image saved elsewhere, # the watchdog not stopped, or the report missing the exit code and the model's report. log = [] - out = tmp_path / "new" / "dir" / "o.png" # Review Focus 5: the directory does not exist yet + out = tmp_path / "new" / "dir" / "o.png" # the directory does not exist yet code, watchdogs = _run( [ "--model", @@ -470,7 +470,7 @@ def failing_watchdog_factory(out_dir, **kw): @pytest.mark.parametrize("text", ["0", "-5", "abc", "inf", "-inf", "nan"]) def test_cache_limit_must_be_a_positive_byte_count(text): - # Bug caught (Review Focus 4): a zero or negative cache limit reaching mx.set_cache_limit. + # Bug caught: a zero or negative cache limit reaching mx.set_cache_limit. # A non-positive value and a non-numeric string raise ValueError with different messages # (the explicit check vs. float()'s own conversion error), so this shared parametrized # case checks the common ValueError contract rather than one specific message. @@ -579,7 +579,7 @@ def _ceiling_for(*extra): (("--tier", "24"), TIER_24_CEILING, 24, False, "CAPPED"), # Bug caught: the 2/3 ratio applied as 3/4 at 16 GB (a 12 GiB budget). (("--tier", "16"), TIER_16_CEILING, 16, False, "CAPPED"), - # Bug caught (Review Focus 2): the host's own tier treated as a cap (the ratio's 24 GiB + # Bug caught: the host's own tier treated as a cap (the ratio's 24 GiB # budget, 22 GiB ceiling) instead of keeping the host's default ceiling. (("--tier", "32"), HOST_CEILING, 32, True, "MEASURED"), ], @@ -597,7 +597,7 @@ def test_ceiling_for_resolves_the_ceiling_the_limits_and_the_label( def test_ceiling_for_refuses_both_flags_and_a_tier_above_the_host(): - # Bug caught (Review Focus 6): --tier and --memory-ceiling together silently picking one of them, + # Bug caught: --tier and --memory-ceiling together silently picking one of them, # or a 48 GB "tier" on a 32 GB host setting limits above its RAM and looking like a pass. with pytest.raises(ValueError, match="--tier") as info: _ceiling_for("--tier", "24", "--memory-ceiling", "12345") @@ -642,7 +642,7 @@ def test_tier_24_applies_the_tier_limits_not_the_host_caps_and_budgets_the_fit_c def test_tier_32_on_a_32_gb_host_keeps_the_host_caps_and_is_measured(tmp_path): - # Bug caught (Review Focus 2): `--tier ` installing a smaller tier's caps instead of the + # Bug caught: `--tier ` installing a smaller tier's caps instead of the # host caps, so MEASURED rows run under other limits than a plain generate. code, watchdogs, calls, report = _limits_run(["--tier", "32"], tmp_path) assert code == 0 @@ -684,7 +684,7 @@ def test_a_plain_run_records_the_host_tiers_limits(tmp_path): def test_both_flags_exit_2_naming_both_before_anything_is_installed(tmp_path, capsys): - # Bug caught (Review Focus 6): argparse or run() silently preferring one flag, or the caps + # Bug caught: argparse or run() silently preferring one flag, or the caps # installed / the watchdog started before the refusal. code, watchdogs, calls, report = _limits_run( ["--tier", "24", "--memory-ceiling", "12345"], tmp_path @@ -699,7 +699,7 @@ def test_both_flags_exit_2_naming_both_before_anything_is_installed(tmp_path, ca def test_a_tier_above_the_host_exits_2_before_anything_is_installed(tmp_path, capsys): - # Bug caught: a 48 GB tier on a 32 GB host installing limits above RAM (Review Focus 2), or the + # Bug caught: a 48 GB tier on a 32 GB host installing limits above RAM, or the # refusal surfacing as a traceback. code, watchdogs, calls, report = _limits_run(["--tier", "48"], tmp_path) assert code == 2 @@ -1201,7 +1201,7 @@ def test_qwen_passes_the_negative_prompt_through(tmp_path, capsys): def test_qwen_guidance_without_a_negative_prompt_warns(tmp_path, capsys): - # Bug caught (Review Focus 4): a Qwen user's --guidance 4 silently running no CFG (mflux's rule: guidance above + # Bug caught: a Qwen user's --guidance 4 silently running no CFG (mflux's rule: guidance above # 1.0 AND a negative prompt, qwen_image_21.py:89), or the call changed instead of warned (the call still runs at # 4.0 without a negative prompt, as mflux would). call = _generate_kwargs("qwen-image-2.1", tmp_path, "--guidance", "4") @@ -1287,7 +1287,7 @@ def test_ernie_steps_follow_mflux_per_variant(): def test_turbo_guidance_other_than_1_is_refused_before_the_build(tmp_path, capsys): - # Bug caught (Review Focus 1, command half): --guidance 4 reaching ERNIE-Image-Turbo (mflux's turbo command errors, + # Bug caught (the command's half): --guidance 4 reaching ERNIE-Image-Turbo (mflux's turbo command errors, # ernie_image_turbo_generate.py:42-45), or the refusal coming after the caps or the build. built = [] args = gen.build_parser().parse_args( @@ -1377,3 +1377,194 @@ def test_an_empty_prompt_is_still_passed_to_the_other_families(tmp_path): ) assert code == 0 assert log[0][1]["prompt"] == "" + + +# --- Krea 2 ------------------------------------------------------------------------------------------------------ + + +def test_krea_negative_prompt_at_guidance_half_does_not_warn_it_has_no_effect(tmp_path, capsys): + # Bug caught: the "runs only above guidance 1.0" warning for Krea 2 at 0.5, where mflux does run + # classifier-free guidance (any guidance other than 1.0, prompt_encoder.py:33), so the negative prompt is used. + call = _generate_kwargs("krea-2", tmp_path, "--guidance", "0.5", "--negative-prompt", "blurry") + assert (call["guidance"], call["negative_prompt"]) == (0.5, "blurry") + assert "has no effect" not in capsys.readouterr().err + + +def test_krea_negative_prompt_at_guidance_1_warns_with_the_default(tmp_path, capsys): + # Bug caught: the warning missing at exactly 1.0 (mflux runs no negative branch there), or worded with the other + # families' "only above 1.0" (wrong for Krea 2, where 0.5 runs it too). + _generate_kwargs("krea-2", tmp_path, "--guidance", "1", "--negative-prompt", "blurry") + err = capsys.readouterr().err + assert ( + "warning: --negative-prompt has no effect: classifier-free guidance runs only for a guidance other than " + "1.0 (the default is 1.0)" + ) in err + assert "only above" not in err + + +def test_flux_dev_negative_prompt_at_guidance_half_still_warns(tmp_path, capsys): + # Bug caught: the Krea rule applied to every family (FLUX.1 dev takes no negative prompt; Z-Image at 0.5 runs no + # CFG, so its warning must stay). + _generate_kwargs("z-image", tmp_path, "--guidance", "0.5", "--negative-prompt", "blurry") + assert ( + "warning: --negative-prompt has no effect: classifier-free guidance runs only above guidance 1.0" + in capsys.readouterr().err + ) + _generate_kwargs("dev", tmp_path, "--guidance", "0.5", "--negative-prompt", "blurry") + assert "warning: --negative-prompt is ignored" in capsys.readouterr().err + + +def test_an_unknown_krea_scheduler_is_refused_before_the_build(tmp_path, capsys): + # Bug caught: flow_match_euler_discrete (another family's scheduler) reaching the model, where + # mflux would raise only after the encoder and the set loaded, or the refusal after the caps. + log = [] + code, watchdogs = _run( + [ + "--model", + "krea-2", + "--prompt", + "p", + "--scheduler", + "flow_match_euler_discrete", + "--output", + str(tmp_path / "o.png"), + ], + log, + tmp_path, + factory=lambda **kw: pytest.fail("built"), + install_caps=lambda: pytest.fail("caps installed before the refusal"), + ) + assert (code, watchdogs, log) == (2, [], []) + assert "error: --scheduler: Krea 2 Turbo runs er_sde, euler, linear" in capsys.readouterr().err + + +@pytest.mark.parametrize("scheduler", ["er_sde", "euler", "linear"]) +def test_the_krea_schedulers_pass_through(scheduler, tmp_path): + # Bug caught: "linear" refused although mflux maps it to er_sde (krea2.py:206-212), or euler refused. + assert ( + _generate_kwargs("krea-2-raw", tmp_path, "--scheduler", scheduler)["scheduler"] == scheduler + ) + + +def test_the_scheduler_allow_list_is_krea_only(): + # Bug caught: another family's --scheduler checked against Krea's list (Qwen's "linear" passes, but Z-Image's + # "euler" or a custom dotted scheduler would be refused). + for model, scheduler in (("z-image", "euler"), ("qwen-image-2.1", "my.module.Scheduler")): + args = gen.build_parser().parse_args( + ["--model", model, "--prompt", "p", "--scheduler", scheduler] + ) + assert gen.scheduler_refusal(args) is None + + +def test_krea_raw_passes_its_defaults(tmp_path, capsys): + # Bug caught: the card's 52 / 3.5 left as Raw's default (a CFG run nobody asked for), or Turbo's 8 steps. Literals: + # 25 steps (mflux defaults.py:19,111-118), guidance 1.0 (krea2.py:58), er_sde (krea2.py:206-212). + call = _generate_kwargs("krea-2-raw", tmp_path, "--negative-prompt", "n") + assert (call["num_inference_steps"], call["guidance"], call["scheduler"]) == (25, 1.0, "er_sde") + assert call["negative_prompt"] == "n" + turbo = _generate_kwargs("krea-2", tmp_path) + assert (turbo["num_inference_steps"], turbo["guidance"], turbo["scheduler"]) == ( + 8, + 1.0, + "er_sde", + ) + + +def test_krea_raw_negative_prompt_at_its_default_guidance_warns(tmp_path, capsys): + # Bug caught: the warning keyed on Raw's former 3.5 (silent at the real default 1.0, where the negative prompt has + # no effect). + _generate_kwargs("krea-2-raw", tmp_path, "--negative-prompt", "n") + err = capsys.readouterr().err + assert "warning: --negative-prompt has no effect" in err + assert "(the default is 1.0)" in err + + +def test_krea_raw_card_recipe_flags_pass_through(tmp_path, capsys): + # Bug caught: the model card's recipe clamped to the defaults, or --guidance 3.5 refused. + call = _generate_kwargs("krea-2-raw", tmp_path, "--steps", "52", "--guidance", "3.5") + assert (call["num_inference_steps"], call["guidance"]) == (52, 3.5) + + +@pytest.mark.parametrize( + ("model", "extra", "note"), + [ + ( + "krea-2", + ["--guidance", "0.5"], + "info: guidance 0.5 on Krea 2 Turbo runs classifier-free guidance: two transformer calls per step " + '(negative prompt " ")', + ), + ( + "krea-2-raw", + ["--guidance", "3.5", "--negative-prompt", "n"], + "info: guidance 3.5 on Krea 2 Raw runs classifier-free guidance: two transformer calls per step", + ), + ("krea-2-raw", [], None), + ("dev", ["--guidance", "3.5"], None), + ("qwen-image-2.1", ["--guidance", "4", "--negative-prompt", "n"], None), + ], +) +def test_cfg_cost_note_names_two_calls_for_krea_off_guidance_one( + model, extra, note, tmp_path, capsys +): + # Bug caught: the note on every family, missing below 1.0 (the rule copied as `> 1.0`), or claiming mflux's " " + # when the user gave a negative prompt. The exit code is unchanged (an info line, not a refusal). + args = gen.build_parser().parse_args(["--model", model, "--prompt", "p", *extra]) + expected = None if note is None else note.removeprefix("info: ") + assert gen.cfg_cost_note(args) == expected + _generate_kwargs(model, tmp_path, *extra) # asserts exit 0 + err = capsys.readouterr().err + if note is None: + assert "info: guidance" not in err + else: + assert note in err + + +@pytest.mark.parametrize( + ("model", "guidance", "expected"), + [ + ("krea-2", None, False), + ("krea-2", 1.0, False), + ("krea-2", 0.5, True), + ("krea-2", 1.0001, True), + ("qwen-image-2.1", 0.5, False), + ("qwen-image-2.1", 1.0, False), + ("qwen-image-2.1", 4.0, True), + ], +) +def test_runs_cfg_follows_each_familys_rule(model, guidance, expected): + # Bug caught: Krea's != 1.0 rule applied to every family, or the other families' > 1.0 rule applied to Krea. + from mlx_dfloat.mflux import families + + assert gen.runs_cfg(families.entry(model), guidance) is expected + + +@pytest.mark.parametrize("model", ["krea-2", "krea-2-raw"]) +@pytest.mark.parametrize("prompt", ["", " "]) +def test_an_empty_krea_prompt_is_refused_before_anything_loads(model, prompt, tmp_path, capsys): + # Bug caught: an empty prompt reaching Krea 2 (its tokenizer gives "" no ids: an empty context through the + # transformer), or the refusal coming after the caps, the watchdog or the build. + log = [] + code, watchdogs = _run( + ["--model", model, "--prompt", prompt, "--output", str(tmp_path / "o.png")], + log, + tmp_path, + factory=lambda **kw: pytest.fail("built"), + install_caps=lambda: pytest.fail("caps installed before the refusal"), + ) + assert (code, watchdogs, log) == (2, [], []) + assert ( + "an empty or blank prompt is refused as a user error (an empty prompt gives Krea 2's tokenizer" + in (capsys.readouterr().err) + ) + + +def test_the_help_names_the_krea_defaults_and_the_cards_recipe(): + # Bug caught: --help still listing only the older families' defaults, or Raw's default shown as the card's 3.5. + text = " ".join(gen.build_parser().format_help().split()) + assert "krea-2-raw 25, krea-2 8" in text + assert ( + "Krea 2 Raw and Turbo 1.0 (Krea 2 Raw's model card uses --steps 52 --guidance 3.5); Krea 2 runs " + "classifier-free guidance, two transformer calls per step, for any value other than 1.0" + ) in text + assert "Krea 2: er_sde (default), euler" in text diff --git a/tests/test_ernie_init.py b/tests/test_ernie_init.py index ff727a6..0121658 100644 --- a/tests/test_ernie_init.py +++ b/tests/test_ernie_init.py @@ -41,7 +41,7 @@ "vae/config.json", "vae/diffusion_pytorch_model.safetensors", ) -# The de-risk prompt and its token count (mflux 0.20.0 TokenizerLoader on both base snapshots, 2026-10-08). +# The calibration prompt and its token count (mflux 0.20.0 TokenizerLoader on both base snapshots, 2026-10-08). PROMPT = ( "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the " "water" @@ -180,7 +180,7 @@ def _write_tiny_encoder(root, *, drop=()): @pytest.mark.mflux def test_an_encoder_tensor_missing_from_the_base_is_refused_by_name(tmp_path): - # Bug caught (Review Focus 5): mflux's applier updates with strict=False, so a base missing one tensor (a lost + # Bug caught: mflux's applier updates with strict=False, so a base missing one tensor (a lost # shard, a renamed key) loads without error and that layer keeps its random float32 init: a wrong image, no error. from mlx_dfloat.mflux.ernie.init import load_text_encoder @@ -239,7 +239,7 @@ def _write_vae(root, *, drop=()): @pytest.mark.mflux @pytest.mark.parametrize("dropped", ["decoder.conv_out.weight", "bn.running_mean"]) def test_a_vae_tensor_missing_from_the_base_is_refused_by_name(tmp_path, dropped): - # Bug caught (Review Focus 5): the VAE loaded through the same strict=False applier with no coverage check (a + # Bug caught: the VAE loaded through the same strict=False applier with no coverage check (a # lost decoder conv keeps its random init and the decode returns noise), or the batch-norm statistics exempted as # computed buffers (they come from the file: flux2_weight_mapping.py:202-209). from mlx_dfloat.mflux.ernie.init import load_vae diff --git a/tests/test_ernie_memory.py b/tests/test_ernie_memory.py index 39d07e6..94f4497 100644 --- a/tests/test_ernie_memory.py +++ b/tests/test_ernie_memory.py @@ -1,6 +1,6 @@ """ERNIE-Image's memory rules: the encoder's resident bytes, the text-token count, the sizes, the measured constants. -Measured inputs (2026-10-08, git a771e58, M1 Max 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0): two one-step de-risk +Measured inputs (2026-10-08, git a771e58, M1 Max 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0): two one-step calibration runs at 1024², seed 42, the 26-token lighthouse prompt, one process each (build, encode, drop, set load, one step at the planner's cache limit for the block-cache probe's allowance, VAE decode on the resident set): ERNIE-Image-Turbo at guidance 1.0 (batch 1) and ERNIE-Image at guidance 4.0 (mflux's " " negative then the prompt: batch 2). Their phase @@ -26,7 +26,7 @@ from mlx_dfloat.mflux._pipeline import plan_call_for from mlx_dfloat.mflux.ernie import memory as emem -# The de-risk prompt's token count (mflux 0.20.0 TokenizerLoader on both base snapshots, 2026-10-08). +# The calibration prompt's token count (mflux 0.20.0 TokenizerLoader on both base snapshots, 2026-10-08). DERISK_TOKENS = 26 @@ -207,7 +207,7 @@ def test_the_encode_warning_is_silent_at_the_bound_and_names_the_prompt_length_o LARGEST = { "layers": 436_207_616 } # one block kind; 36 groups of 218_103_808 elements (4 x 4096² + 3 x 4096 x 12288) -TEXT_TOKENS = 26 # the de-risk prompt's (the base's negative " " is 2) +TEXT_TOKENS = 26 # the calibration prompt's (the base's negative " " is 2) # The block-cache probe's allowances at 4096 + 26 tokens (the smallest limit at which each batch's decoded buffer is # reused: 2.6e9 and 3.6e9 B, minus one decoded block). ALLOWANCE = {1: 2_163_792_384, 2: 3_163_792_384} @@ -249,7 +249,7 @@ def _plan(batch, budget, *, fit_check=True, cache_limit_override=None, policy="p ) -# The constants are calibrated on the two de-risk runs and, for batch 1's denoise term, the CAPPED 24 GB run (the +# The constants are calibrated on the two calibration runs and, for batch 1's denoise term, the CAPPED 24 GB run (the # overall-peak rule), so at their working point the estimate equals the measurement by construction: the in-sample # checks are exact. No ERNIE-Image run away from 1024² and 26 tokens exists yet. @@ -276,7 +276,7 @@ def test_the_calibration_records_hold_the_runs_the_constants_name(): def test_the_record_sizes_count_the_evaluated_language_model(): - # Bug caught: the records carrying every language-model tensor (6_858_012_672 B, Task 0's header sum) while the + # Bug caught: the records carrying every language-model tensor (6_858_012_672 B, the published header's sum) while the # code counts the evaluated layers (the encode term would come out negative: the runs held about 6.63e9 B), or the # other way round. 6_858_012_672 - the last layer's 232_796_160 = 6_625_216_512. assert {b: s.encoders for b, s in SIZES.items()} == {1: 6_625_216_512, 2: 6_625_216_512} @@ -292,7 +292,7 @@ def test_the_overhead_is_the_larger_post_build_gap_of_the_two_runs(): for b in RUNS for label, run in RUNS[b].items() if "after_build" - in run # the one-step de-risk runs; the later samples record a VAE peak only + in run # the one-step calibration runs; the later samples record a VAE peak only } assert gaps == {"derisk-turbo-1": 616_434_178, "derisk-base-1": 548_194_818} assert RECORDS[1]["overhead_from"] == "derisk-turbo-1" @@ -318,10 +318,10 @@ def test_the_estimate_at_1024_equals_each_batchs_calibration_run_phase_for_phase def test_the_vae_term_covers_every_1024_sample_and_equals_the_highest(): - # Bug caught: the VAE transient derived from one run instead of the highest sample (the base's de-risk sample, with + # Bug caught: the VAE transient derived from one run instead of the highest sample (the base's calibration sample, with # its own sizes, is 35_821_552 B lower and puts the Turbo run's VAE phase that much low), or a sample dropped from a - # record. Five samples: the two de-risk runs, the two MEASURED rows and the base identity check's df11 side; the - # highest is still the Turbo de-risk run. One term serves both models, whose sets differ by 59_576 B, so the + # record. Five samples: the two calibration runs, the two MEASURED rows and the base identity check's df11 side; the + # highest is still the Turbo calibration run. One term serves both models, whose sets differ by 59_576 B, so the # estimate equals the sample only for the model whose record holds it (Turbo). assert len(VAE_SAMPLES) == 5 for b in RUNS: @@ -411,7 +411,7 @@ def _tier_budget(tier): [(2, 1024, 450_976_567), (2, 512, 402_653_184), (1, 512, 136_331_547)], ) def test_the_batch_2_denoise_term_is_floored_by_the_float32_weight_copies(batch, size, expected): - # Bug caught (X3): a small CFG call predicted below the float32 copies a batch-2 step holds whatever the image size + # Bug caught: a small CFG call predicted below the float32 copies a batch-2 step holds whatever the image size # (a float32 hidden state times a BF16 weight makes MLX copy the weight to float32: gate_proj and up_proj, 2 x # 12288 x 4096 x 4 = 402_653_184 B), or the floor changing the 1024² term it was calibrated on. 512² + 26 tokens: # 450_976_567 x 1050 / 4122 = 114_877_582 scaled, floored; batch 1 carries no floor (not measured below 1024²): @@ -428,7 +428,7 @@ def test_the_tier_budgets_are_a_real_macs(): def test_turbo_at_24_gb_is_refused_0_39_gib_over_the_budget(): - # Bug caught (X1): the batch-1 denoise term taken from one run instead of the highest denoise sample (the de-risk + # Bug caught: the batch-1 denoise term taken from one run instead of the highest denoise sample (the calibration # run's passed the call 58 MB under the budget; the first CAPPED run's, 37 MB over it; the second identical run # peaked 0.35 GiB higher still), or the dropped-set VAE phase without the extras. forced = _plan(1, _tier_budget(24), fit_check=False) @@ -454,7 +454,7 @@ def test_every_capped_call_is_refused_in_the_denoise_phase(batch, tier): # + the slack. The encode term is per prompt, so the bound does not depend on the batch or on the prompt length. NOT_LOADED = ( 806_610_944 + 33_556_480 -) # the vision tower and the projector in the encoder's file (Task 0's header) +) # the vision tower and the projector in the encoder's file (its published header) def _encode_bound(batch, text_tokens=TEXT_TOKENS): diff --git a/tests/test_ernie_model.py b/tests/test_ernie_model.py index 176694b..79d6aa8 100644 --- a/tests/test_ernie_model.py +++ b/tests/test_ernie_model.py @@ -156,7 +156,7 @@ def test_the_set_load_installs_the_three_nonblock_groups_and_attaches_with_verif def test_a_new_batch_after_a_generation_drops_the_set_and_reinstalls_the_nonblock_weights( tmp_path, monkeypatch ): - # Bug caught (Review Focus 4): a non-block weight left on its placeholder after the reload a new batch forces (a + # Bug caught: a non-block weight left on its placeholder after the reload a new batch forces (a # zero-size matmul on adaln_modulation before the first block), or the encoder reload reading another base. model = fake_model(tmp_path, monkeypatch) loads = [] @@ -337,7 +337,7 @@ def test_a_call_whose_batch_was_not_encoded_is_refused(tmp_path, monkeypatch): def test_blank_and_missing_negatives_share_one_batch(tmp_path, monkeypatch): - # Bug caught (Review Focus 2): None, "" and " " keyed apart (three encodes and set reloads for one mflux batch + # Bug caught: None, "" and " " keyed apart (three encodes and set reloads for one mflux batch # [" ", prompt]), or a blank negative encoded as itself. Real mflux loop, one step each. model = fake_model(tmp_path, monkeypatch) run = {"seed": 1, "prompt": "a", "num_inference_steps": 1, "height": 64, "width": 64} @@ -399,7 +399,7 @@ def _double(allowance): def test_a_cfg_call_plans_with_the_batch_two_constants( tmp_path, monkeypatch, guidance, batch, limit ): - # Bug caught (Review Focus 1, Python-API half): the constants looked up by model name (Turbo planned with batch-1 + # Bug caught (the Python API's half): the constants looked up by model name (Turbo planned with batch-1 # activations for a batch-2 call: 1_000_017_408 at guidance 4), or the batch counted from the model instead of the # prompts. The doubles' reference is 4096 + 1 tokens ("a" and " " are one stub token each), so no scaling, no # floor: the tiny decoded block 17_408 B + the allowance, by hand. @@ -484,7 +484,7 @@ def end(name): def test_an_encode_over_the_bound_warns_through_the_pure_helper(tmp_path, monkeypatch, over, warns): # Bug caught: _check_encode_peak never called after the encode phase, measuring the wrong phase record, or a bound # of its own that drifts from the helper's (`>` at the bound: silent). The bound comes from the plan through the - # helper Task 7 pins with literals. + # helper, whose own tests pin it with literals. import warnings from mlx_dfloat.mflux.ernie import memory as emem @@ -634,7 +634,7 @@ def test_save_model_is_refused_and_freeze_skips_a_dropped_encoder(tmp_path, monk def test_the_constructor_builds_with_the_models_overrides_and_the_language_model_bytes( tmp_path, monkeypatch ): - # Bug caught (S3 review fix 9): __init__ wiring past the resolver untested: the transformer built without the + # Bug caught: __init__ wiring past the resolver untested: the transformer built without the # model's overrides (mflux builds ErnieTransformer(**model_config.transformer_overrides), # ernie_image_initializer.py:63; ModelConfig.ernie_image_turbo() sets rope_axes_dim [32, 48, 48]), the base # components read from another directory, or the sizes taken from the encoder file (vision tower included) instead @@ -841,7 +841,7 @@ def test_a_tiny_generation_gives_stock_mfluxs_latents_bit_for_bit(tmp_path, monk def test_encode_then_generate_on_the_base_reuses_the_batch(tmp_path, monkeypatch): - # Bug caught (C1): encode("p") keyed at guidance 1.0 while generate_image(seed, "p") runs the base's default 4.0 + # Bug caught: encode("p") keyed at guidance 1.0 while generate_image(seed, "p") runs the base's default 4.0 # (CFG: the batch [" ", "p"]), so the call drops the set and reloads the encoder for a batch encode() was meant to # prepare. model = fake_model(tmp_path, monkeypatch, model="ernie-image") @@ -856,7 +856,7 @@ def test_encode_then_generate_on_the_base_reuses_the_batch(tmp_path, monkeypatch @pytest.mark.parametrize("prompt", ["", " "]) def test_encode_refuses_an_empty_or_blank_prompt(tmp_path, monkeypatch, prompt): - # Bug caught (C3): encode() taking a prompt generate_image refuses (an empty batch cached, an encoder load spent). + # Bug caught: encode() taking a prompt generate_image refuses (an empty batch cached, an encoder load spent). model = fake_model(tmp_path, monkeypatch) monkeypatch.setattr(model._lifecycle, "ensure_embeddings", lambda *p: pytest.fail("encoded")) with pytest.raises( @@ -866,7 +866,7 @@ def test_encode_refuses_an_empty_or_blank_prompt(tmp_path, monkeypatch, prompt): def test_encode_runs_under_the_commands_caps_and_restores_mlxs_defaults(tmp_path, monkeypatch): - # Bug caught (X4): encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured + # Bug caught: encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured # phases, sized under the command's caps), or the caps left installed after it returns. from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap diff --git a/tests/test_ernie_names.py b/tests/test_ernie_names.py index e1f7eed..4a276a4 100644 --- a/tests/test_ernie_names.py +++ b/tests/test_ernie_names.py @@ -202,7 +202,7 @@ def test_a_time_embedding_group_with_one_matrix_is_refused(tmp_path): def test_a_zimage_block_under_the_same_group_name_is_refused_naming_the_seven_matrices(tmp_path): - # Bug caught (Review Focus 3): a Z-Image DF11 checkpoint (its main blocks are also named layers.) passed for + # Bug caught: a Z-Image DF11 checkpoint (its main blocks are also named layers.) passed for # ERNIE and failing with a shape error at the first block instead of at build. zimage = ( "attention.to_q", diff --git a/tests/test_flux1_memory.py b/tests/test_flux1_memory.py index c5738a0..56d8f52 100644 --- a/tests/test_flux1_memory.py +++ b/tests/test_flux1_memory.py @@ -22,7 +22,7 @@ GIB = 1024**3 LARGEST = {"transformer_blocks": 679_000_000, "single_transformer_blocks": 283_000_000} -# Measured on the 2026-09-28 de-risk run (schnell DF11 `51a428b9`, schnell base `741f7c3c`). +# Measured on the 2026-09-28 calibration run (schnell DF11 `51a428b9`, schnell base `741f7c3c`). MEASURED_VAE_PEAK_GIB = 23.29 # derisk.json phases.vae.footprint_peak, set resident SCHNELL_SIZES = FluxSizes( compressed=16_195_141_095, extras=113_899_648, encoders=9_770_792_936, vae=167_666_902 @@ -216,8 +216,8 @@ def test_fit_for_uses_the_measured_constants_and_names_the_peak_phase(): def test_the_measured_constants_are_the_recorded_ones(): # Bug caught: a constant edited without a new measurement (a "tidy-up" rounding it, a units slip). - # OVERHEAD: the de-risk run's build-phase footprint minus MLX active and cache. - # VAE_TRANSIENT: the de-risk run's VAE-phase footprint peak (23.29 GiB) minus the footprint at the + # OVERHEAD: the calibration run's build-phase footprint minus MLX active and cache. + # VAE_TRANSIENT: the calibration run's VAE-phase footprint peak (23.29 GiB) minus the footprint at the # end of the set load (15.47 GiB) = 7.82 GiB. The same record's `derisk.json["vae_transient_bytes"]` # (5.47 GiB) is a different quantity: the VAE peak minus the footprint at the end of the denoise # phase, which already carries the retained activations. Re-deriving from the record means the diff --git a/tests/test_flux1_model.py b/tests/test_flux1_model.py index 096346f..546e40d 100644 --- a/tests/test_flux1_model.py +++ b/tests/test_flux1_model.py @@ -213,8 +213,8 @@ def test_save_model_is_refused_and_freeze_skips_dropped_encoders(tmp_path, monke def test_plan_call_rounds_to_multiples_of_16_and_warns_below_the_derived_minimum( tmp_path, monkeypatch, caplog ): - # Bug caught: planning at the requested 1000² while mflux runs 992² (Review Focus 1), or an - # override below the two largest groups accepted silently (Review Focus 4). + # Bug caught: planning at the requested 1000² while mflux runs 992², or an + # override below the two largest groups accepted silently. model = _fake_model(tmp_path, monkeypatch) assert ( model.plan_call(height=1000, width=1000).cache_limit @@ -426,7 +426,7 @@ def recording_load_encoders(root): @pytest.mark.parametrize("where", ["raise_before_loop", "raise_after_loop"]) def test_a_format_error_mid_step_drops_the_set_and_restores_the_limit(tmp_path, monkeypatch, where): - # Bug caught (Review Focus 2): a corrupt block leaving the set attached with stale status words, + # Bug caught: a corrupt block leaving the set attached with stale status words, # so the retry is refused by begin_step, or the cache limit left at the call's value; checked # both before the after-loop guard ran (inside the loop) and after it (during the decode). model = _fake_model(tmp_path, monkeypatch) @@ -475,7 +475,7 @@ def fake(self, **kwargs): def test_an_interrupt_keeps_the_set_resident_and_restores_the_limit(tmp_path, monkeypatch): - # Bug caught (Review Focus 3): Ctrl-C (mflux's StopImageGenerationException) dropping the set + # Bug caught: Ctrl-C (mflux's StopImageGenerationException) dropping the set # (a 24 s reload on the retry) or leaving the limit changed. from mflux.utils.exceptions import StopImageGenerationException @@ -689,7 +689,7 @@ def test_a_python_api_call_runs_under_the_commands_caps_and_restores_mlxs_defaul def test_encode_runs_under_the_commands_caps_and_restores_mlxs_defaults(tmp_path, monkeypatch): - # Bug caught (X4): encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured + # Bug caught: encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured # phases, sized under the command's caps), or the caps left installed after it returns. from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap diff --git a/tests/test_flux2_memory.py b/tests/test_flux2_memory.py index 741f63a..6e7ade9 100644 --- a/tests/test_flux2_memory.py +++ b/tests/test_flux2_memory.py @@ -1,4 +1,4 @@ -"""FLUX.2 Klein's memory constants and its encoder sizing, checked against the 1024² de-risk measurements. +"""FLUX.2 Klein's memory constants and its encoder sizing, checked against the 1024² calibration measurements. Measured inputs (2026-10-08, base 4B and base 9B, guidance 4, 1024², seed 42; one process per size: build, encode, drop, set load, one step, VAE decode on the resident set; M1 Max 32 GB, mlx 0.32.2, mflux 0.20.0): @@ -137,7 +137,7 @@ def _fit_9b(encoders=KLEIN_9B_ENCODER_USED, compressed=12_275_304_729): ) -# The constants are calibrated on the de-risk runs (MEASURED_4B / MEASURED_9B), so there the estimate equals the +# The constants are calibrated on the calibration runs (MEASURED_4B / MEASURED_9B), so there the estimate equals the # measurement by construction: the in-sample checks are exact. Whether the model predicts is decided by the runs below, # none of them used for calibration. diff --git a/tests/test_flux2_model.py b/tests/test_flux2_model.py index a02ae91..c146c41 100644 --- a/tests/test_flux2_model.py +++ b/tests/test_flux2_model.py @@ -204,7 +204,7 @@ def test_the_set_load_installs_the_five_nonblock_weights_and_attaches_with_verif def test_a_new_prompt_drops_the_set_reloads_the_encoder_and_reinstalls_the_nonblock_weights( tmp_path, monkeypatch ): - # Bug caught (Review Focus 5): context_embedder (or any non-block weight) left on its placeholder after the + # Bug caught: context_embedder (or any non-block weight) left on its placeholder after the # reload a new prompt forces, or the reload not reading the base the model was built from. model = fake_model(tmp_path, monkeypatch) loads = [] @@ -491,7 +491,7 @@ def fake(self, **kwargs): [(None, ("p",)), (1.0, ("p",)), (1.0001, ("p", " ")), (4.0, ("p", " "))], ) def test_cfg_prompts_follow_mflux(tmp_path, monkeypatch, guidance, expected): - # Bug caught (Review Focus 4): `>=` for `>` (guidance 1.0 encoding the negative: a wasted encoder reload and a + # Bug caught: `>=` for `>` (guidance 1.0 encoding the negative: a wasted encoder reload and a # second transformer call), or guidance None encoding one. mflux 0.20.0 flux2_klein.py:78-80, 163: the negative is # " " and is encoded only when guidance > 1.0. model = fake_model(tmp_path, monkeypatch, model="flux2-klein-base-4b") @@ -874,7 +874,7 @@ def test_a_python_api_call_runs_under_the_commands_caps_and_restores_mlxs_defaul def test_encode_runs_under_the_commands_caps_and_restores_mlxs_defaults(tmp_path, monkeypatch): - # Bug caught (X4): encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured + # Bug caught: encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured # phases, sized under the command's caps), or the caps left installed after it returns. from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap diff --git a/tests/test_format_arrays.py b/tests/test_format_arrays.py index bde78d4..f92ed19 100644 --- a/tests/test_format_arrays.py +++ b/tests/test_format_arrays.py @@ -85,5 +85,5 @@ def test_pointer_bounds(rows, pointer): def test_tail_block_without_a_code_start_is_accepted(): - # Review focus 5: n_bytes spans 2 blocks, no code starts in the second, 2 entries. + # n_bytes spans 2 blocks, no code starts in the second, 2 entries. validate_group_arrays(_two_blocks(output_positions=np.array([0, 8], np.uint32)), name="g") diff --git a/tests/test_format_row_splits.py b/tests/test_format_row_splits.py index b1ce1fb..115ea50 100644 --- a/tests/test_format_row_splits.py +++ b/tests/test_format_row_splits.py @@ -17,7 +17,7 @@ with_row_splits, ) -# S0 (2026-10-07): every Qwen-Image 2.1 block's stored split_positions and element count. +# Every Qwen-Image 2.1 block's stored split_positions and element count (the published header, read 2026-10-07). QWEN_SPLITS = np.array( [16_777_216, 33_554_432, 50_331_648, 67_108_864, 167_772_160], dtype=np.int64 ) @@ -41,8 +41,8 @@ FUSED = {"img_mlp.gate_up": ("img_mlp.gate_layer", "img_mlp.proj")} -def test_the_gate_up_segment_is_cut_in_half_at_the_s0_seam(): - # Bug caught: the seam taken from the wrong segment (index off by one) or not halved. S0 measured the proj rows +def test_the_gate_up_segment_is_cut_in_half_at_the_measured_seam(): + # Bug caught: the seam taken from the wrong segment (index off by one) or not halved. A 2026-10-07 byte comparison put the proj rows # (gate_up row 12288) starting at element 117,440,512. got = insert_row_splits( QWEN_SPLITS, QWEN_BLOCK_ELEMENTS, ((4, 2),), name="transformer_blocks.0" @@ -118,7 +118,7 @@ def test_a_disordered_stored_split_table_is_refused_before_planning(splits, n): def test_the_plan_expands_gate_up_into_gate_layer_then_proj(): - # Bug caught: the halves swapped (S0: rows 0..12287 are gate_layer), or the fused name left in the list. + # Bug caught: the halves swapped (rows 0..12287 are gate_layer, by the 2026-10-07 byte comparison), or the fused name left in the list. config = with_row_splits(parse_df11_config(QWEN_RAW, source="t"), FUSED, source="t") stored, names, plan = row_split_plan("transformer_blocks.7", config) head = "transformer_blocks.7." diff --git a/tests/test_integrate_blockseam.py b/tests/test_integrate_blockseam.py index af4ff7e..78911eb 100644 --- a/tests/test_integrate_blockseam.py +++ b/tests/test_integrate_blockseam.py @@ -133,3 +133,59 @@ def test_a_raise_inside_a_block_restores_placeholders_and_clears_pending_words() tf(mx.ones((1, 2, 5)), mx.ones((1, 2, 4)), mx.ones((1, 4))) # D=4 blocks given width 5 assert provider.pending == [] assert tf.noise_refiner[0].attention.to_q.weight.size == 0 # the block that raised + + +class _HookedProvider(ResidentProvider): + """A resident provider with the optional ``before_block`` hook: records the hook and the request in order.""" + + def __init__(self, per_block, rec): + super().__init__(per_block) + self.rec = rec + + def before_block(self, block_name, args, kwargs): + self.rec.events.append(("before", block_name, tuple(id(a) for a in args), dict(kwargs))) + + def weights_for(self, block_name, shapes): + self.rec.events.append(("weights", block_name)) + return super().weights_for(block_name, shapes) + + +def test_a_providers_before_block_hook_runs_before_each_request_with_the_blocks_own_inputs(): + # Bug caught: the hook called after the weights are requested (a provider releasing memory there would release + # it after the block's decode, not before), called with another block's inputs, or not called for every block. + # Block 0 of layers is called with (concat(x, cap), t_emb); layers.1 with layers.0's output. + rec = Recorder() + tf, shapes = _seamed(rec) + x, cap, t_emb = zimage_inputs() + tf.attach(_HookedProvider(resident_dicts(shapes), rec), shapes) + mx.eval(tf(x, cap, t_emb)) + runs = [e[1] for e in rec.events if e[0] == "run"] + order = [(e[0], e[1]) for e in rec.events if e[0] != "run"] + names = ["noise_refiner.0", "context_refiner.0", "layers.0", "layers.1"] + assert order == [step for n in names for step in (("before", n), ("weights", n))] + hooks = {e[1]: e[2:] for e in rec.events if e[0] == "before"} + assert hooks["noise_refiner.0"] == ((id(x), id(t_emb)), {}) + assert hooks["context_refiner.0"][0][0] != id( + cap + ) # cap went through the resident cap_embedder first + assert hooks["layers.1"] == ((runs[2], id(t_emb)), {}) + + +def test_without_the_hook_the_seam_requests_the_weights_and_nothing_else(): + # Bug caught: the seam calling something new on a provider that has no hook (every family before Krea 2 runs one + # of these), e.g. a hook looked up with a default that is called, or the inputs evaluated for every provider. + # The provider records every attribute the seam reads during a step. + seen = [] + + class Watching(ResidentProvider): + def __getattribute__(self, name): + if not name.startswith("_"): + seen.append(name) + return super().__getattribute__(name) + + rec = Recorder() + tf, shapes = _seamed(rec) + tf.attach(Watching(resident_dicts(shapes)), shapes) + seen.clear() + mx.eval(tf(*zimage_inputs())) + assert seen == ["pending"] + ["before_block", "weights_for"] * 4 diff --git a/tests/test_integrate_resident.py b/tests/test_integrate_resident.py index 046c965..28379f1 100644 --- a/tests/test_integrate_resident.py +++ b/tests/test_integrate_resident.py @@ -133,3 +133,56 @@ def test_a_group_that_is_not_resident_and_a_size_that_does_not_fit_are_refused() names, decode=REFERENCE, ) + + +def _two_groups(module, names): + """Two one-matrix groups landing on proj.inner and cap_embedder.1 (renamed and plain).""" + rng = np.random.default_rng(9) + a, b = random_bf16(rng, SHAPE), random_bf16(rng, SHAPE) + groups = { + "proj": encoder_group(a.reshape(-1)).to_mx(name="proj"), + "cap_embedder": encoder_group(b.reshape(-1)).to_mx(name="cap_embedder"), + } + matrices = {"proj": ("proj.inner.weight",), "cap_embedder": (MATRIX,)} + shapes = install_nonblock_placeholders(module, ["proj.inner.weight", MATRIX], names) + return groups, matrices, shapes, {"proj.inner.weight": a, MATRIX: b} + + +@pytest.mark.parametrize(("together", "evals"), [(False, [2, 2]), (True, [4])]) +def test_the_nonblock_groups_evaluate_per_group_or_in_one_eval(monkeypatch, together, evals): + # Bug caught: the per-call path (Krea 2 decodes its seven groups at every transformer call) paying one eval per + # group, or the default changed for the families that decode once at set load (it stays one eval per group: + # status word + matrix each). The bits are the same either way. + from mlx_dfloat.integrate import resident + + seen = [] + real = resident._eval + monkeypatch.setattr(resident, "_eval", lambda *a: (seen.append(len(a)), real(*a))) + names = StaticNameMap(TABLE) + module = _module() + groups, matrices, shapes, sources = _two_groups(module, names) + weights = decode_nonblock( + groups, matrices, shapes, names, decode=REFERENCE, eval_together=together + ) + assert seen == evals # one status word and one matrix per group, per eval + for param, source in sources.items(): + assert np.array_equal(np.array(weights[param].view(mx.uint16)), source), param + + +def test_a_decode_error_in_one_eval_still_names_its_group(): + # Bug caught: the one-eval path checking the status words before they are evaluated, or naming the first group + # for an error in the second. + names = StaticNameMap(TABLE) + groups, matrices, shapes, _sources = _two_groups(_module(), names) + + def flag_cap(g): + ok = REFERENCE(g) + if g.name != "cap_embedder": + return ok + status = mx.full(ok.status.shape, STATUS_BROKEN_CHAIN, dtype=mx.uint32) + return DecodeResult( + bits=ok.bits, status=status, backend=ok.backend, direct_blocks=0, threadgroup_bytes=0 + ) + + with pytest.raises(DFloatFormatError, match="cap_embedder"): + decode_nonblock(groups, matrices, shapes, names, decode=flag_cap, eval_together=True) diff --git a/tests/test_integrate_seam.py b/tests/test_integrate_seam.py index ddde22f..61d9269 100644 --- a/tests/test_integrate_seam.py +++ b/tests/test_integrate_seam.py @@ -484,3 +484,23 @@ def test_detach_clears_a_df11_providers_pending_status_words(): assert provider.pending != [] # verify_step was never called tf.detach() assert provider.pending == [] + + +def test_a_seam_that_passes_no_inputs_hands_each_hook_call_its_own_empty_arguments(): + # Bug caught: one shared default argument object for every call of run_block, so a before_block hook that adds a + # key to the kwargs it is given (the FLUX.1 seam passes none) leaks it into every later block's call, and into the + # next step's. + seen = [] + + class Hooked(ResidentProvider): + def before_block(self, block_name, args, kwargs): + seen.append((block_name, args, dict(kwargs))) + kwargs["touched"] = block_name + + rec = Recorder() + tf = FakeSeamTransformer(rec, n_double=2, n_single=1) + shapes = install_placeholders(block_lists(tf), FLUX_TABLE) + tf.attach(Hooked(resident_dicts(shapes)), shapes, eval_policy="per-block") + mx.eval(tf(*inputs())) + mx.eval(tf(*inputs())) + assert [(args, kwargs) for _name, args, kwargs in seen] == [((), {})] * 6 diff --git a/tests/test_krea2_build.py b/tests/test_krea2_build.py new file mode 100644 index 0000000..12ef94d --- /dev/null +++ b/tests/test_krea2_build.py @@ -0,0 +1,425 @@ +"""The Krea 2 build from a config-less DF11 checkpoint: base extras offline, the real tiny build (mflux lane).""" + +from functools import partial +from pathlib import Path +from types import SimpleNamespace + +import mlx.core as mx +import numpy as np +import pytest +from mlx.utils import tree_flatten +from tests._krea2_tiny import TINY, write_tiny_checkpoint + +from mlx_dfloat.decode import decode_group +from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError +from mlx_dfloat.format import open_checkpoint +from mlx_dfloat.mflux.krea2.transformer import base_extras, build_transformer + +# The eight block matrices at TINY size, from mflux 0.20.0 (weights are (out, in)): Krea2Attention's wq / gate / wo +# Linear(32, 32) and wk / wv Linear(32, 16 x 1 kv head) (attention.py:16-21); Krea2SwiGLU's gate / up +# Linear(32, 128) and down Linear(128, 32) (feed_forward.py:7-12, width 84 rounded up to 128). +EIGHT = { + "attn.wq": (32, 32), + "attn.wk": (16, 32), + "attn.wv": (16, 32), + "attn.gate": (32, 32), + "attn.wo": (32, 32), + "mlp.gate": (128, 32), + "mlp.up": (128, 32), + "mlp.down": (32, 128), +} +# A text-fusion block at txtdim 32, 2 heads and 2 kv heads: five Linear(32, 32), the same SwiGLU (text_fusion.py). +TEXT_BLOCK = { + "attn.wq": (32, 32), + "attn.wk": (32, 32), + "attn.wv": (32, 32), + "attn.gate": (32, 32), + "attn.wo": (32, 32), + "mlp.gate": (128, 32), + "mlp.up": (128, 32), + "mlp.down": (32, 128), +} +# The non-block matrices as mflux names them: tmlp Linear(tdim 32, 32) and Linear(32, 32); tproj Linear(32, 6 x 32); +# txtmlp Linear(txtdim 32, 32) and Linear(32, 32) (timestep_embedder.py, text_mlp.py). +NONBLOCK = { + "tmlp.linear_in.weight": (32, 32), + "tmlp.linear_out.weight": (32, 32), + "tproj.linear.weight": (192, 32), + "txtmlp.linear_in.weight": (32, 32), + "txtmlp.linear_out.weight": (32, 32), + **{ + f"txtfusion.{kind}.{i}.{attr}.weight": shape + for kind in ("layerwise_blocks", "refiner_blocks") + for i in (0, 1) + for attr, shape in TEXT_BLOCK.items() + }, +} + + +def test_base_extras_drop_the_block_matrices_and_keep_the_nonblock_ones(): + # Bug caught: the block matrices loaded from the base as extras (a coverage refusal), or a non-block matrix + # dropped (the BF16 side of an identity check would leave tproj on a placeholder). + groups = { + "blocks.0": SimpleNamespace( + matrix_names=("blocks.0.attn.wq.weight", "blocks.0.mlp.down.weight") + ), + "tproj": SimpleNamespace(matrix_names=("tproj.1.weight",)), + "txtfusion.refiner_blocks.1": SimpleNamespace( + matrix_names=("txtfusion.refiner_blocks.1.attn.wq.weight",) + ), + } + index = { + name: (Path(name), None) + for name in ( + "blocks.0.attn.wq.weight", + "blocks.0.mlp.down.weight", + "blocks.0.mod.lin", + "tproj.1.weight", + "txtfusion.refiner_blocks.1.attn.wq.weight", + "first.weight", + ) + } + assert sorted(base_extras(index, SimpleNamespace(groups=groups))) == [ + "blocks.0.mod.lin", + "first.weight", + "tproj.1.weight", + "txtfusion.refiner_blocks.1.attn.wq.weight", + ] + + +def _open(tmp_path, **kwargs): + root = tmp_path / "df11" + matrices, constants, layout = write_tiny_checkpoint(root, np.random.default_rng(3), **kwargs) + return open_checkpoint(root, layouts=(layout,)), matrices, constants + + +def _params(build): + return dict(tree_flatten(build.transformer.parameters())) + + +@pytest.mark.mflux +def test_the_build_installs_placeholders_and_loads_every_extra_by_name(tmp_path): + # Bug caught: a rename lost (txtmlp.0.scale loaded nowhere: a coverage refusal) or an extra on the wrong parameter + # (a constant in the wrong place), a compressed matrix left live (memory), or a non-block matrix not on a + # placeholder (it would load nothing and keep mflux's random init). + ckpt, _matrices, constants = _open(tmp_path) + build = build_transformer(ckpt, transformer_kwargs=TINY) + params = _params(build) + assert build.shapes == {f"blocks.{i}": EIGHT for i in range(2)} + assert build.nonblock == NONBLOCK + for block in build.shapes: + for attr in EIGHT: + assert params[f"{block}.{attr}.weight"].size == 0, (block, attr) + for name in NONBLOCK: + assert params[name].size == 0, name + assert "tproj.linear.bias" in constants + assert "blocks.1.mod.lin" in constants + assert "txtmlp.norm.scale" in constants + for name, value in constants.items(): + assert np.all(np.array(params[name].view(mx.uint16)) == value), name + assert build.counts == {"blocks": 2} + + +@pytest.mark.mflux +def test_nonblock_groups_stored_as_extras_load_as_plain_weights(tmp_path): + # Bug caught: a checkpoint without the non-block groups refused, or tproj left on a placeholder (a zero-size matmul + # before the first block). + ckpt, _m, constants = _open(tmp_path, nonblock_as_groups=False) + assert set(ckpt.groups) == {"blocks.0", "blocks.1"} + assert "tproj.1.weight" in ckpt.extras + build = build_transformer(ckpt, transformer_kwargs=TINY) + weight = _params(build)["tproj.linear.weight"] + assert build.nonblock == {} + assert weight.shape == (192, 32) + assert np.all(np.array(weight.view(mx.uint16)) == constants["tproj.linear.weight"]) + + +@pytest.mark.mflux +def test_a_checkpoint_with_other_block_counts_is_refused_naming_both(tmp_path): + # Bug caught: a 3-block checkpoint built as 2 (block 2 silently dropped) instead of a refusal. + ckpt, _m, _c = _open(tmp_path, n_layers=3) + with pytest.raises(DFloatFormatError, match=r"^checkpoint has 3 blocks; this model builds 2$"): + build_transformer(ckpt, transformer_kwargs=TINY) + + +@pytest.mark.mflux +def test_the_count_check_defaults_to_mfluxs_depth_under_mfluxs_parameter_name(tmp_path): + # Bug caught: the depth read from `num_layers` (Qwen's and ERNIE's name) instead of Krea's `layers`, so the + # published model (no kwargs) or the tiny one would be checked against the wrong depth. + ckpt, _m, _c = _open(tmp_path) + no_depth = {k: v for k, v in TINY.items() if k != "layers"} + with pytest.raises(DFloatFormatError, match=r"^checkpoint has 2 blocks; this model builds 28$"): + build_transformer(ckpt, transformer_kwargs=no_depth) + + +@pytest.mark.mflux +def test_n_layers_builds_fewer_blocks_and_is_range_checked(tmp_path): + # Bug caught: a partial build loading block 1's norms into a transformer without block 1 (an extra without a + # parameter), or a request beyond the checkpoint built from nothing. + ckpt, _m, _c = _open(tmp_path) + build = build_transformer(ckpt, transformer_kwargs=TINY, n_layers=1) + assert build.counts == {"blocks": 1} + assert len(build.transformer.blocks) == 1 + with pytest.raises(DFloatIntegrationError, match=r"asked for 3 blocks; the checkpoint has 2"): + build_transformer(ckpt, transformer_kwargs=TINY, n_layers=3) + + +@pytest.mark.mflux +def test_the_decoded_block_matrices_land_in_pattern_order(tmp_path): + # Bug caught: the order swapped between the layout, the reader and the map (attn.gate and attn.wo, or mlp.up and + # mlp.down: equal shapes or element counts, so only the values show it). + from mlx_dfloat.integrate.coverage import load_resident_set + from mlx_dfloat.integrate.providers import DF11Provider + from mlx_dfloat.mflux.krea2.names import krea2_name_map + + ckpt, matrices, _c = _open(tmp_path) + build = build_transformer(ckpt, transformer_kwargs=TINY) + blocks = ["blocks.0", "blocks.1"] + provider = DF11Provider( + load_resident_set(ckpt, blocks), + {g: ckpt.groups[g].matrix_names for g in blocks}, + krea2_name_map(), + decode=partial(decode_group, backend="reference"), + ) + got = provider.weights_for("blocks.1", build.shapes["blocks.1"]) + provider.verify() + for attr in EIGHT: + want = matrices["blocks.1"][attr] + assert np.array_equal(np.array(got[attr].view(mx.uint16)), want), attr + + +@pytest.mark.mflux +def test_the_raw_model_refuses_the_turbo_layouts_checkpoint_at_build(tmp_path): + # Bug caught: Turbo's distilled weights built and run under Raw's 25-step default (or the card's + # 52 steps at 3.5), and a row labelled Raw measuring Turbo. The refusal comes before any module is built: the + # transformer class is never called (MLX builds modules lazily, so active memory cannot show a stray build). + from mlx_dfloat.mflux.krea2.transformer import seam_transformer_class + + ckpt, _m, _c = _open(tmp_path, key="krea-2-turbo-comfyui") + built = [] + + def recording(**kwargs): + built.append(kwargs) + return seam_transformer_class()(**kwargs) + + with pytest.raises( + DFloatFormatError, match=r"Krea 2 Turbo DF11 checkpoint .*run it with --model krea-2$" + ): + build_transformer( + ckpt, model="krea-2-raw", transformer_kwargs=TINY, transformer_class=recording + ) + assert built == [] + build = build_transformer( + ckpt, model="krea-2", transformer_kwargs=TINY, transformer_class=recording + ) + assert build.counts == {"blocks": 2} + assert len(built) == 1 + + +@pytest.mark.mflux +def test_the_build_refuses_one_that_adds_too_much_active_memory(tmp_path, monkeypatch): + # Bug caught: the active-memory guard not wired (a build that materialised the decoded set would pass). + from mlx_dfloat.mflux.krea2 import transformer as module + + ckpt, _m, _c = _open(tmp_path) + monkeypatch.setattr(module, "MAX_BUILD_ACTIVE_BYTES", 0) + with pytest.raises(DFloatIntegrationError, match="GiB of active memory"): + build_transformer(ckpt, transformer_kwargs=TINY) + + +@pytest.mark.mflux +def test_the_bf16_side_builds_with_the_nonblock_matrices_as_plain_weights(tmp_path): + # Bug caught (the BF16 side of an identity check): the block matrices loaded as extras, or a non-block matrix left + # on a placeholder instead of loading from the base. + from tests._df11_fixtures import write_bf16_original + + from mlx_dfloat.mflux.krea2.transformer import base_transformer_files_index + + ckpt, matrices, _constants = _open(tmp_path) + originals = { + name: np.full(info.shape, 0x4000, dtype=np.uint16) + for name, (_p, info) in ckpt.extras.items() + } + originals.update( + (name, mat) + for group, g in ckpt.groups.items() + for name, mat in zip(g.matrix_names, matrices[group].values(), strict=True) + ) + root = write_bf16_original(tmp_path / "base", originals) + extras = base_extras(base_transformer_files_index(root), ckpt) + assert not any(n.startswith("blocks.") and n.endswith(".weight") for n in extras) + build = build_transformer( + ckpt, transformer_kwargs=TINY, extras=extras, nonblock_from_extras=True + ) + params = _params(build) + assert build.nonblock == {} + tproj = np.array(params["tproj.linear.weight"].view(mx.uint16)) + assert np.array_equal(tproj, matrices["tproj"]["tproj.linear.weight"]) + assert np.all(np.array(params["first.weight"].view(mx.uint16)) == 0x4000) + assert params["blocks.0.mlp.up.weight"].size == 0 + + +# --- the BF16/FP32 base file as the identity check's reference side ----------------------------------------------- + +# FP32 bit patterns and their BF16 by round-to-nearest-even, by hand (the patterns of the parity check's tests): +# a tie with an even neighbour stays, a tie with an odd neighbour rounds up, just above half rounds up, just below +# half rounds down, and the largest finite float32 rounds to +inf. Truncation would give 0x3F80, 0x3F81, 0x3F80, +# 0x3F80, 0x7F7F. +F32_PATTERNS = (0x3F808000, 0x3F818000, 0x3F808001, 0x3F807FFF, 0x7F7FFFFF) +RNE_BF16 = (0x3F80, 0x3F82, 0x3F81, 0x3F80, 0x7F80) + + +def test_mlx_astype_rounds_fp32_to_bf16_nearest_even(): + # Bug caught: if MLX truncated (or rounded half away from zero), the identity's reference side, which casts the + # base's FP32 tensors as mflux does at load (weight_loader.py:448-449), would differ from the checkpoint's own + # BF16 (rounded nearest even when it was made) and report a false mismatch on the five FP32 groups and the extras. + f32 = mx.array(np.array(F32_PATTERNS, dtype=np.uint32)).view(mx.float32) + assert np.array(f32.astype(mx.bfloat16).view(mx.uint16)).tolist() == list(RNE_BF16) + + +def _base_file(path, tensors): + """A safetensors file of ``{name: (numpy bits, mx dtype)}``, written as that dtype.""" + mx.save_safetensors( + str(path), {n: mx.array(bits).view(dtype) for n, (bits, dtype) in tensors.items()} + ) + return path + + +def test_base_bf16_weights_reads_bf16_as_is_and_rounds_fp32_to_nearest_even(tmp_path): + # Bug caught: an FP32 base tensor truncated (two of the five patterns differ) or read as two BF16 halves, a BF16 + # one converted at all, or a name read from another tensor. + from mlx_dfloat.mflux.krea2.transformer import base_bf16_weights + + bf16 = np.array([0x3F80, 0xBF80, 0x0001], dtype=np.uint16) + path = _base_file( + tmp_path / "raw.safetensors", + { + "blocks.0.attn.wq.weight": (bf16, mx.bfloat16), + "tproj.1.bias": (np.array(F32_PATTERNS, dtype=np.uint32), mx.float32), + "unused": (np.array([7], dtype=np.uint16), mx.bfloat16), + }, + ) + got = base_bf16_weights(path, ["tproj.1.bias", "blocks.0.attn.wq.weight"]) + assert sorted(got) == ["blocks.0.attn.wq.weight", "tproj.1.bias"] + assert {n: a.dtype for n, a in got.items()} == { + "blocks.0.attn.wq.weight": mx.bfloat16, + "tproj.1.bias": mx.bfloat16, + } + assert np.array(got["blocks.0.attn.wq.weight"].view(mx.uint16)).tolist() == [ + 0x3F80, + 0xBF80, + 0x0001, + ] + assert np.array(got["tproj.1.bias"].view(mx.uint16)).tolist() == list(RNE_BF16) + + +def test_base_bf16_weights_evaluates_each_cast_as_it_goes(tmp_path, monkeypatch): + # Bug caught: the FP32 casts left lazy, so every FP32 source array (twice its BF16 size) stays alive until the + # caller's first eval: on the real file 174 FP32 tensors held at once instead of one. One eval per FP32 tensor, + # each on its own cast, in read order; a BF16 tensor needs none. + from mlx_dfloat.mflux.krea2 import transformer as ktf + + seen = [] + real = ktf._eval + monkeypatch.setattr(ktf, "_eval", lambda *a: (seen.append(tuple(x.dtype for x in a)), real(*a))) + f32 = (np.array(F32_PATTERNS, dtype=np.uint32), mx.float32) + path = _base_file( + tmp_path / "raw.safetensors", + { + "a.bias": f32, + "b.weight": (np.array([0x3F80], dtype=np.uint16), mx.bfloat16), + "c.bias": f32, + }, + ) + got = ktf.base_bf16_weights(path, ["a.bias", "b.weight", "c.bias"]) + assert seen == [(mx.bfloat16,), (mx.bfloat16,)] + assert np.array(got["c.bias"].view(mx.uint16)).tolist() == list(RNE_BF16) + + +def test_base_bf16_weights_refuses_an_f16_tensor_and_a_missing_name(tmp_path): + # Bug caught: an F16 tensor read as BF16 bits (a silent wrong reference), or a tensor the base lacks skipped (the + # reference side would keep the DF11 file's value for it and agree by construction). + from mlx_dfloat.mflux.krea2.transformer import base_bf16_weights + + path = _base_file( + tmp_path / "raw.safetensors", + {"first.bias": (np.array([0x3C00], dtype=np.uint16), mx.float16)}, + ) + with pytest.raises( + DFloatFormatError, match=r"first\.bias is F16; the base must hold BF16 or F32" + ): + base_bf16_weights(path, ["first.bias"]) + with pytest.raises( + DFloatFormatError, match=r"raw\.safetensors has no tensor 'last\.linear\.bias'" + ): + base_bf16_weights(path, ["last.linear.bias"]) + + +@pytest.mark.mflux +def test_install_base_weights_replaces_every_extra_and_nonblock_matrix_with_the_bases(tmp_path): + # Bug caught: the reference side keeping any DF11 value (an extra or a non-block matrix not replaced: the identity + # would then compare the checkpoint with itself on it), an FP32 base tensor not rounded as mflux loads it, a + # rename lost (tmlp.0 vs tmlp.linear_in), or a shape mismatch accepted. The base holds every tensor of the tiny + # model under the checkpoint's names; extras and non-block matrices as FP32 (random bits), blocks as BF16. + from scripts.verify_remote_group import fp32_to_bf16_rne + from tests._krea2_tiny import write_tiny_checkpoint + + from mlx_dfloat.format import open_checkpoint + from mlx_dfloat.mflux.krea2.names import NONBLOCK_GROUPS, krea2_name_map + from mlx_dfloat.mflux.krea2.transformer import install_base_weights + + rng = np.random.default_rng(13) + _m, _c, layout = write_tiny_checkpoint(tmp_path / "df11", rng) + ckpt = open_checkpoint(tmp_path / "df11", layouts=(layout,)) + build = build_transformer(ckpt, transformer_kwargs=TINY) + names = krea2_name_map() + nonblock = [m for g in NONBLOCK_GROUPS for m in ckpt.groups[g].matrix_names] + tensors, want = {}, {} + built = dict(tree_flatten(build.transformer.parameters())) + for name in [*ckpt.extras, *nonblock]: + param = names.param_name(name) + shape = build.nonblock[param] if param in build.nonblock else tuple(built[param].shape) + bits = rng.integers(0x3F000000, 0x40000000, size=shape, dtype=np.uint64).astype(np.uint32) + tensors[name] = (bits, mx.float32) + want[param] = fp32_to_bf16_rne(bits) # an independent RNE (the parity check's NumPy one) + path = _base_file(tmp_path / "raw.safetensors", tensors) + replaced = install_base_weights(build.transformer, ckpt, path, names, build.nonblock) + assert sorted(replaced) == sorted([*ckpt.extras, *nonblock]) + params = dict(tree_flatten(build.transformer.parameters())) + for param, bits in want.items(): + got = params[param] + assert got.dtype == mx.bfloat16, param + assert np.array_equal(np.array(got.view(mx.uint16)), bits), param + assert params["blocks.0.attn.wq.weight"].size == 0 # the blocks still stream + + +@pytest.mark.mflux +def test_install_base_weights_refuses_a_tensor_of_another_shape(tmp_path): + # Bug caught: a base of another width loaded onto the tiny transformer without a check (a silent wrong reference + # or a matmul error minutes into the run). + from tests._krea2_tiny import write_tiny_checkpoint + + from mlx_dfloat.format import open_checkpoint + from mlx_dfloat.mflux.krea2.names import NONBLOCK_GROUPS, krea2_name_map + from mlx_dfloat.mflux.krea2.transformer import install_base_weights + + _m, _c, layout = write_tiny_checkpoint(tmp_path / "df11", np.random.default_rng(1)) + ckpt = open_checkpoint(tmp_path / "df11", layouts=(layout,)) + build = build_transformer(ckpt, transformer_kwargs=TINY) + names = krea2_name_map() + tensors = { + n: (np.zeros(info.shape, dtype=np.uint16), mx.bfloat16) + for n, (_p, info) in ckpt.extras.items() + } + for g in NONBLOCK_GROUPS: + for m in ckpt.groups[g].matrix_names: + tensors[m] = ( + np.zeros(build.nonblock[names.param_name(m)], dtype=np.uint16), + mx.bfloat16, + ) + tensors["tproj.1.weight"] = (np.zeros((4, 4), dtype=np.uint16), mx.bfloat16) + path = _base_file(tmp_path / "raw.safetensors", tensors) + with pytest.raises( + DFloatFormatError, match=r"tproj\.1\.weight has shape \(4, 4\); the model needs \(192, 32\)" + ): + install_base_weights(build.transformer, ckpt, path, names, build.nonblock) diff --git a/tests/test_krea2_init.py b/tests/test_krea2_init.py new file mode 100644 index 0000000..4472376 --- /dev/null +++ b/tests/test_krea2_init.py @@ -0,0 +1,389 @@ +"""Krea 2 base components (text encoder, VAE, tokenizer) without the transformer.""" + +import os +import re +from fnmatch import fnmatch +from pathlib import Path + +import httpx +import mlx.core as mx +import pytest +from mlx.utils import tree_flatten, tree_unflatten + +from mlx_dfloat.errors import DFloatAccessError, DFloatFormatError +from mlx_dfloat.mflux import _hub +from mlx_dfloat.mflux.krea2.init import BASE_PATTERNS, DF11_PATTERNS, TOKENIZER_PATTERNS + +# krea/Krea-2-Raw @ 6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603, the Hub file listing (2026-10-08); its +# images/ directory (sample pictures, about 40 MB) is represented by its first and last file. +RAW_FILES = ( + ".gitattributes", + "LICENSE.pdf", + "README.md", + "images/00.png", + "images/header.jpg", + "model_index.json", + "raw.safetensors", + "scheduler/scheduler_config.json", + "text_encoder/config.json", + "text_encoder/model.safetensors", + "tokenizer/chat_template.jinja", + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + "transformer/config.json", + "transformer/diffusion_pytorch_model-00001-of-00003.safetensors", + "transformer/diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer/diffusion_pytorch_model-00003-of-00003.safetensors", + "transformer/diffusion_pytorch_model.safetensors.index.json", + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", +) +# krea/Krea-2-Turbo @ 98e0fe11: turbo.safetensors in place of raw.safetensors, no images/. +TURBO_FILES = tuple( + "turbo.safetensors" if f == "raw.safetensors" else f + for f in RAW_FILES + if not f.startswith("images/") +) +WANTED = [ + "text_encoder/config.json", + "text_encoder/model.safetensors", + "tokenizer/chat_template.jinja", + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", +] +# The calibration prompt and its text tokens: 64 ids, the chat-template prefix ends at 34, so 30 reach the transformer +# (mflux 0.20.0 TokenizerLoader on both base snapshots, 2026-10-08). +PROMPT = ( + "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the " + "water" +) + + +def _hits(files): + return sorted(f for f in files if any(fnmatch(f, p) for p in BASE_PATTERNS)) + + +@pytest.mark.parametrize("files", [RAW_FILES, TURBO_FILES]) +def test_base_patterns_fetch_encoder_vae_and_tokenizer_and_never_a_transformer(files): + # Bug caught: 26 GB of transformer pulled on a user's first call (mflux's own Raw patterns fetch + # transformer/*.safetensors, krea2_weight_definition.py:120-124), the root raw.safetensors / turbo.safetensors + # matched by a loose "*.safetensors", the sample images, or a file of the encoder, the VAE or the tokenizer missed. + assert _hits(files) == WANTED + + +def test_the_tokenizer_pattern_fetches_tokenizer_files_only(): + # Bug caught: the tokenizer pattern widened back to the directory (mflux's "tokenizer/**"), fetching whatever a + # later revision drops there. The listing is the snapshot's three files plus files a later revision could add. + listing = ( + "tokenizer/chat_template.jinja", + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + "tokenizer/extra/tokenizer.json", + "tokenizer/model.bin", + "tokenizer/README.md", + "tokenizer/special_tokens_map.json", + ) + assert _hits(listing) == [ + "tokenizer/chat_template.jinja", + "tokenizer/special_tokens_map.json", + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + ] + assert all(not p.endswith("*") for p in TOKENIZER_PATTERNS) + + +def test_df11_patterns_fetch_the_one_comfyui_file(): + # Bug caught: the config-less repository resolved without its only weight file, or with its README. + # mingyi456/Krea-2-Raw-DF11-ComfyUI @ 8320616b holds these three files (Hub, 2026-10-08). + files = (".gitattributes", "README.md", "krea2_raw_bf16-DF11.safetensors") + assert [f for f in files if any(fnmatch(f, p) for p in DF11_PATTERNS)] == [ + "krea2_raw_bf16-DF11.safetensors" + ] + + +def _gated(status=403): + """huggingface_hub's GatedRepoError with the httpx response its __init__ requires.""" + from huggingface_hub.errors import GatedRepoError + + request = httpx.Request("GET", "https://huggingface.co/api/models/krea/Krea-2-Raw") + return GatedRepoError("gated", response=httpx.Response(status, request=request)) + + +def test_a_gated_base_without_access_names_the_licence_page(monkeypatch): + # Bug caught: a user without the licence accepted sees a bare Hub traceback or a generic "on the Hub", not the + # page to accept it on (both bases are gated, auto-approved). + from mlx_dfloat.mflux.krea2 import init + + def refuse(**_kwargs): + raise _gated() + + monkeypatch.setattr(_hub, "_snapshot_download", refuse) + with pytest.raises(DFloatAccessError) as err: + init.resolve("krea/Krea-2-Raw", patterns=BASE_PATTERNS) + assert str(err.value) == ( + "krea/Krea-2-Raw: this repository is gated or this account may not read it; accept the licence at " + "https://huggingface.co/krea/Krea-2-Raw and run `hf auth login`" + ) + + +def test_hub_errors_that_are_not_access_refusals_pass_through_unchanged(monkeypatch): + # Bug caught: every Hub failure turned into a licence message (an offline cache miss or a typo'd repo id would + # send the user to accept a licence). resolve is the one path that names the licence page. + from mlx_dfloat.mflux.krea2 import init + + other = ValueError("offline") + + def offline(**_kwargs): + raise other + + monkeypatch.setattr(_hub, "_snapshot_download", offline) + with pytest.raises(ValueError, match="offline"): + init.resolve("krea/Krea-2-Raw", patterns=BASE_PATTERNS) + + +# --- mflux lane ----------------------------------------------------------------------------------------------- + + +@pytest.mark.mflux +def test_the_definitions_name_only_their_component_and_the_tokenizer_subdir(): + # Bug caught: a subset keeping the transformer component (mflux's loader would read raw.safetensors or + # transformer/), the encoder losing its prefix transform (the vision tower and every language-model tensor would + # miss), or the tokenizer looked up anywhere but tokenizer/ (krea2_weight_definition.py:96-108). + from mflux.models.krea2.weights.krea2_weight_definition import Krea2WeightDefinition + + from mlx_dfloat.mflux.krea2.init import encoder_definition, vae_definition + + (encoder,) = encoder_definition().get_components() + assert encoder.name == "text_encoder" + assert encoder.key_transform is Krea2WeightDefinition.strip_te_prefix + assert [c.name for c in vae_definition().get_components()] == ["vae"] + (tokenizer,) = encoder_definition().get_tokenizers() + assert (tokenizer.name, tokenizer.hf_subdir, tokenizer.max_length) == ( + "qwen3vl", + "tokenizer", + 1024, + ) + + +# A two-layer Qwen3-VL text stack (tiny widths; mflux 0.20.0 Krea2TextEncoder's constructor arguments). +TINY_ENCODER = { + "vocab_size": 100, + "hidden_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "intermediate_size": 128, + "head_dim": 32, +} + + +@pytest.mark.mflux +def test_the_computed_exemptions_are_exactly_what_the_base_files_do_not_hold(): + # Bug caught: a computed buffer missing from COMPUTED (every real base refused), or a loaded weight listed there (a + # missing tensor never noticed). mflux 0.20.0: the encoder's rotary_emb.inv_freq is computed in + # Qwen3TextRotaryEmbedding.__init__ and the base file holds no such tensor (its 713 tensors are language_model.* + # and visual.*, 2026-10-08 header); every VAE parameter is a target of mflux's VAE mapping. + from mflux.models.krea2.model.krea2_text_encoder.text_encoder import Krea2TextEncoder + from mflux.models.krea2.weights.krea2_weight_mapping import Krea2WeightMapping + from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE + + from mlx_dfloat.mflux.krea2.init import COMPUTED + + encoder = {n for n, _ in tree_flatten(Krea2TextEncoder(**TINY_ENCODER).parameters())} + held = {n for n in encoder if n.startswith(("embed_tokens.", "layers.", "norm."))} + assert encoder - held == {"rotary_emb.inv_freq"} + targets = [ + re.compile(re.sub(r"\\\{[a-z_]+\\\}", r"[0-9]+", re.escape(t.to_pattern)) + "$") + for t in Krea2WeightMapping.get_vae_mapping() + ] + vae = {n for n, _ in tree_flatten(QwenVAE().parameters())} + assert {n for n in vae if not any(t.match(n) for t in targets)} == set() + assert sorted(COMPUTED["text_encoder"]) == ["rotary_emb.inv_freq"] + assert sorted(COMPUTED["vae"]) == [] + + +def _write_tiny_encoder(root, *, drop=()): + """The tiny encoder under the base file's names (``language_model.`` + mflux's path), plus a vision tower tensor + as the real file carries (``visual.*``, which mflux's key transform drops), in one file. + + ``rotary_emb.inv_freq`` is left out: no base holds it (computed at construction). + """ + from mflux.models.krea2.model.krea2_text_encoder.text_encoder import Krea2TextEncoder + + params = dict(tree_flatten(Krea2TextEncoder(**TINY_ENCODER).parameters())) + flat = { + f"language_model.{n}": mx.full(p.shape, 0.5, dtype=mx.bfloat16) + for n, p in params.items() + if n != "rotary_emb.inv_freq" and f"language_model.{n}" not in drop + } + flat["visual.blocks.0.attn.qkv.weight"] = mx.zeros((4, 4), dtype=mx.bfloat16) + (root / "text_encoder").mkdir(parents=True, exist_ok=True) + mx.save_safetensors(str(root / "text_encoder" / "model.safetensors"), flat) + + +@pytest.mark.mflux +def test_a_missing_encoder_tensor_is_refused(tmp_path): + # Bug caught: mflux's applier updates with strict=False, so a base missing one tensor (a lost + # shard, a renamed key) loads without error and that layer keeps its random float32 init: a wrong image, no error. + from mlx_dfloat.mflux.krea2.init import load_text_encoder + + _write_tiny_encoder(tmp_path, drop={"language_model.layers.1.self_attn.q_proj.weight"}) + with pytest.raises( + DFloatFormatError, match=r"layers\.1\.self_attn\.q_proj\.weight not in the base" + ): + load_text_encoder(tmp_path, dims=TINY_ENCODER) + + +@pytest.mark.mflux +def test_a_complete_encoder_loads_from_the_language_model_tensors_only(tmp_path): + # Bug caught: the coverage check demanding rotary_emb.inv_freq (every real base refused), the vision tower in the + # same file breaking the load, or the file's weights not landing (still the float32 init). + from mlx_dfloat.mflux.krea2.init import load_text_encoder + + _write_tiny_encoder(tmp_path) + encoder = load_text_encoder(tmp_path, dims=TINY_ENCODER) + q = encoder.layers[1].self_attn.q_proj.weight + assert q.dtype == mx.bfloat16 + assert float(q[0, 0]) == 0.5 + + +@pytest.mark.mflux +def test_a_base_with_another_encoder_width_is_refused_before_any_encode(monkeypatch, tmp_path): + # Bug caught: Qwen-Image 2.1's base (a Qwen3-VL encoder 4096 wide) passed as --base: mflux's applier assigns its + # weights without a shape check, so the failure would be a matmul error at the first encode, not a refusal. + from mlx_dfloat.mflux.krea2.init import load_text_encoder + + def wrong_width(root, definition, models, *, what): + models["text_encoder"].embed_tokens.weight = mx.zeros((151936, 4096), dtype=mx.bfloat16) + + monkeypatch.setattr(_hub, "load_into", wrong_width) + with pytest.raises( + DFloatFormatError, + match=r"embed_tokens\.weight has shape \(151936, 4096\), this model needs \(151936, 2560\)", + ): + load_text_encoder(tmp_path) + + +def _replace_all_but(component, missing): + """A ``load_into`` that replaces every parameter of ``component`` but ``missing`` with lazy zeros (no allocation).""" + + def fake(root, definition, models, *, what): + module = models[component] + flat = dict(tree_flatten(module.parameters())) + new = {n: mx.zeros(p.shape, dtype=mx.bfloat16) for n, p in flat.items() if n != missing} + module.update(tree_unflatten(list(new.items())), strict=False) + + return fake + + +@pytest.mark.mflux +def test_a_missing_vae_tensor_is_refused(monkeypatch, tmp_path): + # Bug caught: the VAE loaded through the same strict=False applier with no coverage check: a lost + # decoder conv keeps its random init and the decode returns noise. + from mlx_dfloat.mflux.krea2.init import load_vae + + missing = "decoder.conv_out.conv3d.weight" + monkeypatch.setattr(_hub, "load_into", _replace_all_but("vae", missing)) + with pytest.raises(DFloatFormatError, match=re.escape(missing) + " not in the base"): + load_vae(tmp_path) + + +@pytest.mark.mflux +def test_a_complete_vae_loads_as_qwen_image_1s_vae(monkeypatch, tmp_path): + # Bug caught: a complete VAE refused, or Krea 2 given Qwen-Image 2.1's VAE class (a different module: Krea 2 uses + # Qwen-Image 1's QwenVAE, krea2.py:19 and krea2_weight_mapping.py get_vae_mapping). + from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE + + from mlx_dfloat.mflux.krea2.init import load_vae + + monkeypatch.setattr(_hub, "load_into", _replace_all_but("vae", None)) + vae = load_vae(tmp_path) + assert type(vae) is QwenVAE + assert vae.decoder.conv_out.conv3d.weight.dtype == mx.bfloat16 + + +@pytest.mark.mflux +def test_a_directory_without_the_text_encoder_or_the_vae_is_a_format_error(tmp_path): + # Bug caught: a DF11 checkpoint passed as --base crashing inside mflux's loader with a bare FileNotFoundError. + from mlx_dfloat.mflux.krea2.init import load_text_encoder, load_vae + + with pytest.raises(DFloatFormatError, match="Krea 2 base"): + load_text_encoder(tmp_path) + with pytest.raises(DFloatFormatError, match="Krea 2 base"): + load_vae(tmp_path) + + +@pytest.mark.mflux +def test_missing_tokenizer_files_are_a_format_error(tmp_path): + # Bug caught: an empty base directory surfacing mflux's own tokenizer error instead of the package's. + from mlx_dfloat.mflux.krea2.init import load_tokenizers + + with pytest.raises(DFloatFormatError, match="no usable Krea 2 tokenizer"): + load_tokenizers(tmp_path) + + +@pytest.mark.mflux +def test_a_partial_tokenizer_directory_is_a_format_error(monkeypatch, tmp_path): + # Bug caught: transformers' plain OSError for a tokenizer/ missing one of its files escaping as a raw error + # instead of the package's format error naming the base. + from mflux.models.common.tokenizer.tokenizer_loader import TokenizerLoader + + from mlx_dfloat.mflux.krea2.init import load_tokenizers + + def partial(**_kwargs): + raise OSError("Can't load tokenizer for 'tokenizer'") + + monkeypatch.setattr(TokenizerLoader, "load_all", staticmethod(partial)) + with pytest.raises( + DFloatFormatError, match=r"no usable Krea 2 tokenizer.*Can't load tokenizer" + ): + load_tokenizers(tmp_path) + + +def test_load_base_returns_the_three_components(monkeypatch, tmp_path): + # Bug caught: load_base wiring the encoder into the vae slot (or the reverse), or the tokenizers dropped. + from mlx_dfloat.mflux.krea2 import init as kinit + + monkeypatch.setattr(kinit, "load_vae", lambda root: ("vae", root)) + monkeypatch.setattr(kinit, "load_text_encoder", lambda root: ("encoder", root)) + monkeypatch.setattr(kinit, "load_tokenizers", lambda root: {"qwen3vl": root}) + parts = kinit.load_base(tmp_path) + assert (parts.vae, parts.text_encoder, parts.tokenizers) == ( + ("vae", tmp_path), + ("encoder", tmp_path), + {"qwen3vl": tmp_path}, + ) + + +@pytest.mark.mflux +@pytest.mark.network +def test_the_real_tokenizer_loads_from_the_pattern_files_alone(tmp_path): + # Bug caught: a narrowed tokenizer pattern that leaves out a file mflux's TokenizerLoader (transformers' + # AutoTokenizer) needs, so a fresh download loads no tokenizer; or the template not applied (the 34-token prefix + # mflux strips would be missing). Links only the snapshot's tokenizer/ files a pattern selects. Marked network as + # the Qwen-Image 2.1 twin is, but it reads only a local base snapshot (MLX_DFLOAT_KREA_BASE): no download, no + # Hub request (it passes with HF_HUB_OFFLINE=1). + base = os.environ.get("MLX_DFLOAT_KREA_BASE") + if not base: + pytest.skip("MLX_DFLOAT_KREA_BASE is not set (a Krea 2 base snapshot with tokenizer/)") + from mflux.models.krea2.model.krea2_text_encoder.text_encoder import Krea2TextEncoder + + from mlx_dfloat.mflux.krea2.init import load_tokenizers + + (tmp_path / "tokenizer").mkdir() + for path in sorted(Path(base, "tokenizer").iterdir()): + name = f"tokenizer/{path.name}" + if any(fnmatch(name, p) for p in BASE_PATTERNS): + (tmp_path / name).symlink_to(path.resolve()) + linked = sorted(f"tokenizer/{p.name}" for p in (tmp_path / "tokenizer").iterdir()) + assert linked == [ + "tokenizer/chat_template.jinja", + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + ] + tokenizer = load_tokenizers(tmp_path)["qwen3vl"] + assert tokenizer.tokenize("a").input_ids.shape == (1, 40) + ids = tokenizer.tokenize(PROMPT).input_ids + assert (int(ids.shape[1]), Krea2TextEncoder._template_end(ids)) == (64, 34) diff --git a/tests/test_krea2_memory.py b/tests/test_krea2_memory.py new file mode 100644 index 0000000..4a6ac8c --- /dev/null +++ b/tests/test_krea2_memory.py @@ -0,0 +1,668 @@ +"""Krea 2's memory rules: the encoder's resident bytes, the text-token count, the sizes, the encode warning, and the +constants measured at 1024² against the committed calibration records.""" + +import dataclasses +import hashlib +import json +import logging +from pathlib import Path +from types import SimpleNamespace + +import mlx.core as mx +import pytest + +from mlx_dfloat.bench.capped import tier_limits +from mlx_dfloat.errors import DFloatResourceError +from mlx_dfloat.mflux._phases import FamilySizes, PhaseConstants, activation_allowance, fit_for +from mlx_dfloat.mflux._pipeline import plan_call_for +from mlx_dfloat.mflux.krea2 import memory as kmem + +# The calibration prompt's text tokens (mflux 0.20.0 TokenizerLoader on both base snapshots, 2026-10-08: 64 ids, the +# chat-template prefix ends at 34). +DERISK_TOKENS = 30 + + +def _bf16(*shape): + return mx.zeros(shape, dtype=mx.bfloat16) + + +def _write_encoder(root, *, prefix="language_model."): + """A Krea 2 text encoder file the way the base stores it: the language model under ``language_model.`` (two + numbered layers) and the vision tower beside it under ``visual.`` (the real file's prefixes, 2026-10-08 header). + + bf16 bytes: embed_tokens 10 x 4 = 80, one up_proj 4 x 4 = 32 per layer, the final norm (4,) = 8, the vision tensor + 8 x 4 = 64. + """ + enc = root / "text_encoder" + enc.mkdir(parents=True) + mx.save_safetensors( + str(enc / "model.safetensors"), + { + f"{prefix}embed_tokens.weight": _bf16(10, 4), + f"{prefix}layers.0.mlp.up_proj.weight": _bf16(4, 4), + f"{prefix}layers.1.mlp.up_proj.weight": _bf16(4, 4), + f"{prefix}norm.weight": _bf16(4), + "visual.blocks.0.attn.qkv.weight": _bf16(8, 4), + }, + ) + + +@pytest.mark.parametrize("prefix", ["language_model.", "model.language_model."]) +def test_encoder_bytes_count_the_evaluated_language_model_only(tmp_path, prefix): + # Bug caught: the vision tower counted (+64 here, +0.77 GiB on the real file), or the last layer counted although + # the deepest tap (hidden state 35 = the output of layer 34) means layer 35 never runs (+32 here, +0.19 GiB real). + # Both of mflux's prefixes (krea2_weight_definition.py:14) are the language model. By hand: 80 + 32 + 8 = 120. + _write_encoder(tmp_path, prefix=prefix) + assert kmem.encoder_bytes_used(tmp_path) == 120 + # On the published header the same rule gives 8_044_936_192 - layer 35's 201_861_632 = 7_843_074_560 B, the + # encoder bytes the calibration records carry (test_the_record_sizes_charge_no_resident_nonblock_in_the_shipped_mode). + + +class _Tokenizer: + """mflux's tokenizer protocol as the prompt encoder uses it: ``tokenize(prompt).input_ids`` of shape (1, L).""" + + def __init__(self, ids=None, lengths=None): + self.ids = ids + self.lengths = dict(lengths or {}) + + def tokenize(self, prompt): + ids = self.ids if self.ids is not None else [list(range(1, self.lengths[prompt] + 1))] + return SimpleNamespace( + input_ids=mx.array(ids, dtype=mx.int32), attention_mask=mx.ones((1, len(ids[0]))) + ) + + +@pytest.mark.mflux +def test_prompt_tokens_are_the_ids_after_the_template_prefix(): + # Bug caught: the template prefix counted (the plan sized by about 35 extra tokens) or the attention mask counted + # instead. By hand from text_encoder.py:117-127: the second <|im_start|> (151644) is at 3, followed by "user" + # (872) and a newline (198), so the prefix ends at 6 and 3 ids remain; without the template nothing is stripped. + templated = _Tokenizer(ids=[[151644, 9, 10, 151644, 872, 198, 5, 6, 7]]) + assert kmem.prompt_tokens(templated, "p") == 3 + assert kmem.prompt_tokens(_Tokenizer(ids=[[1, 2, 3]]), "p") == 3 + + +@pytest.mark.mflux +def test_a_cfg_step_is_sized_by_its_longer_prompt(): + # Bug caught: a CFG step sized by its first prompt (the negative " " first would give 6) or by the sum; the two + # calls run one after the other, so the step's activations are the longer call's. + tokenizer = _Tokenizer(lengths={" ": 6, "long": 40}) + assert kmem.text_tokens_for(tokenizer, [" ", "long"]) == 40 + assert kmem.text_tokens_for(tokenizer, ["long", " "]) == 40 + + +def test_the_encode_bound_is_the_phase_without_overhead_plus_the_slack(): + # Bug caught: the overhead kept in the bound, or the slack missing. 10e9 - 0.4e9 + 268_435_456, by hand. + c = PhaseConstants( + overhead_bytes=400_000_000, + vae_transient_bytes=0, + denoise_activation_at_reference=0, + reference_tokens=1, + ) + assert kmem.encode_peak_bound(10_000_000_000, c) == 9_868_435_456 + + +def test_the_encode_warning_is_silent_at_the_bound_and_names_the_prompt_length_over_it(): + # Bug caught: `>=` for `>` at the bound (a normal encode warning), or a warning that blames mflux's loading alone + # when the prompt is longer than the one the encode term was measured on (30 tokens). + bound = 8_000_000_000 + assert kmem.encode_peak_warning(bound, bound, text_tokens=1024) is None + message = kmem.encode_peak_warning(bound + 1, bound, text_tokens=1024) + assert message is not None + assert "for a 1024-token prompt" in message + assert f"measured on one {DERISK_TOKENS}-token prompt" in message + assert "the vision tower stored in the text encoder's file" in message + + +def test_the_sizes_of_the_published_raw_file_are_the_arithmetic_tables(tmp_path): + # Bug caught: a non-block group left out of the resident set (the four text-fusion blocks are 0.64 GiB of the + # 1.234), a block counted as non-block, the extras counted in the compressed set, or the encoder counted by file + # size. Literals from the published header's matrix shapes: non-block 4 x 171_704_320 + 78_643_200 + 452_984_832 + + # 106_954_752 = 1_325_400_064; extras 4_560_024 (header sum); compressed 17_545_490_584 - 4_560_024. + from mlx_dfloat._layouts import KREA_2_RAW_COMFYUI + from mlx_dfloat.format import open_checkpoint + + fixture = Path(__file__).parent / "fixtures" / "krea-2-raw-df11-header.bin" + df11 = tmp_path / "df11" + df11.mkdir() + with (df11 / "model.safetensors").open( + "wb" + ) as fh: # sparse: the header, then zeros to the file size + fh.write(fixture.read_bytes()) + fh.truncate(17_545_490_584) + zero = hashlib.sha256(bytes(4096)).hexdigest() # the sparse file reads as zeros at every probe + layout = dataclasses.replace( + KREA_2_RAW_COMFYUI, + probes=tuple(dataclasses.replace(p, sha256=zero) for p in KREA_2_RAW_COMFYUI.probes), + ) + ckpt = open_checkpoint(df11, layouts=(layout,)) + base = tmp_path / "base" + _write_encoder(base) + (base / "vae").mkdir() + (base / "vae" / "diffusion_pytorch_model.safetensors").write_bytes(b"x" * 1000) + sizes = kmem.sizes_for(ckpt, base) + assert sizes.nonblock == 1_325_400_064 + assert sizes.extras == 4_560_024 + assert sizes.compressed == 17_540_930_560 + assert sizes.encoders == 120 + assert sizes.vae == 1000 + # Per-call decode (the shipped mode): no non-block weights stay decoded with the set. Bug caught: + # the per-call path still charging 1.234 GiB, so the fit check refuses calls the measured runs passed. + per_call = kmem.sizes_for(ckpt, base, nonblock_per_call=True) + assert per_call.nonblock == 0 + assert (per_call.compressed, per_call.extras, per_call.encoders) == ( + 17_540_930_560, + 4_560_024, + 120, + ) + + +# --- the measured constants, against the committed calibration records ----------------------------------------- +# Measured 2026-10-08 (M1 Max 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0), the 30-token lighthouse prompt, 1024², +# seed 42, at the planner's limit for the block-cache probe's allowance (4_500_000_000 B), the set dropped before the +# VAE decode. The shipped mechanism (per-call non-block decode, MLX's cache emptied and drained at each call's start +# and before block 0, git c555f50): Krea 2 Raw, one step at guidance 3.5 with mflux's " " negative (two batch-1 calls, +# derisk-raw-2), and Krea 2 Turbo, two steps at guidance 1.0 (derisk-turbo-2). Also recorded: Raw with the non-block +# groups resident (derisk-raw-1, git d1aa9df; a VAE sample), and four runs of superseded per-call code (history-*), +# never a constant's source. The tests read them from the records. + +GIB = 1024**3 +REPO = Path(__file__).resolve().parents[1] +RECORDS = { + "krea-2-raw": json.loads((REPO / "bench/results/calibration/krea-2-raw-1024.json").read_text()), + "krea-2": json.loads((REPO / "bench/results/calibration/krea-2-1024.json").read_text()), +} +RUNS = {m: {run["label"]: run for run in rec["runs"]} for m, rec in RECORDS.items()} +ALL_RUNS = {label: run for runs in RUNS.values() for label, run in runs.items()} +# The samples the constants rest on: the shipped mechanism's runs (every term), plus the resident run for the VAE term +# (its VAE phase ran with the set dropped, like every mechanism's). +SHIPPED = {label: run for label, run in ALL_RUNS.items() if run["mechanism"].startswith("shipped")} +# The one-process calibration runs: they alone record every phase's peak, the post-build counters and the encode hold. +DERISK = {label: run for label, run in ALL_RUNS.items() if "after_build" in run} +SHIPPED_DERISK = {label: run for label, run in SHIPPED.items() if label in DERISK} +VAE_RUNS = {**SHIPPED, "derisk-raw-1": ALL_RUNS["derisk-raw-1"]} +# Which run each term came from (one set of constants for both models, each phase named exactly once). +SOURCE = { + phase: label for rec in RECORDS.values() for phase, label in rec["constants_from"].items() +} +# The planner's sizes: the shipped mode (per-call), so no decoded non-block bytes stay with the set. +SIZES = { + m: FamilySizes(**RUNS[m][f"derisk-{'raw' if m == 'krea-2-raw' else 'turbo'}-2"]["sizes"]) + for m in RECORDS +} +LARGEST = { + "blocks": 868_220_928 +} # one block kind: 434_110_464 elements (the header's eight block matrix shapes) +TEXT_TOKENS = 30 # the calibration prompt's (the negative " " is 6) +ALLOWANCE = ( + 3_631_779_072 # the block-cache probe's, at 4096 + 30 tokens (limit 4_500_000_000 - one block) +) +BUDGET_32 = 24_653_119_488 # this Mac's fit budget: the recommended working set 26_800_603_136 - the 2 GiB reserve +RECOMMENDED_32 = 26_800_603_136 +RAM_32 = 34_359_738_368 + + +def _plan( + model, budget=BUDGET_32, *, fit_check=True, cache_limit_override=None, policy="per-block" +): + return plan_call_for( + constants=kmem.CONSTANTS, + sizes=SIZES[model], + largest=LARGEST, + policy=policy, + cache_limit_override=cache_limit_override, + fit_check=fit_check, + budget=budget, + height=1024, + width=1024, + text_tokens=TEXT_TOKENS, + log=logging.getLogger("test"), + ) + + +def _fit(model, sizes=None, *, budget=BUDGET_32, vae_with_set=False): + c = kmem.CONSTANTS + return fit_for( + c, + sizes=SIZES[model] if sizes is None else sizes, + largest=LARGEST, + policy="per-block", + cache_limit=868_220_928 + ALLOWANCE, + allowance=activation_allowance(c, height=1024, width=1024, text_tokens=TEXT_TOKENS), + budget=budget, + height=1024, + width=1024, + text_tokens=TEXT_TOKENS, + vae_with_set=vae_with_set, + ) + + +def test_the_calibration_records_name_the_pinned_repos_and_hold_the_runs_the_constants_name(): + # Bug caught: a constant re-derived from a run that is not recorded, a record of the wrong model or snapshot, a + # phase sourced twice (or not at all), or the calls per step of the CFG run mislabelled. + assert ( + RECORDS["krea-2-raw"]["df11"] + == "mingyi456/Krea-2-Raw-DF11-ComfyUI@8320616b25ac9340a830a7fb21f1b0237e160e66" + ) + assert ( + RECORDS["krea-2-raw"]["base"] == "krea/Krea-2-Raw@6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603" + ) + assert ( + RECORDS["krea-2"]["df11"] + == "mingyi456/Krea-2-Turbo-DF11-ComfyUI@978da5fb7647bd222d33125993abd8fdc2840cfc" + ) + assert RECORDS["krea-2"]["base"] == "krea/Krea-2-Turbo@98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7" + assert (RECORDS["krea-2-raw"]["calls_per_step"], RECORDS["krea-2"]["calls_per_step"]) == (2, 1) + phases = [p for rec in RECORDS.values() for p in rec["constants_from"]] + assert sorted(phases) == ["denoise", "encode", "vae"] + for m, rec in RECORDS.items(): + assert set(rec["constants_from"].values()) <= set(RUNS[m]), m + assert rec["text_tokens"] == TEXT_TOKENS + assert RECORDS["krea-2"]["overhead_from"] == "derisk-turbo-2" + # Every constant comes from a run of the shipped code, or (the VAE term) from the resident run. + assert set(SOURCE.values()) | {RECORDS["krea-2"]["overhead_from"]} <= set(VAE_RUNS) + assert { + label: (run["git"][:7], run["mechanism"].split(":")[0]) for label, run in ALL_RUNS.items() + } == { + "derisk-raw-2": ("c555f50", "shipped"), + "confirm-raw-24711fb": ("24711fb", "shipped"), + "derisk-raw-1": ("d1aa9df", "resident"), + "history-v1-raw": ("59c961b", "superseded (history)"), + "history-mm1-raw": ("bedba36", "superseded (history)"), + "measured-25-steps": ("f0f348f", "shipped"), + "card-52-steps-cfg": ("f0f348f", "shipped"), + "identity-4-steps-cfg": ("e3a508e", "shipped"), + "derisk-turbo-2": ("c555f50", "shipped"), + "confirm-turbo-24711fb": ("24711fb", "shipped"), + "measured-8-steps": ("f0f348f", "shipped"), + "history-v1-turbo": ("81e79f3", "superseded (history)"), + "history-mm1-turbo": ("bedba36", "superseded (history)"), + } + assert (SHIPPED["derisk-raw-2"]["steps"], SHIPPED["derisk-turbo-2"]["steps"]) == (1, 2) + + +def test_the_record_sizes_charge_no_resident_nonblock_in_the_shipped_mode(): + # Bug caught: the per-call runs recorded with the 1.234 GiB of resident non-block weights (the estimate would + # charge them twice: once as nonblock, once inside the measured denoise term), or the encoder counted by file. + assert {label: run["sizes"]["nonblock"] for label, run in ALL_RUNS.items()} == { + "derisk-raw-1": 1_325_400_064, + "derisk-raw-2": 0, + "derisk-turbo-2": 0, + "history-v1-raw": 0, + "history-mm1-raw": 0, + "history-v1-turbo": 0, + "history-mm1-turbo": 0, + "measured-8-steps": 0, + "measured-25-steps": 0, + "confirm-raw-24711fb": 0, + "confirm-turbo-24711fb": 0, + "card-52-steps-cfg": 0, + "identity-4-steps-cfg": 0, + } + assert {s.encoders for s in SIZES.values()} == {8_044_936_192 - 201_861_632} # minus layer 35 + + +def test_the_overhead_is_the_larger_post_build_gap_of_the_shipped_runs(): + # Bug caught: the overhead taken from one run (Raw's 522_673_362 would put the Turbo run's phases 1_687_552 B low), + # from a superseded run (the first per-call Turbo run's 529_374_466), or computed otherwise than footprint - MLX + # active - MLX cache after the build. + gaps = { + label: run["after_build"]["footprint"] + - run["after_build"]["mlx_active"] + - run["after_build"]["mlx_cache"] + for label, run in SHIPPED_DERISK.items() + } + assert gaps == { + "derisk-raw-2": 522_673_362, + "confirm-raw-24711fb": 523_590_866, + "derisk-turbo-2": 524_360_914, + "confirm-turbo-24711fb": 523_787_474, + } + assert kmem.CONSTANTS.overhead_bytes == 524_360_914 == max(gaps.values()) + + +def test_the_estimate_at_1024_equals_the_calibration_runs_phases(): + # Bug caught: a term dropped from a phase or counted twice, REFERENCE_TOKENS left at another prompt's count, the + # resident non-block bytes charged in the per-call mode, or a term derived at another cache limit than the + # planner's. In-sample, so exact for encode and VAE: + # - encode: the evaluated language model 7_843_074_560 + the footprint encode term 230_829_150 + overhead + # 524_360_914 = 8_598_264_624, confirm-raw-24711fb's encode peak (the highest of the shipped runs); + # - VAE (set dropped): extras 4_560_024 + VAE file 507_591_892 + transient 10_832_964_778 + overhead 524_360_914 = + # 11_869_477_608, derisk-raw-1's VAE peak; + # - denoise (Raw): compressed 17_540_930_560 + extras 4_560_024 + nonblock 0 + one block 868_220_928 + cache limit + # 4_500_000_000 + activation 805_306_368 (the floor: the measured term is 0) + overhead 524_360_914 = + # 24_243_378_794, 1_541_677_514 B over derisk-raw-2's 22_701_701_280 (the drained cache never fills to the + # cache-limit term at the peak). + fit = _fit("krea-2-raw") + assert ( + fit.phases["encode"] + == 8_598_264_624 + == ALL_RUNS[SOURCE["encode"]]["footprint_peaks"]["encode"] + == max(run["footprint_peaks"]["encode"] for run in SHIPPED_DERISK.values()) + ) + assert fit.phases["vae"] == 11_869_477_608 == ALL_RUNS[SOURCE["vae"]]["footprint_peaks"]["vae"] + assert fit.phases["denoise"] == 24_243_378_794 + assert ( + fit.phases["denoise"] - SHIPPED["derisk-raw-2"]["footprint_peaks"]["denoise"] + == 1_541_677_514 + ) + assert (SOURCE["encode"], SOURCE["denoise"], SOURCE["vae"]) == ( + "confirm-raw-24711fb", + "measured-8-steps", + "derisk-raw-1", + ) + + +def test_the_denoise_term_is_the_highest_shipped_sample_floored_at_0(): + # Bug caught: the term taken from a superseded run (the second per-call Raw run's denoise sat at the + # memory limit, 24_135_923_728 B: a term of 698_850_888 B that the shipped code no longer needs), from one shipped run + # instead of the highest (Raw's residual -736_371_146 vs Turbo's -726_734_444), or left negative. Per shipped run, + # the denoise peak minus the phase's other terms at that run's own sizes: the highest is Turbo's, below 0, so the + # term is 0 and the float32-copy floor alone carries the activation. + residual = { + label: run["footprint_peaks"]["denoise"] + - ( + run["sizes"]["compressed"] + + run["sizes"]["extras"] + + run["sizes"]["nonblock"] + + 868_220_928 + + run["cache_limit"] + + kmem.CONSTANTS.overhead_bytes + ) + for label, run in SHIPPED_DERISK.items() + } + assert residual == { + "derisk-raw-2": -736_371_146, + "confirm-raw-24711fb": -733_766_258, + "derisk-turbo-2": -726_734_444, + "confirm-turbo-24711fb": -733_730_340, + } + assert kmem.CONSTANTS.denoise_activation_at_reference == max(0, *residual.values()) == 0 + for label, run in SHIPPED.items(): + assert ( + run["footprint_peaks"]["denoise"] + <= _fit("krea-2-raw", FamilySizes(**run["sizes"])).phases["denoise"] + ), label + + +def test_the_vae_term_covers_every_1024_sample_and_equals_the_highest(): + # Bug caught: the VAE transient derived from one run instead of the highest sample (the shipped runs alone would + # put the resident run's decode 1_646_167_480 B over its estimate). Every run decoded with the set dropped, so each + # sample, the superseded ones too, is checked against the dropped-set VAE phase with the run's own sizes. + peaks = {label: run["footprint_peaks"]["vae"] for label, run in DERISK.items()} + for label, run in DERISK.items(): + assert peaks[label] <= _fit("krea-2-raw", FamilySizes(**run["sizes"])).phases["vae"], label + assert _fit("krea-2-raw").phases["vae"] == max(peaks.values()) == peaks["derisk-raw-1"] + assert max(run["footprint_peaks"]["vae"] for run in SHIPPED_DERISK.values()) == 10_525_591_448 + # The later runs record one watched peak, reached in the denoise phase: no VAE-phase sample beyond these. + assert all( + "vae" not in run["footprint_peaks"] + for label, run in ALL_RUNS.items() + if label not in DERISK + ) + + +def test_the_planners_1024_cache_limit_is_one_block_plus_the_allowance(): + # Bug caught: the shared 1.5e9 allowance, two blocks budgeted for a one-kind family, or the allowance scaled from + # another reference token count. 868_220_928 + 3_631_779_072. + assert _plan("krea-2-raw").cache_limit == 4_500_000_000 + assert (kmem.CONSTANTS.allowance_at_reference, kmem.CONSTANTS.allowance_reference_tokens) == ( + 3_631_779_072, + 4096 + 30, + ) + + +def test_a_cache_limit_override_under_the_one_kind_minimum_warns_for_both_policies(caplog): + # Bug caught: the depth2 minimum missing its look-ahead group, or the per-block minimum without the + # allowance. Per-block 868_220_928 + 3_631_779_072; depth2 1_736_441_856 + 3_631_779_072. + per_block, depth2 = 4_500_000_000, 5_368_220_928 + with caplog.at_level("WARNING", logger="test"): + _plan("krea-2-raw", cache_limit_override=per_block) + assert "below the derived minimum" not in caplog.text + with caplog.at_level("WARNING", logger="test"): + _plan("krea-2-raw", cache_limit_override=per_block - 1) + assert f"below the derived minimum {per_block}" in caplog.text + caplog.clear() + with caplog.at_level("WARNING", logger="test"): + assert _plan("krea-2-raw", policy="depth2", fit_check=False).cache_limit == depth2 + assert "below the derived minimum" not in caplog.text + with caplog.at_level("WARNING", logger="test"): + _plan("krea-2-raw", policy="depth2", cache_limit_override=per_block, fit_check=False) + assert f"below the derived minimum {depth2}" in caplog.text + + +@pytest.mark.parametrize( + ("model", "denoise"), [("krea-2-raw", 24_243_378_794), ("krea-2", 24_242_523_532)] +) +def test_a_32_gb_mac_drops_the_set_before_the_vae_at_1024_and_fits_in_the_denoise_phase( + model, denoise +): + # Bug caught: the set-drop decision inverted (a decode on the resident set runs into the watchdog), or the 32 GB + # verdict lost. With the set the VAE phase would be compressed + extras + VAE file + transient + overhead: Raw + # 17_540_930_560 + 4_560_024 + 507_591_892 + 10_832_964_778 + 524_360_914 = 29_410_408_168 B, Turbo 855_262 B less + # (its compressed set is 17_540_075_298), both over the 24_653_119_488 budget. Without it the call peaks in the + # denoise phase (the activation at the 805_306_368 floor): Raw 24_243_378_794 (409_740_694 B under the budget), + # Turbo 24_242_523_532 (410_595_956 B under). + plan = _plan(model) + assert plan.drop_set_before_vae is True + assert (plan.estimate.peak_phase, plan.estimate.peak_bytes) == ("denoise", denoise) + assert ( + _fit(model, vae_with_set=True).phases["vae"] + == {"krea-2-raw": 29_410_408_168, "krea-2": 29_409_552_906}[model] + ) + assert BUDGET_32 - denoise == {"krea-2-raw": 409_740_694, "krea-2": 410_595_956}[model] + + +@pytest.mark.parametrize( + ("size", "expected"), [(512, 805_306_368), (1024, 805_306_368), (2048, 805_306_368)] +) +def test_the_denoise_term_is_floored_by_the_float32_mlp_weight_copies(size, expected): + # Bug caught: a call predicted below the float32 copies every step holds whatever the image size (the float32 + # residual stream times a BF16 weight makes MLX copy the weight to float32: mlp.gate and mlp.up, 2 x 24576 x 4096 + # x 4 = 805_306_368 B). The measured term is 0 (the shipped runs' highest residual is negative), so the floor is + # the whole activation at every size, 2048² included; whether it scales above 1024² is unmeasured (sizes above + # 1024² are refused without --no-fit-check). + from mlx_dfloat.mflux._phases import denoise_activation_bytes + + assert kmem.CONSTANTS.denoise_activation_floor_bytes == 805_306_368 == 2 * 24576 * 4096 * 4 + got = denoise_activation_bytes(kmem.CONSTANTS, height=size, width=size, text_tokens=30) + assert got == expected + + +def _tier_budget(tier): + return tier_limits( + tier, host_ram_bytes=RAM_32, host_recommended_bytes=RECOMMENDED_32 + ).fit_budget_bytes + + +def test_the_tier_budgets_are_a_real_macs(): + # Bug caught: a tier judged against its watchdog ceiling or the host's budget instead of the fit budget. + assert (_tier_budget(16), _tier_budget(24)) == (9_305_762_474, 15_032_385_536) + + +# --- CAPPED predictions: each model's default call at 1024² (Raw 25 steps at 1.0, Turbo 8 at 1.0: one call, the +# 30-token prompt), against a 24 GB Mac's fit budget (15_032_385_536 B) and a 16 GB Mac's (9_305_762_474 B). The +# compressed set alone (17_540_930_560 B, 16.34 GiB) is over both, so the denoise phase decides: refused there at both +# tiers. Forced (fit_check off) the set is dropped before the decode, and the dropped-set VAE phase keeps the extras +# (they stay on the transformer): 4_560_024 + 507_591_892 + 10_832_964_778 + 524_360_914 = 11_869_477_608 B. + + +@pytest.mark.parametrize("model", ["krea-2-raw", "krea-2"]) +@pytest.mark.parametrize("tier", [16, 24]) +def test_every_capped_default_call_is_refused_in_the_denoise_phase(model, tier): + # Bug caught: the refusal naming the VAE phase (the set drop not applied), a call that cannot hold the compressed + # set planned to fit, or the dropped-set VAE phase without the extras. + budget = _tier_budget(tier) + with pytest.raises(DFloatResourceError, match="in the denoise phase"): + _plan(model, budget) + forced = _plan(model, budget, fit_check=False) + assert forced.drop_set_before_vae is True + assert forced.estimate.peak_phase == "denoise" + assert ( + forced.estimate.peak_bytes + == {"krea-2-raw": 24_243_378_794, "krea-2": 24_242_523_532}[model] + ) + assert forced.estimate.phases["vae"] == 11_869_477_608 + + +# What the prompt encode may hold before the model warns (model._check_encode_peak): MLX's encode peak minus what was +# active when the phase began, against the encode phase without its overhead (evaluated language model + encode +# term) + the slack: 7_843_074_560 + 230_829_150 + 268_435_456 = 8_342_339_166 B. +ENCODE_BOUND = 8_342_339_166 +VISION_TOWER = 830_695_424 # the visual.* tensors in the encoder's file (its published header) + + +def test_the_encode_bound_holds_each_measured_encode_with_its_spare(): + # Bug caught: a bound under what a normal encode holds (every call would warn). + bound = kmem.encode_peak_bound(_fit("krea-2-raw").phases["encode"], kmem.CONSTANTS) + assert bound == ENCODE_BOUND + # Every recorded encode, the superseded runs' too: the encode path did not change between them. + spare = {label: bound - run["mlx_encode_held"] for label, run in DERISK.items()} + assert spare == { + "derisk-raw-2": 396_039_902, + "confirm-raw-24711fb": 389_145_310, + "derisk-raw-1": 399_350_774, + "history-v1-raw": 402_987_390, + "history-mm1-raw": 396_047_070, + "derisk-turbo-2": 434_814_750, + "confirm-turbo-24711fb": 405_439_478, + "history-v1-turbo": 399_323_126, + "history-mm1-turbo": 474_185_294, + } + + +@pytest.mark.parametrize("text_tokens", [1, TEXT_TOKENS, 1024]) +def test_an_encode_that_also_loads_the_vision_tower_is_over_the_bound_at_any_prompt(text_tokens): + # Bug caught: the encode bound carrying the denoise allowance (it grows with the prompt), or a slack wide enough + # that an mflux change loading the vision tower goes unnoticed. Over by 7_953_193_856 + 830_695_424 - 8_342_339_166 + # = 441_550_114 B at every prompt length, for the largest measured hold. + c = kmem.CONSTANTS + fit = fit_for( + c, + sizes=SIZES["krea-2-raw"], + largest=LARGEST, + policy="per-block", + cache_limit=4_500_000_000, + allowance=activation_allowance(c, height=1024, width=1024, text_tokens=text_tokens), + budget=BUDGET_32, + height=1024, + width=1024, + text_tokens=text_tokens, + ) + bound = kmem.encode_peak_bound(fit.phases["encode"], c) + held = max(run["mlx_encode_held"] for run in DERISK.values()) + VISION_TOWER + assert held - bound == 441_550_114 + assert kmem.encode_peak_warning(held, bound, text_tokens=text_tokens) is not None + + +# --- the committed 1024² records of the shipped code, beyond the calibration runs ------------------------------------- +# Each is a run's own report (mlx-dfloat generate --report, or verify_image's df11 side), never hand-copied; the +# calibration records file them under the label given here. Their watched peak is reached in the denoise phase (every +# other phase's boundary samples stay at or under 18.1e9 B; the VAE decode runs with the set dropped). +OUT_OF_SAMPLE = [ + ("bench/results/tiers/krea-2-1024.json", "krea-2", "measured-8-steps"), + ("bench/results/tiers/krea-2-raw-1024.json", "krea-2-raw", "measured-25-steps"), + ( + "bench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.json", + "krea-2-raw", + "card-52-steps-cfg", + ), + ( + "bench/results/identity/krea-2-raw-1024/df11-result.json", + "krea-2-raw", + "identity-4-steps-cfg", + ), +] + + +def _record_run(path): + """A committed record's report, its height and width, and its watched (footprint) peak.""" + record = json.loads((REPO / path).read_text()) + if ( + "report" in record + ): # verify_image's df11 side: the model's report under "report", the size in the key + size = record["key"]["size"] + return record["report"], size, size, record["footprint_peak_bytes"] + return record, record["height"], record["width"], record["watched_peak_bytes"] + + +@pytest.mark.parametrize(("path", "model", "label"), OUT_OF_SAMPLE) +def test_the_later_denoise_samples_are_the_committed_runs_watched_peaks(path, model, label): + # Bug caught: a calibration sample that is not the committed run's peak (a typo, or a re-run whose record was not + # carried over), or a sample filed under the other model or the wrong mechanism. + _report, _h, _w, watched = _record_run(path) + run = RUNS[model][label] + assert run["footprint_peaks"] == {"denoise": watched} + assert run["mechanism"].startswith("shipped") + + +def test_the_denoise_term_stays_0_over_every_shipped_sample(): + # Bug caught: the term re-derived from the calibration runs only, so a longer run whose peak rose above its + # estimate would not move it. Every shipped sample's residual (peak minus the phase's other terms) is still negative; + # the highest is the Turbo MEASURED row's, 22_736_402_568 - 23_437_217_164 = -700_814_596. + def residual(run): + s = run["sizes"] + return run["footprint_peaks"]["denoise"] - ( + s["compressed"] + + s["extras"] + + s["nonblock"] + + 868_220_928 + + run["cache_limit"] + + kmem.CONSTANTS.overhead_bytes + ) + + residuals = {label: residual(run) for label, run in SHIPPED.items()} + assert len(residuals) == 8 + assert max(residuals.values()) == residuals["measured-8-steps"] == -700_814_596 + assert kmem.CONSTANTS.denoise_activation_at_reference == 0 + + +# The committed records re-estimated with the current constants from their own reports (sizes, cache limit in force, +# text tokens, policy, VAE strategy). The check is one-sided (a run under the estimate is not the bug); the margins are +# pinned too, because the estimate sits about 1.53 GB over every shipped peak (the drained cache never fills to the +# cache-limit term at block 0), so a 0.25 GiB tolerance alone would let a 1.7 GB rise pass unnoticed. +RECORD_MARGINS = [ + pytest.param( + "bench/results/tiers/krea-2-1024.json", 1_506_120_964, id="turbo-measured-8-steps" + ), + pytest.param( + "bench/results/tiers/krea-2-raw-1024.json", 1_529_635_490, id="raw-measured-25-steps" + ), + pytest.param( + "bench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.json", + 1_528_865_130, + id="raw-card-52-steps-cfg", + ), + pytest.param( + "bench/results/identity/krea-2-raw-1024/df11-result.json", + 1_531_863_714, + id="raw-identity-df11-4-steps-cfg", + ), +] + + +@pytest.mark.parametrize(("path", "margin"), RECORD_MARGINS) +def test_a_committed_1024_record_peaks_within_the_current_estimate(path, margin): + # Bug caught: a 1024² peak that grows with the step count or with CFG while the estimate stays at the one-step + # value, the estimate peaking in another phase, or a record re-estimated from other sizes than its own. The + # estimate is the report's own prediction (Turbo 24_242_523_532, Raw 24_243_378_794) minus the watched peak. + report, height, width, watched = _record_run(path) + c = kmem.CONSTANTS + tokens = report["text_tokens"] + fit = fit_for( + c, + sizes=FamilySizes(**report["sizes"]), + largest=LARGEST, + policy=report["eval_policy"], + cache_limit=report["cache_limit_in_force"], + allowance=activation_allowance(c, height=height, width=width, text_tokens=tokens), + budget=report["fit"]["budget_bytes"], + height=height, + width=width, + text_tokens=tokens, + vae_with_set=not report["drop_set_before_vae"], + ) + assert fit.peak_phase == "denoise" + assert fit.peak_bytes == report["fit"]["peak_bytes"] + assert watched <= fit.peak_bytes + GIB // 4 + assert fit.peak_bytes - watched == margin diff --git a/tests/test_krea2_model.py b/tests/test_krea2_model.py new file mode 100644 index 0000000..7d26d93 --- /dev/null +++ b/tests/test_krea2_model.py @@ -0,0 +1,1296 @@ +"""Tests for `DFloatKrea2`: construction, refusals, the lifecycle and the set, the prelude, CFG and the report. + +Every test is `@pytest.mark.mflux`; the module imports mflux only inside helpers, so it collects without it. +""" + +import gc +import inspect +import re +from types import SimpleNamespace + +import mlx.core as mx +import numpy as np +import pytest +from tests._krea2_tiny import ( + TINY_SEED, + StubTextEncoder, + StubTokenizer, + fake_model, + write_tiny_checkpoint, +) + +from mlx_dfloat.errors import DFloatIntegrationError, DFloatUnsupportedError +from mlx_dfloat.mflux.krea2 import init as kinit + +pytestmark = pytest.mark.mflux + +# The non-block matrices as mflux names them (krea2_weight_mapping.py: the six timestep / text-MLP renames, the text +# fusion by concrete index): 2 + 1 + 2 + 4 x 8 = 37 matrices in the seven groups. +_FUSION = ( + "attn.wq", + "attn.wk", + "attn.wv", + "attn.gate", + "attn.wo", + "mlp.gate", + "mlp.up", + "mlp.down", +) +NONBLOCK = { + "tmlp": ("tmlp.linear_in.weight", "tmlp.linear_out.weight"), + "tproj": ("tproj.linear.weight",), + "txtmlp": ("txtmlp.linear_in.weight", "txtmlp.linear_out.weight"), + **{ + f"txtfusion.{kind}.{i}": tuple(f"txtfusion.{kind}.{i}.{s}.weight" for s in _FUSION) + for kind in ("layerwise_blocks", "refiner_blocks") + for i in (0, 1) + }, +} + + +def _nonblock_sources(tmp_path): + """The non-block matrices ``fake_model`` compressed: the same seed gives the same arrays.""" + matrices, _c, _layout = write_tiny_checkpoint( + tmp_path / "again", np.random.default_rng(TINY_SEED) + ) + return {p: matrices[group][p] for group, params in NONBLOCK.items() for p in params} + + +def _param(model, name): + module = model.transformer + for part in name.split(".")[:-1]: + module = module[int(part)] if part.isdigit() else getattr(module, part) + return module.weight + + +def test_every_attribute_upstreams_generate_image_reads_exists_on_the_model(tmp_path, monkeypatch): + # Bug caught: an mflux point release reading a new self. in Krea2.generate_image our assembly never sets + # (mflux's own initializer is not run). + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + + names = set(re.findall(r"self\.(\w+)", inspect.getsource(Krea2.generate_image))) + assert names >= { + "model_config", + "callbacks", + "transformer", + "bits", + "lora_paths", + "lora_scales", + } + model = fake_model(tmp_path, monkeypatch) + assert [n for n in sorted(names) if not hasattr(model, n)] == [] + assert isinstance(model, Krea2) + assert (model.bits, model.prompt_cache, model.lora_paths, model.lora_scales) == ( + None, + {}, + [], + [], + ) + assert model.tiling_config is None # krea2_initializer.py:47-54: the VAE is never tiled + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"quantize": 8}, "quantize"), + ({"lora_paths": ["a"]}, "lora_paths"), + ({"lora_scales": [1.0]}, "lora_scales"), + ({"bake_lora": False}, "bake_lora"), + ({"eval_policy": "none"}, "none"), + ({"model": "krea-dev"}, r"'krea-dev'.*\['krea-2', 'krea-2-raw'\]"), + ], +) +def test_construction_refuses_quantize_lora_and_unknown_names(monkeypatch, kwargs, match): + # Bug caught: a quantize / LoRA argument or the "none" policy silently honoured (mflux's applier would quantize the + # decoded weights), FLUX.1 Krea-dev (another family) accepted by this class, or any of them resolved (a 17.5 GB + # download) before the refusal. + from mlx_dfloat.mflux.krea2.model import DFloatKrea2 + + monkeypatch.setattr( + kinit, "resolve", lambda *a, **k: pytest.fail("resolved before the refusal") + ) + with pytest.raises(DFloatUnsupportedError, match=match): + DFloatKrea2(**{"model": "krea-2", **kwargs}) + + +@pytest.mark.parametrize( + ("model", "df11", "base"), + [ + ( + "krea-2", + ("mingyi456/Krea-2-Turbo-DF11-ComfyUI", "978da5fb7647bd222d33125993abd8fdc2840cfc"), + ("krea/Krea-2-Turbo", "98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7"), + ), + ( + "krea-2-raw", + ("mingyi456/Krea-2-Raw-DF11-ComfyUI", "8320616b25ac9340a830a7fb21f1b0237e160e66"), + ("krea/Krea-2-Raw", "6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603"), + ), + ], +) +def test_the_default_checkpoint_and_base_resolve_at_their_pins_and_user_paths_unpinned( + tmp_path, monkeypatch, model, df11, base +): + # Bug caught: a pin not reaching the resolver (the default checkpoint, one file on one person's Hub account, or the + # base following whatever lands on main instead of the revisions whose bytes and runs were checked), the Raw and + # Turbo pins swapped, or a pin applied to a user's own --df11 / --base. + from mlx_dfloat import _metal_decode + from mlx_dfloat.mflux._hub import ResolvedRepo + from mlx_dfloat.mflux.krea2 import model as model_module + from mlx_dfloat.mflux.krea2.model import DFloatKrea2 + + class StopError(Exception): + pass + + calls = [] + + def fake_resolve(spec, *, patterns, revision=None): + calls.append((spec, revision)) + if patterns == kinit.DF11_PATTERNS: + return ResolvedRepo(root=tmp_path, repo_id=spec, revision=revision) + raise StopError + + monkeypatch.setattr(_metal_decode, "ensure_canary", lambda **kwargs: None) + monkeypatch.setattr( + model_module, + "open_checkpoint", + lambda root: SimpleNamespace(root=root, config_source="config.json"), + ) + monkeypatch.setattr(kinit, "resolve", fake_resolve) + for kwargs in ({}, {"df11_path": "me/fork", "base_path": "me/base"}): + with pytest.raises(StopError): + DFloatKrea2(model, **kwargs) + assert calls == [df11, base, ("me/fork", None), ("me/base", None)] + + +def test_the_other_models_checkpoint_is_refused_before_the_base_resolves(tmp_path, monkeypatch): + # Bug caught: the variant check run only at the build, after the base is resolved, so asking for Krea 2 Raw + # with Turbo's DF11 file downloads the whole Raw base (about 16 GB of encoder and VAE) before the refusal. + from mlx_dfloat import _metal_decode + from mlx_dfloat.errors import DFloatFormatError + from mlx_dfloat.mflux._hub import ResolvedRepo + from mlx_dfloat.mflux.krea2 import model as model_module + + calls = [] + + def fake_resolve(spec, *, patterns, revision=None): + calls.append(patterns == kinit.DF11_PATTERNS) + return ResolvedRepo(root=tmp_path, repo_id=spec, revision=revision) + + turbo = "header and spot checks match layout krea-2-turbo-comfyui (mingyi456/Krea-2-Turbo-DF11-ComfyUI@978da5f)" + monkeypatch.setattr(_metal_decode, "ensure_canary", lambda **kwargs: None) + monkeypatch.setattr( + model_module, + "open_checkpoint", + lambda root: SimpleNamespace(root=root, config_source=turbo), + ) + monkeypatch.setattr(kinit, "resolve", fake_resolve) + with pytest.raises(DFloatFormatError, match=r"Krea 2 Turbo DF11 checkpoint.*--model krea-2$"): + model_module.DFloatKrea2("krea-2-raw", df11_path="me/turbo") + assert calls == [True] # the checkpoint only + + +@pytest.mark.parametrize(("nonblock", "bytes_"), [("per-call", 10), ("resident", 0)]) +def test_sizes_that_disagree_with_the_builds_nonblock_mode_are_refused( + tmp_path, monkeypatch, nonblock, bytes_ +): + # Bug caught: a per-call build planned with the resident non-block bytes (every call over-estimated by 1.23 + # GiB on the published model, refused on a 32 GB Mac) or a resident build planned without them (the fit check + # 1.23 GiB short). The checkpoint holds all seven non-block groups. + from mlx_dfloat.errors import DFloatIntegrationError + from mlx_dfloat.mflux._phases import FamilySizes + + sizes = FamilySizes(compressed=1_000, extras=100, nonblock=bytes_, encoders=1_000, vae=10) + with pytest.raises( + DFloatIntegrationError, match=f"{nonblock} non-block decode.*nonblock={bytes_}" + ): + fake_model(tmp_path, monkeypatch, nonblock=nonblock, sizes=sizes) + + +def test_the_report_lists_only_the_nonblock_groups_the_checkpoint_compresses(tmp_path, monkeypatch): + # Bug caught: the report naming all seven non-block groups whatever the checkpoint holds, so a checkpoint + # storing them as plain BF16 extras reads as if it compressed them. + from tests._krea2_tiny import TINY + + from mlx_dfloat.format import open_checkpoint + from mlx_dfloat.mflux.krea2.transformer import build_transformer, seam_transformer_class + + _m, _c, layout = write_tiny_checkpoint( + tmp_path / "plain", np.random.default_rng(TINY_SEED), nonblock_as_groups=False + ) + ckpt = open_checkpoint(tmp_path / "plain", layouts=(layout,)) + build = build_transformer( + ckpt, transformer_kwargs=TINY, transformer_class=seam_transformer_class(per_call=True) + ) + model = fake_model(tmp_path, monkeypatch, ckpt=ckpt, build=build) + assert model.report()["nonblock_groups"] == [] + + +def _weights_at_first_block(model, monkeypatch): + """Run one uncompiled guided call and return every non-block weight as it stood when the first release ran (at + block 0, before its weights are requested: the weights the pre-block graph read).""" + from tests._krea2_tiny import tiny_inputs, tiny_predict + + from mlx_dfloat.mflux.krea2 import transformer as ktf + + seen = {} + real = ktf.clear_nonblock + + def recording(module, shapes): + if not seen: + seen.update( + {p: np.array(_param(model, p).view(mx.uint16)) for p in _nonblock_sources_names()} + ) + real(module, shapes) + + monkeypatch.setattr(ktf, "clear_nonblock", recording) + mx.eval(tiny_predict(model.transformer, tiny_inputs(), guidance=3.5)) + monkeypatch.setattr(ktf, "clear_nonblock", real) + return seen + + +def _nonblock_sources_names(): + return [p for params in NONBLOCK.values() for p in params] + + +def test_the_model_decodes_the_nonblock_groups_per_call_by_default(tmp_path, monkeypatch): + # Bug caught: the shipped model keeping the seven groups decoded with the set (the layout whose one-step + # calibration run peaked 0.095 GiB over the 32 GB line on Krea 2 Raw), a group bound but not installed for the call (a placeholder on + # tproj.linear is a zero-size matmul before the first block), a group's matrices on each other's parameters, or + # the set attached without verify_in_call (a guided step's second call would be refused). + model = fake_model(tmp_path, monkeypatch) + model._lifecycle.ensure_set() + assert model.transformer.nonblock_bound() is True + assert _param(model, "tproj.linear.weight").size == 0 # nothing decoded at set load + assert model.transformer.seam_cell.state.verify_in_call is True + assert model.report()["nonblock"] == "per-call" + seen = _weights_at_first_block(model, monkeypatch) + sources = _nonblock_sources(tmp_path) + assert len(sources) == 37 + for name, source in sources.items(): + assert np.array_equal(seen[name], source), name + assert _param(model, "tproj.linear.weight").size == 0 # released after the call + + +def test_the_resident_set_load_installs_all_seven_nonblock_groups(tmp_path, monkeypatch): + # Bug caught: on the resident build, a non-block group decoded but not installed, or installed on another + # parameter. + model = fake_model(tmp_path, monkeypatch, nonblock="resident") + model._lifecycle.ensure_set() + assert model.report()["nonblock"] == "resident" + sources = _nonblock_sources(tmp_path) + assert len(sources) == 37 + for name, source in sources.items(): + got = _param(model, name) + assert np.array_equal(np.array(got.view(mx.uint16)), source), name + assert model.transformer.seam_cell.state.verify_in_call is True + + +@pytest.mark.parametrize("nonblock", ["per-call", "resident"]) +def test_a_new_prompt_after_a_generation_drops_the_set_and_reinstalls_the_nonblock_weights( + tmp_path, monkeypatch, nonblock +): + # Bug caught: a non-block weight left on its placeholder after the reload a new prompt forces + # (tproj.linear before the first block, a refiner block's wo in the text fusion), the per-call groups not bound + # again, or the encoder reload reading another base than the model was built from. + model = fake_model(tmp_path, monkeypatch, nonblock=nonblock) + loads = [] + monkeypatch.setattr( + kinit, + "load_text_encoder", + lambda root, **kwargs: (loads.append(root), StubTextEncoder())[1], + ) + model.encode("a") # the encoder came with the assembly: no reload + model._lifecycle.ensure_set() + model.encode("b") # the set must go before the encoder comes back + counters = model._lifecycle.counters + assert (counters.forced_set_drops, counters.encoder_loads) == (1, 1) + assert loads == [model._base.root] + assert not model._lifecycle.set_resident + tproj, wo = "tproj.linear.weight", "txtfusion.refiner_blocks.1.attn.wo.weight" + assert (_param(model, tproj).size, _param(model, wo).size) == (0, 0) + if nonblock == "per-call": + assert model.transformer.nonblock_bound() is False # the drop forgot the compressed groups + model._lifecycle.ensure_set() + assert counters.set_loads == 2 + sources = _nonblock_sources(tmp_path) + if nonblock == "per-call": + seen = _weights_at_first_block(model, monkeypatch) + for name in (tproj, wo): + assert np.array_equal(seen[name], sources[name]), name + else: + assert (_param(model, tproj).shape, _param(model, wo).shape) == ((192, 32), (32, 32)) + for name in (tproj, wo): + assert np.array_equal(np.array(_param(model, name).view(mx.uint16)), sources[name]), ( + name + ) + + +def test_embeddings_are_cached_per_prompt_and_mfluxs_cache_stays_empty(tmp_path, monkeypatch): + # Bug caught: the embeddings put into mflux's own prompt_cache (keyed by (prompt, negative, guidance) there: + # prompt_encoder.py:25, so ours would never match and mflux would call the dropped encoder), or one prompt encoded + # for another. The stub tokenizer gives 1 id per character after its 5-id template prefix mflux strips. + model = fake_model(tmp_path, monkeypatch) + model.encode("a", " ", "abc") + assert set(model._embeds) == {"a", " ", "abc"} + assert model.prompt_cache == {} + assert [model._embeds[p][0].shape for p in ("a", " ", "abc")] == [ + (1, 1, 384), + (1, 1, 384), + (1, 3, 384), + ] + a, blank = (np.array(model._embeds[p][0].view(mx.uint16)) for p in ("a", " ")) + assert not np.array_equal(a, blank) + + +def test_the_empty_prompt_is_refused_by_encode_as_it_gives_no_tokens(tmp_path, monkeypatch): + # Bug caught: "" encoded (Krea 2's tokenizer gives it no ids at all, 2026-10-08: the transformer would get an + # empty context), while " " (mflux's own blank negative) must still encode. + model = fake_model(tmp_path, monkeypatch) + with pytest.raises(DFloatUnsupportedError, match="empty prompt"): + model.encode("a", "") + assert model._embeds == {} + model.encode(" ") + assert set(model._embeds) == {" "} + + +def test_a_dropped_encoder_refuses_with_a_named_error(tmp_path, monkeypatch): + # Bug caught: None left in the encoder's place (mflux would call it on a cache miss and fail with a TypeError far + # from the cause). + model = fake_model(tmp_path, monkeypatch) + model.encode("a") + with pytest.raises(DFloatIntegrationError, match="not encoded"): + model.text_encoder.get_prompt_embeds(mx.ones((1, 4), dtype=mx.int32)) + with pytest.raises(DFloatIntegrationError, match="not encoded"): + model.text_encoder(mx.ones((1, 4), dtype=mx.int32)) + + +def test_a_failed_set_load_leaves_nothing_installed(tmp_path, monkeypatch): + # Bug caught: a raise late in the set load (here at attach) leaving the decoded non-block weights installed and + # the provider holding the compressed set while the lifecycle believes no set is resident; the retry's + # active-memory check would then refuse to load a second copy. + model = fake_model(tmp_path, monkeypatch) + + def broken_attach(*args, **kwargs): + raise RuntimeError("attach") + + monkeypatch.setattr(model.transformer, "attach", broken_attach) + with pytest.raises(RuntimeError, match="attach"): + model._lifecycle.ensure_set() + assert model._provider is None + assert _param(model, "tproj.linear.weight").size == 0 + assert not model._lifecycle.set_resident + + +def _corrupting(group_name): + """A reference decode whose status words report an invalid code for ``group_name`` while ``on`` is set.""" + from dataclasses import replace + from functools import partial + + from mlx_dfloat.decode import STATUS_INVALID_CODE, decode_group + + reference = partial(decode_group, backend="reference") + state = {"on": True} + + def decode(group): + result = reference(group) + if state["on"] and group.name == group_name: + return replace(result, status=result.status | STATUS_INVALID_CODE) + return result + + return decode, state + + +def _frames(tb): + while tb is not None: + yield tb.tb_frame + tb = tb.tb_next + + +def test_a_corrupt_nonblock_group_at_set_load_releases_the_set(tmp_path, monkeypatch): + # Bug caught: a status failure while the set loads (the non-block groups decode once at the load) leaving the + # compressed set reachable from the exception's frames (_install_set's `resident`) while the caller holds the + # exception, so the retry refuses with "a previous set is still resident". + from mlx_dfloat.errors import DFloatFormatError + + decode, state = _corrupting("tproj") + model = fake_model(tmp_path, monkeypatch, decode=decode, nonblock="resident") + model.encode("p") + gc.collect() + mx.clear_cache() + before = int(mx.get_active_memory()) + with pytest.raises(DFloatFormatError, match="tproj") as info: + model._lifecycle.ensure_set() + assert not model._lifecycle.set_resident + install = [f for f in _frames(info.value.__traceback__) if f.f_code.co_name == "_install_set"] + assert len(install) == 1 + assert "resident" not in install[0].f_locals + monkeypatch.setattr(model._lifecycle, "retained_bound", lambda: before) + state["on"] = False + model._lifecycle.ensure_set() # the retry, with the exception still held + assert model._lifecycle.counters.set_loads == 1 + del info + + +def test_a_corrupt_nonblock_group_under_per_call_decode_fails_the_call_and_drops_the_set( + tmp_path, monkeypatch +): + # Bug caught: a status failure in the per-call decode (at the call, not at set load) surfacing as anything but the + # format error naming the group, leaving decoded weights installed, or keeping the set for a retry that would hit + # the same corrupt group again without a clean reload. + from mlx_dfloat.errors import DFloatFormatError + + decode, _state = _corrupting("tproj") + model = fake_model(tmp_path, monkeypatch, decode=decode) + with pytest.raises(DFloatFormatError, match="tproj"): + model.generate_image(seed=1, prompt="p", num_inference_steps=1, height=64, width=64) + assert not model._lifecycle.set_resident + for name in _nonblock_sources_names(): + assert _param(model, name).size == 0, name + + +def test_the_retained_bound_counts_the_cached_embeddings_and_the_vae(tmp_path, monkeypatch): + # Bug caught: the bound reading mflux's prompt_cache (empty: every drop after an encode would refuse) instead of + # our per-prompt cache, where the embeddings live. + model = fake_model(tmp_path, monkeypatch) + empty = model._retained_bound() + model.encode("abc") + # One prompt: 3 rows of 384 bf16 = 2_304 B. + assert model._retained_bound() - empty == 3 * 384 * 2 + + +def test_save_model_is_refused_and_freeze_skips_a_dropped_encoder(tmp_path, monkeypatch): + # Bug caught: freeze() calling .freeze() on the encoder stand-in (it is not a module), or save_model writing a + # checkpoint with zero-size placeholders as the transformer. + model = fake_model(tmp_path, monkeypatch) + model.encode("p") # drops the encoder + model.freeze() + with pytest.raises(DFloatUnsupportedError, match="save"): + model.save_model(str(tmp_path)) + + +def _innermost_locals(tb): + while tb.tb_next is not None: + tb = tb.tb_next + return tb.tb_frame.f_locals + + +def test_clear_frames_walks_the_whole_chain_of_causes_and_contexts(): + # Bug caught: only the raised exception and its __context__ cleared, so a __cause__ one link further (raise ... from + # an earlier error) keeps its frame's locals, in production a block's decoded weights, alive while the caller holds + # the exception; or a chain that loops back on itself never ending. + from mlx_dfloat.mflux.krea2.model import _clear_frames + + def fail(kind, message): + big = mx.zeros(1024) # noqa: F841 # stands in for a block's decoded weights + raise kind(message) + + try: + try: + try: + fail(ValueError, "first") + except ValueError as first: + raise KeyError("second") from first + except KeyError: + fail(RuntimeError, "third") + except RuntimeError as third: + caught = third + second = caught.__context__ + first = second.__cause__ + first.__context__ = caught # a loop: the walk must still end + _clear_frames(caught) + for exc in (caught, second, first): + assert "big" not in _innermost_locals(exc.__traceback__), exc + + +def _users_own_error(): + """An exception a caller is handling, raised in a finished frame that holds a local.""" + mine = "the caller's state" # noqa: F841 # must survive our cleanup + raise LookupError("the caller's own") + + +def test_clear_frames_stops_at_the_exception_the_caller_was_already_handling(): + # Bug caught: the walk following __context__ out of our call into the exception the caller was handling when + # it called us, and clearing the caller's own frames (their locals gone from a traceback they still use). + from mlx_dfloat.mflux.krea2.model import _clear_frames + + def ours(): + big = mx.zeros(1024) # noqa: F841 + raise RuntimeError("ours") + + users = caught = None + try: + _users_own_error() + except LookupError as exc: + users = exc + try: + ours() + except RuntimeError as inner: + caught = inner + _clear_frames(caught, outer=users) + assert caught.__context__ is users + assert "big" not in _innermost_locals(caught.__traceback__) + assert _innermost_locals(users.__traceback__)["mine"] == "the caller's state" + + +def test_a_failed_generation_inside_a_callers_except_leaves_the_callers_frames( + tmp_path, monkeypatch +): + # Bug caught: generate_image's cleanup clearing the frames of the exception its caller was handling (the + # failure's __context__), not only its own. + from mlx_dfloat.errors import DFloatFormatError + + decode, _state = _corrupting("tproj") + model = fake_model(tmp_path, monkeypatch, decode=decode) + users = failure = None + try: + _users_own_error() + except LookupError as exc: + users = exc + with pytest.raises(DFloatFormatError, match="tproj") as info: + model.generate_image(seed=1, prompt="p", num_inference_steps=1, height=64, width=64) + failure = info.value + chain, link = [], failure + while link is not None: + chain.append(link) + link = link.__context__ + assert users in chain + assert _innermost_locals(users.__traceback__)["mine"] == "the caller's state" + + +# --- the generate prelude, CFG, the plan ------------------------------------------------------------------------ + +GIB = 1024**3 + + +def _cache_limit_in_force(): + previous = mx.set_cache_limit(0) + mx.set_cache_limit(previous) + return previous + + +def _patch_upstream_generate(monkeypatch, *, raise_after_loop=None, raise_before_loop=None): + """Replace Krea2.generate_image with a probe: records the cache limit and MLX limits in force and the kwargs, + fires our after-loop subscriber the way mflux's GenerationContext would, returns a sentinel.""" + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + from tests._mlx_limits import current_limits + + seen = {} + + def fake(self, **kwargs): + seen["limit_at_entry"] = _cache_limit_in_force() + seen["limits_at_entry"] = current_limits() + seen["set_resident"] = self._lifecycle.set_resident + seen["kwargs"] = kwargs + if raise_before_loop is not None: + raise raise_before_loop + for subscriber in self.callbacks.after_loop_callbacks(): + subscriber.call_after_loop( + seed=kwargs["seed"], prompt=kwargs["prompt"], latents=mx.zeros(1), config=None + ) + seen["limit_after_loop"] = _cache_limit_in_force() + seen["set_resident_after_loop"] = self._lifecycle.set_resident + if raise_after_loop is not None: + raise raise_after_loop + return "image" + + monkeypatch.setattr(Krea2, "generate_image", fake) + return seen + + +@pytest.mark.parametrize( + ("guidance", "negative", "expected"), + [ + (1.0, "n", ("p",)), + (3.5, "n", ("p", "n")), + (0.5, None, ("p", " ")), + (1.0001, "", ("p", " ")), + (3.5, " ", ("p", " ")), + (3.5, " x ", ("p", " x ")), + ], +) +def test_cfg_prompts_follow_mflux(guidance, negative, expected): + # Bug caught: `> 1.0` copied from the other families (guidance 0.5 would skip mflux's negative branch: a different + # image), or a blank negative not replaced by mflux's " ". Literal table from prompt_encoder.py:33-35. + from mlx_dfloat.mflux.krea2.model import DFloatKrea2 + + assert DFloatKrea2.cfg_prompts("p", negative_prompt=negative, guidance=guidance) == expected + + +@pytest.mark.parametrize(("guidance", "negative"), [(3.5, "n"), (1.0, "n"), (0.5, None)]) +def test_the_override_returns_mfluxs_own_pair(tmp_path, monkeypatch, guidance, negative): + # Bug caught: the copied rule drifting from mflux's (the override answering another pair than mflux's own + # _encode_prompts would build), or the pair swapped. The oracle is mflux's method on a throwaway object. + from types import SimpleNamespace + + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + + model = fake_model(tmp_path, monkeypatch, tokenizer=StubTokenizer({"n": 12})) + oracle = SimpleNamespace( + prompt_cache={}, tokenizers=model.tokenizers, text_encoder=StubTextEncoder() + ) + want = Krea2._encode_prompts(oracle, prompt="p", negative_prompt=negative, guidance=guidance) + model.encode(*model.cfg_prompts("p", negative_prompt=negative, guidance=guidance)) + got = model._encode_prompts(prompt="p", negative_prompt=negative, guidance=guidance) + assert np.array_equal(np.array(got[0].view(mx.uint16)), np.array(want[0].view(mx.uint16))) + if guidance == 1.0: + assert got[1] is None + assert want[1] is None + else: + assert np.array_equal(np.array(got[1].view(mx.uint16)), np.array(want[1].view(mx.uint16))) + + +def test_a_prompt_not_encoded_first_is_refused_by_the_override(tmp_path, monkeypatch): + # Bug caught: a cache miss falling through to mflux's encoder (dropped: a refusal deep in the loop, or an encoder + # reload next to the set). + model = fake_model(tmp_path, monkeypatch) + model.encode("p") + with pytest.raises(DFloatIntegrationError, match=r"\[' '\] were not encoded"): + model._encode_prompts(prompt="p", negative_prompt=None, guidance=3.5) + + +def test_an_unknown_scheduler_is_refused_before_anything_loads(tmp_path, monkeypatch): + # Bug caught: mflux raising only inside its own generate_image, after the prelude spent the + # encoder and the set load (about 45 s on the real model); or "linear" refused although mflux maps it to er_sde. + from mlx_dfloat.errors import DFloatUnsupportedError as Unsupported + + model = fake_model(tmp_path, monkeypatch) + monkeypatch.setattr(model._lifecycle, "ensure_embeddings", lambda *p: pytest.fail("encoded")) + with pytest.raises(Unsupported, match=r"'ddim'.*er_sde, euler, linear"): + model.generate_image(seed=1, prompt="p", height=64, width=64, scheduler="ddim") + counters = model._lifecycle.counters + assert (counters.encoder_loads, counters.set_loads) == (0, 0) + monkeypatch.undo() + model = fake_model(tmp_path / "again", monkeypatch) + seen = _patch_upstream_generate(monkeypatch) + model.generate_image(seed=1, prompt="p", height=64, width=64, scheduler="linear") + assert seen["kwargs"]["scheduler"] == "linear" + + +@pytest.mark.parametrize( + ("model_name", "kwargs", "want"), + [ + # mflux: Raw has no step-table entry, so the fallback 25 (defaults.py:19,111-118); guidance 1.0 (krea2.py:58). + ("krea-2-raw", {}, (1.0, 25, "er_sde", 1)), + # mflux: Turbo 8 steps (defaults.py:47), guidance 1.0 (krea2_generate.py:17). + ("krea-2", {}, (1.0, 8, "er_sde", 1)), + # The model card's recipe for Raw passes through: two calls per step. + ("krea-2-raw", {"num_inference_steps": 52, "guidance": 3.5}, (3.5, 52, "er_sde", 2)), + ("krea-2", {"guidance": 2.0}, (2.0, 8, "er_sde", 2)), + ], +) +def test_the_python_api_defaults_follow_the_registry_per_model( + tmp_path, monkeypatch, model_name, kwargs, want +): + # Bug caught: Raw's entry still carrying the card's 52 / 3.5 (a CFG run nobody asked for, about 4x the time), the + # per-model lookup taking Turbo's 8 for Raw, or the scheduler left None for mflux (it resolves None itself, but the + # run would not record it). + model = fake_model(tmp_path, monkeypatch, model=model_name) + seen = _patch_upstream_generate(monkeypatch) + model.generate_image(seed=1, prompt="p", height=64, width=64, **kwargs) + got = seen["kwargs"] + assert (got["guidance"], got["num_inference_steps"], got["scheduler"]) == want[:3] + assert model.report()["cfg_calls_per_step"] == want[3] + + +def test_the_generate_signature_leaves_the_recipe_to_the_registry(): + # Bug caught: a default baked into the signature (mflux's class default of 8 steps would then run Raw at 8). + from mlx_dfloat.mflux.krea2.model import DFloatKrea2 + + defaults = { + n: p.default + for n, p in inspect.signature(DFloatKrea2.generate_image).parameters.items() + if p.default is not inspect.Parameter.empty + } + assert defaults == { + "num_inference_steps": None, + "height": 1024, + "width": 1024, + "guidance": None, + "negative_prompt": None, + "image_path": None, + "image_strength": None, + "scheduler": None, + "pid_decode": False, + "pid_degrade_sigma": 0.0, + } + + +@pytest.mark.parametrize("prompt", ["", " "]) +def test_an_empty_or_blank_prompt_is_refused_before_anything_loads(tmp_path, monkeypatch, prompt): + # Bug caught: "" sent on (Krea 2's tokenizer gives it no ids: the transformer gets an empty context), or a blank + # prompt treated as a real one; either way after an encoder load. + model = fake_model(tmp_path, monkeypatch) + monkeypatch.setattr(model._lifecycle, "ensure_embeddings", lambda *p: pytest.fail("encoded")) + with pytest.raises(DFloatUnsupportedError, match="empty or blank prompt"): + model.generate_image(seed=1, prompt=prompt, height=64, width=64) + + +@pytest.mark.parametrize( + "kwargs", [{"image_path": "in.png"}, {"image_strength": 0.5}, {"pid_decode": True}] +) +def test_img2img_and_pid_decode_are_refused_before_anything_loads(tmp_path, monkeypatch, kwargs): + # Bug caught: encoding the prompt (an encoder load) or loading the set before refusing img2img or mflux's PiD + # decoder (which loads its own caption encoder next to the set). + model = fake_model(tmp_path, monkeypatch) + monkeypatch.setattr(model._lifecycle, "ensure_embeddings", lambda *p: pytest.fail("encoded")) + with pytest.raises(DFloatUnsupportedError, match=next(iter(kwargs))): + model.generate_image(seed=1, prompt="p", **kwargs) + assert model._lifecycle.counters.set_loads == 0 + + +def test_text_tokens_follow_the_longer_prompt_of_the_call(tmp_path, monkeypatch): + # Bug caught: the plan sized by the prompt only (a longer negative prompt's activations unplanned), or the + # template prefix counted. Stub lengths: the prompt 7 ids, the negative 12, after the 5-id prefix mflux strips. + model = fake_model(tmp_path, monkeypatch, tokenizer=StubTokenizer({"p": 7, "n": 12})) + planned = [] + real = model.plan_call + monkeypatch.setattr( + model, "plan_call", lambda **kw: (planned.append(kw["text_tokens"]), real(**kw))[1] + ) + _patch_upstream_generate(monkeypatch) + model.generate_image(seed=1, prompt="p", height=64, width=64, guidance=3.5, negative_prompt="n") + assert planned == [12] + assert model.report()["text_tokens"] == 12 + + +def test_plan_call_reads_the_familys_sizes_and_the_one_block_kind(tmp_path, monkeypatch): + # Bug caught: plan_call not wired to the shared planner with this model's own largest block group (the cache limit + # would miss the decoded group), the non-block groups counted as the largest block, or the plan not kept for the + # VAE guard. Injected constants (literals): allowance 3e9 at 4096 + 30 tokens. + from mlx_dfloat.mflux._phases import PhaseConstants + + model = fake_model(tmp_path, monkeypatch) + model._constants = PhaseConstants( + overhead_bytes=1, + vae_transient_bytes=1, + denoise_activation_at_reference=1, + reference_tokens=4096 + 30, + allowance_at_reference=3_000_000_000, + allowance_reference_tokens=4096 + 30, + ) + plan = model.plan_call(height=1030, width=1030, text_tokens=30) # rounded to 1024^2 first + # The tiny block group: 1024 + 512 + 512 + 1024 + 1024 + 3 x 4096 = 16_384 elements, 32_768 B (TINY's shapes). + assert model._largest == {"blocks": 32_768} + assert plan.cache_limit == 32_768 + 3_000_000_000 + assert model._plan is plan + + +def test_a_dropped_set_vae_phase_keeps_the_extras(tmp_path, monkeypatch): + # Bug caught: the VAE phase planned after the set drop leaving out the extras, which stay loaded (they are + # the transformer's own BF16 parameters, not part of the set): the fit estimate would under-predict the VAE phase. + # Literals: VAE phase with the set 12e9 + 777 + 100 + 11e9 + 400 > the 16e9 budget (per-call: no resident + # non-block bytes), so the set drops; without it the phase is 777 + 100 + 11e9 + 400. + from mlx_dfloat.mflux._phases import FamilySizes, PhaseConstants + + sizes = FamilySizes(compressed=12_000_000_000, extras=777, nonblock=0, encoders=0, vae=100) + model = fake_model(tmp_path, monkeypatch, sizes=sizes, budget_bytes=16_000_000_000) + model._constants = PhaseConstants( + overhead_bytes=400, + vae_transient_bytes=11_000_000_000, + denoise_activation_at_reference=0, + reference_tokens=4096, + allowance_at_reference=0, + allowance_reference_tokens=4096, + allowance_floor=0, + ) + plan = model.plan_call(height=1024, width=1024, text_tokens=0) + assert plan.drop_set_before_vae is True + assert plan.estimate.phases["vae"] == 777 + 100 + 11_000_000_000 + 400 + + +def test_the_vae_guard_follows_the_plan_and_zeroes_the_cache_limit(tmp_path, monkeypatch): + # Bug caught: the guard not reading this call's plan (the set kept through a decode the budget cannot hold), or a + # transformer-sized cache limit left for the VAE decode. + # Injected constants (literals): the VAE phase with the set is 12e9 + 11e9 + 400 > 16e9, without it 11e9 + 400; + # the denoise phase stays near 12e9 (no allowance, no activation term). + from mlx_dfloat.mflux._phases import FamilySizes, PhaseConstants + + constants = PhaseConstants( + overhead_bytes=400, + vae_transient_bytes=11_000_000_000, + denoise_activation_at_reference=0, + reference_tokens=4096, + allowance_at_reference=0, + allowance_reference_tokens=4096, + allowance_floor=0, + ) + big = FamilySizes(compressed=12_000_000_000, extras=0, nonblock=0, encoders=0, vae=0) + tight = fake_model(tmp_path / "t", monkeypatch, sizes=big, budget_bytes=16_000_000_000) + tight._constants = constants + seen = _patch_upstream_generate(monkeypatch) + assert tight.generate_image(seed=1, prompt="p", height=1024, width=1024) == "image" + assert seen["set_resident"] + assert seen["limit_at_entry"] == tight._plan.cache_limit + assert seen["limit_after_loop"] == 0 + assert seen["set_resident_after_loop"] is False + roomy = fake_model(tmp_path / "r", monkeypatch, sizes=big, budget_bytes=40 * GIB) + roomy._constants = constants + seen = _patch_upstream_generate(monkeypatch) + roomy.generate_image(seed=1, prompt="p", height=1024, width=1024) + assert seen["limit_after_loop"] == 0 + assert seen["set_resident_after_loop"] is True + + +def test_the_prelude_encodes_then_drops_the_encoder_then_loads_the_set_then_runs_mflux( + tmp_path, monkeypatch +): + # Bug caught: the set loaded while the encoder is resident, mflux's loop entered before the embeddings exist, or + # the negative prompt not passed through (mflux reads it for the image metadata). + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + + model = fake_model(tmp_path, monkeypatch) + events = [] + life = model._lifecycle + for attr in ("_encode", "_unload_encoders", "_load_set"): + real = getattr(life, attr) + + def wrapped(*args, _real=real, _name=attr, **kwargs): + events.append(_name + (f":{args[0]!r}" if args else "")) + return _real(*args, **kwargs) + + monkeypatch.setattr(life, attr, wrapped) + seen = _patch_upstream_generate(monkeypatch) + upstream = Krea2.generate_image + + def recording(self, **kwargs): + events.append("upstream") + return upstream(self, **kwargs) + + monkeypatch.setattr(Krea2, "generate_image", recording) + image = model.generate_image( + seed=7, + prompt="p", + num_inference_steps=1, + height=256, + width=256, + guidance=3.5, + negative_prompt="n", + ) + assert image == "image" + assert events == ["_encode:'p'", "_encode:'n'", "_unload_encoders", "_load_set", "upstream"] + kwargs = seen["kwargs"] + assert (kwargs["image_path"], kwargs["image_strength"], kwargs["pid_decode"]) == ( + None, + None, + False, + ) + assert (kwargs["guidance"], kwargs["negative_prompt"], kwargs["num_inference_steps"]) == ( + 3.5, + "n", + 1, + ) + + +@pytest.mark.parametrize("where", ["raise_before_loop", "raise_after_loop"]) +def test_a_format_error_mid_step_drops_the_set_and_restores_the_limit(tmp_path, monkeypatch, where): + # Bug caught: a corrupt block's retry reusing stale resident state, or the process cache limit left at the call's. + from mlx_dfloat.errors import DFloatFormatError + + model = fake_model(tmp_path, monkeypatch) + before = _cache_limit_in_force() + _patch_upstream_generate(monkeypatch, **{where: DFloatFormatError("block 1: invalid code")}) + with pytest.raises(DFloatFormatError): + model.generate_image(seed=1, prompt="p", height=256, width=256) + assert not model._lifecycle.set_resident + assert _cache_limit_in_force() == before + assert _param(model, "tproj.linear.weight").size == 0 + assert model.open_phase is None + + +def test_keyboard_interrupt_restores_the_cache_limit_and_leaves_no_pending_words( + tmp_path, monkeypatch +): + # Bug caught: mflux's StopImageGenerationException (raised for a Ctrl-C in its loop, krea2.py:103-107) leaving + # status words pending (begin_step would refuse the next call), the cache limit not restored, or the interrupt + # dropping the set (a reload per Ctrl-C). Real mflux loop; the interrupt fires at the second block of the step's + # second transformer call (CFG: 2 blocks per call). + from mflux.utils.exceptions import StopImageGenerationException + + model = fake_model(tmp_path, monkeypatch) + before = _cache_limit_in_force() + model.encode("p", "n") + model._lifecycle.ensure_set() + provider = model._provider + real = provider.weights_for + calls = [] + interrupt = {"on": True} + + def interrupting(name, shapes): + calls.append(name) + if interrupt["on"] and len(calls) == 4: + raise KeyboardInterrupt + return real(name, shapes) + + monkeypatch.setattr(provider, "weights_for", interrupting) + run = { + "seed": 1, + "prompt": "p", + "num_inference_steps": 1, + "height": 64, + "width": 64, + "guidance": 3.5, + "negative_prompt": "n", + } + with pytest.raises(StopImageGenerationException): + model.generate_image(**run) + assert len(calls) == 4 + assert _cache_limit_in_force() == before + assert provider.pending == [] + assert model._lifecycle.set_resident + interrupt["on"] = False + assert model.generate_image(**run) is not None + assert model._lifecycle.counters.set_loads == 1 + + +@pytest.mark.parametrize("chained", [False, True]) +def test_any_exception_from_the_loop_releases_its_frames_locals(tmp_path, monkeypatch, chained): + # Bug caught: only DFloatFormatError clearing the traceback's frames, so a RuntimeError (or an error raised while + # handling another) keeps the raising frame's locals, in production the seam's decoded weights, alive for as long + # as the caller holds the exception. + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + + model = fake_model(tmp_path, monkeypatch) + + def failing_step(): + big = mx.zeros(1024) # noqa: F841 # stands in for a block's decoded weights + raise ValueError("inner") + + def fake(self, **kwargs): + if not chained: + failing_step() + try: + failing_step() + except ValueError as exc: + raise RuntimeError("outer") from exc + + monkeypatch.setattr(Krea2, "generate_image", fake) + with pytest.raises((ValueError, RuntimeError)) as info: + model.generate_image(seed=1, prompt="p", height=256, width=256) + exc = info.value + assert "big" not in _innermost_locals(exc.__traceback__) + if chained: + assert isinstance(exc.__cause__, ValueError) + assert "big" not in _innermost_locals(exc.__cause__.__traceback__) + assert model._lifecycle.set_resident # only a format error drops the set + + +def test_open_phase_names_the_phase_running_now_and_none_outside_a_call(tmp_path, monkeypatch): + # Bug caught: the watchdog's abort context naming no phase, or a stale one. + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + + model = fake_model(tmp_path, monkeypatch) + assert model.open_phase is None + seen = [] + life = model._lifecycle + for attr in ("_encode", "_load_set"): + real = getattr(life, attr) + + def wrapped(*args, _real=real, **kwargs): + seen.append(model.open_phase) + return _real(*args, **kwargs) + + monkeypatch.setattr(life, attr, wrapped) + + def fake(self, **kwargs): + seen.append(self.open_phase) + for subscriber in self.callbacks.after_loop_callbacks(): + subscriber.call_after_loop(seed=1, prompt="p", latents=mx.zeros(1), config=None) + seen.append(self.open_phase) + return "image" + + monkeypatch.setattr(Krea2, "generate_image", fake) + model.generate_image(seed=1, prompt="p", height=256, width=256) + assert seen == ["encode", "set_load", "denoise", "vae"] + assert model.open_phase is None + + +def _fake_encode_peak(monkeypatch, model, above_bound): + """Make the tracker record an encode whose MLX hold is ``above_bound`` bytes over the call's warning bound.""" + from mlx_dfloat.mflux.krea2 import memory as kmem + + real_end = model._tracker.end + + def end(name): + real_end(name) + if name == "encode": + record = model._tracker.peaks["encode"] + bound = kmem.encode_peak_bound(model._plan.estimate.phases["encode"], model._constants) + record["mlx_peak"] = record["active_at_start"] + bound + above_bound + + monkeypatch.setattr(model._tracker, "end", end) + + +def test_an_encode_over_the_bound_warns_through_the_pure_helper(tmp_path, monkeypatch): + # Bug caught: the encode hold never compared with the bound (an mflux change that loads the vision tower, +0.77 + # GiB on the real file, going unnoticed), `>=` for `>` at the bound, or the warning missing the call's prompt + # length. The bound itself is pinned by the helper's own tests; this pins the wiring. The stub gives "p" 1 token. + model = fake_model(tmp_path, monkeypatch) + _patch_upstream_generate(monkeypatch) + _fake_encode_peak(monkeypatch, model, 0) + assert ( + model.generate_image(seed=1, prompt="p", height=256, width=256) == "image" + ) # at the bound: silent + model = fake_model(tmp_path / "again", monkeypatch) + _patch_upstream_generate(monkeypatch) + _fake_encode_peak(monkeypatch, model, 1) + with pytest.warns(UserWarning, match=r"for a 1-token prompt.*vision tower"): + model.generate_image(seed=1, prompt="p", height=256, width=256) + + +def test_the_budget_is_the_constructors_else_the_devices_read_at_each_call(tmp_path, monkeypatch): + # Bug caught: the constructor's budget ignored (a --tier run planned against the whole Mac), or the device's read + # once at assembly instead of at each call. + from mlx_dfloat.mflux.krea2 import model as model_module + + assert fake_model(tmp_path / "a", monkeypatch, budget_bytes=9 * GIB)._budget() == 9 * GIB + model = fake_model(tmp_path / "b", monkeypatch) + assert model._budget() == 23 * GIB # fake_model's fixed device budget + monkeypatch.setattr(model_module, "budget_bytes", lambda: 10 * GIB) + assert model._budget() == 10 * GIB + + +# --- the report ----------------------------------------------------------------------------------------------- + + +def test_the_report_names_the_family_the_calls_the_layout_and_the_uncompiled_predict( + tmp_path, monkeypatch +): + # Bug caught: the family, the call count, the predict mode or the checkpoint's layout missing from the report (the + # CLI and the bench table read them), a non-block group missing from the list, or the estimate labelled as a + # measurement. + model = fake_model(tmp_path, monkeypatch) + report = model.report() + assert (report["family"], report["model"]) == ("krea2", "krea-2") + assert ( + report["cfg_calls_per_step"], + report["text_tokens"], + report["predict"], + report["fit"], + ) == ( + 1, + None, + None, + None, + ) + assert report["nonblock_groups"] == [ + "tmlp", + "tproj", + "txtfusion.layerwise_blocks.0", + "txtfusion.layerwise_blocks.1", + "txtfusion.refiner_blocks.0", + "txtfusion.refiner_blocks.1", + "txtmlp", + ] + assert report["layout"].startswith("header and spot checks match layout tiny-krea") + assert report["nonblock"] == "per-call" + assert report["sizes"] == { + "compressed": 1_000, + "extras": 100, + "nonblock": 0, + "encoders": 1_000, + "vae": 10, + } + assert report["eval_policy_status"] == "measured" + model.generate_image( + seed=1, + prompt="a", + num_inference_steps=1, + height=64, + width=64, + guidance=3.5, + negative_prompt="n", + ) + report = model.report() + assert (report["cfg_calls_per_step"], report["text_tokens"], report["predict"]) == ( + 2, + 1, + "uncompiled", + ) + assert report["fit"]["label"] == "predicted" + assert report["decode_launches"] == 1 * 2 * 2 + + +def test_the_report_labels_depth2_unmeasured(tmp_path, monkeypatch): + # Bug caught: a depth2 run reported like the measured per-block path. + report = fake_model(tmp_path, monkeypatch, eval_policy="depth2").report() + assert (report["eval_policy"], report["eval_policy_status"]) == ("depth2", "unmeasured") + + +# --- the constructor, end to end past the resolver ------------------------------------------------------------ + + +def test_the_constructor_builds_from_local_dirs_with_the_variant_check_and_the_language_model_bytes( + tmp_path, monkeypatch +): + # Bug caught: __init__ wiring past the resolver untested: the config-less checkpoint not opened through + # the layout table, the build without the variant check (the model name not passed), the base components read + # from another directory, or the sizes taken from the encoder file (vision tower and last layer included) instead + # of the evaluated language model. + from tests._krea2_tiny import TINY, stub_components + + from mlx_dfloat import _layouts, _metal_decode + from mlx_dfloat.mflux.krea2 import model as model_module + + _m, _c, layout = write_tiny_checkpoint(tmp_path / "df11", np.random.default_rng(TINY_SEED)) + monkeypatch.setattr(_layouts, "KNOWN_LAYOUTS", (layout,)) + base = tmp_path / "base" + (base / "text_encoder").mkdir(parents=True) + # The language model: embed_tokens (10, 4) 80 B, two layers' up_proj (4, 4) 32 B each, final norm (4,) 8 B; the + # last layer never runs. Never loaded: a vision tensor (8, 4) 64 B. 80 + 32 + 8 = 120. + mx.save_safetensors( + str(base / "text_encoder" / "model.safetensors"), + { + "language_model.embed_tokens.weight": mx.zeros((10, 4), dtype=mx.bfloat16), + "language_model.layers.0.mlp.up_proj.weight": mx.zeros((4, 4), dtype=mx.bfloat16), + "language_model.layers.1.mlp.up_proj.weight": mx.zeros((4, 4), dtype=mx.bfloat16), + "language_model.norm.weight": mx.zeros((4,), dtype=mx.bfloat16), + "visual.blocks.0.attn.qkv.weight": mx.zeros((8, 4), dtype=mx.bfloat16), + }, + ) + (base / "vae").mkdir() + (base / "vae" / "diffusion_pytorch_model.safetensors").write_bytes(b"\0" * 100) + + monkeypatch.setattr(_metal_decode, "ensure_canary", lambda **kwargs: None) + loaded = [] + monkeypatch.setattr( + kinit, "load_base", lambda root: (loaded.append(root), stub_components())[1] + ) + built = [] + real_build = model_module.build_transformer + + def tiny_build(ckpt, **kwargs): + built.append(kwargs) + return real_build(ckpt, transformer_kwargs=TINY, **kwargs) + + from mlx_dfloat.mflux.krea2.transformer import seam_transformer_class + + monkeypatch.setattr(model_module, "build_transformer", tiny_build) + model = model_module.DFloatKrea2( + "krea-2-raw", df11_path=str(tmp_path / "df11"), base_path=str(base) + ) + # The per-call class: the shipped build. + assert built == [ + {"model": "krea-2-raw", "transformer_class": seam_transformer_class(per_call=True)} + ] + assert loaded == [base] + report = model.report() + assert report["sizes"]["encoders"] == 120 + assert report["sizes"]["vae"] == 100 + # Per-call decode: the seven non-block groups are not kept decoded with the set. + assert report["sizes"]["nonblock"] == 0 + assert report["nonblock"] == "per-call" + assert report["layout"].startswith("header and spot checks match layout tiny-krea") + + +# --- per-call caps ---------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("raised", [None, RuntimeError("in the loop")], ids=["returns", "raises"]) +def test_a_python_api_call_runs_under_the_commands_caps_and_restores_mlxs_defaults( + tmp_path, monkeypatch, raised +): + # Bug caught: generate_image not run under the per-call caps (a Python-API process sits at MLX's default wired + # limit 0, while every VAE term is measured under the command's caps), or the caps left installed after a call + # that returns or raises. + from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap + + model = fake_model(tmp_path, monkeypatch) + seen = _patch_upstream_generate(monkeypatch, raise_before_loop=raised) + wired, memory = command_caps() + assert wired > 0 + with mlx_without_wired_cap() as start: + if raised is None: + model.generate_image(seed=1, prompt="p", height=256, width=256) + else: + with pytest.raises(RuntimeError, match="in the loop"): + model.generate_image(seed=1, prompt="p", height=256, width=256) + after = current_limits() + assert (seen["limits_at_entry"]["wired"], seen["limits_at_entry"]["memory"]) == (wired, memory) + assert after == start + + +def test_encode_runs_under_the_commands_caps_and_restores_mlxs_defaults(tmp_path, monkeypatch): + # Bug caught: encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured + # phases, sized under the command's caps), or the caps left installed after it returns. + from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap + + model = fake_model(tmp_path, monkeypatch) + seen = [] + monkeypatch.setattr( + model._lifecycle, "ensure_embeddings", lambda *prompts: seen.append(current_limits()) + ) + wired, memory = command_caps() + assert wired > 0 + with mlx_without_wired_cap() as start: + model.encode("p") + after = current_limits() + assert (seen[0]["wired"], seen[0]["memory"]) == (wired, memory) + assert after == start + + +# --- the tiny generation against stock mflux ---------------------------------------------------------------- + + +@pytest.mark.parametrize("nonblock", ["resident", "per-call"]) +@pytest.mark.parametrize( + ("guidance", "negative"), [(1.0, None), (3.5, "n")], ids=["guidance-1", "cfg"] +) +def test_the_tiny_generation_matches_stock_krea2_latents( + tmp_path, monkeypatch, guidance, negative, nonblock +): + # Bug caught: anything between the checkpoint and mflux's loop changing the result: (embeds, neg_embeds) swapped by + # the _encode_prompts override (the negative is 12 tokens, the prompt 1, so a swap changes the latents), a decoded + # matrix on the wrong parameter, a non-block group or rename lost, an extra mis-routed, the CFG formula or the + # er_sde seeded noise stream not mflux's. The reference is stock mflux over the same BF16 weights; mflux's predict + # factory returns the plain closure when the chip check says M1/M2, so it is patched True for both sides (ours + # bypasses compile on any chip). Equality is bit for bit on the float32 latents; the extras are random because + # constant ones blind the RMSNorm-fed forward. The per-call build (the shipped one) releases the non-block weights + # and clears MLX's cache before block 0 of every call: the latents must not move. + from mflux.utils.apple_silicon import AppleSiliconUtil + from tests._krea2_tiny import TINY, StubVAE, stock_krea2, stub_components + + from mlx_dfloat.format import open_checkpoint + from mlx_dfloat.mflux.krea2.transformer import build_transformer, seam_transformer_class + + monkeypatch.setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: True)) + rng = np.random.default_rng(11) + matrices, extras, layout = write_tiny_checkpoint(tmp_path / "random", rng, random_extras=rng) + ckpt = open_checkpoint(tmp_path / "random", layouts=(layout,)) + tokenizer = StubTokenizer({"n": 12}) + ours_vae, stock_vae = StubVAE(), StubVAE() + components = stub_components(tokenizer) + components.vae, components.text_encoder = ours_vae, StubTextEncoder() + model = fake_model( + tmp_path, + monkeypatch, + model="krea-2-raw", + nonblock=nonblock, + ckpt=ckpt, + build=build_transformer( + ckpt, + transformer_kwargs=TINY, + transformer_class=seam_transformer_class(per_call=nonblock == "per-call"), + ), + components=components, + ) + stock = stock_krea2(tokenizer, StubTextEncoder(), stock_vae, matrices, extras) + kwargs = { + "seed": 3, + "prompt": "p", + "num_inference_steps": 2, + "height": 64, + "width": 64, + "guidance": guidance, + "negative_prompt": negative, + "scheduler": "er_sde", + } + model.generate_image(**kwargs) + stock.generate_image(**kwargs) + (ours,), (theirs,) = ours_vae.seen, stock_vae.seen + assert model.report()["cfg_calls_per_step"] == (2 if negative else 1) + assert model.report()["predict"] == "uncompiled" + assert bool(mx.all(mx.isfinite(theirs))) + assert ours.dtype == theirs.dtype == mx.float32 + assert ours.shape == theirs.shape == (1, 16, 8, 8) + assert np.array_equal(np.array(ours.view(mx.uint32)), np.array(theirs.view(mx.uint32))) diff --git a/tests/test_krea2_names.py b/tests/test_krea2_names.py new file mode 100644 index 0000000..1610a5e --- /dev/null +++ b/tests/test_krea2_names.py @@ -0,0 +1,347 @@ +"""Krea 2 naming: the map derived from mflux's weight mapping, the DF11 group check and the variant check.""" + +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +from tests._df11_fixtures import random_bf16, write_checkpoint + +from mlx_dfloat._layouts import KREA_2_RAW_COMFYUI +from mlx_dfloat.errors import DFloatFormatError +from mlx_dfloat.format import matrix_names_for, open_checkpoint +from mlx_dfloat.mflux._mapping import name_map_from_mapping +from mlx_dfloat.mflux.krea2.names import ( + KIND, + LAYOUT_FOR_MODEL, + MATRIX_SUBS, + NONBLOCK_GROUPS, + TEXT_FUSION_GROUPS, + check_krea2_groups, + check_variant, +) + +# The stored order a byte comparison with the originals confirmed in both published files (registration order of +# mflux's modules; see the Krea 2 layout comment in mlx_dfloat/_layouts.py). +EIGHT = ("attn.wq", "attn.wk", "attn.wv", "attn.gate", "attn.wo", "mlp.gate", "mlp.up", "mlp.down") + + +def _t(to, frm): + return SimpleNamespace(to_pattern=to, from_pattern=[frm], transform=None) + + +def _ids(*keys): + return [_t(k, k) for k in keys] + + +def _attn(p): + return _ids(*(f"{p}.attn.{s}.weight" for s in ("wq", "wk", "wv", "wo", "gate"))) + _ids( + f"{p}.attn.qknorm.qnorm.scale", f"{p}.attn.qknorm.knorm.scale" + ) + + +def _mlp(p): + return _ids(f"{p}.mlp.gate.weight", f"{p}.mlp.up.weight", f"{p}.mlp.down.weight") + + +def _norms(p): + return _ids(f"{p}.prenorm.scale", f"{p}.postnorm.scale") + + +def _mapping(): + # A copy of Krea2WeightMapping.get_transformer_mapping() (mflux 0.20.0 krea2_weight_mapping.py:41-95), in its + # order: patch embed; the {layer}-templated block targets; the six timestep renames; the text fusion by concrete + # index (two layerwise blocks, the projector, two refiner blocks); the five text-MLP renames; the final layer. + # test_the_real_mapping_matches_the_copy_and_every_matrix_lands_on_a_linear compares it with mflux's. + targets = _ids("first.weight", "first.bias") + b = "blocks.{layer}" + targets += _norms(b) + _attn(b) + _mlp(b) + _ids(f"{b}.mod.lin") + targets += [ + _t("tmlp.linear_in.weight", "tmlp.0.weight"), + _t("tmlp.linear_in.bias", "tmlp.0.bias"), + _t("tmlp.linear_out.weight", "tmlp.2.weight"), + _t("tmlp.linear_out.bias", "tmlp.2.bias"), + _t("tproj.linear.weight", "tproj.1.weight"), + _t("tproj.linear.bias", "tproj.1.bias"), + ] + for i in range(2): + p = f"txtfusion.layerwise_blocks.{i}" + targets += _norms(p) + _attn(p) + _mlp(p) + targets += _ids("txtfusion.projector.weight") + for i in range(2): + p = f"txtfusion.refiner_blocks.{i}" + targets += _norms(p) + _attn(p) + _mlp(p) + targets += [ + _t("txtmlp.norm.scale", "txtmlp.0.scale"), + _t("txtmlp.linear_in.weight", "txtmlp.1.weight"), + _t("txtmlp.linear_in.bias", "txtmlp.1.bias"), + _t("txtmlp.linear_out.weight", "txtmlp.3.weight"), + _t("txtmlp.linear_out.bias", "txtmlp.3.bias"), + ] + targets += _ids( + "last.norm.scale", "last.linear.weight", "last.linear.bias", "last.modulation.lin" + ) + return targets + + +def test_the_map_places_every_block_matrix_on_its_own_path_in_decoded_order(): + # Bug caught: attn.gate and attn.wo swapped against the checkpoint (equal shapes: the block runs, the image is + # noise), or a block sub-path renamed on its way to the module. + m = name_map_from_mapping(_mapping(), matrix_subs=MATRIX_SUBS) + assert m.attrs_of("blocks") == EIGHT + assert m.place("blocks.27.mlp.down.weight").attr == "mlp.down" + assert m.place("blocks.27.mlp.down.weight").block == "blocks.27" + assert m.param_name("blocks.3.mod.lin") == "blocks.3.mod.lin" + + +def test_the_nonblock_renames_are_mfluxs(): + # Bug caught: a rename lost (tmlp.0.weight has no parameter of that name in mflux: an extras-coverage refusal at + # build), or a text-fusion name swallowed by the block table. + m = name_map_from_mapping(_mapping(), matrix_subs=MATRIX_SUBS) + assert m.param_name("tmlp.0.weight") == "tmlp.linear_in.weight" + assert m.param_name("tmlp.2.bias") == "tmlp.linear_out.bias" + assert m.param_name("tproj.1.weight") == "tproj.linear.weight" + assert m.param_name("txtmlp.0.scale") == "txtmlp.norm.scale" + assert m.param_name("txtmlp.3.weight") == "txtmlp.linear_out.weight" + name = "txtfusion.refiner_blocks.1.attn.wq.weight" + assert m.param_name(name) == name + + +def test_the_tables_are_the_published_layouts_matrix_names(): + # Bug caught: the adapter's order or group names drifting from the layout's (the group check would then refuse + # both published files, or install a matrix under another's name). + pattern = KREA_2_RAW_COMFYUI.raw_config["pattern_dict"] + assert matrix_names_for("blocks.9", pattern) == tuple( + f"blocks.9.{s}.weight" for s in MATRIX_SUBS[KIND] + ) + assert MATRIX_SUBS == {"blocks": EIGHT} + assert TEXT_FUSION_GROUPS == ( + "txtfusion.layerwise_blocks.0", + "txtfusion.layerwise_blocks.1", + "txtfusion.refiner_blocks.0", + "txtfusion.refiner_blocks.1", + ) + assert { + "tmlp": ("tmlp.0.weight", "tmlp.2.weight"), + "tproj": ("tproj.1.weight",), + "txtmlp": ("txtmlp.1.weight", "txtmlp.3.weight"), + **{g: tuple(f"{g}.{s}.weight" for s in EIGHT) for g in TEXT_FUSION_GROUPS}, + } == NONBLOCK_GROUPS + for group, matrices in NONBLOCK_GROUPS.items(): + assert matrix_names_for(group, pattern) == matrices + + +# --- the group check, on written tiny checkpoints (the published layout's pattern_dict, 4 x 4 matrices) ------------ + +_PATTERNS = KREA_2_RAW_COMFYUI.raw_config["pattern_dict"] + + +def _write(root, *, blocks=(0, 1), nonblock=True, extra_groups=(), subs=EIGHT): + rng = np.random.default_rng(3) + groups = {f"blocks.{b}": [random_bf16(rng, (4, 4)) for _ in subs] for b in blocks} + patterns = {r"blocks\.\d+": subs} + if nonblock: + for name, matrices in NONBLOCK_GROUPS.items(): + groups[name] = [random_bf16(rng, (4, 4)) for _ in matrices] + patterns.update({k: v for k, v in _PATTERNS.items() if k != r"blocks\.\d+"}) + for name in extra_groups: + groups[name] = [random_bf16(rng, (4, 4))] + patterns[name.replace(".", r"\.")] = () + write_checkpoint(root, groups=groups, patterns=patterns) + return open_checkpoint(root) + + +@pytest.mark.parametrize("nonblock", [False, True]) +def test_the_block_count_is_returned_with_and_without_the_nonblock_groups(tmp_path, nonblock): + # Bug caught: an optional non-block group (a text-fusion block is named like a block list) counted as a block or + # refused, or the count off by one. + assert check_krea2_groups(_write(tmp_path, nonblock=nonblock)) == {"blocks": 2} + + +def test_a_hole_in_the_block_indices_is_refused(tmp_path): + # Bug caught: blocks.0 and blocks.2 accepted, blocks.1 left on placeholders (a zero-size matmul at run time). + with pytest.raises(DFloatFormatError, match=r"^blocks groups are not contiguous: \[0, 2\]$"): + check_krea2_groups(_write(tmp_path, blocks=(0, 2))) + + +def test_a_zero_padded_block_index_is_refused(tmp_path): + # Bug caught: blocks.01 counted as block 1 (int("01") == 1); the seam then asks for blocks.1, a group the + # checkpoint does not have, after the encoder and the set have loaded. + with pytest.raises(DFloatFormatError, match=r"^blocks\.01: not a Krea 2 group"): + check_krea2_groups(_write(tmp_path, blocks=("0", "01"))) + + +def test_a_checkpoint_without_blocks_is_refused(tmp_path): + # Bug caught: a non-block-only checkpoint accepted as zero blocks. + with pytest.raises(DFloatFormatError, match=r"^no blocks groups in the checkpoint$"): + check_krea2_groups(_write(tmp_path, blocks=())) + + +def test_a_stray_group_is_refused_naming_the_accepted_nonblock_groups(tmp_path): + # Bug caught: an unknown group (another model's) silently ignored, its matrices never placed. + with pytest.raises(DFloatFormatError) as err: + check_krea2_groups(_write(tmp_path, extra_groups=("cap_embedder",))) + assert str(err.value) == ( + "cap_embedder: not a Krea 2 group (blocks: blocks.; non-block groups accepted: tmlp, tproj, txtmlp, " + "txtfusion.layerwise_blocks.0, txtfusion.layerwise_blocks.1, txtfusion.refiner_blocks.0, " + "txtfusion.refiner_blocks.1)" + ) + + +def test_a_tmlp_group_with_one_matrix_is_refused(tmp_path): + # Bug caught: a non-block group of the wrong matrices accepted; tmlp.linear_out never installed (a zero-size + # placeholder in the timestep MLP before the first block). + rng = np.random.default_rng(5) + write_checkpoint( + tmp_path, + groups={ + "blocks.0": [random_bf16(rng, (4, 4)) for _ in EIGHT], + "tmlp": [random_bf16(rng, (4, 4))], + }, + patterns={r"blocks\.\d+": EIGHT, "tmlp": ("0",)}, + ) + with pytest.raises(DFloatFormatError, match=r"^tmlp: holds \['tmlp\.0\.weight'\], expected"): + check_krea2_groups(open_checkpoint(tmp_path)) + + +def test_a_block_of_other_matrices_is_refused_naming_the_eight(tmp_path): + # Bug caught: a block whose stored matrices are not Krea's (another model's group under the same name) accepted + # and failing with a shape error at the first block instead of at build. + seven = EIGHT[:-1] + with pytest.raises(DFloatFormatError) as err: + check_krea2_groups(_write(tmp_path, nonblock=False, subs=seven)) + assert str(err.value) == ( + "blocks.0: holds 7 matrices " + + str([f"blocks.0.{s}.weight" for s in seven]) + + "; a Krea 2 block holds attn.wq, attn.wk, attn.wv, attn.gate, attn.wo, mlp.gate, mlp.up, mlp.down" + ) + + +@pytest.mark.mflux +def test_a_qwen_image_21_config_less_checkpoint_is_refused_by_the_group_check(tmp_path): + # Bug caught: another family's config-less file accepted for Krea 2 and failing with a shape + # error at the first block (or silently decoding into the wrong modules) instead of at build. + from tests._qwen21_tiny import write_tiny_checkpoint + + _m, _c, layout = write_tiny_checkpoint( + tmp_path, np.random.default_rng(2), modulation_as_group=False + ) + ckpt = open_checkpoint(tmp_path, layouts=(layout,)) + with pytest.raises( + DFloatFormatError, + match=r"^transformer_blocks\.0: not a Krea 2 group \(blocks: blocks\.;", + ): + check_krea2_groups(ckpt) + + +def _ckpt(source): + return SimpleNamespace(root=Path("snap"), config_source=source) + + +def test_the_turbo_file_is_refused_for_raw_and_the_raw_file_for_turbo(): + # Bug caught: Raw's 25-step default (or the card's 52 steps at 3.5) run silently on Turbo's + # distilled weights, and a README row labelled Raw measuring Turbo; or the reverse. + turbo = _ckpt( + "header and spot checks match layout krea-2-turbo-comfyui (mingyi456/Krea-2-Turbo-DF11-ComfyUI@978d)" + ) + raw = _ckpt( + "header and spot checks match layout krea-2-raw-comfyui (mingyi456/Krea-2-Raw-DF11-ComfyUI@8320)" + ) + with pytest.raises(DFloatFormatError) as err: + check_variant("krea-2-raw", turbo) + assert str(err.value) == ( + "snap: this is the Krea 2 Turbo DF11 checkpoint (layout krea-2-turbo-comfyui); run it with --model krea-2" + ) + with pytest.raises(DFloatFormatError) as err: + check_variant("krea-2", raw) + assert str(err.value) == ( + "snap: this is the Krea 2 Raw DF11 checkpoint (layout krea-2-raw-comfyui); run it with --model krea-2-raw" + ) + check_variant("krea-2-raw", raw) + check_variant("krea-2", turbo) + + +@pytest.mark.parametrize("model", ["krea-2-raw", "krea-2"]) +@pytest.mark.parametrize( + "source", + [ + "config.json", + "header and spot checks match layout qwen-image-2.1-comfyui (mingyi456/Qwen-Image-2.1-DF11-ComfyUI@1b22)", + "header and spot checks match layout tiny (t/t@0000)", + ], +) +def test_a_checkpoint_of_neither_krea_layout_passes_the_variant_check(model, source): + # Bug caught: a user's own config.json export (or a test's pinned tiny file) refused as "the other model", or any + # layout but the model's own read as the other Krea model's. + check_variant(model, _ckpt(source)) + + +def test_each_model_names_its_own_layout(): + # Bug caught: the two keys swapped (every published file refused for its own model) or a layout key that no + # layout carries. + from mlx_dfloat._layouts import KREA_2_TURBO_COMFYUI + + assert { + "krea-2-raw": KREA_2_RAW_COMFYUI.key, + "krea-2": KREA_2_TURBO_COMFYUI.key, + } == LAYOUT_FOR_MODEL + assert LAYOUT_FOR_MODEL == { + "krea-2-raw": "krea-2-raw-comfyui", + "krea-2": "krea-2-turbo-comfyui", + } + + +@pytest.mark.mflux +def test_the_real_mapping_matches_the_copy_and_every_matrix_lands_on_a_linear(): + # Bug caught: mflux's real mapping no longer matching the copy (a rename or transform added upstream), or a + # decoded matrix with no nn.Linear to land on in a real tiny Krea2Transformer. + import mlx.nn as nn + from mflux.models.krea2.model.krea2_transformer.transformer import Krea2Transformer + from mflux.models.krea2.weights.krea2_weight_mapping import Krea2WeightMapping + from tests._krea2_tiny import TINY + + from mlx_dfloat.integrate.placeholders import get_attr_path + from mlx_dfloat.mflux.krea2.names import krea2_name_map + + real = [ + (t.to_pattern, list(t.from_pattern), t.transform is not None) + for t in Krea2WeightMapping.get_transformer_mapping() + ] + copy = [(t.to_pattern, list(t.from_pattern), t.transform is not None) for t in _mapping()] + assert sorted(real) == sorted(copy) + m = krea2_name_map() + assert m.attrs_of(KIND) == EIGHT + assert m.param_name("tproj.1.weight") == "tproj.linear.weight" + tf = Krea2Transformer(**TINY) + for attr in m.attrs_of(KIND): + assert isinstance(get_attr_path(tf.blocks[0], attr), nn.Linear), attr + for matrices in NONBLOCK_GROUPS.values(): + for name in matrices: + module = get_attr_path(tf, m.param_name(name).removesuffix(".weight")) + assert isinstance(module, nn.Linear), name + + +@pytest.mark.mflux +def test_the_default_layer_count_is_mfluxs_constructor_default(): + # Bug caught: an mflux bump changing Krea2Transformer's default depth while the count check keeps 28: every real + # checkpoint refused, or one of another depth accepted. + import inspect + + from mflux.models.krea2.model.krea2_transformer.transformer import Krea2Transformer + + from mlx_dfloat.mflux.krea2.names import DEFAULT_LAYERS + + assert inspect.signature(Krea2Transformer.__init__).parameters["layers"].default == 28 + assert DEFAULT_LAYERS == 28 + + +@pytest.mark.parametrize("module", ["mlx_dfloat.mflux.krea2", "mlx_dfloat.mflux.krea2.names"]) +def test_importing_the_adapter_does_not_import_mflux(module): + # Bug caught: a module-level mflux import in the adapter (a user without the extra gets a bare ImportError on + # import instead of the dependency error naming it, and every import pays for mflux). + import subprocess + import sys + + code = f"import sys, {module}; print('mflux' in sys.modules)" + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert out.stdout.strip() == "False" diff --git a/tests/test_krea2_real_blocks.py b/tests/test_krea2_real_blocks.py new file mode 100644 index 0000000..30464d4 --- /dev/null +++ b/tests/test_krea2_real_blocks.py @@ -0,0 +1,183 @@ +"""The real Krea 2 Raw DF11 checkpoint: its layout and groups, Metal decode against the CPU reference, a one-block build. + +Env-gated: every test skips unless ``MLX_DFLOAT_KREA_DF11`` names the ``mingyi456/Krea-2-Raw-DF11-ComfyUI`` snapshot +(one config-less file, ``krea2_raw_bf16-DF11.safetensors``); the one-block comparison also needs +``MLX_DFLOAT_KREA_RAW_FILE``, a local ``raw.safetensors`` of ``krea/Krea-2-Raw`` (26 GB, kept only for the image +identity check). Expected counts are literals from the checkpoint's header (35 groups: ``blocks.0..27``, four +text-fusion blocks, ``tmlp``, ``tproj``, ``txtmlp``; 169 BF16 extras; read 2026-10-08) and from mflux 0.20.0 +(``Krea2Transformer`` builds 28 blocks, transformer.py:21); the expected bits come from the CPU reference decoder or the +base's tensors, never from the code under test. Krea 2 Turbo's file shares the layout and is not tested here. +""" + +import gc +import os +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import mlx.core as mx +import numpy as np +import pytest + +from mlx_dfloat.decode import check, decode_group, split_matrices +from mlx_dfloat.format import DF11Checkpoint, load_group_mx, open_checkpoint +from mlx_dfloat.mflux.krea2.names import KIND, NONBLOCK_GROUPS, check_krea2_groups + +DF11_ENV, RAW_ENV = "MLX_DFLOAT_KREA_DF11", "MLX_DFLOAT_KREA_RAW_FILE" +# The split positions, by hand from the matrix shapes in the published header: a block stores wq 6144^2, wk and +# wv 1536 x 6144, gate and wo 6144^2, mlp.gate and mlp.up 16384 x 6144, mlp.down 6144 x 16384 (434_110_464 in all); a +# text-fusion block five 2560^2 and three 6912 x 2560 (85_852_160); tmlp 256 x 6144 then 6144^2; tproj one 36864 x +# 6144; txtmlp 2560 x 6144 then 6144^2. +SPLITS = { + "blocks": (37748736, 47185920, 56623104, 94371840, 132120576, 232783872, 333447168), + "txtfusion": (6553600, 13107200, 19660800, 26214400, 32768000, 50462720, 68157440), + "tmlp": (1572864,), + "tproj": (), + "txtmlp": (15728640,), +} + + +def _env(name: str) -> Path: + value = os.environ.get(name) + if not value: + pytest.skip(f"{name} not set") + return Path(value) + + +def _bits(array: mx.array) -> np.ndarray: + """The array's raw bits through a dtype-sized unsigned view (equality of bits, never of floats).""" + view = {2: mx.uint16, 4: mx.uint32}[array.dtype.size] + return np.array(array.view(view)) + + +@pytest.mark.slow +def test_the_real_checkpoint_opens_through_its_layout() -> None: + # Bug caught: the published file not matched by its pinned layout (every user's first call refused), a group or + # extra the adapter does not expect (a group count other than the header's 35, a non-block group missing, an + # extra count other than 169), or a split table other than the shapes give. + ckpt = open_checkpoint(_env(DF11_ENV)) + assert ckpt.config_source.startswith( + "header and spot checks match layout krea-2-raw-comfyui " + "(mingyi456/Krea-2-Raw-DF11-ComfyUI@8320616b25ac9340a830a7fb21f1b0237e160e66)" + ) + assert check_krea2_groups(ckpt) == {KIND: 28} + assert len(ckpt.groups) == 35 + assert sorted(g for g in ckpt.groups if not g.startswith("blocks.")) == sorted(NONBLOCK_GROUPS) + assert len(ckpt.extras) == 169 + assert sorted({info.dtype for _path, info in ckpt.extras.values()}) == ["BF16"] + for name in ( + "blocks.0", + "blocks.27", + "txtfusion.layerwise_blocks.0", + "tmlp", + "tproj", + "txtmlp", + ): + key = name.split(".")[0] + assert load_group_mx(ckpt.groups[name]).split_positions == SPLITS[key], name + + +@pytest.mark.slow +@pytest.mark.metal +def test_four_real_groups_decode_on_metal_equal_to_the_cpu_reference() -> None: + # Bug caught: the Metal decoder differing from the reference on this file's codecs (the first and last block, a + # text-fusion block, the largest non-block group tproj at 226_492_416 elements), or a group cut into matrices at + # other points than its split table. The reference decoder is the oracle; equality is by bits. + ckpt = open_checkpoint(_env(DF11_ENV)) + for name in ("blocks.0", "blocks.27", "txtfusion.refiner_blocks.1", "tproj"): + group = ckpt.groups[name] + loaded = load_group_mx(group) + metal = decode_group(loaded, backend="metal") + check(metal, name=name) + reference = decode_group(loaded, backend="reference") + check(reference, name=name) + assert np.array_equal(_bits(metal.bits), _bits(reference.bits)), name + parts = split_matrices(metal.bits, loaded.split_positions) + assert len(parts) == len(group.matrix_names), name + del loaded, metal, reference, parts + gc.collect() + mx.clear_cache() + + +def _inputs(*, seed: int = 1) -> dict[str, Any]: + """One call's inputs at 512²: float32 latents (1, 16, 64, 64) as create_noise makes them, 16 text tokens of the + 30720-wide context (12 taps x 2560), bf16 as the text encoder returns them, one sigma.""" + keys = mx.random.split(mx.random.key(seed), 2) + return { + "latents": mx.random.normal((1, 16, 64, 64), key=keys[0]), + "timestep": mx.array([0.5]), + "embeds": mx.random.normal((1, 16, 30720), key=keys[1]).astype(mx.bfloat16), + } + + +def _run(transformer: Any, inputs: Mapping[str, Any]) -> mx.array: + out = transformer(inputs["latents"], inputs["timestep"], inputs["embeds"]) + mx.eval(out) + return out + + +def df11_one_block(ckpt: DF11Checkpoint, inputs: Mapping[str, Any]) -> mx.array: + """Build one block the way the model does (non-block groups decoded per call), attach the DF11 decode, run.""" + from mlx_dfloat.integrate.coverage import load_resident_set + from mlx_dfloat.integrate.providers import DF11Provider + from mlx_dfloat.mflux.krea2.names import krea2_name_map + from mlx_dfloat.mflux.krea2.transformer import build_transformer, seam_transformer_class + + names = krea2_name_map() + build = build_transformer( + ckpt, + model="krea-2-raw", + n_layers=1, + transformer_class=seam_transformer_class(per_call=True), + ) + blocks = list(build.shapes) + assert blocks == ["blocks.0"] + groups = load_resident_set(ckpt, [*blocks, *NONBLOCK_GROUPS]) + matrix_names = {n: ckpt.groups[n].matrix_names for n in groups} + build.transformer.bind_nonblock({g: groups[g] for g in NONBLOCK_GROUPS}, build.nonblock, names) + provider = DF11Provider( + {n: groups[n] for n in blocks}, {n: matrix_names[n] for n in blocks}, names + ) + build.transformer.attach(provider, build.shapes, verify_in_call=True) + return _run(build.transformer, inputs) + + +def bf16_one_block(ckpt: DF11Checkpoint, raw_file: Path, inputs: Mapping[str, Any]) -> mx.array: + """The same one-block build reading nothing of the DF11 file but its shapes: every extra and every non-block + matrix replaced by the base's (FP32 cast to BF16 as mflux casts at load), block 0 streamed from the base.""" + from mlx_dfloat._safetensors import read_header + from mlx_dfloat.integrate.providers import StreamingBF16Provider + from mlx_dfloat.mflux.krea2.names import krea2_name_map + from mlx_dfloat.mflux.krea2.transformer import build_transformer, install_base_weights + + names = krea2_name_map() + build = build_transformer(ckpt, model="krea-2-raw", n_layers=1) + replaced = install_base_weights(build.transformer, ckpt, raw_file, names, build.nonblock) + # The built block's five extras, the 49 non-block extras and the 37 non-block matrices: 169 - 27 x 5 + 37. + assert len(replaced) == 169 - 27 * 5 + 37 + index = {name: (raw_file, info) for name, info in read_header(raw_file).items()} + provider = StreamingBF16Provider( + index, {n: ckpt.groups[n].matrix_names for n in build.shapes}, names + ) + build.transformer.attach(provider, build.shapes, verify_in_call=True) + return _run(build.transformer, inputs) + + +@pytest.mark.slow +@pytest.mark.metal +@pytest.mark.mflux +def test_a_one_block_build_matches_the_bf16_stream_bit_for_bit() -> None: + # Bug caught: a non-block rename or a block matrix misplaced on one side (tmlp.0 -> tmlp.linear_in, a text-fusion + # matrix on another's parameter, a non-block group left on its placeholder: a zero-size matmul before the first + # block), the per-call weights released before the forward starts, or any seam difference between the DF11 and + # the streamed BF16 paths (the decoder itself is the test above's). The five FP32 base matrices are + # BF16-representable, so the cast cannot show a rounding difference here. The output is float32 (the residual stream follows the float32 latents). Needs the 26 GB + # raw.safetensors of krea/Krea-2-Raw @ 6b0ece7f (MLX_DFLOAT_KREA_RAW_FILE). + ckpt = open_checkpoint(_env(DF11_ENV)) + raw_file = _env(RAW_ENV) + inputs = _inputs() + out_a = df11_one_block(ckpt, inputs) + out_b = bf16_one_block(ckpt, raw_file, inputs) + assert out_a.shape == out_b.shape == (1, 16, 64, 64) + assert out_a.dtype == out_b.dtype == mx.float32 + assert np.array_equal(_bits(out_a), _bits(out_b)) diff --git a/tests/test_krea2_tiny.py b/tests/test_krea2_tiny.py new file mode 100644 index 0000000..8876fff --- /dev/null +++ b/tests/test_krea2_tiny.py @@ -0,0 +1,65 @@ +"""One real tiny Krea 2 generation through mflux's own loop: our prelude, the uncompiled predict, the per-prompt cache, +the seam, the VAE guard and mflux's image assembly composing. `@pytest.mark.mflux`.""" + +import mlx.core as mx +import numpy as np +import pytest +from tests._krea2_tiny import fake_model + +pytestmark = pytest.mark.mflux + + +def _cache_limit_in_force(): + previous = mx.set_cache_limit(0) + mx.set_cache_limit(previous) + return previous + + +@pytest.mark.parametrize( + ("model_name", "guidance", "calls"), + [("krea-2", 3.5, 2), ("krea-2", 1.0, 1), ("krea-2", 0.5, 2), ("krea-2-raw", None, 1)], + ids=["cfg", "guidance-1", "guidance-half", "raw-default"], +) +def test_one_tiny_generation_decodes_once_per_block_per_transformer_call( + tmp_path, monkeypatch, model_name, guidance, calls +): + # Bug caught: anything between our prelude and mflux's loop not composing (the uncompiled predict, the per-prompt + # cache answering mflux's _encode_prompts, the VAE guard), the negative branch skipped below guidance 1.0 (mflux + # runs CFG for any guidance other than 1.0), or Raw's default running CFG. Tiny transformer: 2 + # blocks; 2 steps: launches == 2 steps x calls x 2 blocks == 8 with CFG, 4 without. The chip check says "not + # M1/M2", so stock mflux would compile the predict here. + from mflux.utils.apple_silicon import AppleSiliconUtil + + monkeypatch.setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: False)) + model = fake_model(tmp_path, monkeypatch, model=model_name) + seen = {} + + class Probe: + def call_in_loop(self, t, seed, prompt, latents, config, time_steps): + seen.setdefault("limits", []).append(_cache_limit_in_force()) + + def call_after_loop(self, seed, prompt, latents, config): + seen["latents"] = np.array(latents) + + model.callbacks.register(Probe()) + before = _cache_limit_in_force() + image = model.generate_image( + seed=1, prompt="a", num_inference_steps=2, height=64, width=64, guidance=guidance + ) + assert np.asarray(image.image).shape == (64, 64, 3) + # The denoised latents: 16 channels at 64 / 8 = 8 x 8, float32 (krea2_latent_creator.py:6-10), finite. + assert seen["latents"].shape == (1, 16, 8, 8) + assert seen["latents"].dtype == np.float32 + assert np.isfinite(seen["latents"]).all() + report = model.report() + assert report["decode_launches"] == 2 * calls * 2 + assert report["cfg_calls_per_step"] == calls + assert report["predict"] == "uncompiled" + assert model.prompt_cache == {} + # The tiny block group (16_384 elements, 32_768 B) next to the allowance's floor (500_000_000 B: 16 image tokens + # and a few text tokens are far below the reference). + assert report["cache_limit_in_force"] == 32_768 + 500_000_000 + assert seen["limits"] == [report["cache_limit_in_force"]] * 2 # in force in the loop + assert set(report["peaks"]) >= {"encode", "set_load", "denoise", "vae"} + assert report["lifecycle"]["set_loads"] == 1 + assert _cache_limit_in_force() == before diff --git a/tests/test_krea2_transformer.py b/tests/test_krea2_transformer.py new file mode 100644 index 0000000..7a37326 --- /dev/null +++ b/tests/test_krea2_transformer.py @@ -0,0 +1,877 @@ +"""The Krea 2 transformer: the class-swap seam, the compile bypass and the float32 residual, on a real tiny mflux one.""" + +import importlib.util +import types +from types import SimpleNamespace + +import mlx.core as mx +import numpy as np +import pytest +from mlx.utils import tree_flatten, tree_map, tree_unflatten +from tests._krea2_tiny import TINY, tiny_groups, tiny_inputs, tiny_predict, tiny_seamed + +from mlx_dfloat.decode import decode_group +from mlx_dfloat.errors import DFloatDependencyError +from mlx_dfloat.integrate import seam +from mlx_dfloat.integrate.placeholders import get_attr_path +from mlx_dfloat.integrate.providers import DF11Provider, ResidentProvider +from mlx_dfloat.mflux._compile import uncompiled +from mlx_dfloat.mflux.krea2.transformer import block_lists, seam_transformer_class + +COMPILED = type( + mx.compile(lambda x: x) +) # mlx.gc_func, a FunctionType subclass: only `type(f) is` tells them apart +BLOCKS = ["blocks.0", "blocks.1"] + + +@pytest.fixture +def not_m1(monkeypatch): + """mflux's chip check answers "not a base or Pro M1/M2": stock mflux compiles Krea 2's predict there. + + On this M1 Max the unpatched check is already False ("max" in the chip name, apple_silicon.py:15-16); CI's base + M1/M2 runners answer True, where mflux hands back the plain function without our override and a bypass test + would pass vacuously. The fixture keeps every chip on the compiling side. + """ + from mflux.utils.apple_silicon import AppleSiliconUtil + + monkeypatch.setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: False)) + return AppleSiliconUtil + + +def _stock_args(inputs, guidance=3.5): + from mflux.models.krea2.model.krea2_transformer.transformer import Krea2Transformer + + return (Krea2Transformer(**TINY), inputs["embeds"], inputs["neg_embeds"], guidance) + + +def _random_weights(shapes, seed=11): + """Every block matrix its own random bfloat16 values (small, so the tiny forward stays finite).""" + rng = np.random.default_rng(seed) + return { + block: { + attr: mx.array((rng.standard_normal(shape) * 0.05).astype(np.float32)).astype( + mx.bfloat16 + ) + for attr, shape in per.items() + } + for block, per in shapes.items() + } + + +def _constant_weights(shapes): + """Block i's matrices filled with the constant i + 1.""" + return { + block: { + attr: mx.full(shape, float(i + 1), dtype=mx.bfloat16) for attr, shape in per.items() + } + for i, (block, per) in enumerate(shapes.items()) + } + + +def _eager_reference(tf, shapes, weights): + """Stock mflux holding tf's non-block parameters and ``weights`` on its blocks.""" + from mflux.models.krea2.model.krea2_transformer.transformer import Krea2Transformer + + plain = Krea2Transformer(**TINY) + matrices = {f"{b}.{a}.weight" for b, per in shapes.items() for a in per} + plain.update( + tree_unflatten([(k, v) for k, v in tree_flatten(tf.parameters()) if k not in matrices]) + ) + for block, per in weights.items(): + idx = int(block.partition(".")[2]) + for attr, w in per.items(): + get_attr_path(plain.blocks[idx], attr).weight = w + return plain + + +def _call(inputs): + """``Krea2Transformer.__call__`` positional inputs (latents, timestep, context), as mflux's predict passes them.""" + return (inputs["latents"], inputs["timestep"], inputs["embeds"]) + + +def test_block_lists_is_the_one_list_the_forward_runs(): + # Bug caught: another attribute than the list __call__ iterates (transformer.py:89-92: `for block in + # self.blocks`), so no block would run through the seam. + tf = SimpleNamespace(blocks=[1, 2], txtfusion=SimpleNamespace(layerwise_blocks=[3])) + assert block_lists(tf) == [("blocks", [1, 2])] + + +def test_the_seamed_class_needs_the_mflux_extra(monkeypatch): + # Bug caught: the adapter importing mflux without the guard (a bare ImportError instead of the package's own + # dependency error naming the extra), for both classes. The check runs on every call, before the cached class. + real = importlib.util.find_spec + monkeypatch.setattr( + importlib.util, + "find_spec", + lambda name, *a, **k: None if name == "mflux" else real(name, *a, **k), + ) + with pytest.raises(DFloatDependencyError, match=r"mlx-dfloat\[mflux\]"): + seam_transformer_class() + with pytest.raises(DFloatDependencyError, match=r"mlx-dfloat\[mflux\]"): + seam_transformer_class(per_call=True) + + +@pytest.mark.mflux +def test_off_m1_m2_the_stock_factory_compiles_and_on_them_it_does_not(not_m1, monkeypatch): + # Bug caught (the control): the reason for the bypass gone (an mflux change: krea2.py:202-204 no longer compiling + # off base and Pro M1/M2), or a bypass test that passes vacuously on a base M1/M2 CI runner (where the unpatched + # chip check is True). On this M1 Max the unpatched check is already False: stock mflux compiles here. + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + + args = _stock_args(tiny_inputs()) + assert type(Krea2._predict(*args)) is COMPILED + monkeypatch.setattr(not_m1, "is_m1_or_m2", classmethod(lambda cls: True)) + assert type(Krea2._predict(*args)) is types.FunctionType + + +@pytest.mark.mflux +def test_uncompiled_returns_mfluxs_plain_predict_on_any_chip(not_m1): + # Bug caught: the chip override not reaching Krea2._predict (the seam's per-block mx.eval would then run inside + # mx.compile), or leaking past the factory call (every later mflux factory would run uncompiled). + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + + fn = uncompiled(Krea2._predict, *_stock_args(tiny_inputs())) + assert type(fn) is types.FunctionType + assert not_m1.is_m1_or_m2() is False + + +@pytest.mark.mflux +def test_the_compiled_predict_cannot_run_the_seam(): + # Bug caught: none in our code; it pins why the bypass exists (if the compiled predict could run per-block evals, + # the bypass and its cost on Max/Ultra and M3+ chips would be unjustified). + tf, shapes = tiny_seamed() + tf.attach(ResidentProvider(_constant_weights(shapes)), shapes) + with pytest.raises(ValueError, match=r"\[eval\] Attempting to eval"): + mx.eval(tiny_predict(tf, tiny_inputs(), guidance=3.5, compiled=True)) + + +@pytest.mark.mflux +def test_the_swap_keeps_parameter_paths_and_nothing_is_computed(): + # Bug caught: the class swap breaking on SingleStreamBlock or moving parameters under another path; a parameter + # mflux computes at construction added upstream (no checkpoint holds it: every build would be refused by the + # extras coverage), or a stale exemption hiding a parameter no checkpoint fills. + from mflux.models.krea2.model.krea2_transformer.transformer import Krea2Transformer + from mflux.models.krea2.model.krea2_transformer.transformer_block import SingleStreamBlock + from mflux.models.krea2.weights.krea2_weight_mapping import Krea2WeightMapping + + from mlx_dfloat.integrate.blockseam import seam_blocks + from mlx_dfloat.mflux.krea2.transformer import COMPUTED_PARAMS + + tf = seam_transformer_class()(**TINY) + before = sorted(k for k, _v in tree_flatten(tf.parameters())) + assert seam_blocks(block_lists(tf), tf.seam_cell) == ("blocks.0", "blocks.1") + after = sorted(k for k, _v in tree_flatten(tf.parameters())) + assert after == before + assert isinstance(tf.blocks[1], SingleStreamBlock) + + params = {k for k, _v in tree_flatten(Krea2Transformer(**TINY).parameters())} + targets = set() + for t in Krea2WeightMapping.get_transformer_mapping(): + if "{layer}" in t.to_pattern: + targets |= {t.to_pattern.replace("{layer}", str(i)) for i in range(2)} + else: + targets.add(t.to_pattern) + assert frozenset(params - targets) == COMPUTED_PARAMS + assert len(COMPUTED_PARAMS) == 0 # mflux 0.20.0: every parameter is a mapping target + + +@pytest.mark.mflux +def test_the_residual_stream_runs_float32_from_float32_latents(): + # Bug caught (in our model of the world): activation terms calibrated or predicted on a bf16 residual (half the + # true size). With every weight in bf16, as mflux loads the published checkpoint (weight_loader.py:448-449), the + # float32 latents (create_noise, krea2_latent_creator.py:6-10) still make the stream every block sees float32: + # `first` promotes the patches and the bf16 text joins them by concatenation. The control: bf16 latents give a + # bf16 stream, so the latents' dtype is what decides. + from mflux.models.krea2.model.krea2_transformer.transformer import Krea2Transformer + + def run(latents): + tf = Krea2Transformer(**TINY) + tf.update(tree_map(lambda p: p.astype(mx.bfloat16), tf.parameters())) + seen = [] + block = tf.blocks[0] + + def recording(x, *args): + seen.append(x.dtype) + return block(x, *args) + + tf.blocks[0] = recording + inputs = tiny_inputs() + out = tf(latents, inputs["timestep"], inputs["embeds"]) + mx.eval(out) + return seen, out.dtype + + latents = tiny_inputs()["latents"] + assert latents.dtype == mx.float32 + assert run(latents) == ([mx.float32], mx.float32) + assert run(latents.astype(mx.bfloat16)) == ([mx.bfloat16], mx.bfloat16) + + +@pytest.mark.mflux +def test_a_forward_runs_each_block_with_its_own_weights_and_ends_on_placeholders(monkeypatch): + # Bug caught: a block run on another block's weights (blocks.1 skipped or fed blocks.0's matrices: the output + # differs from stock mflux holding the same weights), a block output not evaluated at its boundary, or the + # placeholders not restored afterwards. + tf, shapes = tiny_seamed() + weights = _constant_weights(shapes) + evals = [] + real_eval = seam._eval + + def recording_eval(*args): + evals.append(args[0]) + real_eval(*args) + + monkeypatch.setattr(seam, "_eval", recording_eval) + tf.attach(ResidentProvider(weights), shapes) + inputs = tiny_inputs() + out = tf(*_call(inputs)) + mx.eval(out) + tf.verify_step() + want = _eager_reference(tf, shapes, weights)(*_call(inputs)) + mx.eval(want) + + assert len(evals) == 2 + # Each block returns the joint [6 text | 16 image] sequence at width 32, float32 like the latents. + assert [(tuple(e.shape), e.dtype) for e in evals] == [((1, 22, 32), mx.float32)] * 2 + assert out.shape == (1, 16, 8, 8) + assert bool(mx.isfinite(out).all().item()) + assert out.dtype == want.dtype == mx.float32 + assert mx.array_equal(out.view(mx.uint32), want.view(mx.uint32)).item() + for block, per in shapes.items(): + idx = int(block.partition(".")[2]) + for attr in per: + assert get_attr_path(tf.blocks[idx], attr).weight.size == 0, (block, attr) + + +@pytest.mark.mflux +def test_depth2_drains_the_last_block(monkeypatch): + # Bug caught: end_step not draining the tail (the last block's output never evaluated inside the step), or the + # look-ahead evaluating in the wrong order. Expected (seam._seam_eval + end_step): async r0, async r1, eval r0, + # eval r1. + tf, shapes = tiny_seamed() + events = [] + real_eval, real_async = seam._eval, seam._async_eval + + def recording_eval(*args): + events.append(("eval", args[0])) + real_eval(*args) + + def recording_async(*args): + events.append(("async", args[0])) + real_async(*args) + + monkeypatch.setattr(seam, "_eval", recording_eval) + monkeypatch.setattr(seam, "_async_eval", recording_async) + tf.attach(ResidentProvider(_constant_weights(shapes)), shapes, eval_policy="depth2") + mx.eval(tf(*_call(tiny_inputs()))) + tf.verify_step() + + outputs = [] + for _kind, obj in events: + if not any(obj is seen for seen in outputs): + outputs.append(obj) + order = [(kind, next(i for i, o in enumerate(outputs) if o is obj)) for kind, obj in events] + assert order == [("async", 0), ("async", 1), ("eval", 0), ("eval", 1)] + + +@pytest.mark.mflux +@pytest.mark.parametrize( + ("guidance", "negative", "expected"), + [(3.5, True, BLOCKS + BLOCKS), (1.0, False, BLOCKS)], +) +def test_cfg_runs_two_calls_and_decodes_each_block_once_per_call(guidance, negative, expected): + # Bug caught: a decode cached across the two CFG calls without verification, or the second call skipped. mflux + # runs a CFG step as two batch-1 transformer calls (krea2.py:195-200): 4 decodes per step at 2 blocks; at + # guidance 1.0 with no negative, one call and 2 decodes. + from mlx_dfloat.mflux.krea2.names import krea2_name_map + + tf, shapes = tiny_seamed() + groups, names, _src = tiny_groups(shapes, np.random.default_rng(9)) + calls = [] + + def count(group): + calls.append(group.name) + return decode_group(group, backend="reference") + + tf.attach( + DF11Provider(groups, names, krea2_name_map(), decode=count), shapes, verify_in_call=True + ) + inputs = tiny_inputs() + if not negative: + inputs["neg_embeds"] = None + out = tiny_predict(tf, inputs, guidance=guidance) + mx.eval(out) + assert calls == expected + assert out.shape == (1, 16, 8, 8) + assert bool(mx.isfinite(out).all().item()) + + +@pytest.mark.mflux +def test_a_cfg_step_through_the_seam_gives_stock_mfluxs_prediction_bit_for_bit(monkeypatch): + # Bug caught: the two CFG calls wrong through the seam where one call is right: the second call run on the first + # call's placeholders-restored blocks, the two prompts' contexts (6 and 3 tokens) mixed, or the guidance + # combination fed the wrong call. The reference is mflux's own predict over a stock transformer holding the same + # weights, both uncompiled; equality is bit for bit. + from mflux.models.krea2.variants.txt2img.krea2 import Krea2 + from mflux.utils.apple_silicon import AppleSiliconUtil + + tf, shapes = tiny_seamed() + weights = _random_weights(shapes) + tf.attach(ResidentProvider(weights), shapes, verify_in_call=True) + inputs = tiny_inputs() + ours = tiny_predict(tf, inputs, guidance=3.5) + mx.eval(ours) + plain = _eager_reference(tf, shapes, weights) + monkeypatch.setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: True)) + predict = Krea2._predict(plain, inputs["embeds"], inputs["neg_embeds"], 3.5) + assert type(predict) is types.FunctionType # stock, uncompiled like ours + theirs = predict(latents=inputs["latents"], timestep=inputs["timestep"]) + mx.eval(theirs) + assert ours.shape == theirs.shape == (1, 16, 8, 8) + assert ours.dtype == theirs.dtype == mx.float32 + assert mx.array_equal(ours.view(mx.uint32), theirs.view(mx.uint32)).item() + + +@pytest.mark.mflux +def test_the_krea_schedulers_are_registered_by_the_adapters_own_import(): + # Bug caught: the adapter's import path no longer registering er_sde / euler (mflux 0.20.0 registers them in + # mflux/models/krea2/__init__.py, which any mflux.models.krea2 import runs), so the user's first call fails with + # NotImplementedError when Config resolves the name (krea2.py:78 reads config.scheduler.sigmas). A fresh interpreter: before any Krea import the scheduler is unknown (the + # control), after the adapter builds its transformer class it constructs. + import subprocess + import sys + + code = """ +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig + +def make(): # Config.scheduler is a lazy property (config.py:145-159): reading it is what resolves the name + return Config(model_config=ModelConfig.krea2_raw(), num_inference_steps=2, height=64, width=64, + guidance=3.5, scheduler="er_sde").scheduler + +try: + make() + print("registered-before") +except NotImplementedError: + print("unknown-before") +from mlx_dfloat.mflux.krea2.transformer import seam_transformer_class +seam_transformer_class() +print(type(make()).__name__) +""" + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert out.stdout.split() == ["unknown-before", "Krea2FlowScheduler"] + + +def test_importing_the_transformer_module_does_not_import_mflux(): + # Bug caught: a module-level mflux import in the transformer module (a user without the extra gets a bare + # ImportError on import). + import subprocess + import sys + + code = "import sys, mlx_dfloat.mflux.krea2.transformer; print('mflux' in sys.modules)" + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert out.stdout.strip() == "False" + + +# --- per-call non-block decode (the shipped mode) ------------------------------------------------------------------ + +TPROJ = "tproj.linear.weight" + + +def _per_call_build(tmp_path, *, per_call, n_layers=2, seed=4): + """A tiny config-less checkpoint (random extras, so the forward sees every weight), built resident or per-call; + returns (build, groups, names): the block and non-block DF11 groups loaded, nothing decoded or bound yet.""" + from tests._krea2_tiny import write_tiny_checkpoint + + from mlx_dfloat.format import open_checkpoint + from mlx_dfloat.integrate.coverage import load_resident_set + from mlx_dfloat.mflux.krea2.names import krea2_name_map + from mlx_dfloat.mflux.krea2.transformer import build_transformer + + rng = np.random.default_rng(seed) + root = tmp_path / ("per-call" if per_call else "resident") + _m, _c, layout = write_tiny_checkpoint(root, rng, n_layers=n_layers, random_extras=rng) + ckpt = open_checkpoint(root, layouts=(layout,)) + tf_class = seam_transformer_class(per_call=True) if per_call else None + build = build_transformer( + ckpt, transformer_kwargs={**TINY, "layers": n_layers}, transformer_class=tf_class + ) + return build, load_resident_set(ckpt), krea2_name_map(), ckpt + + +def _split(groups): + from mlx_dfloat.mflux.krea2.names import NONBLOCK_GROUPS + + nonblock = {g: v for g, v in groups.items() if g in NONBLOCK_GROUPS} + blocks = {g: v for g, v in groups.items() if g not in NONBLOCK_GROUPS} + return nonblock, blocks + + +def _reference(): + from functools import partial + + return partial(decode_group, backend="reference") + + +def _attach_blocks(build, blocks, ckpt, names, *, eval_policy="per-block", provider_class=None): + provider = DF11Provider( + blocks, {n: ckpt.groups[n].matrix_names for n in blocks}, names, decode=_reference() + ) + if provider_class is not None: + provider = provider_class(provider) + build.transformer.attach(provider, build.shapes, eval_policy=eval_policy, verify_in_call=True) + return provider + + +@pytest.mark.mflux +def test_per_call_decode_gives_the_resident_paths_output_bit_for_bit(tmp_path): + # Bug caught: the non-block weights released before the forward starts (the pre-block graph, text fusion, tmlp, + # tproj, txtmlp, would read a placeholder: a zero-size matmul error), a non-block weight installed on another + # parameter, or the second CFG call run on cleared weights. The release point itself is the size-recording tests' + # (a release later in the call changes no output: the graph already holds the arrays). One uncompiled predict at + # guidance 3.5 (two calls) on each path. + from mlx_dfloat.integrate.resident import decode_nonblock, install_nonblock + from mlx_dfloat.mflux.krea2.names import NONBLOCK_GROUPS + + outs = [] + for per_call in (False, True): + build, groups, names, ckpt = _per_call_build(tmp_path, per_call=per_call) + nonblock, blocks = _split(groups) + if per_call: + build.transformer.bind_nonblock(nonblock, build.nonblock, names, decode=_reference()) + else: + matrices = {g: NONBLOCK_GROUPS[g] for g in nonblock} + weights = decode_nonblock( + nonblock, matrices, build.nonblock, names, decode=_reference() + ) + install_nonblock(build.transformer, weights) + _attach_blocks(build, blocks, ckpt, names) + out = tiny_predict(build.transformer, tiny_inputs(), guidance=3.5) + mx.eval(out) + outs.append(out) + resident, per_call = outs + assert per_call.dtype == resident.dtype == mx.float32 + assert bool(mx.isfinite(per_call).all().item()) + assert mx.array_equal(per_call.view(mx.uint32), resident.view(mx.uint32)).item() + + +@pytest.mark.mflux +@pytest.mark.parametrize(("guidance", "neg", "per_group"), [(3.5, True, 2), (1.0, False, 1)]) +def test_per_call_decode_runs_once_per_transformer_call(tmp_path, guidance, neg, per_group): + # Bug caught: the decode cached across the CFG pair (the second call reading cleared placeholders: a zero-size + # matmul) or run twice in one call. Two calls per CFG step: each of the seven groups decoded twice; one at 1.0. + from mlx_dfloat.mflux.krea2.names import NONBLOCK_GROUPS + + build, groups, names, ckpt = _per_call_build(tmp_path, per_call=True) + nonblock, blocks = _split(groups) + seen = [] + reference = _reference() + + def counting(group): + seen.append(group.name) + return reference(group) + + build.transformer.bind_nonblock(nonblock, build.nonblock, names, decode=counting) + _attach_blocks(build, blocks, ckpt, names) + inputs = tiny_inputs() + if not neg: + inputs["neg_embeds"] = None + mx.eval(tiny_predict(build.transformer, inputs, guidance=guidance)) + assert sorted(seen) == sorted(list(NONBLOCK_GROUPS) * per_group) + + +# What MLX may still hold in its buffer cache right after a drain: a few tiny buffers (the drain's own one-element +# kick, a scalar) whose free lands after the final clear when the GPU's completion handler runs late, which a CI +# runner's paravirtual GPU showed as 4 to 16 B. The tests' released sets are far larger (the tiny model's activations +# and non-block weights, or a 4 MiB array), so a missed clear or drain still fails these bounds. +TINY_CACHE_BYTES = 1024 + + +class _Recording: + """A provider that records, at each block request, the size of the transformer's tproj weight and the MLX cache, + then delegates.""" + + def __init__(self, inner, transformer, *, raise_at=None, events=None): + self.inner, self.transformer, self.raise_at = inner, transformer, raise_at + self.sizes, self.cache = [], [] + self.events = [] if events is None else events + self.launches, self.launching, self.policies = 0, inner.launching, inner.policies + + def weights_for(self, name, shapes): + self.sizes.append((name, int(self.transformer.tproj.linear.weight.size))) + self.cache.append(int(mx.get_cache_memory())) + self.events.append(("weights", name)) + if name == self.raise_at: + raise RuntimeError(f"stop at {name}") + return self.inner.weights_for(name, shapes) + + def verify(self): + self.inner.verify() + + def reset(self): + self.inner.reset() + + +@pytest.mark.mflux +@pytest.mark.parametrize(("policy", "n_layers"), [("per-block", 2), ("depth2", 3)]) +def test_the_nonblock_weights_are_released_before_block_0_requests_its_weights( + tmp_path, policy, n_layers +): + # Bug caught: the non-block weights released after block 0's decode (at blocks.1, the first design, which + # measured 0.14 GiB saved: the denoise peak is block 0's eval), released never, or released with their buffers + # left in MLX's cache (the footprint keeps them). At every block request tproj (192 x 32 at TINY size) is a + # placeholder, and at block 0's request the cache is empty. + build, groups, names, ckpt = _per_call_build(tmp_path, per_call=True, n_layers=n_layers) + nonblock, blocks = _split(groups) + build.transformer.bind_nonblock(nonblock, build.nonblock, names, decode=_reference()) + recorder = _attach_blocks( + build, + blocks, + ckpt, + names, + eval_policy=policy, + provider_class=lambda inner: _Recording(inner, build.transformer), + ) + inputs = tiny_inputs() + inputs["neg_embeds"] = None + mx.eval(tiny_predict(build.transformer, inputs, guidance=1.0)) + want = [(f"blocks.{i}", 0) for i in range(n_layers)] + assert recorder.sizes == want + assert recorder.cache[0] < TINY_CACHE_BYTES + assert build.transformer.tproj.linear.weight.size == 0 # and after the call + # The attached provider is back after the call: a wrapper left in place would be wrapped again by the next call, + # one level deeper per call (about 100 deep over a 50-step guided run). + assert build.transformer.seam_cell.state.provider is recorder + mx.eval(tiny_predict(build.transformer, inputs, guidance=1.0)) + assert recorder.sizes == want + want # the second call releases at the same point + + +@pytest.mark.mflux +def test_each_call_drains_the_cache_then_block_0s_inputs_are_evaluated_released_and_drained( + tmp_path, monkeypatch +): + # Bug caught: the call-start drain missing or placed after the non-block decode (call 2 then decodes and + # builds its pre-block graph on top of call 1's cached activations: the two-call peak sat at the memory limit), the + # release before block 0's inputs are evaluated (the lazy pre-block graph still holds the decoded arrays, so + # nothing is freed), the drain before the release (the released buffers land in the cache and stay in the + # footprint), or the inputs of a later block evaluated instead. Block 0 is called with (combined, tvec, freqs, + # attention_mask): three arrays and None at guidance 1.0. The decode marks where the call's non-block decode runs. + from mlx_dfloat.mflux.krea2 import transformer as ktf + + events = [] + real_eval, real_clear, real_drain = ktf._eval, ktf.clear_nonblock, ktf._release_and_drain + + def eval_(*arrays): + events.append(("eval", len(arrays))) + real_eval(*arrays) + + def clear(module, shapes): + events.append(("clear", int(module.tproj.linear.weight.size))) + real_clear(module, shapes) + + def drain(): + events.append(("drain",)) + return real_drain() + + reference = _reference() + + def decode(group): + if not events or events[-1] != ("decode",): + events.append(("decode",)) + return reference(group) + + monkeypatch.setattr(ktf, "_eval", eval_) + monkeypatch.setattr(ktf, "clear_nonblock", clear) + monkeypatch.setattr(ktf, "_release_and_drain", drain) + build, groups, names, ckpt = _per_call_build(tmp_path, per_call=True) + nonblock, blocks = _split(groups) + build.transformer.bind_nonblock(nonblock, build.nonblock, names, decode=decode) + _attach_blocks( + build, + blocks, + ckpt, + names, + provider_class=lambda inner: _Recording(inner, build.transformer, events=events), + ) + inputs = tiny_inputs() + inputs["neg_embeds"] = None + events.clear() + out = tiny_predict(build.transformer, inputs, guidance=1.0) + assert events == [ + ("drain",), + ("decode",), + ("eval", 3), + ("clear", 192 * 32), + ("drain",), + ("weights", "blocks.0"), + ("weights", "blocks.1"), + # The call ends with its output evaluated and the placeholders put back again (already placeholders). + ("eval", 1), + ("clear", 0), + ] + assert bool(mx.isfinite(out).all().item()) + + +@pytest.mark.mflux +def test_a_guided_steps_second_call_starts_with_mlxs_cache_empty(tmp_path): + # Bug caught: call 2 of a guided step starting with call 1's freed activations still in MLX's cache, under + # a cache limit large enough to keep them (the planner's 4.5e9 B at 1024²; 1 GiB here). Recorded at each call's + # first non-block decode: call 1 leaves its activations cached (checked), and both calls start from 0 B. + build, groups, names, ckpt = _per_call_build(tmp_path, per_call=True) + nonblock, blocks = _split(groups) + reference = _reference() + at_decode, seen = [], set() + + def decode(group): + if group.name in seen: + seen.clear() + if not seen: + at_decode.append(int(mx.get_cache_memory())) + seen.add(group.name) + return reference(group) + + build.transformer.bind_nonblock(nonblock, build.nonblock, names, decode=decode) + _attach_blocks(build, blocks, ckpt, names) + previous = mx.set_cache_limit(1024**3) + try: + mx.eval(tiny_predict(build.transformer, tiny_inputs(), guidance=3.5)) + left = int(mx.get_cache_memory()) + finally: + mx.set_cache_limit(previous) + assert all(cached < TINY_CACHE_BYTES for cached in at_decode), at_decode + assert ( + left > 0 + ) # the step's last call did leave buffers in the cache: the drain had something to clear + + +class _Drain: + """The drain's inputs, scripted: MLX's cache bytes, a footprint sequence and a clock sequence; records calls.""" + + def __init__(self, monkeypatch, *, cached, footprints, clock): + from mlx_dfloat.mflux.krea2 import transformer as ktf + + self.events = [] + self._fp, self._clock = list(footprints), list(clock) + monkeypatch.setattr(ktf, "_synchronize", lambda: self.events.append("sync")) + monkeypatch.setattr(ktf, "_cache_memory", lambda: cached) + monkeypatch.setattr(ktf, "_clear_cache", lambda: self.events.append("clear")) + monkeypatch.setattr(ktf, "_kick", lambda: self.events.append("kick")) + monkeypatch.setattr(ktf, "_footprint", self._footprint) + monkeypatch.setattr(ktf, "_clock", lambda: self._clock.pop(0)) + monkeypatch.setattr(ktf, "_sleep", lambda s: self.events.append(("sleep", s))) + self.drain = ktf._release_and_drain + + def _footprint(self): + self.events.append("footprint") + return self._fp.pop(0) + + +def test_the_drain_waits_until_the_footprint_has_dropped_by_90_percent_of_the_released_bytes( + monkeypatch, +): + # Bug caught: the drain returning at the clear (the released pages stay in the footprint for ~20 ms to 0.7 s, and + # the next call's 1.33 GB decode stacks on them), waiting for the whole release, or reading the footprint after + # the clear (the drop is measured from before it). Released 1e9 B, so 9e8 B of drop ends the wait: 5e8 and 8.5e8 + # wait, 9e8 stops. + d = _Drain( + monkeypatch, + cached=1_000_000_000, + footprints=[20_000_000_000, 19_500_000_000, 19_150_000_000, 19_100_000_000], + clock=[0.0, 0.01, 0.02], + ) + assert d.drain() is True + assert d.events == [ + "sync", + "footprint", + "clear", + "kick", + "footprint", + ("sleep", 0.0005), + "footprint", + ("sleep", 0.0005), + "footprint", + "sync", + "clear", + ] + + +def test_the_drain_gives_up_at_a_quarter_second(monkeypatch): + # Bug caught: a drain that never returns when the footprint does not fall (another allocation in flight), or a cap + # other than 0.25 s: 0.24 s still waits, 0.26 s stops. + d = _Drain( + monkeypatch, cached=1_000_000_000, footprints=[20_000_000_000] * 4, clock=[0.0, 0.24, 0.26] + ) + assert d.drain() is False + assert d.events == [ + "sync", + "footprint", + "clear", + "kick", + "footprint", + ("sleep", 0.0005), + "footprint", + "sync", + "clear", + ] + + +def test_the_drain_empties_the_cache_of_buffers_whose_free_was_still_pending(): + # Bug caught: the drain clearing the cache before MLX's pending frees have landed in it (a buffer freed after its + # evaluation reaches the cache from the GPU's completion handler, after the Python reference is gone), so the call + # starts with them cached after all. Without the synchronisation first, a freed 4 MiB array was still in the cache + # after the clear in 1980 of 2000 tries on an M1 Max (mlx 0.32.2); with it, in none. 50 tries here. + from mlx_dfloat.mflux.krea2 import transformer as ktf + + left = [] + for _ in range(50): + big = mx.ones((1024, 1024)) + mx.eval(big) + del big + ktf._release_and_drain() + left.append(int(mx.get_cache_memory())) + assert all(cached < TINY_CACHE_BYTES for cached in left), left + + +def test_a_release_under_64_mib_is_cleared_and_kicked_without_a_wait(monkeypatch): + # Bug caught: every block boundary of a small model paying the poll (up to 0.25 s when a few MB cannot show in + # the footprint), the clear and the driver kick skipped for a small cache, or the kick's own buffers left behind: + # they reach the cache once its command completes (16 B on an M1 Max), so the drain waits for that and clears + # again. 64 MiB - 1 B released. + d = _Drain(monkeypatch, cached=64 * 1024**2 - 1, footprints=[], clock=[]) + assert d.drain() is True + assert d.events == ["sync", "clear", "kick", "sync", "clear"] + + +@pytest.mark.mflux +def test_the_none_policy_keeps_the_nonblock_weights_until_the_call_ends(tmp_path): + # Bug caught: a release under "none" (no per-block eval: evaluating block 0's inputs there would force the + # pre-block graph mid-call, which that policy promises not to do). Two blocks, the most "none" allows with a + # launching provider, and verify_step() after the step instead of verify_in_call. + build, groups, names, ckpt = _per_call_build(tmp_path, per_call=True) + nonblock, blocks = _split(groups) + build.transformer.bind_nonblock(nonblock, build.nonblock, names, decode=_reference()) + inner = DF11Provider( + blocks, {n: ckpt.groups[n].matrix_names for n in blocks}, names, decode=_reference() + ) + recorder = _Recording(inner, build.transformer) + build.transformer.attach(recorder, build.shapes, eval_policy="none") + inputs = tiny_inputs() + inputs["neg_embeds"] = None + mx.eval(tiny_predict(build.transformer, inputs, guidance=1.0)) + build.transformer.verify_step() + assert recorder.sizes == [("blocks.0", 192 * 32), ("blocks.1", 192 * 32)] + assert build.transformer.tproj.linear.weight.size == 0 + + +@pytest.mark.mflux +def test_an_exception_mid_call_leaves_placeholders(tmp_path, monkeypatch): + # Bug caught: a call that raises leaving the decoded non-block weights installed (1.23 GiB on the real model kept + # past the call, and a retry would see them). The evaluation of block 0's inputs raises, before the release, so + # only the call's own cleanup can put the placeholders back. + from mlx_dfloat.mflux.krea2 import transformer as ktf + + def boom(*_arrays): + raise RuntimeError("stop before the release") + + build, groups, names, ckpt = _per_call_build(tmp_path, per_call=True) + nonblock, blocks = _split(groups) + build.transformer.bind_nonblock(nonblock, build.nonblock, names, decode=_reference()) + _attach_blocks(build, blocks, ckpt, names) + monkeypatch.setattr(ktf, "_eval", boom) + inputs = tiny_inputs() + inputs["neg_embeds"] = None + with pytest.raises(RuntimeError, match="stop before the release"): + mx.eval(tiny_predict(build.transformer, inputs, guidance=1.0)) + for param in build.nonblock: + assert get_attr_path(build.transformer, param.removesuffix(".weight")).weight.size == 0, ( + param + ) + + +@pytest.mark.mflux +def test_a_per_call_decode_evaluates_the_seven_groups_in_one_eval(tmp_path, monkeypatch): + # Bug caught: the per-call decode paying one synchronisation per non-block group (seven per transformer call, 14 + # per guided step) where one eval of all seven status words and 37 matrices does. Two calls at guidance 3.5. + from mlx_dfloat.integrate import resident + + seen = [] + real = resident._eval + monkeypatch.setattr(resident, "_eval", lambda *a: (seen.append(len(a)), real(*a))) + build, groups, names, ckpt = _per_call_build(tmp_path, per_call=True) + nonblock, blocks = _split(groups) + build.transformer.bind_nonblock(nonblock, build.nonblock, names, decode=_reference()) + _attach_blocks(build, blocks, ckpt, names) + mx.eval(tiny_predict(build.transformer, tiny_inputs(), guidance=3.5)) + assert seen == [7 + 37, 7 + 37] + + +@pytest.mark.mflux +def test_the_bound_groups_live_outside_the_module_tree_and_detach_forgets_them(tmp_path): + # Bug caught: the binding stored as a dict attribute, which nn.Module.__setattr__ routes into the parameter tree + # (the compressed groups would show up as parameters: the extras coverage and load_weights break); or detach + # keeping the compressed non-block groups alive after the set drops (about 0.9 GB on the real model, through the + # VAE phase). A call without a binding is refused naming the method to call. + from mlx_dfloat.errors import DFloatIntegrationError + + build, groups, names, ckpt = _per_call_build(tmp_path, per_call=True) + nonblock, blocks = _split(groups) + before = sorted(k for k, _v in tree_flatten(build.transformer.parameters())) + build.transformer.bind_nonblock(nonblock, build.nonblock, names, decode=_reference()) + assert sorted(k for k, _v in tree_flatten(build.transformer.parameters())) == before + _attach_blocks(build, blocks, ckpt, names) + build.transformer.detach() + assert build.transformer.nonblock_bound() is False + _attach_blocks(build, blocks, ckpt, names) + with pytest.raises(DFloatIntegrationError, match="bind_nonblock"): + tiny_predict(build.transformer, tiny_inputs(), guidance=1.0) + + +@pytest.mark.mflux +def test_the_release_wrapper_carries_the_providers_attributes(tmp_path): + # Bug caught: the wrapper exposing the Protocol's methods but not its attributes (`launches`, `launching`, + # `policies`): the seam and the model's launch count read them through the state's provider during a call. + from mlx_dfloat.mflux.krea2.transformer import ReleaseAtFirstBlock + + build, groups, names, ckpt = _per_call_build(tmp_path, per_call=True) + _nonblock, blocks = _split(groups) + inner = DF11Provider( + blocks, {n: ckpt.groups[n].matrix_names for n in blocks}, names, decode=_reference() + ) + wrapper = ReleaseAtFirstBlock(inner, build.transformer, build.nonblock, policy="per-block") + assert (wrapper.launching, wrapper.policies) == (inner.launching, inner.policies) + wrapper.weights_for("blocks.0", build.shapes["blocks.0"]) + wrapper.verify() + assert wrapper.launches == inner.launches == 1 + # The Protocol's attribute is writable; the count lives in the wrapped provider. + wrapper.launches = 0 + assert inner.launches == 0 + + +@pytest.mark.mflux +def test_a_per_call_step_is_refused_while_the_last_steps_status_words_are_unchecked(tmp_path): + # Bug caught: the release wrapper hiding the provider's `pending` list, so the seam's unchecked-status + # refusal (begin_step reads state.provider.pending) never fires on the per-call class: a caller attached without + # verify_in_call who forgets verify_step() would pile status words up unnoticed. + from mlx_dfloat.errors import DFloatIntegrationError + from mlx_dfloat.mflux.krea2.transformer import ReleaseAtFirstBlock + + build, groups, names, ckpt = _per_call_build(tmp_path, per_call=True) + nonblock, blocks = _split(groups) + build.transformer.bind_nonblock(nonblock, build.nonblock, names, decode=_reference()) + inner = DF11Provider( + blocks, {n: ckpt.groups[n].matrix_names for n in blocks}, names, decode=_reference() + ) + assert ( + ReleaseAtFirstBlock(inner, build.transformer, build.nonblock, policy="per-block").pending + is inner.pending + ) + build.transformer.attach(inner, build.shapes, verify_in_call=False) + inputs = tiny_inputs() + inputs["neg_embeds"] = None + mx.eval(tiny_predict(build.transformer, inputs, guidance=1.0)) + assert len(inner.pending) == 2 # one status word per block, unchecked + with pytest.raises( + DFloatIntegrationError, match="2 decode status words from an earlier step are unchecked" + ): + tiny_predict(build.transformer, inputs, guidance=1.0) + build.transformer.verify_step() + mx.eval(tiny_predict(build.transformer, inputs, guidance=1.0)) diff --git a/tests/test_layouts.py b/tests/test_layouts.py index 7df633d..cff168a 100644 --- a/tests/test_layouts.py +++ b/tests/test_layouts.py @@ -4,7 +4,7 @@ ``qwen_image_2.1_bf16-DF11.safetensors`` in ``mingyi456/Qwen-Image-2.1-DF11-ComfyUI`` at revision ``1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc``, exactly as on disk: the u64 little-endian header length (28,408) followed by the JSON header. It was range-read on 2026-10-08 (file bytes 0..28415) -and is byte-identical to the header saved by the 2026-10-07 layout spike. +and is byte-identical to the header saved by the 2026-10-07 range reads. The spot-check digests were range-read at the same revision on 2026-10-08, 4096 bytes each (data starts at file byte 28,416; ``transformer_blocks.0.sign_mantissa`` at tensor data offset @@ -23,7 +23,7 @@ - block 31 element 0 (``attn.to_q``): 7,730,137,195..7,730,141,290. - ``modulation.1`` element 0: 163,724,295..163,728,390. -The 2026-10-07 spike matched the sign/mantissa bytes at each of block 0's offsets against the BF16 +A 2026-10-07 byte comparison matched the sign/mantissa bytes at each of block 0's offsets against the BF16 base (``Qwen/Qwen-Image-2.1``): 4096 of 4096 equal for the expected tensor, at most 29 for any other. """ @@ -52,7 +52,7 @@ FIXTURE = Path(__file__).parent / "fixtures" / "qwen-image-2.1-df11-header.bin" QWEN_REVISION = "1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc" QWEN_FILE = "qwen_image_2.1_bf16-DF11.safetensors" -QWEN_FILE_SIZE = 9_724_491_411 # the 2026-10-07 spike, and the Hub's file size +QWEN_FILE_SIZE = 9_724_491_411 # the 2026-10-07 range reads, and the Hub's file size QWEN_HEADER_SHA256 = "fa5b70707d9d5e08d1c92e40dcdeedd64f72508f88ca942217c6fd419ff5ba5e" QWEN_FILE_SHA256 = ( "fc07ef38e3609b1a38dc8f6ec585ef9921a4d4b5e6f74b744899e4dea91b4e2e" # the Hub's LFS oid @@ -137,7 +137,7 @@ def test_the_fixture_is_the_published_header_as_on_disk(): assert len(data) == 28_416 assert struct.unpack(" Path: + """The published header at full file size; every data byte reads as zero (a sparse file, ~36 KB on disk).""" + root = tmp_path / path.stem + root.mkdir() + with (root / "model.safetensors").open("wb") as fh: + fh.write(path.read_bytes()) + fh.truncate(size) + return root + + +@pytest.mark.parametrize(("layout", "path", "size", "_probes"), LAYOUTS) +def test_the_published_header_opens_from_the_default_table_when_the_content_matches( + tmp_path, monkeypatch, layout, path, size, _probes +): + # Bug caught (end to end through open_checkpoint and the default table): a Krea layout missing from + # KNOWN_LAYOUTS, or its counts, names or source wrong. In the table, the layout's probes are swapped for the + # digest of a zero page because the sparse file reads as zeros; a layout absent from the table swaps nothing, and + # the open then fails on the real probes. + from mlx_dfloat import _layouts + + zeros = tuple(dataclasses.replace(p, sha256=ZERO_PAGE_SHA256) for p in layout.probes) + table = tuple( + dataclasses.replace(known, probes=zeros) if known.key == layout.key else known + for known in _layouts.KNOWN_LAYOUTS + ) + monkeypatch.setattr(_layouts, "KNOWN_LAYOUTS", table) + ckpt = open_checkpoint(_sparse(tmp_path, path, size)) + assert (len(ckpt.groups), len(ckpt.extras)) == (35, 169) + assert ckpt.groups["blocks.27"].matrix_names == tuple(f"blocks.27.{s}.weight" for s in BLOCK) + assert ckpt.config_source == ( + f"header and spot checks match layout {layout.key} ({layout.repo_id}@{layout.revision})" + ) + + +@pytest.mark.parametrize(("_layout", "path", "size", "_probes"), LAYOUTS) +def test_the_published_header_with_other_content_is_refused_by_the_default_table( + tmp_path, _layout, path, size, _probes +): + # Bug caught: the default table's Krea layout read without its content probes (the real refusal path). + with pytest.raises( + DFloatFormatError, + match=rf"differs .*\(sha256 {ZERO_PAGE_SHA256} at byte 0 of blocks\.0\.sign_mantissa\)", + ): + open_checkpoint(_sparse(tmp_path, path, size)) + + +@pytest.mark.network +@pytest.mark.parametrize(("layout", "path", "size", "_probes"), LAYOUTS) +def test_the_pinned_header_size_lfs_and_probes_match_the_published_file( + layout, path, size, _probes +): + # Bug caught: a digest copied wrong from the range reads, or a pin of a revision other than the one recorded. + # Reads the header, the 25 probe ranges (100 KiB) and the Hub's size + LFS oid; no download. + from huggingface_hub import HfApi + from scripts.verify_remote_group import HfRangeSource + + src = HfRangeSource(layout.repo_id, layout.revision) + assert src.size(layout.file_name) == size + (entry,) = HfApi().get_paths_info(layout.repo_id, [layout.file_name], revision=layout.revision) + assert entry.lfs is not None + assert entry.lfs.sha256 == layout.file_sha256 + assert src.read(layout.file_name, 0, 35_976) == path.read_bytes() + infos = parse_header(path.read_bytes()[8:], data_start=35_976, file_size=size, source=path.name) + for probe in layout.probes: + start = infos[probe.tensor].offset + probe.offset + assert ( + hashlib.sha256(src.read(layout.file_name, start, probe.length)).hexdigest() + == probe.sha256 + ) diff --git a/tests/test_mflux_families.py b/tests/test_mflux_families.py index 956ecb9..cc95c4d 100644 --- a/tests/test_mflux_families.py +++ b/tests/test_mflux_families.py @@ -12,7 +12,7 @@ def test_importing_the_registry_does_not_import_mflux(): # Bug caught: a family module imported at registry import time (the CLI's --help would need mflux). code = "import sys, mlx_dfloat.mflux.families as f; assert 'mflux' not in sys.modules; print(len(f.MODELS))" out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) - assert out.stdout.strip() == "12" + assert out.stdout.strip() == "14" def test_resolving_an_adapter_module_and_importing_the_package_loads_no_mlx(): @@ -52,6 +52,19 @@ def test_resolving_the_qwen21_coverage_specs_loads_no_mlx(): assert out.stdout.strip() == "False" +def test_resolving_the_krea2_coverage_specs_loads_no_mlx(): + # Bug caught: an eager import in mlx_dfloat/mflux/krea2/__init__.py. The mflux lane's + # --cov=mlx_dfloat.mflux.krea2.model / .init / .transformer resolve those specs (importing the krea2 package) when + # coverage starts; MLX loaded there registers its native types twice and the lane aborts (exit 134). + code = ( + "import importlib.util, sys; " + "[importlib.util.find_spec(f'mlx_dfloat.mflux.krea2.{m}') for m in ('model', 'init', 'transformer')]; " + "import mlx_dfloat.mflux; print('mlx.core' in sys.modules)" + ) + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert out.stdout.strip() == "False" + + def test_dfloat_model_is_still_importable_from_the_package(): # Bug caught: the lazy attribute dropping DFloatModel (the README's `from mlx_dfloat.mflux import DFloatModel`). from mlx_dfloat.mflux import DFloatModel @@ -203,6 +216,9 @@ def test_only_the_distilled_klein_models_fix_their_guidance(): # mflux 0.20.0 ernie_image_turbo_generate.py:42-45: any guidance but 1.0 is an error for Turbo. "ernie-image-turbo": 1.0, "ernie-image": None, + # mflux 0.20.0: Krea 2 Turbo and Raw take any guidance (krea2.py:58; CFG for any value other than 1.0). + "krea-2": None, + "krea-2-raw": None, } @@ -227,7 +243,7 @@ def test_the_qwen_entry_names_its_published_repos_and_pinned_checkpoint(): def test_pinned_bases_are_the_snapshots_the_recorded_runs_used_and_no_other_base_is(): # Bug caught: a base (text encoder, VAE, tokenizer) following whatever lands on main instead of the snapshot its # recorded runs used (Qwen-Image 2.1: d26bb61, bench/results/tiers/qwen-image-2.1-1024.json; ERNIE-Image and - # ERNIE-Image-Turbo: the snapshots their parity and de-risk runs read, 2026-10-08), or a pin set on a family whose + # ERNIE-Image-Turbo: the snapshots their parity and calibration runs read, 2026-10-08), or a pin set on a family whose # base no recorded run names. assert { n: e.base_revision for n, e in families.MODELS.items() if e.base_revision is not None @@ -235,6 +251,9 @@ def test_pinned_bases_are_the_snapshots_the_recorded_runs_used_and_no_other_base "qwen-image-2.1": "d26bb61231c349cf6b7896fa83353113880e1ba3", "ernie-image": "5346b31d68c9c23758ba56ef8be5e9dc174c7f99", "ernie-image-turbo": "bc68c81e2a1730a394d5fc9fae70713dee940140", + # Krea 2: the gated base snapshots the range-read parity reads and every Krea run uses (2026-10-08). + "krea-2": "98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7", + "krea-2-raw": "6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", } @@ -361,3 +380,75 @@ def test_the_ernie_family_loads_its_class_and_one_kind_map(): "ernie-image": ("mingyi456/ERNIE-Image-DF11", "baidu/ERNIE-Image"), "ernie-image-turbo": ("mingyi456/ERNIE-Image-Turbo-DF11", "baidu/ERNIE-Image-Turbo"), } + + +# --- Krea 2 ---------------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("name", "label", "steps", "df11", "revision", "base", "base_revision"), + [ + # mflux 0.20.0: 8 steps (cli/defaults/defaults.py:47), guidance 1.0 (krea2_generate.py:17; krea2.py:58). + ( + "krea-2", + "Krea 2 Turbo", + 8, + "mingyi456/Krea-2-Turbo-DF11-ComfyUI", + "978da5fb7647bd222d33125993abd8fdc2840cfc", + "krea/Krea-2-Turbo", + "98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7", + ), + # mflux's Krea command refuses krea-2-raw (krea2_generate.py:46-49): the step table's fallback 25 + # (defaults.py:19,111-118) and generate_image's guidance 1.0 (krea2.py:58) are its only defaults. + ( + "krea-2-raw", + "Krea 2 Raw", + 25, + "mingyi456/Krea-2-Raw-DF11-ComfyUI", + "8320616b25ac9340a830a7fb21f1b0237e160e66", + "krea/Krea-2-Raw", + "6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", + ), + ], +) +def test_the_krea_entries_carry_mfluxs_defaults_and_their_pins( + name, label, steps, df11, revision, base, base_revision +): + # Bug caught: Raw carrying the model card's 52 / 3.5 (a CFG run nobody asked for, about 4x the time) or Turbo's 8 + # steps, a pin or repository swapped between the two models, or the CFG rule copied as `> 1.0`. + e = families.entry(name) + assert (e.family, e.label, e.df11_repo, e.df11_revision, e.base_repo, e.base_revision) == ( + "krea2", + label, + df11, + revision, + base, + base_revision, + ) + assert (e.default_steps, e.default_guidance, e.default_scheduler) == (steps, 1.0, "er_sde") + assert (e.cfg_two_calls, e.uses_negative_prompt, e.cfg_needs_negative) == (True, True, False) + assert (e.cfg_unless_guidance_one, e.schedulers) == (True, ("er_sde", "euler", "linear")) + assert e.fixed_guidance is None + + +def test_new_fields_default_off_for_every_other_family(): + # Bug caught: another family's warnings or scheduler handling changed by the two Krea fields. + others = [e for e in families.MODELS.values() if e.family != "krea2"] + assert len(others) == 12 + assert all(e.cfg_unless_guidance_one is False for e in others) + assert all(e.schedulers is None for e in others) + + +@pytest.mark.mflux +def test_the_krea2_family_loads_its_class_and_one_block_kind(): + # Bug caught: the family's loaders pointing at another family's class or map. + from mlx_dfloat.mflux.krea2.model import MODELS as KREA + from mlx_dfloat.mflux.krea2.model import DFloatKrea2 + + assert families.FAMILIES["krea2"].load_model_class() is DFloatKrea2 + # mflux 0.20.0 Krea2Transformer.__call__ (transformer.py:89-92): one block list. + assert families.FAMILIES["krea2"].load_name_map().kinds == ("blocks",) + assert KREA == { + "krea-2": ("mingyi456/Krea-2-Turbo-DF11-ComfyUI", "krea/Krea-2-Turbo"), + "krea-2-raw": ("mingyi456/Krea-2-Raw-DF11-ComfyUI", "krea/Krea-2-Raw"), + } diff --git a/tests/test_mflux_phases.py b/tests/test_mflux_phases.py index 54c1e89..626c0cd 100644 --- a/tests/test_mflux_phases.py +++ b/tests/test_mflux_phases.py @@ -17,7 +17,7 @@ vae_transient_bytes, ) -# Klein 4B decoded group sizes (S3 plan arithmetic table: double block 490_733_568 B, single block 245_366_784 B). +# Klein 4B decoded group sizes (from the header's matrix shapes: double block 490_733_568 B, single block 245_366_784 B). KLEIN_4B_LARGEST = {"transformer_blocks": 490_733_568, "single_transformer_blocks": 245_366_784} @@ -71,7 +71,7 @@ def test_the_denoise_activation_scales_with_tokens_from_the_reference(): def test_a_denoise_activation_floor_holds_below_the_reference_and_is_off_by_default(): - # Bug caught (X3): a family whose step holds a size-independent buffer (ERNIE-Image's float32 copies of a block's + # Bug caught: a family whose step holds a size-independent buffer (ERNIE-Image's float32 copies of a block's # weights) predicted below it at a small image, or the floor applied to every family (the default must leave the # scaled term as it was). 512^2 + 512: 333_333 scaled, floored at 400_000; 1024^2 + 512: 1_000_000 above the floor. floored = _constants( diff --git a/tests/test_mflux_pipeline.py b/tests/test_mflux_pipeline.py index 965b1be..9dc3078 100644 --- a/tests/test_mflux_pipeline.py +++ b/tests/test_mflux_pipeline.py @@ -14,7 +14,7 @@ from mlx_dfloat.mflux._phases import FamilySizes, PhaseConstants GIB = 1024**3 -# Klein 4B decoded group sizes (S3 plan arithmetic table): double block 490_733_568 B, single block 245_366_784 B. +# Klein 4B decoded group sizes (from the header's matrix shapes): double block 490_733_568 B, single block 245_366_784 B. LARGEST = {"transformer_blocks": 490_733_568, "single_transformer_blocks": 245_366_784} LOG = logging.getLogger("tests.pipeline") diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 87b83e3..cb1ab00 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -101,3 +101,23 @@ def test_dfloat_ernie_image_is_reachable_lazily(): from mlx_dfloat.mflux.ernie.model import DFloatErnieImage assert adapters.DFloatErnieImage is DFloatErnieImage + + +@pytest.mark.mflux +def test_dfloat_krea2_is_reachable_lazily(): + # Bug caught: DFloatKrea2 missing from the lazy attribute hook (AttributeError on first use), or added to __all__ + # (a star-import would then need mflux; the star-import test pins __all__). + import mlx_dfloat.mflux as adapters + from mlx_dfloat.mflux.krea2.model import DFloatKrea2 + + assert adapters.DFloatKrea2 is DFloatKrea2 + + +def test_importing_the_mflux_package_loads_no_mlx_core(): + # Bug caught: an eager import of the Krea 2 model (or any model) in mlx_dfloat/mflux/__init__.py. + import subprocess + import sys + + code = "import sys, mlx_dfloat.mflux; print('mlx.core' in sys.modules, 'mflux' in sys.modules)" + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert out.stdout.strip() == "False False" diff --git a/tests/test_qwen21_build.py b/tests/test_qwen21_build.py index 6fa7a1f..f5db582 100644 --- a/tests/test_qwen21_build.py +++ b/tests/test_qwen21_build.py @@ -83,7 +83,7 @@ def test_the_build_installs_placeholders_and_loads_every_extra_by_name(tmp_path) for block in build.shapes: for attr in SEVEN: assert params[f"{block}.{attr}.weight"].size == 0, (block, attr) - # Every parameter a checkpoint holds outside the compressed groups (S0: 72 extras for 32 blocks = 8 + 2 x 32). + # Every parameter a checkpoint holds outside the compressed groups (the published header: 72 extras for 32 blocks = 8 + 2 x 32). assert sorted(constants) == [ "img_in.weight", "norm_out.linear.weight", diff --git a/tests/test_qwen21_memory.py b/tests/test_qwen21_memory.py index 2d84cb7..b33de37 100644 --- a/tests/test_qwen21_memory.py +++ b/tests/test_qwen21_memory.py @@ -1,6 +1,6 @@ """Qwen-Image 2.1's memory rules: the encoder's resident bytes, the text-token count, the sizes, the measured constants. -Measured inputs (2026-10-08, the first one-step de-risk run: Qwen-Image-2.1 DF11 +Measured inputs (2026-10-08, the first one-step calibration run: Qwen-Image-2.1 DF11 (mingyi456/Qwen-Image-2.1-DF11-ComfyUI @ 1b22a3a) over the Qwen/Qwen-Image-2.1 base (@ d26bb61), 1024², seed 42, guidance 4 with the negative prompt " " (two transformer calls per step), text tokens 33 (prompt) and 9 (negative); one process: build, encode, drop, set load, one step, VAE decode on the resident set; M1 Max 32 GB, macOS 27.0.1, @@ -13,10 +13,10 @@ ``bench/results/tiers/qwen-image-2.1-1024.json`` carries the same five sizes (``sizes``), the 33 text tokens, the 2_731_761_634 cache limit and this Mac's working set and RAM (``limits.tier``). -The denoise term is re-derived at the planner's limit (2026-10-08, the first one-step de-risk run at that limit, in +The denoise term is re-derived at the planner's limit (2026-10-08, the first one-step calibration run at that limit, in the cache-limit A/B; same model, prompt, size, seed and guidance): denoise footprint peak 13_939_106_072 at the 2_731_761_634 limit. The VAE term covers the highest of eleven 1024² VAE-phase footprint peaks (21.76-22.79 GB; MLX's -VAE peak was the same in every de-risk run): 22_788_580_400, from the second one-step run of the A/B at the old limit. +VAE peak was the same in every calibration run): 22_788_580_400, from the second one-step run of the A/B at the old limit. The calibration runs' phase peaks are committed in ``bench/results/calibration/qwen-image-2.1-1024.json`` (each run's label, date, cache limit and footprint peaks; ``constants_from`` names the run each term was derived from), and the @@ -54,21 +54,21 @@ # The planner's 1024² limit: the largest decoded block 436_207_616 + Qwen's own allowance 2_295_554_018 at 4096 + 33 # tokens (the block-cache probe: a decoded buffer is reused at a 2605 MiB limit, released at 1773 and 2189 MiB). CACHE_LIMIT = 2_731_761_634 -# The de-risk run's limit: the block + the shared 1.5e9 allowance at 4096 + 256 tokens, rescaled to 4096 + 33. +# The calibration run's limit: the block + the shared 1.5e9 allowance at 4096 + 256 tokens, rescaled to 4096 + 33. CALIBRATION_CACHE_LIMIT = 1_859_346_402 -TEXT_TOKENS = 33 # the de-risk prompt's (the negative " " is 9) +TEXT_TOKENS = 33 # the calibration prompt's (the negative " " is 9) REPO = Path(__file__).resolve().parents[1] CALIBRATION = json.loads((REPO / "bench/results/calibration/qwen-image-2.1-1024.json").read_text()) CALIBRATION_RUNS = {run["label"]: run for run in CALIBRATION["runs"]} # The phase peaks the constants were derived from: the denoise peak is the first one-step A/B run's at CACHE_LIMIT; # the VAE peak the highest of the eleven 1024² samples (the second A/B run at the old limit); the encode peak the first -# de-risk run's. +# calibration run's. MEASURED = { phase: CALIBRATION_RUNS[label]["footprint_peaks"][phase] for phase, label in CALIBRATION["constants_from"].items() } VAE_SAMPLES = {label: run["footprint_peaks"]["vae"] for label, run in CALIBRATION_RUNS.items()} -# This Mac (M1 Max 32 GB): mx.device_info()["max_recommended_working_set_size"] and the RAM, from the de-risk record. +# This Mac (M1 Max 32 GB): mx.device_info()["max_recommended_working_set_size"] and the RAM, from the calibration record. RECOMMENDED_32 = 26_800_603_136 RAM_32 = 34_359_738_368 @@ -183,7 +183,7 @@ def test_the_token_cap_is_mfluxs_tokenizer_max_length(): assert qmem.TEXT_MAX_LENGTH == 2048 -# The de-risk prompt (the run record's env.sh, $PROMPT): 33 text tokens by the real tokenizer. +# The calibration prompt (the run record's env.sh, $PROMPT): 33 text tokens by the real tokenizer. DERISK_PROMPT = ( "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the " "water" @@ -199,7 +199,7 @@ def test_the_token_cap_is_mfluxs_tokenizer_max_length(): ) def test_the_real_tokenizer_counts_the_text_tokens_the_calibration_used(prompt, tokens): # Bug caught: prompt_tokens disagreeing with the real tokenizer where the constants were calibrated (33 for the - # de-risk prompt, 9 for its negative " "): the template or the system prefix miscounted, or "" not encoded as " " + # calibration prompt, 9 for its negative " "): the template or the system prefix miscounted, or "" not encoded as " " # (the stub tests cannot see either). Needs the base's processor/ files (a few MB, CPU only). base = os.environ.get("MLX_DFLOAT_QWEN21_BASE") if not base: @@ -242,7 +242,7 @@ def _plan(budget, *, fit_check=True, cache_limit_override=None): ) -# The constants are calibrated on the de-risk run, so at its working point the estimate equals the measurement by +# The constants are calibrated on the calibration run, so at its working point the estimate equals the measurement by # construction: the in-sample checks are exact. No Qwen-Image 2.1 run outside the calibration exists yet (another size # or prompt length), so whether the model predicts away from 1024² and 33 tokens is not tested. @@ -261,7 +261,7 @@ def test_the_calibration_record_holds_the_runs_the_constants_name(): def test_the_vae_term_covers_every_1024_sample_and_equals_the_highest(): # Bug caught: the VAE transient derived from one run instead of the highest of the eleven samples (a fit check that - # under-predicts the runs it was calibrated on: the first de-risk run alone puts the phase 648_937_376 B low), or a + # under-predicts the runs it was calibrated on: the first calibration run alone puts the phase 648_937_376 B low), or a # sample dropped from the record. assert len(VAE_SAMPLES) == 11 fit = _fit() @@ -271,14 +271,14 @@ def test_the_vae_term_covers_every_1024_sample_and_equals_the_highest(): def test_the_estimate_at_1024_equals_the_calibration_run_phase_for_phase(): # Bug caught: a term dropped from a phase or counted twice; the VAE transient taken from one run instead of the - # highest sample (the first de-risk's 22_139_643_024 puts the VAE phase 648_937_376 B under it); + # highest sample (the first calibration run's 22_139_643_024 puts the VAE phase 648_937_376 B under it); # REFERENCE_TOKENS left at another prompt's count (the activation scaled off the calibration point, e.g. 4096 + 512 # puts denoise 94_905_646 B low); or a borrowed family's constants (Klein 9B's put the VAE phase 4.2 GiB low). # The encode phase: the language model 15_136_811_008 + the measured encode term 66_567_800 + overhead # 459_445_098 (by footprint the run's encode peaked 195_114_382 B higher: memory MLX does not count, which the # overhead measured after the build does not cover). # The denoise phase is calibrated at the planner's own limit (CACHE_LIMIT), so a cache-limit term taken from - # another limit (the de-risk's 1_859_346_402 puts it 872_415_232 B low) also shows here. + # another limit (the calibration run's 1_859_346_402 puts it 872_415_232 B low) also shows here. fit = _fit() assert fit.peak_phase == "vae" assert fit.phases["vae"] == MEASURED["vae"] @@ -287,14 +287,14 @@ def test_the_estimate_at_1024_equals_the_calibration_run_phase_for_phase(): def test_the_planners_1024_cache_limit_is_one_block_plus_qwens_own_allowance(): - # Bug caught: Qwen planned on the shared 1.5e9 allowance (the de-risk's 1_859_346_402: the probe released the + # Bug caught: Qwen planned on the shared 1.5e9 allowance (the calibration run's 1_859_346_402: the probe released the # decoded buffer at that limit, so every block's output was allocated fresh), two blocks budgeted for the one # block kind (3_167_969_250), or the allowance's reference tokens left at the shared 4096 + 256 (2_659_262_210). assert _plan(24_653_119_488).cache_limit == CACHE_LIMIT def test_a_cache_limit_override_under_qwens_derived_minimum_warns(caplog): - # Bug caught: an override of one block or more passing silently for this one-kind family (the de-risk's own + # Bug caught: an override of one block or more passing silently for this one-kind family (the calibration run's own # 1_859_346_402 would have): the minimum is the block + the allowance, the planner's own limit. with caplog.at_level("WARNING", logger="test"): _plan(24_653_119_488, cache_limit_override=CACHE_LIMIT) @@ -327,7 +327,7 @@ def test_a_32_gb_mac_keeps_the_set_through_the_vae_decode(monkeypatch): @pytest.mark.parametrize("tier", [16, 24]) def test_a_16_or_24_gb_tier_is_refused_in_the_encode_phase(tier): - # The prediction for the CAPPED tiers at 1024² (the de-risk prompt, 33 tokens): the VAE phase with the set + # The prediction for the CAPPED tiers at 1024² (the calibration prompt, 33 tokens): the VAE phase with the set # (22_788_580_400 B) is over either tier's budget, so the set is dropped before the decode: the VAE phase is then # the extras (they stay on the transformer) + the VAE + its transient + overhead, 137_388_032 + 1_350_989_512 + # 11_119_436_651 + 459_445_098 = 13_067_259_293 B; the encode phase (the language model 15_136_811_008 + the encode term 66_567_800 + overhead 459_445_098 = @@ -353,7 +353,7 @@ def test_a_16_or_24_gb_tier_is_refused_in_the_encode_phase(tier): # What the prompt encode may hold before the model warns (model._check_encode_peak): MLX's encode peak minus what was # active when the phase began, against the encode phase without its overhead (language model + encode term) + the slack. -ENCODE_HELD = 15_350_204_544 - 146_825_736 # de-risk: MLX encode peak - MLX active after the build +ENCODE_HELD = 15_350_204_544 - 146_825_736 # calibration: MLX encode peak - active after the build # The tensors in the encoder's shards mflux does not load (the vision tower and the lm_head), counted as the shards' # file size minus the language model's tensor bytes (the headers' few KB included: an upper bound by those bytes). NOT_LOADED = 17_534_339_488 - 15_136_811_008 @@ -364,7 +364,7 @@ def _encode_bound(text_tokens): def test_the_encode_warning_bound_holds_the_measured_encode_with_the_slack_to_spare(): - # Bug caught: a bound under what a normal encode holds (every call would warn): the de-risk encode held + # Bug caught: a bound under what a normal encode holds (every call would warn): the calibration encode held # 66_567_800 B over the language model, exactly the encode term, so the whole 256 MiB slack is left. assert _encode_bound(TEXT_TOKENS) - ENCODE_HELD == 268_435_456 @@ -393,7 +393,7 @@ def test_the_encode_warning_is_silent_at_the_bound_and_names_the_prompt_length_o assert "more of the encoder than its language model" in message -# The two committed 1024² records with more steps than the one-step calibration runs, both on the 33-token de-risk +# The two committed 1024² records with more steps than the one-step calibration runs, both on the 33-token calibration # prompt: the MEASURED row (`mlx-dfloat generate --tier 32`, 40 steps, no CFG) and the identity check's df11 side # (4 steps, guidance 4 with CFG). Each is re-estimated with the current constants from its own report (its sizes, the # cache limit in force, its text tokens, its policy and its VAE strategy), so a re-run that changes the record is diff --git a/tests/test_qwen21_model.py b/tests/test_qwen21_model.py index 1aa2220..9519193 100644 --- a/tests/test_qwen21_model.py +++ b/tests/test_qwen21_model.py @@ -96,7 +96,7 @@ def test_the_set_load_installs_modulation_and_attaches_with_verify_in_call(tmp_p def test_a_new_prompt_drops_the_set_and_reinstalls_modulation(tmp_path, monkeypatch): - # Bug caught (Review Focus 5): modulation.layers.1.weight left on its placeholder after the reload a new prompt + # Bug caught: modulation.layers.1.weight left on its placeholder after the reload a new prompt # forces, or the encoder reload reading another base than the model was built from. model = fake_model(tmp_path, monkeypatch) loads = [] @@ -361,7 +361,7 @@ def test_a_missing_guidance_and_scheduler_become_mflux_defaults(tmp_path, monkey def test_an_empty_prompt_and_a_negative_equal_to_the_prompt_run_from_the_cache( tmp_path, monkeypatch ): - # Bug caught (Review Focus 3): the empty prompt cached under " " (mflux's loop looks it up as "" and calls the + # Bug caught: the empty prompt cached under " " (mflux's loop looks it up as "" and calls the # dropped encoder, which raises), an empty negative encoded or run as CFG, or a negative equal to the prompt # encoded twice. Real mflux loop, one step. model = fake_model(tmp_path, monkeypatch) @@ -840,7 +840,7 @@ def fake_resolve(spec, *, patterns, revision=None): def test_the_constructor_builds_from_local_dirs_with_mfluxs_defaults_and_the_language_model_bytes( tmp_path, monkeypatch ): - # Bug caught (S3 review fix 9 for Klein): __init__ wiring past the resolver untested: the config-less checkpoint + # Bug caught: __init__ wiring past the resolver untested: the config-less checkpoint # not opened through the layout table, the transformer built with arguments mflux's initializer does not use # (Qwen21Initializer._init_models builds Qwen21Transformer() bare, qwen21_initializer.py:52), the base components # read from another directory, or the sizes taken from the shared rule (the encoder shards' file size, vision @@ -1040,7 +1040,7 @@ def test_a_tiny_generation_gives_stock_mfluxs_latents_bit_for_bit( def test_encode_runs_under_the_commands_caps_and_restores_mlxs_defaults(tmp_path, monkeypatch): - # Bug caught (X4): encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured + # Bug caught: encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured # phases, sized under the command's caps), or the caps left installed after it returns. from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap diff --git a/tests/test_qwen21_names.py b/tests/test_qwen21_names.py index b70f577..c16bf61 100644 --- a/tests/test_qwen21_names.py +++ b/tests/test_qwen21_names.py @@ -17,7 +17,7 @@ templated_targets, ) -# The seven block matrices in the order a block decodes into after the gate_up split (S0: the stored order is +# The seven block matrices in the order a block decodes into after the gate_up split (the stored order, by a 2026-10-07 byte comparison, is # to_q, to_k, to_v, to_out.0, gate_up = [gate_layer; proj], out). Literal, so a reordered MATRIX_SUBS shows. SEVEN = ( "attn.to_q", @@ -76,7 +76,7 @@ def _folded(targets=None): def test_the_folded_map_places_every_block_matrix_on_its_own_path_in_decoded_order(): # Bug caught: a concrete-index mapping not folded (no {block} target: "a compressed matrix with no mflux target"), - # or the table order not the decoded order (gate_layer before proj, S0). + # or the table order not the decoded order (gate_layer before proj). m = _folded() assert m.attrs_of("transformer_blocks") == SEVEN assert m.place("transformer_blocks.2.img_mlp.proj.weight").attr == "img_mlp.proj" @@ -203,7 +203,7 @@ def test_a_checkpoint_without_blocks_is_refused(tmp_path): def test_a_qwen_image_1_block_is_refused_naming_the_group_and_the_seven_matrices(tmp_path): - # Bug caught (Review Focus 2): the Qwen-Image 1 DF11 checkpoint (dual-stream blocks, 14 matrices, same group + # Bug caught: the Qwen-Image 1 DF11 checkpoint (dual-stream blocks, 14 matrices, same group # names) passed for 2.1 and failing with a shape error at the first block instead of at build. qwen_image_1 = ( "img_mod.1", diff --git a/tests/test_reference_roundtrip.py b/tests/test_reference_roundtrip.py index 8720503..9907c8a 100644 --- a/tests/test_reference_roundtrip.py +++ b/tests/test_reference_roundtrip.py @@ -121,7 +121,7 @@ def test_skewed_one_bit_code_fills_a_block_with_32k_elements(): def test_byte_aligned_end_without_eof_roundtrips(): - # Review focus 5a. Pick a 1-bit symbol, then 8 of it = 8 bits exactly: no EOF emitted. + # Pick a 1-bit symbol, then 8 of it = 8 bits exactly: no EOF emitted. counter = {127: 1000, 126: 2, 125: 1} _, table, _ = codec_for(counter) one_bit = next(k for k, v in table.items() if isinstance(k, int) and v[0] == 1) @@ -133,7 +133,7 @@ def test_byte_aligned_end_without_eof_roundtrips(): def test_last_code_straddling_into_a_code_free_tail_block_roundtrips(): - # Review focus 5b: 32767 one-bit codes + one 9-bit code at bit 32767 -> 32776 bits = 4097 bytes, + # 32767 one-bit codes + one 9-bit code at bit 32767 -> 32776 bits = 4097 bytes, # byte-aligned, no EOF; block 1 has no code start (2 positions for 2 blocks). counter = {127 - i: 2 ** (20 - i) for i in range(12)} _, table, _ = codec_for(counter) @@ -147,7 +147,7 @@ def test_last_code_straddling_into_a_code_free_tail_block_roundtrips(): def test_dropped_entry_for_a_block_that_starts_codes_is_refused(): - # Review focus 5c: remove the last block's start from a real 4-block group. + # Remove the last block's start from a real 4-block group. rng = np.random.default_rng(6) arrays = _roundtrip([random_bf16(rng, (20_000,))]) assert arrays.n_blocks >= 3 diff --git a/tests/test_safetensors_reader.py b/tests/test_safetensors_reader.py index a833620..b0ecd2c 100644 --- a/tests/test_safetensors_reader.py +++ b/tests/test_safetensors_reader.py @@ -38,7 +38,7 @@ def test_bf16_is_returned_as_uint16_bits(tmp_path): def test_reads_through_a_symlink(tmp_path): - # Review focus 1: HF snapshots are symlinks into blobs/. + # HF snapshots are symlinks into blobs/. blob = tmp_path / "blob.safetensors" _write(blob, {"t": np.arange(4, dtype=np.uint8)}) link = tmp_path / "model.safetensors" @@ -56,7 +56,6 @@ def test_dangling_symlink_is_a_format_error(tmp_path): def test_truncated_file_is_a_format_error(tmp_path): - # Review focus 2. path = tmp_path / "c.safetensors" _write(path, {"t": np.zeros(1000, dtype=np.uint8)}) path.write_bytes(path.read_bytes()[:-10]) diff --git a/tests/test_verify_image.py b/tests/test_verify_image.py index 0d7d93a..d802fec 100644 --- a/tests/test_verify_image.py +++ b/tests/test_verify_image.py @@ -865,7 +865,7 @@ def test_ernie_negative_prompt_is_accepted_and_keyed_with_mfluxs_guidance(tmp_pa @pytest.mark.parametrize("guidance", ["4", "1.5"]) def test_ernie_turbo_guidance_other_than_1_is_refused(tmp_path, capsys, guidance): - # Bug caught (C2): an identity run of Turbo at a guidance its command refuses (mflux's turbo command errors on any + # Bug caught: an identity run of Turbo at a guidance its command refuses (mflux's turbo command errors on any # but 1.0, ernie_image_turbo_generate.py:42-45), so the check would cover a path no user runs. with pytest.raises(SystemExit) as info: _parse_ernie(tmp_path, tmp_path, "--guidance", guidance, model=ERNIE_TURBO) @@ -877,3 +877,249 @@ def test_ernie_turbo_guidance_1_is_accepted(tmp_path): # Bug caught: the refusal firing on Turbo's own 1.0 (or on the default, which resolves to it). assert _parse_ernie(tmp_path, tmp_path, "--guidance", "1", model=ERNIE_TURBO).guidance == 1.0 assert _parse_ernie(tmp_path, tmp_path, model=ERNIE_TURBO).guidance == 1.0 + + +# --- Krea 2 ------------------------------------------------------------------------------------------------------ + +KREA_RAW, KREA_TURBO = "krea-2-raw", "krea-2" + + +def _krea_base(tmp_path, native="raw.safetensors"): + """A base snapshot with the native single-file transformer at its root (a link into a blob, as hf lays it out).""" + base = tmp_path / "kbase" + (base / "blobs").mkdir(parents=True) + blob = base / "blobs" / "f99bb0ff" + blob.write_bytes(b"\0") + (base / native).symlink_to(blob) + return base + + +def _parse_krea(tmp_path, base, *extra, model=KREA_RAW, side=("--side", "df11")): + return vi.parse_args( + [ + *side, + "--model", + model, + "--df11", + str(tmp_path), + "--base", + str(base), + "--out", + str(tmp_path / "out"), + *extra, + ] + ) + + +@pytest.mark.parametrize("model", [KREA_RAW, KREA_TURBO]) +def test_runner_for_dispatches_krea2_to_its_sides(model): + # Bug caught: Krea 2 routed to another family's sides (a wrong-model identity verdict), or no krea2 row (a + # KeyError after the df11 side already ran for minutes). + assert vi.runner_for(model, "df11") is vi.run_df11_krea2 + assert vi.runner_for(model, "bf16") is vi.run_bf16_krea2 + + +@pytest.mark.parametrize( + ("model", "native", "repo"), + [ + (KREA_RAW, "raw.safetensors", "krea/Krea-2-Raw"), + (KREA_TURBO, "turbo.safetensors", "krea/Krea-2-Turbo"), + ], +) +@pytest.mark.parametrize("layout", ["missing", "dangling", "directory", "other_model"]) +def test_krea2_base_problem_names_the_missing_native_file(tmp_path, model, native, repo, layout): + # Bug caught: the identity launched on a base without the model's own native transformer file (the bf16 side would + # fail after the df11 side's minutes), a dangling link or a directory taken for the file, or Turbo's file accepted + # as Raw's reference (the two bases share their encoder and VAE, not the transformer). + base = tmp_path / "kbase" + base.mkdir() + if layout == "dangling": + (base / native).symlink_to(base / "gone") + elif layout == "directory": + (base / native).mkdir() + elif layout == "other_model": + other = "turbo.safetensors" if native == "raw.safetensors" else "raw.safetensors" + (base / other).write_bytes(b"\0") + problem = vi.krea2_base_problem(model, base) + assert problem == ( + f"--base {base}: the bf16 side streams the transformer from {native} at the base's root (download " + f"{repo}'s {native} first)" + ) + + +def test_krea2_base_problem_is_none_with_the_native_file_and_for_the_other_families(tmp_path): + # Bug caught: a complete base refused (the identity check could never run), or the Krea check applied to another + # family's base. + assert vi.krea2_base_problem(KREA_RAW, _krea_base(tmp_path)) is None + assert vi.krea2_base_problem("qwen-image-2.1", tmp_path) is None + assert vi.krea2_base_problem("schnell", tmp_path) is None + + +@pytest.mark.parametrize("side", [("--side", "bf16"), ("--orchestrate",)]) +def test_a_krea_run_that_streams_the_bf16_side_is_refused_without_its_native_file( + tmp_path, capsys, side +): + # Bug caught: the check not wired into parse_args for the runs that read the base's transformer. + base = tmp_path / "kbase" + base.mkdir() + with pytest.raises(SystemExit) as info: + _parse_krea(tmp_path, base, side=side) + assert info.value.code == 2 + assert "raw.safetensors" in capsys.readouterr().err + + +def test_a_krea_df11_side_needs_no_native_file(tmp_path): + # Bug caught: the df11 side refused for a base that holds only the encoder, the VAE and the tokenizer (all it + # reads). + base = tmp_path / "kbase" + base.mkdir() + assert _parse_krea(tmp_path, base).model == KREA_RAW + + +def test_krea_raw_cfg_run_is_keyed_with_its_guidance_and_negative_prompt(tmp_path, monkeypatch): + # Bug caught: the identity's CFG run (guidance 3.5, two batch-1 calls per step, mflux's " " negative) keyed or + # passed to the children without the guidance or the negative prompt, so a stored side of another run is reused; + # or the default guidance left None instead of mflux's 1.0 for Raw. + monkeypatch.setattr(vi, "source_hash", lambda: "src") + base = _krea_base(tmp_path) + args = _parse_krea( + tmp_path, base, "--guidance", "3.5", "--negative-prompt", " ", side=("--orchestrate",) + ) + key = vi.run_key(args) + assert (key["guidance"], key["negative_prompt"]) == (3.5, " ") + command = vi.child_command(args, "bf16") + assert command[command.index("--guidance") + 1] == "3.5" + assert command[command.index("--negative-prompt") + 1] == " " + assert _parse_krea(tmp_path, base).guidance == 1.0 + + +@pytest.mark.parametrize("guidance", ["3.5", "0.5"]) +def test_krea_turbo_guidance_other_than_1_is_refused(tmp_path, capsys, guidance): + # Bug caught: an identity run of Turbo at a guidance off its distilled recipe (CFG at any value but 1.0, 0.5 + # included), which the check does not cover. + with pytest.raises(SystemExit) as info: + _parse_krea(tmp_path, tmp_path, "--guidance", guidance, model=KREA_TURBO) + assert info.value.code == 2 + assert "krea-2 runs at guidance 1.0 only in this check" in capsys.readouterr().err + + +def test_krea_turbo_negative_prompt_is_refused_and_guidance_1_accepted(tmp_path, capsys): + # Bug caught: a negative prompt keyed for Turbo at 1.0 (no negative branch there), or the refusal firing on 1.0. + assert _parse_krea(tmp_path, tmp_path, "--guidance", "1", model=KREA_TURBO).guidance == 1.0 + with pytest.raises(SystemExit) as info: + _parse_krea(tmp_path, tmp_path, "--negative-prompt", " ", model=KREA_TURBO) + assert info.value.code == 2 + assert "there is no negative branch" in capsys.readouterr().err + + +def _tiny_krea_native(path, ckpt, matrices, extras): + """The tiny checkpoint's source as a base's native file, under the checkpoint's names: the block matrices BF16, + the non-block matrices FP32 holding the same BF16 values (the published file stores five of them that way), the + extras BF16.""" + import numpy as np + + from mlx_dfloat.mflux.krea2.names import NONBLOCK_GROUPS, krea2_name_map + + names = krea2_name_map() + tensors = {} + for group, per in matrices.items(): + for name in ckpt.groups[group].matrix_names: + if group in NONBLOCK_GROUPS: + bits = per[names.param_name(name)].astype(np.uint32) << np.uint32(16) + tensors[name] = mx.array(bits).view(mx.float32) + else: + sub = name.removeprefix(f"{group}.").removesuffix(".weight") + tensors[name] = mx.array(per[sub]).view(mx.bfloat16) + for name in ckpt.extras: + tensors[name] = mx.array(extras[names.param_name(name)]).view(mx.bfloat16) + mx.save_safetensors(str(path), tensors) + + +@pytest.mark.mflux +@pytest.mark.parametrize( + ("guidance", "negative"), [(1.0, None), (3.5, "n")], ids=["guidance-1", "cfg"] +) +def test_the_krea2_bf16_side_gives_stock_mfluxs_latents_on_a_tiny_checkpoint( + tmp_path, monkeypatch, guidance, negative +): + # Bug caught: the bf16 side's copy of mflux's loop body drifting from mflux's own (the CFG formula, the + # stepper's seed, the sigma passed as the timestep, the embeddings swapped), a native tensor not replacing the DF11 + # value it stands for, or the streamed blocks read under other names. Stock Krea2 over the same BF16 weights, the + # same embeddings and seed; equality is bit for bit on the float32 latents, before the heavy identity run. + import numpy as np + from mflux.models.krea2.model.krea2_text_encoder.prompt_encoder import Krea2PromptEncoder + from mflux.utils.apple_silicon import AppleSiliconUtil + from tests._krea2_tiny import ( + TINY, + StubTextEncoder, + StubTokenizer, + StubVAE, + stock_krea2, + write_tiny_checkpoint, + ) + + from mlx_dfloat import _layouts + from mlx_dfloat.format import open_checkpoint + from mlx_dfloat.mflux.krea2 import init as kinit + from mlx_dfloat.mflux.krea2 import transformer as ktf + + monkeypatch.setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: True)) + rng = np.random.default_rng(11) + df11 = tmp_path / "df11" + matrices, extras, layout = write_tiny_checkpoint(df11, rng, random_extras=rng) + monkeypatch.setattr(_layouts, "KNOWN_LAYOUTS", (layout,)) + ckpt = open_checkpoint(df11) + base = tmp_path / "base" + base.mkdir() + _tiny_krea_native(base / "raw.safetensors", ckpt, matrices, extras) + real_build = ktf.build_transformer + monkeypatch.setattr( + ktf, "build_transformer", lambda c, **kw: real_build(c, transformer_kwargs=TINY, **kw) + ) + ours_vae, stock_vae = StubVAE(), StubVAE() + monkeypatch.setattr(kinit, "load_vae", lambda root: ours_vae) + tokenizer = StubTokenizer({"n": 12}) + embeds, neg = Krea2PromptEncoder.encode_prompt_pair( + prompt="p", + negative_prompt=negative, + guidance=guidance, + tokenizer=tokenizer, + text_encoder=StubTextEncoder(), + prompt_cache={}, + ) + out = tmp_path / "out" + (out / "df11").mkdir(parents=True) + saved = {"embeds": embeds} if neg is None else {"embeds": embeds, "neg_embeds": neg} + mx.save_safetensors(str(out / "df11" / "embeds.safetensors"), saved) + args = vi.argparse.Namespace( + out=out, + df11=str(df11), + base=str(base), + model=KREA_RAW, + eval_policy="per-block", + seed=3, + steps=2, + size=64, + guidance=guidance, + ) + result = vi.run_bf16_krea2(args, vi.argparse.Namespace(peak_footprint=0)) + ours = mx.load(str(out / "bf16" / "latents.safetensors"))["latents"] + stock = stock_krea2(tokenizer, StubTextEncoder(), stock_vae, matrices, extras) + stock.generate_image( + seed=3, + prompt="p", + num_inference_steps=2, + height=64, + width=64, + guidance=guidance, + negative_prompt=negative, + ) + (theirs,) = stock_vae.seen + assert result["calls_per_step"] == (2 if negative else 1) + assert result["replaced_from_base"] == len(ckpt.extras) + 37 # every extra and non-block matrix + assert ours.dtype == theirs.dtype == mx.float32 + assert ours.shape == theirs.shape == (1, 16, 8, 8) + assert bool(mx.all(mx.isfinite(theirs))) + assert np.array_equal(np.array(ours.view(mx.uint32)), np.array(theirs.view(mx.uint32))) + (decoded,) = ours_vae.seen # the VAE decodes the side's own final latents + assert np.array_equal(np.array(decoded.view(mx.uint32)), np.array(ours.view(mx.uint32))) diff --git a/tests/test_verify_remote.py b/tests/test_verify_remote.py index 3c9b5b6..4b3f9d5 100644 --- a/tests/test_verify_remote.py +++ b/tests/test_verify_remote.py @@ -7,6 +7,7 @@ import pytest from scripts.verify_checkpoint import VerifyError from scripts.verify_remote_group import ( + GROUP_FIELDS, HfRangeSource, LocalRangeSource, bf16_index, @@ -928,3 +929,222 @@ def test_a_layout_whose_spot_checks_cannot_run_is_refused(tmp_path, probes, mess config_from_source( LocalRangeSource(df11), layouts=(dataclasses.replace(layout, probes=probes),) ) + + +# --- extras: the uncompressed tensors compared by name with their originals ------------------------------------------ + +# Three BF16 extras with literal bit patterns (1.0, 2.0, -1.0, the smallest denormal; 0.25, -3.0; 1.0078125 ...). +_EXTRAS = { + "first.bias": np.array([0x3F80, 0x4000, 0xBF80, 0x0001], dtype=np.uint16), + "norm.scale": np.array([0x3E80, 0xC040], dtype=np.uint16), + "mod.lin": np.array([[0x3F81, 0x3F82], [0x3F83, 0x8000]], dtype=np.uint16), +} +# An FP32 original of one extra: the two ties and the two near-ties where nearest-even and truncation differ in two +# of four places. By hand: 0x3F80_8000 (tie, even 0x3F80 below) -> 0x3F80; 0x3F81_8000 (tie, odd 0x3F81) -> 0x3F82; +# 0x3F80_8001 (just above half) -> 0x3F81; 0x3F80_7FFF (just below half) -> 0x3F80. Truncation: 0x3F80, 0x3F81, +# 0x3F80, 0x3F80. +_F32_ORIGINAL = np.array([0x3F808000, 0x3F818000, 0x3F808001, 0x3F807FFF], dtype=np.uint32) +_F32_AS_DF11 = np.array([0x3F80, 0x3F82, 0x3F81, 0x3F80], dtype=np.uint16) + + +def _extras_pair(tmp_path, *, df11_extras=None, originals=None, single_file=True): + """A DF11 checkpoint (one group, `_EXTRAS`) and one original file holding the group's matrix and the extras.""" + rng = np.random.default_rng(41) + q = random_bf16(rng, (4, 6)) + df11 = write_checkpoint( + tmp_path / "d", + groups={"blocks.0": [q]}, + pattern=r"blocks\.\d+", + sub_paths=(), + extras=_EXTRAS if df11_extras is None else df11_extras, + single_file=single_file, + ) + root = tmp_path / "b" + root.mkdir() + tensors = {"blocks.0.weight": q, **(_EXTRAS if originals is None else originals)} + mx.save_safetensors( + str(root / "model.safetensors"), + { + n: mx.array(v).view(mx.float32 if v.dtype == np.uint32 else mx.bfloat16) + for n, v in tensors.items() + }, + ) + return df11, root + + +def _extras_run(tmp_path, df11, bf16, *extra_args): + out = tmp_path / "out" / "RESULT.json" + code = main([*_argv(df11, bf16, out), *extra_args], source_factory=_local_factory) + return code, json.loads(out.read_text()) + + +@pytest.mark.parametrize("single_file", [True, False]) +def test_extras_equal_to_bf16_originals_pass_and_are_counted(tmp_path, single_file): + # Bug caught: extras never read (a pass with 0 compared), or an extras-only shard of a sharded checkpoint (no group + # in it, so absent from the group index) skipped. + df11, bf16 = _extras_pair(tmp_path, single_file=single_file) + code, result = _extras_run(tmp_path, df11, bf16, "--groups", "none", "--extras") + assert code == 0, result + assert (result["compared"], result["extras_compared"], result["extras_mismatched"]) == (0, 3, 0) + got = [(e["name"], e["n"], e["equal"], e["original_dtype"]) for e in result["extras"]["extras"]] + assert got == [ + ("first.bias", 4, True, "BF16"), + ("mod.lin", 4, True, "BF16"), + ("norm.scale", 2, True, "BF16"), + ] + assert result["control"] == "bf16" + + +def test_an_fp32_extra_is_compared_rounded_to_nearest_even(tmp_path): + # Bug caught: truncation instead of RNE (two of four differ), or a silent cast without the flag. The control label + # must name fp32 although no group matrix was compared (the extras count in it). + df11, bf16 = _extras_pair( + tmp_path, + df11_extras={**_EXTRAS, "mod.lin": _F32_AS_DF11}, + originals={**_EXTRAS, "mod.lin": _F32_ORIGINAL}, + ) + code, result = _extras_run( + tmp_path, df11, bf16, "--groups", "none", "--extras", "--cast-fp32-to-bf16" + ) + assert code == 0, result + assert result["control"] == "fp32 rounded to bf16 (nearest even)" + dtypes = {e["name"]: e["original_dtype"] for e in result["extras"]["extras"]} + assert dtypes == {"first.bias": "BF16", "mod.lin": "F32", "norm.scale": "BF16"} + code, result = _extras_run(tmp_path, df11, bf16, "--groups", "none", "--extras") + assert code == 2 + assert result["extras"]["status"] == "error" + assert "--cast-fp32-to-bf16" in result["extras"]["error"] + + +def test_one_flipped_extra_bit_is_a_mismatch_exit_1(tmp_path): + # Bug caught: any / all swapped, or the exit code computed before the extras (the group matrix is equal, so only + # the extras can make this run exit 1). + flipped = {**_EXTRAS, "norm.scale": _EXTRAS["norm.scale"] ^ np.uint16(0x0001)} + df11, bf16 = _extras_pair(tmp_path, df11_extras=flipped) + code, result = _extras_run(tmp_path, df11, bf16, "--groups", "first", "--extras") + assert code == 1, result + assert (result["compared"], result["mismatched"]) == (1, 0) + assert (result["extras_compared"], result["extras_mismatched"]) == (3, 1) + assert result["extras"]["status"] == "mismatch" + assert [e["name"] for e in result["extras"]["extras"] if not e["equal"]] == ["norm.scale"] + + +def test_an_extra_without_an_original_is_an_error_not_a_pass(tmp_path): + # Bug caught: an extra with no original skipped (the run passes on the extras that happen to exist). + df11, bf16 = _extras_pair( + tmp_path, originals={k: v for k, v in _EXTRAS.items() if k != "mod.lin"} + ) + code, result = _extras_run(tmp_path, df11, bf16, "--groups", "none", "--extras") + assert code == 2 + assert result["extras"]["status"] == "error" + assert result["extras"]["error"] == "no original for extra mod.lin" + + +def test_an_extra_whose_original_holds_another_element_count_is_an_error_not_a_mismatch(tmp_path): + # Bug caught: a size disagreement (a mapping problem) reported as a bit mismatch, or compared by a prefix. + longer = {**_EXTRAS, "norm.scale": np.array([0x3E80, 0xC040, 0x3F80], dtype=np.uint16)} + df11, bf16 = _extras_pair(tmp_path, originals=longer) + code, result = _extras_run(tmp_path, df11, bf16, "--groups", "none", "--extras") + assert code == 2 + assert result["extras"]["error"] == "extra norm.scale: shape (2,), original has (3,)" + + +def test_a_df11_extra_that_is_not_bf16_is_an_error(tmp_path): + # Bug caught: an F32 extra in the DF11 file read as two-byte elements and reported as a mismatch (exit 1), or + # passed through. + df11, bf16 = _extras_pair(tmp_path, df11_extras={}, single_file=False) + mx.save_safetensors( + str(df11 / "extra.safetensors"), {"first.bias": mx.array([1.0, 2.0], dtype=mx.float32)} + ) + code, result = _extras_run(tmp_path, df11, bf16, "--groups", "none", "--extras") + assert code == 2 + assert result["extras"]["error"] == "extra first.bias is F32, not BF16" + + +def test_without_the_extras_flag_the_result_keeps_its_old_shape(tmp_path): + # Bug caught: extras keys written on every run (a merge reading them would count a run that compared none). + df11, bf16 = _extras_pair(tmp_path) + code, result = _extras_run(tmp_path, df11, bf16, "--groups", "first") + assert code == 0, result + assert {"extras", "extras_compared", "extras_mismatched"}.isdisjoint(result) + + +def test_extras_with_structural_only_is_refused(tmp_path, capsys): + # Bug caught: a structural run accepting --extras and claiming a comparison it cannot make (no original). + df11, _ = _extras_pair(tmp_path) + argv = [ + "--df11-repo", str(df11), "--df11-revision", "x", "--structural-only", + "--extras", "--out", str(tmp_path / "o.json"), + ] # fmt: skip + with pytest.raises(SystemExit) as info: + main(argv, source_factory=_local_factory) + assert info.value.code == 2 + assert "--extras" in capsys.readouterr().err + + +def test_groups_none_without_extras_is_refused(tmp_path, capsys): + # Bug caught: a run that compares nothing at all exiting 0 (a vacuous pass recorded as parity). + df11, bf16 = _extras_pair(tmp_path) + with pytest.raises(SystemExit) as info: + main( + [*_argv(df11, bf16, tmp_path / "o.json"), "--groups", "none"], + source_factory=_local_factory, + ) + assert info.value.code == 2 + assert "--groups none" in capsys.readouterr().err + + +def test_groups_none_selects_nothing_and_mixed_with_a_group_is_refused(): + # Bug caught: `none` silently dropping the other selectors, or read as a group name. + stats = {"blocks.0": {"elements_per_block": 5, "max_code_length": 30}} + assert pick_groups(stats, "none") == [] + for selector in ("none,first", "first,none"): + with pytest.raises(KeyError, match="none"): + pick_groups(stats, selector) + + +def test_an_extras_only_run_on_a_repository_without_extras_is_an_error(tmp_path): + # Bug caught: `--groups none --extras` on a checkpoint with no extras exiting 0 having compared nothing. + df11, bf16 = _extras_pair(tmp_path, df11_extras={}, single_file=False) + code, result = _extras_run(tmp_path, df11, bf16, "--groups", "none", "--extras") + assert code == 2 + assert result["extras"]["status"] == "error" + assert result["extras"]["error"] == ( + "the checkpoint holds no extras: --groups none --extras compares nothing" + ) + # With a group selected the same checkpoint passes: the group is compared, and no extras is not an error then. + code, result = _extras_run(tmp_path, df11, bf16, "--groups", "first", "--extras") + assert code == 0, result + assert (result["compared"], result["extras_compared"]) == (1, 0) + + +def test_an_extra_whose_original_has_another_shape_is_an_error_before_its_data_is_read(tmp_path): + # Bug caught: a transposed or reshaped original (same element count) compared byte for byte, a mapping + # problem passed off as a pass or a mismatch. mod.lin is 2 x 2 in the checkpoint; its original here is 1 x 4. + reshaped = {**_EXTRAS, "mod.lin": _EXTRAS["mod.lin"].reshape(1, 4)} + df11, bf16 = _extras_pair(tmp_path, originals=reshaped) + code, result = _extras_run(tmp_path, df11, bf16, "--groups", "none", "--extras") + assert code == 2 + assert result["extras"]["error"] == "extra mod.lin: shape (2, 2), original has (1, 4)" + + +def test_a_shard_outside_the_checkpoints_index_is_not_an_extra(tmp_path): + # Bug caught: every safetensors file in the repository read as the checkpoint's (a DF11 repository that also + # hosts an encoder file would fail with "no original for extra ..."). The index names the group shard and the + # extras shard; vae.safetensors sits next to them, unlisted. + df11, bf16 = _extras_pair(tmp_path, single_file=False) + weight_map = { + **{f"blocks.0.{f}": "blocks_0.safetensors" for f in GROUP_FIELDS}, + **dict.fromkeys(_EXTRAS, "model.safetensors"), + } + (df11 / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + mx.save_safetensors( + str(df11 / "vae.safetensors"), {"decoder.conv.weight": mx.ones((2,), dtype=mx.bfloat16)} + ) + code, result = _extras_run(tmp_path, df11, bf16, "--groups", "none", "--extras") + assert code == 0, result + assert [e["name"] for e in result["extras"]["extras"]] == [ + "first.bias", + "mod.lin", + "norm.scale", + ] diff --git a/tests/test_zimage_init.py b/tests/test_zimage_init.py index 6d66482..b917fb5 100644 --- a/tests/test_zimage_init.py +++ b/tests/test_zimage_init.py @@ -10,7 +10,10 @@ from mlx_dfloat.mflux.zimage.init import ( BASE_PATTERNS, encoder_definition, + load_base, load_text_encoder, + load_tokenizers, + load_vae, vae_definition, ) @@ -70,6 +73,30 @@ def test_a_directory_without_the_text_encoder_is_a_format_error(tmp_path): load_text_encoder(tmp_path) +@pytest.mark.mflux +def test_a_directory_without_the_vae_is_a_format_error(tmp_path): + # Bug caught: a DF11 checkpoint passed as --base crashing inside mflux's VAE loader with a bare error instead of + # naming what the directory lacks. + with pytest.raises(DFloatFormatError, match=r"not a Z-Image base with \['vae'\]"): + load_vae(tmp_path) + + +@pytest.mark.mflux +def test_a_directory_without_the_tokenizers_is_a_format_error(tmp_path): + # Bug caught: mflux's tokenizer loader error (FileNotFoundError or RuntimeError) escaping unwrapped, so the user + # sees an mflux traceback instead of the base that lacks the files. + with pytest.raises(DFloatFormatError, match="no usable Z-Image tokenizers"): + load_tokenizers(tmp_path) + + +@pytest.mark.mflux +def test_load_base_reads_the_vae_first_and_refuses_an_empty_directory(tmp_path): + # Bug caught: load_base skipping a component (the refusal would come later, at the first call), or swallowing the + # loader's format error. + with pytest.raises(DFloatFormatError, match=r"\['vae'\]"): + load_base(tmp_path) + + @pytest.mark.network @pytest.mark.parametrize("repo", ["Tongyi-MAI/Z-Image", "Tongyi-MAI/Z-Image-Turbo"]) def test_the_base_patterns_match_text_encoder_vae_and_tokenizer_files_on_the_hub(repo): diff --git a/tests/test_zimage_model.py b/tests/test_zimage_model.py index 2c1c7a7..1fb61f8 100644 --- a/tests/test_zimage_model.py +++ b/tests/test_zimage_model.py @@ -135,7 +135,7 @@ def test_the_set_load_installs_the_nonblock_weight_and_attaches_with_verify_in_c def test_a_drop_restores_every_placeholder_and_a_reload_reinstalls_the_nonblock_weight( tmp_path, monkeypatch ): - # Bug caught (Review Focus 2): the second load after a new prompt finding cap_embedder.1 still a placeholder, + # Bug caught: the second load after a new prompt finding cap_embedder.1 still a placeholder, # or the dropped set still referenced by the module (the active-memory check would refuse the reload). model = fake_model(tmp_path, monkeypatch) expected = _source(tmp_path)["cap_embedder"][0] @@ -269,7 +269,7 @@ def fake(self, **kwargs): ], ) def test_cfg_prompts_follow_mflux_rules(tmp_path, monkeypatch, model, guidance, negative, expected): - # Bug caught (Review Focus 1): a negative prompt encoded at guidance <= 1 (a wasted encoder reload), not encoded + # Bug caught: a negative prompt encoded at guidance <= 1 (a wasted encoder reload), not encoded # at > 1 (the CFG call would find nothing cached), or Turbo running CFG although mflux forces guidance 0 # (z_image.py:65-67, 177-179). instance = fake_model(tmp_path, monkeypatch, model=model) @@ -351,7 +351,7 @@ def test_generate_refusals_come_before_encoding_or_loading(tmp_path, monkeypatch def test_plan_call_uses_the_rounded_size_and_refuses_above_1024_squared( tmp_path, monkeypatch, caplog ): - # Bug caught (Review Focus 3): planning 1000x1000 while mflux runs 992x992; a 1040x1040 request accepted without + # Bug caught: planning 1000x1000 while mflux runs 992x992; a 1040x1040 request accepted without # a measurement. 1040x1040 with fit_check=True -> DFloatResourceError; fit_check=False -> a warning. from mlx_dfloat.errors import DFloatResourceError @@ -498,7 +498,7 @@ def test_a_format_error_mid_step_drops_the_set_and_restores_the_limit(tmp_path, def test_a_keyboard_interrupt_keeps_the_set_restores_the_limit_and_the_next_call_runs( tmp_path, monkeypatch ): - # Bug caught (Review Focus 4): StopImageGenerationException leaving status words pending (begin_step refuses + # Bug caught: StopImageGenerationException leaving status words pending (begin_step refuses # the next call) or the cache limit not restored; or the interrupt dropping the set (a reload per Ctrl-C). from mflux.models.z_image.variants.z_image import ZImage from mflux.utils.exceptions import StopImageGenerationException @@ -760,7 +760,7 @@ def test_a_python_api_call_runs_under_the_commands_caps_and_restores_mlxs_defaul def test_encode_runs_under_the_commands_caps_and_restores_mlxs_defaults(tmp_path, monkeypatch): - # Bug caught (X4): encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured + # Bug caught: encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured # phases, sized under the command's caps), or the caps left installed after it returns. from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap