diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 547f73b..98d163f 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 three steps gate each model.py on its own. + # No combined threshold: init.py's tests run in the default lane (its live ones skip without a token); the next four steps gate each model.py 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,13 +62,15 @@ 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-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.flux2.model --cov=mlx_dfloat.mflux.qwen21.model --cov=mlx_dfloat.mflux.qwen21.init --cov=mlx_dfloat.mflux.qwen21.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) run: uv run --extra mflux --group dev coverage report --include='*/mflux/zimage/model.py' --fail-under=80 - name: flux2/model.py coverage on its own (mflux lane) run: uv run --extra mflux --group dev coverage report --include='*/mflux/flux2/model.py' --fail-under=80 + - name: qwen21/model.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/qwen21/model.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 0367837..6b84ddc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,23 @@ All notable changes to this project are documented here. The format follows Mac's caps the fit check refuses the base 9B (`bench/results/refusals/flux2-klein-base-9b-1024-tier24.json`) and a forced run was stopped by the watchdog, so that tier is at risk. A base model and the distilled model of the same size have the same text-encoder and VAE files on the Hub, and the 4B pair the same tokenizer files too. +- Qwen-Image-2.1 generation through mflux: `mlx-dfloat generate --model qwen-image-2.1` and the `DFloatQwenImage21` + class, also returned by `DFloatModel`. The text encoder, VAE and tokenizer come from `Qwen/Qwen-Image-2.1`, pinned + to the revision the measured runs used. The defaults follow mflux: 40 steps, guidance 1.0, the `linear` scheduler. + Classifier-free guidance runs only with a guidance above 1.0 and a negative prompt, and the command warns when one + comes without the other. +- The default Qwen-Image-2.1 checkpoint, `mingyi456/Qwen-Image-2.1-DF11-ComfyUI`, is one ComfyUI file with no + `config.json`. The reader opens it through a layout pinned to that file's revision and refuses a file without a + `config.json` that it does not know. +- Every compressed matrix and stored tensor of the Qwen-Image-2.1 transformer (225 and 72) equals its BF16 original, + and the GPU decoder matches the CPU reference on all 225 (both records in `bench/results/parity/qwen-image-2.1/`). + The final latents match the BF16 transformer's bit for bit with classifier-free guidance on + (`bench/results/identity/qwen-image-2.1-1024/compare.json`). +- The measured 32 GB Qwen-Image-2.1 run peaks at 20.42 GiB at 1024² over 40 steps + (`bench/results/tiers/qwen-image-2.1-1024.json`) and keeps the compressed transformer loaded through the VAE + decode. Under a 16 GB and a 24 GB Mac's caps the fit check refuses it, predicting 14.6 GiB while the prompt is + encoded (`bench/results/refusals/qwen-image-2.1-1024-tier16.json` and `qwen-image-2.1-1024-tier24.json`), and forced + runs were stopped by the watchdog in that phase, so both sizes are at risk. - A FLUX.2 Klein base that lacks a text-encoder or VAE tensor is refused with `DFloatFormatError` instead of leaving that weight at its random initial value. - A FLUX.2 Klein call warns when the prompt encode holds more memory than the fit estimate allows for it. The @@ -63,6 +80,12 @@ All notable changes to this project are documented here. The format follows column is now headed "Working set − reserve", because on those rows it is the watchdog ceiling, not the fit budget. - The README starts with what the library is for, a figure of a FLUX.1 run, a comparison with mflux's BF16 and 8-bit paths, and a Python example. +- From Python, every model's `generate_image` runs under the memory caps `mlx-dfloat generate` installs and restores + MLX's limits when it returns. Before, a Python caller ran at MLX's default limits, with no wired cap, while the fit + estimates' VAE terms were measured under the caps. A wired limit the process already set is left alone. +- 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 `mlx-dfloat generate --help` text no longer prints reST markup, explains the `--eval-policy` choices, and says @@ -77,8 +100,8 @@ All notable changes to this project are documented here. The format follows - `mlx-dfloat generate --tier 16` and `--tier 24` checked the fit against the watchdog ceiling, the working set minus 1.5 GiB, which is looser than the budget `generate` applies on a real Mac of that size (the working set minus 2 GiB). The fit check now uses the real Mac's budget: 8.67 GiB at 16 GB and 14.00 GiB at 24 GB. The watchdog ceiling stays - at 9.17 and 14.50 GiB. The committed 16 GB and 24 GB results were recorded under the older budget, and none of - their outcomes changes under the new one. + at 9.17 and 14.50 GiB. The committed 16 GB and 24 GB FLUX.2 Klein and Z-Image-Turbo results were recorded under the + older budget, and none of their outcomes changes under the new one. ## [0.1.0] - 2026-10-01 diff --git a/README.md b/README.md index 766b971..67e74f5 100644 --- a/README.md +++ b/README.md @@ -7,15 +7,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 and FLUX.2 Klein arrive in the next release. +Z-Image, FLUX.2 Klein and Qwen-Image-2.1 arrive in the next release. 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 and FLUX.2 Klein follow in the -next release. +on the GPU as it runs. That is how it generates FLUX.1 images through mflux; Z-Image, FLUX.2 Klein and Qwen-Image-2.1 +follow in the next release. ``` one BF16 weight, 16 bits: s eeeeeeee mmmmmmm @@ -46,9 +46,11 @@ of the three FLUX.1 models; no other Mac has been measured. The numbers are unde | 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 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 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, Qwen-Image, 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 covers FLUX.1 only; Z-Image and FLUX.2 Klein arrive in the next release. | +| 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 and Qwen-Image-2.1. | ## Status @@ -94,6 +96,17 @@ In the next release (not yet on PyPI): FLUX.2-klein-base-9B and 30 for FLUX.2-klein-9B. Each model's two records are `kernel-vs-reference.json` and `sampled-vs-bf16.json` in `bench/results/parity//` (`flux2-klein-4b`, `flux2-klein-base-9b`, `flux2-klein-9b`). Their latents have not been compared with a reference. +- Qwen-Image-2.1 makes 1024² images through mflux, from the command line or from Python. Its DFloat11 checkpoint + (`mingyi456/Qwen-Image-2.1-DF11-ComfyUI`) is a single file made for ComfyUI, with no `config.json` next to it. + mlx-dfloat carries the missing layout for that one published file, pinned to one revision on the Hub, and refuses a + file without a `config.json` that it does not recognise. Every one of the 225 compressed matrices and 72 other tensors + equals the BF16 original in `Qwen/Qwen-Image-2.1` (`bench/results/parity/qwen-image-2.1/full-vs-bf16.json`), and the + GPU decoder gives the same bits as the CPU reference on all 225 + (`bench/results/parity/qwen-image-2.1/kernel-vs-reference.json`). Its final latents match, bit for bit, those of + the same transformer streaming its BF16 weights one block at a time, and the images match pixel for pixel. That check + ran at 1024² for 4 steps at guidance 4.0. The negative prompt was a single space, so classifier-free guidance ran + (record: `bench/results/identity/qwen-image-2.1-1024/compare.json`). `scripts/verify_image.py` runs both sides + without mflux's step compilation. - `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. @@ -110,7 +123,7 @@ 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 and FLUX.2 Klein generation. +`mlx-dfloat generate` command needs for FLUX.1, Z-Image, FLUX.2 Klein and Qwen-Image-2.1 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/): @@ -162,10 +175,10 @@ 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 or FLUX.2 Klein 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 or Qwen-Image-2.1 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. ### Generate a FLUX.1 image @@ -206,9 +219,10 @@ image.save("lighthouse.png") ``` FLUX.1-dev and FLUX.1-Krea-dev want more steps; the measured runs below used 20 steps and `guidance=3.5`. -`base_path=` takes the `--base` value, and `fit_check=False` replaces `--no-fit-check` (see below). From Python -nothing installs memory caps or a memory watchdog; `mlx-dfloat generate` does both, so prefer the command on a Mac -near its memory limit. +`base_path=` takes the `--base` value, and `fit_check=False` replaces `--no-fit-check` (see below). From Python, each +`generate_image` call runs under the same memory caps as `mlx-dfloat generate` and puts MLX's limits back when it +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, @@ -384,6 +398,68 @@ FLUX.1. FLUX.2 Klein's image-editing and KV-cache variants are not on this path. a warning: a base model above guidance 1.0 runs its negative branch on mflux's blank prompt and takes no custom one. Sizes above 1024² are refused unless you pass `--no-fit-check`, because none has been measured. +### Generate a Qwen-Image-2.1 image + +Qwen-Image-2.1 runs through the same command. Its BF16 repository and the DFloat11 checkpoint are both ungated on the +Hub, so no Hub login is needed. Both come under the Qwen Research License Agreement +([license](https://huggingface.co/Qwen/Qwen-Image-2.1/blob/main/LICENSE)). + +``` +mlx-dfloat generate --model qwen-image-2.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" \ + --seed 42 --height 1024 --width 1024 --report report.json +``` + +The defaults are mflux's own for this model: 40 steps, guidance 1.0 and the `linear` scheduler. At guidance 1.0 each +step calls the transformer once, with no classifier-free guidance. As in mflux, classifier-free guidance needs both a +`--guidance` above 1 and a `--negative-prompt`. Pass a guidance above 1 without a negative prompt and the command warns +that no classifier-free guidance runs; pass a negative prompt at guidance 1 or below and it warns that the prompt has no +effect. + +The command takes the DFloat11 transformer from that checkpoint and the text encoder, VAE and tokenizer from +`Qwen/Qwen-Image-2.1`, without downloading that repository's BF16 transformer. Both default 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("qwen-image-2.1") +image = model.generate_image( + seed=42, + prompt="A stone lighthouse on a rocky shore at dawn", + height=1024, + width=1024, +) +image.save("lighthouse.png") +``` + +`num_inference_steps` defaults to 40 here too. For classifier-free guidance pass a `guidance` above 1.0 and a +`negative_prompt`. + +mflux compiles Qwen-Image-2.1's denoising step. As with Z-Image and FLUX.2 Klein, this path runs the step +without that compilation (see "Generate a Z-Image image" for why), and what that costs against stock mflux was not +measured. + +Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0) at 1024², seed 42, 40 steps at guidance 1.0, +with the command under "Measured numbers": the run peaked at 20.42 GiB and took 6.9 minutes (the `elapsed_seconds` of +`bench/results/tiers/qwen-image-2.1-1024.json`). MLX's own memory peaked during the VAE decode. On this 32 GB Mac the +compressed transformer stays loaded through the VAE decode: the fit estimate predicts 21.2 GiB for that step (the +report's `fit.phases.vae`), under the 22.96 GiB budget. + +On a 16 GB or 24 GB Mac, Qwen-Image-2.1 is 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. It predicts a 14.6 GiB peak while the prompt is encoded, against a budget of 8.67 GiB at 16 GB and +14.00 GiB at 24 GB; the text encoder alone holds 14.1 GiB. Forced runs (`--no-fit-check`) were stopped by the +watchdog while the prompt was being encoded, before the transformer ran. At 16 GB that came after 4.4 s, when MLX's +active plus cached memory reached 9.28 GiB against the 9.17 GiB ceiling. At 24 GB it came after 4.8 s, with a +footprint of 14.59 GiB against the 14.50 GiB ceiling. + +Quantisation, LoRA, img2img, mflux's alternative image decoder (`--pid-decode`) and ControlNet are refused, as for +FLUX.1. Only text-to-image is on this path; Qwen-Image-2.1's image editing is not. 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 @@ -395,21 +471,24 @@ and 16 GiB), so it sees the limits a user of that Mac runs under. PROOF marks a ceiling, made only to show that the watchdog works; it never appears as a tier row. 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, while FLUX.2-klein-base-9B under a 24 GB Mac's caps and Z-Image-Turbo under a -16 GB Mac's 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, 50 with guidance 4.0 for Z-Image and the two FLUX.2 Klein -base models), with `--tier` and `--report` writing the file in the last column. +FLUX.2-klein-base-4B passed under both, while FLUX.2-klein-base-9B under a 24 GB Mac's caps, Z-Image-Turbo under a +16 GB Mac's and Qwen-Image-2.1 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, 50 with guidance 4.0 for Z-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 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 +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). "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 @@ -468,6 +547,24 @@ 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. +The Qwen-Image-2.1 rows come from these commands, with the same prompt: + +``` +mlx-dfloat generate --model qwen-image-2.1 --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/qwen-image-2.1-1024.json +mlx-dfloat generate --model qwen-image-2.1 --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 qwen-image-2.1 --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 +``` + +In the last two 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` both are refused with exit 2 +(`bench/results/refusals/qwen-image-2.1-1024-tier16.json` and `qwen-image-2.1-1024-tier24.json`). + Measured on an Apple M1 Max, 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0. @@ -481,10 +578,13 @@ Measured on an Apple M1 Max, 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0. | 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 | 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` | | 32 GB | 22.96 GiB | Z-Image | 7.80 GiB | 13.15 GiB | 13.15 GiB | 12.65 GiB | MEASURED | target | host caps | `bench/results/tiers/z-image-1024.json` | | 32 GB | 22.96 GiB | Z-Image-Turbo | 7.80 GiB | 13.11 GiB | 13.11 GiB | 12.20 GiB | MEASURED | target | host caps | `bench/results/tiers/z-image-turbo-1024.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 | 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` | diff --git a/bench/results/calibration/qwen-image-2.1-1024.json b/bench/results/calibration/qwen-image-2.1-1024.json new file mode 100644 index 0000000..1309719 --- /dev/null +++ b/bench/results/calibration/qwen-image-2.1-1024.json @@ -0,0 +1,144 @@ +{ + "model": "qwen-image-2.1", + "purpose": "the 1024² runs the Qwen-Image 2.1 memory constants were calibrated on (mlx_dfloat.mflux.qwen21.memory)", + "df11": "mingyi456/Qwen-Image-2.1-DF11-ComfyUI@1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "base": "Qwen/Qwen-Image-2.1@d26bb61231c349cf6b7896fa83353113880e1ba3", + "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": 33, + "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; for the 40-step and identity runs the run's watched peak, reached in the VAE phase", + "constants_from": { + "encode": "derisk-1", + "denoise": "ab-new-limit-1", + "vae": "ab-old-limit-2" + }, + "runs": [ + { + "label": "derisk-1", + "date": "2026-10-08", + "run": "the first de-risk run: one step at guidance 4 with the negative prompt \" \" (two transformer calls), the VAE decode on the resident set", + "steps": 1, + "cache_limit": 1859346402, + "footprint_peaks": { + "encode": 15857938288, + "denoise": 13526704528, + "vae": 22139643024 + } + }, + { + "label": "ab-new-limit-1", + "date": "2026-10-08", + "run": "cache-limit A/B, the planner's limit: one step at guidance 4 with the negative prompt \" \" (two transformer calls), the VAE decode on the resident set", + "steps": 1, + "cache_limit": 2731761634, + "footprint_peaks": { + "encode": 15857561480, + "denoise": 13939106072, + "vae": 21755487280 + } + }, + { + "label": "ab-old-limit-1", + "date": "2026-10-08", + "run": "cache-limit A/B, the earlier limit: one step at guidance 4 with the negative prompt \" \" (two transformer calls), the VAE decode on the resident set", + "steps": 1, + "cache_limit": 1859346402, + "footprint_peaks": { + "encode": 15868063600, + "denoise": 13943660944, + "vae": 22586942584 + } + }, + { + "label": "ab-new-limit-2", + "date": "2026-10-08", + "run": "cache-limit A/B, the planner's limit: one step at guidance 4 with the negative prompt \" \" (two transformer calls), the VAE decode on the resident set", + "steps": 1, + "cache_limit": 2731761634, + "footprint_peaks": { + "encode": 15871209352, + "denoise": 13506126128, + "vae": 22074123432 + } + }, + { + "label": "ab-old-limit-2", + "date": "2026-10-08", + "run": "cache-limit A/B, the earlier limit: one step at guidance 4 with the negative prompt \" \" (two transformer calls), the VAE decode on the resident set", + "steps": 1, + "cache_limit": 1859346402, + "footprint_peaks": { + "encode": 15866245000, + "denoise": 13633331456, + "vae": 22788580400 + } + }, + { + "label": "measured-40-steps-1", + "date": "2026-10-08", + "run": "mlx-dfloat generate --tier 32: 40 steps, guidance 1.0 (one transformer call per step)", + "steps": 40, + "cache_limit": 1859346402, + "footprint_peaks": { + "vae": 21942805432 + } + }, + { + "label": "measured-40-steps-2", + "date": "2026-10-08", + "run": "mlx-dfloat generate --tier 32: 40 steps, guidance 1.0 (one transformer call per step)", + "steps": 40, + "cache_limit": 2731761634, + "footprint_peaks": { + "vae": 21832147968 + } + }, + { + "label": "identity-4-steps-cfg-1", + "date": "2026-10-08", + "run": "the identity check's df11 side: 4 steps, guidance 4 with the negative prompt \" \"", + "steps": 4, + "cache_limit": 1859346402, + "footprint_peaks": { + "vae": 21977228264 + } + }, + { + "label": "identity-4-steps-cfg-2", + "date": "2026-10-08", + "run": "the identity check's df11 side: 4 steps, guidance 4 with the negative prompt \" \"", + "steps": 4, + "cache_limit": 2731761634, + "footprint_peaks": { + "vae": 22065292384 + } + }, + { + "label": "measured-40-steps-final", + "date": "2026-10-08", + "run": "mlx-dfloat generate --tier 32: 40 steps, guidance 1.0 (one transformer call per step)", + "steps": 40, + "cache_limit": 2731761634, + "footprint_peaks": { + "vae": 21929337856 + } + }, + { + "label": "identity-4-steps-cfg-final", + "date": "2026-10-08", + "run": "the identity check's df11 side: 4 steps, guidance 4 with the negative prompt \" \"", + "steps": 4, + "cache_limit": 2731761634, + "footprint_peaks": { + "vae": 21935465568 + } + } + ] +} diff --git a/bench/results/identity/qwen-image-2.1-1024/bf16-result.json b/bench/results/identity/qwen-image-2.1-1024/bf16-result.json new file mode 100644 index 0000000..2e48dc0 --- /dev/null +++ b/bench/results/identity/qwen-image-2.1-1024/bf16-result.json @@ -0,0 +1,57 @@ +{ + "exit_code": 0, + "reads": 1792, + "cfg_calls_per_step": 2, + "output_shape": [ + 1, + 4096, + 64 + ], + "output_dtype": "mlx.core.bfloat16", + "footprint_peak_bytes": 13110216840, + "mlx_peak_bytes_process": 12887490430, + "key": { + "model": "qwen-image-2.1", + "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": 4.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "base": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", + "source": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "mlx": "0.32.2", + "negative_prompt": " " + }, + "memory_caps_gb": [ + 20, + 22 + ], + "provenance": { + "git": "341c865e4b5ab1f6f611d489c9c40e1e640dcd3c", + "source_hash": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "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": 100 + }, + "date": "2026-10-08" + } +} diff --git a/bench/results/identity/qwen-image-2.1-1024/compare.json b/bench/results/identity/qwen-image-2.1-1024/compare.json new file mode 100644 index 0000000..be738da --- /dev/null +++ b/bench/results/identity/qwen-image-2.1-1024/compare.json @@ -0,0 +1,36 @@ +{ + "exit_code": 0, + "equal": true, + "degenerate": [], + "df11_key": { + "model": "qwen-image-2.1", + "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": 4.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "base": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", + "source": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "mlx": "0.32.2", + "negative_prompt": " " + }, + "bf16_key": { + "model": "qwen-image-2.1", + "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": 4.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "base": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", + "source": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "mlx": "0.32.2", + "negative_prompt": " " + }, + "pixels_identical": true, + "df11_png_sha256": "172f2c760124b619ba1656b939585899bec3050808cfcf8d8ae06d322206773e", + "bf16_png_sha256": "7522a96aa16f2687f0df52a53f7f417b09b55cd88ef068580d821ba03d505829" +} diff --git a/bench/results/identity/qwen-image-2.1-1024/df11-result.json b/bench/results/identity/qwen-image-2.1-1024/df11-result.json new file mode 100644 index 0000000..d030316 --- /dev/null +++ b/bench/results/identity/qwen-image-2.1-1024/df11-result.json @@ -0,0 +1,139 @@ +{ + "exit_code": 0, + "degenerate": [], + "report": { + "family": "qwen21", + "model": "qwen-image-2.1", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "repo_id": null, + "revision": "1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "config_source": "header and spot checks match layout qwen-image-2.1-comfyui (mingyi456/Qwen-Image-2.1-DF11-ComfyUI@1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc)" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", + "repo_id": null, + "revision": "d26bb61231c349cf6b7896fa83353113880e1ba3" + }, + "eval_policy": "per-block", + "eval_policy_status": "measured", + "cfg_calls_per_step": 2, + "text_tokens": 33, + "predict": "uncompiled", + "nonblock_groups": [ + "modulation.1" + ], + "cache_limit_in_force": 2731761634, + "drop_set_before_vae": false, + "fit": { + "label": "predicted", + "phases": { + "encode": 15662823906, + "denoise": 13939106072, + "vae": 22788580400 + }, + "peak_phase": "vae", + "peak_bytes": 22788580400, + "budget_bytes": 24653119488, + "fits": true + }, + "sizes": { + "compressed": 9587103379, + "extras": 137388032, + "nonblock": 134217728, + "encoders": 15136811008, + "vae": 1350989512 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 610518288, + "mlx_peak": 15353362024, + "footprint_end": 13757443256, + "active_at_start": 146825876 + }, + "set_load": { + "footprint_start": 13757443256, + "mlx_peak": 9869647680, + "footprint_end": 10344735096, + "active_at_start": 147416384 + }, + "denoise": { + "footprint_start": 10344735096, + "mlx_peak": 11532343934, + "footprint_end": 13060727112, + "active_at_start": 9776175680 + }, + "vae": { + "footprint_start": 12856713544, + "mlx_peak": 21174302250, + "footprint_end": 15545099360, + "active_at_start": 9782073950 + } + }, + "decode_launches": 256, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 14.079782124958001, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + } + }, + "output_shape": [ + 1, + 4096, + 64 + ], + "output_dtype": "mlx.core.bfloat16", + "footprint_peak_bytes": 21935465568, + "mlx_peak_bytes_max_over_phases": 21174302250, + "key": { + "model": "qwen-image-2.1", + "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": 4.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "base": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", + "source": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "mlx": "0.32.2", + "negative_prompt": " " + }, + "memory_caps_gb": [ + 20, + 22 + ], + "provenance": { + "git": "341c865e4b5ab1f6f611d489c9c40e1e640dcd3c", + "source_hash": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "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": 3507384, + "power": { + "ac": true, + "battery_pct": 100 + }, + "date": "2026-10-08" + } +} diff --git a/bench/results/parity/qwen-image-2.1/full-vs-bf16.json b/bench/results/parity/qwen-image-2.1/full-vs-bf16.json new file mode 100644 index 0000000..8629401 --- /dev/null +++ b/bench/results/parity/qwen-image-2.1/full-vs-bf16.json @@ -0,0 +1,61 @@ +{ + "key": { + "df11_revision": "1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "bf16_revision": "d26bb61231c349cf6b7896fa83353113880e1ba3", + "mode": "parity", + "decoder": "metal", + "source": "b59bf7889068d0726a90be25f1d32b58323efb1eff7d88edf95334998724eac2", + "git": "8d3c2c5350c32e6353bb5be1366ee2d034c45895", + "mlx": "0.32.2" + }, + "exit_code": 0, + "groups": 33, + "partial": false, + "selected_groups": [ + "modulation.1", + "transformer_blocks.0", + "transformer_blocks.1", + "transformer_blocks.2", + "transformer_blocks.3", + "transformer_blocks.4", + "transformer_blocks.5", + "transformer_blocks.6", + "transformer_blocks.7", + "transformer_blocks.8", + "transformer_blocks.9", + "transformer_blocks.10", + "transformer_blocks.11", + "transformer_blocks.12", + "transformer_blocks.13", + "transformer_blocks.14", + "transformer_blocks.15", + "transformer_blocks.16", + "transformer_blocks.17", + "transformer_blocks.18", + "transformer_blocks.19", + "transformer_blocks.20", + "transformer_blocks.21", + "transformer_blocks.22", + "transformer_blocks.23", + "transformer_blocks.24", + "transformer_blocks.25", + "transformer_blocks.26", + "transformer_blocks.27", + "transformer_blocks.28", + "transformer_blocks.29", + "transformer_blocks.30", + "transformer_blocks.31" + ], + "compared": 225, + "expected": 225, + "extras_compared": 72, + "mismatched": 0, + "missing_originals": [], + "uncovered_originals": [], + "ignored_originals": [], + "memory_caps_gb": [ + 20, + 22 + ], + "seconds": 20.16619499999797 +} diff --git a/bench/results/parity/qwen-image-2.1/kernel-vs-reference.json b/bench/results/parity/qwen-image-2.1/kernel-vs-reference.json new file mode 100644 index 0000000..ccce5d0 --- /dev/null +++ b/bench/results/parity/qwen-image-2.1/kernel-vs-reference.json @@ -0,0 +1,61 @@ +{ + "key": { + "df11_revision": "1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "bf16_revision": "none", + "mode": "kernel-vs-reference", + "decoder": "metal", + "source": "b59bf7889068d0726a90be25f1d32b58323efb1eff7d88edf95334998724eac2", + "git": "8d3c2c5350c32e6353bb5be1366ee2d034c45895", + "mlx": "0.32.2" + }, + "exit_code": 0, + "groups": 33, + "partial": false, + "selected_groups": [ + "modulation.1", + "transformer_blocks.0", + "transformer_blocks.1", + "transformer_blocks.2", + "transformer_blocks.3", + "transformer_blocks.4", + "transformer_blocks.5", + "transformer_blocks.6", + "transformer_blocks.7", + "transformer_blocks.8", + "transformer_blocks.9", + "transformer_blocks.10", + "transformer_blocks.11", + "transformer_blocks.12", + "transformer_blocks.13", + "transformer_blocks.14", + "transformer_blocks.15", + "transformer_blocks.16", + "transformer_blocks.17", + "transformer_blocks.18", + "transformer_blocks.19", + "transformer_blocks.20", + "transformer_blocks.21", + "transformer_blocks.22", + "transformer_blocks.23", + "transformer_blocks.24", + "transformer_blocks.25", + "transformer_blocks.26", + "transformer_blocks.27", + "transformer_blocks.28", + "transformer_blocks.29", + "transformer_blocks.30", + "transformer_blocks.31" + ], + "compared": 225, + "expected": 225, + "extras_compared": 0, + "mismatched": 0, + "missing_originals": [], + "uncovered_originals": [], + "ignored_originals": [], + "memory_caps_gb": [ + 20, + 22 + ], + "seconds": 422.43284683302045 +} diff --git a/bench/results/refusals/qwen-image-2.1-1024-tier16.json b/bench/results/refusals/qwen-image-2.1-1024-tier16.json new file mode 100644 index 0000000..454eb41 --- /dev/null +++ b/bench/results/refusals/qwen-image-2.1-1024-tier16.json @@ -0,0 +1,49 @@ +{ + "exit_code": 2, + "output": "q.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "qwen-image-2.1", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "base_path": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", + "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 14.6 GiB in the encode phase exceeds the budget 8.7 GiB (encode 14.6 GiB, denoise 13.0 GiB, vae 12.2 GiB); pass fit_check=False to run anyway", + "footprint_peak_bytes": 610862376, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 610862376, + "mlx_peak_bytes": 147444978, + "elapsed_seconds": 2.905 +} diff --git a/bench/results/refusals/qwen-image-2.1-1024-tier24.json b/bench/results/refusals/qwen-image-2.1-1024-tier24.json new file mode 100644 index 0000000..be90ae9 --- /dev/null +++ b/bench/results/refusals/qwen-image-2.1-1024-tier24.json @@ -0,0 +1,49 @@ +{ + "exit_code": 2, + "output": "q.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "qwen-image-2.1", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "base_path": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", + "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 14.6 GiB in the encode phase exceeds the budget 14.0 GiB (encode 14.6 GiB, denoise 13.0 GiB, vae 12.2 GiB); pass fit_check=False to run anyway", + "footprint_peak_bytes": 613549328, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 613549328, + "mlx_peak_bytes": 147446258, + "elapsed_seconds": 3.289 +} diff --git a/bench/results/tiers/aborts/qwen-image-2.1-1024-tier16.json b/bench/results/tiers/aborts/qwen-image-2.1-1024-tier16.json new file mode 100644 index 0000000..78d78b9 --- /dev/null +++ b/bench/results/tiers/aborts/qwen-image-2.1-1024-tier16.json @@ -0,0 +1,30 @@ +{ + "reason": "memory", + "footprint": 9558646368, + "peak_footprint": 9558646368, + "ceiling": 9842633386, + "elapsed": 4.377274958009366, + "budget": 1800.0, + "rss": 9139568640, + "mlx_active": 9949674348, + "mlx_cache": 14371993, + "verdict_memory": 9964046341, + "verdict_counter": "mlx", + "peak_watched": 9964046341, + "peak_mlx": 9964046341, + "context": { + "model": "qwen-image-2.1", + "height": 1024, + "width": 1024, + "seed": 42, + "steps": 40, + "tier_gb": 16, + "label": "CAPPED", + "limits": { + "memory": 10737418240, + "cache": 10737418240, + "wired": 8589934592 + }, + "phase": "encode" + } +} diff --git a/bench/results/tiers/aborts/qwen-image-2.1-1024-tier24.json b/bench/results/tiers/aborts/qwen-image-2.1-1024-tier24.json new file mode 100644 index 0000000..6d8979f --- /dev/null +++ b/bench/results/tiers/aborts/qwen-image-2.1-1024-tier24.json @@ -0,0 +1,30 @@ +{ + "reason": "memory", + "footprint": 15666704168, + "peak_footprint": 15666704168, + "ceiling": 15569256448, + "elapsed": 4.753989499993622, + "budget": 1800.0, + "rss": 15378350080, + "mlx_active": 15301874200, + "mlx_cache": 66220525, + "verdict_memory": 15666704168, + "verdict_counter": "footprint", + "peak_watched": 15666704168, + "peak_mlx": 15368094725, + "context": { + "model": "qwen-image-2.1", + "height": 1024, + "width": 1024, + "seed": 42, + "steps": 40, + "tier_gb": 24, + "label": "CAPPED", + "limits": { + "memory": 17179869184, + "cache": 16320875724, + "wired": 15032385536 + }, + "phase": "encode" + } +} diff --git a/bench/results/tiers/qwen-image-2.1-1024.json b/bench/results/tiers/qwen-image-2.1-1024.json new file mode 100644 index 0000000..3e161d7 --- /dev/null +++ b/bench/results/tiers/qwen-image-2.1-1024.json @@ -0,0 +1,131 @@ +{ + "exit_code": 0, + "output": "qwen-image-2.1.png", + "memory_caps_gb": [ + 20, + 22 + ], + "model_kwargs": { + "model": "qwen-image-2.1", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "base_path": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", + "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": "qwen21", + "model": "qwen-image-2.1", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "repo_id": null, + "revision": "1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "config_source": "header and spot checks match layout qwen-image-2.1-comfyui (mingyi456/Qwen-Image-2.1-DF11-ComfyUI@1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc)" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", + "repo_id": null, + "revision": "d26bb61231c349cf6b7896fa83353113880e1ba3" + }, + "eval_policy": "per-block", + "eval_policy_status": "measured", + "cfg_calls_per_step": 1, + "text_tokens": 33, + "predict": "uncompiled", + "nonblock_groups": [ + "modulation.1" + ], + "cache_limit_in_force": 2731761634, + "drop_set_before_vae": false, + "fit": { + "label": "predicted", + "phases": { + "encode": 15662823906, + "denoise": 13939106072, + "vae": 22788580400 + }, + "peak_phase": "vae", + "peak_bytes": 22788580400, + "budget_bytes": 24653119488, + "fits": true + }, + "sizes": { + "compressed": 9587103379, + "extras": 137388032, + "nonblock": 134217728, + "encoders": 15136811008, + "vae": 1350989512 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 611714344, + "mlx_peak": 15353361928, + "footprint_end": 13554035800, + "active_at_start": 146825876 + }, + "set_load": { + "footprint_start": 13554035800, + "mlx_peak": 9869450888, + "footprint_end": 10343883224, + "active_at_start": 147219592 + }, + "denoise": { + "footprint_start": 10343883224, + "mlx_peak": 11503639112, + "footprint_end": 12186051032, + "active_at_start": 9775978888 + }, + "vae": { + "footprint_start": 12172665304, + "mlx_peak": 21171435010, + "footprint_end": 15541412888, + "active_at_start": 9779206710 + } + }, + "decode_launches": 1280, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 13.601246499980334, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "footprint_peak_bytes": 21929337856, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 21929337856, + "mlx_peak_bytes": 19963475458, + "elapsed_seconds": 413.966 +} diff --git a/pyproject.toml b/pyproject.toml index 7324230..543f1aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -174,6 +174,8 @@ omit = [ "src/mlx_dfloat/mflux/zimage/model.py", # the same for FLUX.2 Klein's model class: every test is @pytest.mark.mflux, gated in the integration lane. "src/mlx_dfloat/mflux/flux2/model.py", + # the same for Qwen-Image 2.1's model class: every test is @pytest.mark.mflux, gated in the integration lane. + "src/mlx_dfloat/mflux/qwen21/model.py", ] [tool.coverage.report] diff --git a/scripts/verify_image.py b/scripts/verify_image.py index ba4650a..fd07d52 100644 --- a/scripts/verify_image.py +++ b/scripts/verify_image.py @@ -1,8 +1,8 @@ -"""Image identity: DFloat11 latents (FLUX.1, Z-Image, FLUX.2 Klein) against the same seam streaming the BF16 weights. +"""Image identity: DFloat11 latents against the same seam streaming the BF16 weights, for every BF16 family. -The ``df11`` side generates through ``DFloatFlux1``, ``DFloatZImage`` or ``DFloatFlux2Klein`` (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 +The ``df11`` side generates through ``DFloatFlux1``, ``DFloatZImage``, ``DFloatFlux2Klein`` or ``DFloatQwenImage21`` +(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 the shards as the block runs, never all resident) through exactly mflux's loop body on the saved embeddings and the same seed, then decodes with the base's VAE. ``compare`` checks the two latent @@ -15,6 +15,10 @@ repository (a distilled model's transformer lives in the distilled repository, not in the base one). The Klein bf16 side's copy of mflux's loop body is measurement glue, run by the identity check and not unit-tested: the df11 side runs mflux's own loop, so a drift between the two shows as a false mismatch (exit 1), never as a false pass. +For Qwen-Image 2.1, ``--base`` must hold the BF16 transformer (``transformer/``: its index and both shards) for the +bf16 side and ``--orchestrate``; the df11 side needs only the text encoder, the VAE and the DF11 file. Both Qwen sides +run the forward pass uncompiled, and the bf16 side's copy of mflux's Qwen loop body is measurement glue in the same +sense as Klein's. 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] @@ -120,6 +124,38 @@ def klein_base_problem(model: str, base: Path) -> str | None: return None +def qwen21_base_problem(model: str, base: Path) -> str | None: + """Why ``base`` cannot feed a Qwen-Image 2.1 bf16 side, or None (and None for every other family). + + The bf16 side streams the BF16 transformer from ``base/transformer``: its ``*.safetensors.index.json`` and every + shard the index names must be there, each named by a plain file name (as the streaming reader requires). The df11 + side never reads it. + """ + e = entry(model) + if e.family != "qwen21": + return None + root = base / "transformer" + need = ( + f"--base {base}: the bf16 side streams the BF16 transformer from transformer/ " + f"(download {e.base_repo}'s transformer/* first)" + ) + indexes = sorted(root.glob("*.safetensors.index.json")) if root.is_dir() else [] + if len(indexes) != 1: + return f"{need}; found {len(indexes)} weight indexes in {root}" + try: + weight_map = json.loads(indexes[0].read_text())["weight_map"] + shards = sorted({str(v) for v in weight_map.values()}) + except (OSError, ValueError, KeyError, TypeError, AttributeError) as exc: + return f"{indexes[0]}: unreadable weight index ({type(exc).__name__}); {need}" + unsafe = [s for s in shards if Path(s).name != s] + if unsafe: + return f"{indexes[0]}: a shard name that is not a plain file name: {unsafe[0]!r}" + missing = [s for s in shards if not (root / s).is_file()] + if not shards or missing: + return f"{need}; the index names {len(shards)} shards, missing {missing}" + return None + + 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( @@ -143,12 +179,12 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: type=float, 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)", + "FLUX.2 Klein 1.0, base models with CFG at 4; Qwen-Image 2.1 1.0, CFG above 1 with --negative-prompt)", ) p.add_argument( "--negative-prompt", default=None, - help="the negative prompt (used by z-image, a base model)", + help="the negative prompt (used by z-image, a base model, and qwen-image-2.1 above guidance 1)", ) p.add_argument("--eval-policy", choices=("per-block", "depth2"), default="per-block") p.add_argument( @@ -164,6 +200,11 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: problem = klein_base_problem(args.model, Path(args.base).expanduser()) if problem is not None: p.error(f"--model {args.model}: {problem}") + # 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": + problem = qwen21_base_problem(args.model, Path(args.base).expanduser()) + 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 args.guidance = entry(args.model).default_guidance # Absolute paths: the children run with the repository root as their cwd. @@ -802,10 +843,188 @@ def run_bf16_flux2(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, An } +# --- the Qwen-Image 2.1 sides (mflux-touching; exercised by the identity run, not unit-tested here) ----------- + + +def qwen21_embeds( + prompt_cache: Mapping[str, tuple[mx.array, mx.array]], prompts: tuple[str, ...] +) -> dict[str, mx.array]: + """The embeddings file of a Qwen df11 side: the prompt's pair, and the negative's when CFG ran (two prompts). + + ``prompts`` is ``DFloatQwenImage21.cfg_prompts``'s answer for the call; the pairs are ``(embeds, mask)`` as + mflux's prompt cache holds them. + """ + embeds, mask = prompt_cache[prompts[0]] + out = {"prompt_embeds": embeds, "prompt_mask": mask} + if len(prompts) > 1: + out["negative_prompt_embeds"], out["negative_prompt_mask"] = prompt_cache[prompts[1]] + return out + + +def run_df11_qwen21(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, Any]: + """Generate through ``DFloatQwenImage21``, capture the final latents, and save the image, latents and embeddings. + + The embeddings file holds ``prompt_embeds`` and ``prompt_mask`` and, when classifier-free guidance ran (a + guidance above 1.0 with a non-empty negative prompt), ``negative_prompt_embeds`` and ``negative_prompt_mask``. + + Raises: + VerifyImageError: The after-loop callback never ran (mflux's own loop never called it). + """ + from mflux.models.qwen21.latent_creator.qwen21_latent_creator import Qwen21LatentCreator + + from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 + + out_dir = args.out / "df11" + out_dir.mkdir(parents=True, exist_ok=True) + model = DFloatQwenImage21( + 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}) + wanted = model.cfg_prompts( + args.prompt, negative_prompt=args.negative_prompt, guidance=args.guidance + ) + mx.save_safetensors( + str(out_dir / "embeds.safetensors"), qwen21_embeds(model.prompt_cache, wanted) + ) + noise = Qwen21LatentCreator.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_qwen21(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, Any]: + """Stream the base's BF16 transformer shards through the seam and mflux's Qwen-Image 2.1 loop body, then the VAE. + + ``modulation.1`` loads from the base as a plain weight; the blocks stream as they run, the forward pass + uncompiled as on the df11 side (``EagerForward`` sets mflux's ``_step_fn`` to the plain method). + + Raises: + VerifyImageError: The df11 side's saved embeddings are not present yet. + """ + 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.qwen21.latent_creator.qwen21_latent_creator import Qwen21LatentCreator + from mflux.utils.image_util import ImageUtil + + from mlx_dfloat.mflux.qwen21 import init as qinit + from mlx_dfloat.mflux.qwen21 import transformer as qtf + from mlx_dfloat.mflux.qwen21.names import qwen21_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() + ckpt = open_checkpoint(df11_root) + index = qtf.base_transformer_files_index(base_root / "transformer") + build = qtf.build_transformer( + ckpt, extras=qtf.base_extras(index, ckpt), nonblock_from_extras=True + ) + provider = StreamingBF16Provider( + index, {n: ckpt.groups[n].matrix_names for n in build.shapes}, qwen21_name_map() + ) + transformer = build.transformer + transformer.attach(provider, build.shapes, eval_policy=args.eval_policy, verify_in_call=True) + + embeds = mx.load(str(embeds_path)) + prompt_embeds, prompt_mask = embeds["prompt_embeds"], embeds["prompt_mask"] + negative_prompt_embeds = embeds.get("negative_prompt_embeds") # present only when CFG ran + negative_prompt_mask = embeds.get("negative_prompt_mask") + do_true_cfg = negative_prompt_embeds is not None + 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="linear", + ) + # mflux qwen_image_21.py: txt2img noise, cast to the model precision. + latents = Qwen21LatentCreator.create_noise(args.seed, args.size, args.size).astype( + ModelConfig.precision + ) + mx.eval(latents, prompt_embeds, prompt_mask) + + for t in config.time_steps: + # mflux qwen_image_21.py's step body: scale, positive call, negative call when CFG ran, combine, step, eval. + latents = config.scheduler.scale_model_input(latents, t) + noise = transformer( + t=t, + config=config, + hidden_states=latents, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_mask, + ) + if do_true_cfg: + noise_negative = transformer( + t=t, + config=config, + hidden_states=latents, + encoder_hidden_states=negative_prompt_embeds, + encoder_hidden_states_mask=negative_prompt_mask, + ) + noise = noise_negative + config.guidance * (noise - noise_negative) + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) + mx.eval(latents) + # No verify_step(): attach(..., verify_in_call=True) checks inside each call, and this provider defers nothing. + + if not qtf.is_eager(transformer): + raise VerifyImageError("the bf16 side's forward pass was compiled; the df11 side's is not") + mx.save_safetensors(str(out_dir / "latents.safetensors"), {"latents": latents}) + + vae = qinit.load_vae(base_root) + mx.clear_cache() + mx.set_cache_limit(0) # the transformer-shaped buffers cannot serve the decoder + unpacked = Qwen21LatentCreator.unpack_latents( + latents=latents, height=args.size, width=args.size + ) + decoded = VAEUtil.decode(vae=vae, latent=unpacked, tiling_config=None) + ImageUtil.to_pil(decoded).save(str(out_dir / "image.png")) + + return { + "exit_code": EXIT_OK, + "reads": provider.reads, + "cfg_calls_per_step": 2 if do_true_cfg else 1, + "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), } diff --git a/scripts/verify_remote_group.py b/scripts/verify_remote_group.py index db81845..d883b5e 100644 --- a/scripts/verify_remote_group.py +++ b/scripts/verify_remote_group.py @@ -9,17 +9,22 @@ --groups first,last,max-block,max-code --out RESULT.json ``uv run python scripts/verify_remote_group.py ...`` works too. ``--cast-fp32-to-bf16`` compares against an FP32 original (Z-Image-Turbo's) rounded to BF16, nearest even. +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. """ import argparse +import dataclasses +import hashlib import json import math import re import struct import sys import traceback -from collections.abc import Callable +from collections.abc import Callable, Sequence from pathlib import Path from typing import Protocol @@ -34,6 +39,8 @@ from scripts._watchdog import Watchdog, default_ceiling from scripts.verify_checkpoint import VerifyError, natural_key + from mlx_dfloat import _layouts + from mlx_dfloat._layouts import SynthesizedLayout, identify_layout from mlx_dfloat._memory_caps import install_memory_caps from mlx_dfloat._safetensors import MAX_HEADER_BYTES, TensorInfo, parse_header from mlx_dfloat.errors import DFloatError @@ -43,9 +50,12 @@ MAX_LUT_ROWS, DF11Config, GroupArrays, - matrix_names_for, + config_for_layout, + group_headers, + insert_row_splits, n_blocks_for, parse_df11_config, + row_split_plan, validate_group_name, ) from mlx_dfloat.reference import ( @@ -204,6 +214,80 @@ def original_range(meta: object, name: str, *, cast_fp32: bool) -> tuple[int, in return offsets[0], offsets[1] - offsets[0], math.prod(shape), dtype +def _check_spot_checks( + src: RangeSource, name: str, infos: dict[str, TensorInfo], layout: SynthesizedLayout +) -> None: + """Fetch each of the layout's spot checks by a range read and compare its sha256.""" + if not layout.probes: + raise VerifyError(f"{name}: layout {layout.key} pins no spot check; refused") + for probe in layout.probes: + info = infos.get(probe.tensor) + if info is None: + raise VerifyError( + f"{name}: layout {layout.key} spot-checks {probe.tensor}, which the file lacks" + ) + if probe.length <= 0 or not 0 <= probe.offset <= info.nbytes - probe.length: + raise VerifyError( + f"{name}: layout {layout.key} spot check of {probe.tensor} at byte {probe.offset} " + f"({probe.length} bytes) lies outside its {info.nbytes} bytes" + ) + digest = hashlib.sha256( + src.read(name, info.offset + probe.offset, probe.length) + ).hexdigest() + if digest != probe.sha256: + raise VerifyError( + f"{name}: header matches layout {layout.key} ({layout.repo_id}@{layout.revision}), but " + f"the spot check of {probe.tensor} at byte {probe.offset} differs (sha256 {digest}); " + "the stored data is not the copy the layout was checked against" + ) + + +def config_from_source( + src: RangeSource, *, layouts: Sequence[SynthesizedLayout] +) -> tuple[DF11Config, str]: + """The repository's DF11 config and where it came from (``"config.json"`` or ``"layout "``). + + A ``config.json`` is read as it always was and no layout is consulted. Without one, the repository + must hold exactly one safetensors file whose raw header sha256, spot checks (range reads) and group + and extra counts match one of ``layouts``. + + Raises: + VerifyError: No config.json and not exactly one safetensors file; a header that matches no + layout; a spot check that differs or lies outside its tensor; other group or extra counts. + DFloatError: The header, the config or the layout is malformed. + KeyError: A config.json without a ``dfloat11_config`` block. + """ + files = src.files() + if "config.json" in files: + raw_config = json.loads(src.text("config.json"))["dfloat11_config"] + return parse_df11_config(raw_config, source="config.json"), "config.json" + shards = [n for n in files if n.endswith(".safetensors")] + if len(shards) != 1: + raise VerifyError( + f"no config.json and {len(shards)} safetensors files; a config-less checkpoint must be " + "a single file" + ) + name = shards[0] + length = _header_length(src, name) + raw = src.read(name, 8, length) + layout = identify_layout(raw, known=layouts) + if layout is None: + digest = hashlib.sha256(raw).hexdigest() + raise VerifyError( + f"{name}: no config.json and the header (sha256 {digest[:16]}...) matches no known layout" + ) + infos = parse_header(raw, data_start=8 + length, file_size=src.size(name), source=name) + config = config_for_layout(layout) + groups, extras = group_headers({Path(name): infos}, config) + if len(groups) != layout.groups or len(extras) != layout.extras: + raise VerifyError( + f"{name}: {len(groups)} groups and {len(extras)} extras; layout {layout.key} expects " + f"{layout.groups} groups and {layout.extras} extras" + ) + _check_spot_checks(src, name, infos, layout) + return config, f"layout {layout.key}" + + def check_small_fields(header: dict[str, TensorInfo], group: str) -> None: """Refuse luts / output_positions larger than any valid group has, before they are read.""" luts = header[f"{group}.luts"] @@ -340,7 +424,18 @@ def verify_group( arrays = _group_arrays(df11, file, header, group) record["max_code_length"] = max_code_length(arrays.luts) record["n_blocks"] = arrays.n_blocks - names = matrix_names_for(group, config.pattern_dict) + stored, names, plan = row_split_plan(group, config) + if arrays.split_positions.size + 1 != len(stored): + raise VerifyError( + f"{group}: stores {arrays.split_positions.size + 1} matrices but its pattern names " + f"{len(stored)}" + ) + arrays = dataclasses.replace( + arrays, + split_positions=insert_row_splits( + arrays.split_positions, int(arrays.sign_mantissa.size), plan, name=group + ), + ) parts = split_matrices(decode_group(arrays, name=group), arrays.split_positions) if bf16 is None: record["status"] = "structural-ok" @@ -383,13 +478,16 @@ def main( argv: list[str] | None = None, *, source_factory: Callable[[str, str, str], RangeSource] | None = None, + layouts: Sequence[SynthesizedLayout] | None = None, ) -> int: """CLI entry point; any failure other than a real mismatch exits 2. ``source_factory(repo, revision, subdir)`` builds the ``RangeSource`` for the DF11 and BF16 repos (subdir is always ``""`` for the DF11 side); it defaults to :class:`HfRangeSource`. Tests pass a factory that returns :class:`LocalRangeSource` over on-disk fixtures with - explicit ``--df11-revision``/``--bf16-revision``, so no network is touched. + explicit ``--df11-revision``/``--bf16-revision``, so no network is touched. ``layouts`` are the + pinned layouts a config-less DF11 repository may match (default: the package's table, read at + call time). """ factory = source_factory or ( lambda repo, revision, subdir: HfRangeSource(repo, revision, subdir) @@ -445,8 +543,8 @@ def main( api = HfApi() df11_rev = args.df11_revision or api.model_info(args.df11_repo).sha df11 = factory(args.df11_repo, df11_rev, "") - config = parse_df11_config( - json.loads(df11.text("config.json"))["dfloat11_config"], source="config.json" + config, config_source = config_from_source( + df11, layouts=_layouts.KNOWN_LAYOUTS if layouts is None else layouts ) bf16, bf16_rev, bindex = None, None, {} if mode == "parity": @@ -492,6 +590,7 @@ def main( "bf16_repo": args.bf16_repo, "bf16_revision": bf16_rev, "mode": mode, + "config_source": config_source, "control": "fp32 rounded to bf16 (nearest even)" if any(m.get("original_dtype") == "F32" for m in matrices) else "bf16", diff --git a/src/mlx_dfloat/mflux/__init__.py b/src/mlx_dfloat/mflux/__init__.py index 07cc76c..9fcf365 100644 --- a/src/mlx_dfloat/mflux/__init__.py +++ b/src/mlx_dfloat/mflux/__init__.py @@ -19,9 +19,10 @@ def require_mflux() -> None: def __getattr__(name: str) -> Any: - """``DFloatModel``, ``DFloatFlux1``, ``DFloatZImage`` and ``DFloatFlux2Klein`` are imported on first use. + """``DFloatModel`` and the model classes are imported on first use. - So ``import mlx_dfloat.mflux`` needs neither mflux nor MLX. + The classes are ``DFloatFlux1``, ``DFloatZImage``, ``DFloatFlux2Klein`` and ``DFloatQwenImage21``. So ``import + mlx_dfloat.mflux`` needs neither mflux nor MLX. """ if name == "DFloatModel": from mlx_dfloat.mflux.families import DFloatModel @@ -39,9 +40,14 @@ def __getattr__(name: str) -> Any: from mlx_dfloat.mflux.flux2.model import DFloatFlux2Klein return DFloatFlux2Klein + if name == "DFloatQwenImage21": + from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 + + return DFloatQwenImage21 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` and `mlx_dfloat.mflux.DFloatFlux2Klein` still work. +# `mlx_dfloat.mflux.DFloatZImage`, `mlx_dfloat.mflux.DFloatFlux2Klein` and `mlx_dfloat.mflux.DFloatQwenImage21` still +# work. __all__ = ["DFloatModel", "require_mflux"] diff --git a/src/mlx_dfloat/mflux/_phases.py b/src/mlx_dfloat/mflux/_phases.py index 825d87c..7664b79 100644 --- a/src/mlx_dfloat/mflux/_phases.py +++ b/src/mlx_dfloat/mflux/_phases.py @@ -20,7 +20,9 @@ class PhaseConstants: at 1024²; ``denoise_activation_at_reference`` the activation volume a denoise step holds beyond the cache limit at ``reference_tokens`` (image + text tokens). The allowance fields set the cache room for freed activation buffers: ``allowance_at_reference`` at ``allowance_reference_tokens``, never below ``allowance_floor``. - ``max_measured_pixels`` is the largest image the constants were measured at. + ``max_measured_pixels`` is the largest image the constants were measured at. ``encode_activation_bytes``, when + set, is what a prompt encode measured to hold beyond the encoder's weights, and sizes the encode phase's + activations; unset (None), the encode phase carries the denoise allowance instead. """ overhead_bytes: int @@ -31,14 +33,16 @@ class PhaseConstants: allowance_reference_tokens: int = 4096 + 256 allowance_floor: int = 500_000_000 max_measured_pixels: int = 1024 * 1024 + encode_activation_bytes: int | None = None @dataclass(frozen=True, slots=True, kw_only=True) class FamilySizes: """Resident bytes of the components: file sizes (exact), the extras' tensor sizes, the decoded non-block groups. - ``encoders`` is the text encoders' file size for FLUX.1 and Z-Image; FLUX.2 Klein replaces it with the bytes of - the encoder layers its prompt embedding reads (``flux2.memory.encoder_bytes_used``), which is less than the file. + ``encoders`` is the text encoders' file size for FLUX.1 and Z-Image. FLUX.2 Klein and Qwen-Image 2.1 replace it + with the bytes an encode makes resident, less than the files: Klein's the encoder layers its prompt embedding + reads (``flux2.memory.encoder_bytes_used``), Qwen's the language model's tensors (``qwen21.memory.encoder_bytes_used``). """ compressed: int @@ -85,9 +89,11 @@ def cache_limit_for( text_tokens: int, override: int | None = None, ) -> int: - """The MLX buffer-cache limit for one generate call: the two largest decoded groups next to the activations. + """The MLX buffer-cache limit for one generate call: the decoded groups next to the activations. - ``depth2`` adds one more of the largest decoded group (the look-ahead buffer); ``override`` wins outright. + ``largest`` holds the largest decoded group per block kind; the limit holds the two largest of them, so one group + for a one-kind family (per-block evaluation frees one block-sized buffer at a time). ``depth2`` adds one more of + the largest decoded group (the look-ahead buffer); ``override`` wins outright. """ if override is not None: return override @@ -131,13 +137,18 @@ def family_phases( vae_transient_bytes: int, overhead_bytes: int, denoise_activation_bytes: int = 0, + encode_activation_bytes: int | None = None, ) -> dict[str, dict[str, int]]: - """The three phases of a generation (``encode``, ``denoise``, ``vae``), each term once.""" + """The three phases of a generation (``encode``, ``denoise``, ``vae``), each term once. + + The encode phase's activations are ``encode_activation_bytes`` when given, else the denoise ``allowance``. + """ in_flight = max(largest.values()) * (2 if policy == "depth2" else 1) + encode_activations = allowance if encode_activation_bytes is None else encode_activation_bytes return { "encode": { "encoders": encoders_bytes, - "activations": allowance, + "activations": encode_activations, "overhead": overhead_bytes, }, "denoise": { @@ -176,8 +187,9 @@ def fit_for( ) -> FitEstimate: """The phase estimate for one generate call against ``budget`` (a prediction, labelled as such by the caller). - ``vae_with_set=False`` plans the VAE phase after the compressed set has been dropped (compressed, extras and - the decoded non-block weights all leave with it). + ``vae_with_set=False`` plans the VAE phase after the compressed set has been dropped: the compressed groups and + the decoded non-block weights leave with it; the extras stay (they are the transformer's own BF16 parameters, + loaded once at construction and not part of the set). """ phases = family_phases( compressed_bytes=sizes.compressed, @@ -194,8 +206,9 @@ def fit_for( denoise_activation_bytes=denoise_activation_bytes( c, height=height, width=width, text_tokens=text_tokens ), + encode_activation_bytes=c.encode_activation_bytes, ) if not vae_with_set: - for term in ("compressed", "extras", "nonblock"): + for term in ("compressed", "nonblock"): phases["vae"].pop(term) return fit_estimate(phases, budget_bytes=budget) diff --git a/src/mlx_dfloat/mflux/_pipeline.py b/src/mlx_dfloat/mflux/_pipeline.py index fcb603b..fe7dd42 100644 --- a/src/mlx_dfloat/mflux/_pipeline.py +++ b/src/mlx_dfloat/mflux/_pipeline.py @@ -5,13 +5,16 @@ decode. The family supplies its measured numbers (``PhaseConstants``) and sizes; nothing here imports mflux. """ +import functools import logging import math -from collections.abc import Callable, Mapping -from typing import Any, NoReturn +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from typing import Any, NoReturn, ParamSpec, TypeVar import mlx.core as mx +from mlx_dfloat._memory_caps import caps_for_recommended_bytes from mlx_dfloat._watchdog import phys_footprint from mlx_dfloat.errors import DFloatResourceError, DFloatUnsupportedError from mlx_dfloat.integrate.memory import CallPlan @@ -68,6 +71,56 @@ def refuse_construction_args( refuse("bake_lora", "LoRA of any kind") +@contextmanager +def call_caps(limits: Any = mx) -> Iterator[tuple[int, int]]: + """Run the block under the ``mlx-dfloat`` command's memory caps when no wired cap is in force; restore after. + + The Python API starts at MLX's default wired limit 0 (no cap), while every VAE term was measured under the caps + the command installs (``install_memory_caps``: with wired 0 a FLUX.2 Klein VAE decode's MLX peak rose from 6.89 + to 9.49 GiB). A wired limit already set (the command's own caps, or a CAPPED tier's) is left untouched. mlx 0.32.2 + has no getters: the wired limit is read by setting it to 0 and back (``bench.capped.current_limits``), which + leaves nothing behind. ``limits`` is the object holding the setters and the device report (``mlx.core``). + Yields the (wired, memory) caps in bytes installed for the block, (0, 0) when none were. + """ + wired_now = int(limits.set_wired_limit(0)) + limits.set_wired_limit(wired_now) + if wired_now != 0: + yield (0, 0) + return + try: + recommended = int(limits.device_info().get("max_recommended_working_set_size", 0)) + except Exception: # a device or backend that reports nothing: no caps, as install_memory_caps + recommended = 0 + wired, memory = caps_for_recommended_bytes(recommended) + if wired == 0: + yield (0, 0) + return + previous_memory = int(limits.set_memory_limit(memory)) + try: + limits.set_wired_limit(wired) + try: + yield (wired, memory) + finally: + limits.set_wired_limit(wired_now) + finally: + limits.set_memory_limit(previous_memory) + + +_P = ParamSpec("_P") +_R = TypeVar("_R") + + +def with_call_caps(method: Callable[_P, _R]) -> Callable[_P, _R]: + """Wrap a model's ``generate_image`` in ``call_caps``: the whole call, prelude to VAE decode, runs capped.""" + + @functools.wraps(method) + def capped(*args: _P.args, **kwargs: _P.kwargs) -> _R: + with call_caps(): + return method(*args, **kwargs) + + return capped + + class PhaseTracker: """Footprint and MLX peak per phase, sampled at the phase boundaries; one phase open at a time.""" @@ -171,7 +224,23 @@ def plan_call_for( side, side, ) - derived_minimum = sum(sorted(largest.values(), reverse=True)[:2]) + allowance = activation_allowance(constants, height=height, width=width, text_tokens=text_tokens) + if len(largest) == 1: + # One block kind: the limit holds one decoded group (two under depth2, the look-ahead's too), and the + # activations freed beside it must fit as well, or they may evict the decoded buffer and each block's output + # be allocated fresh. A prediction: Qwen-Image 2.1's A/B of real steps below this limit measured no slowdown. + groups = 2 if policy == "depth2" else 1 + derived_minimum = groups * max(largest.values()) + allowance + minimum_is = ( + "the largest decoded group twice, for depth2's look-ahead, and the activation allowance" + if groups == 2 + else "the largest decoded group and the activation allowance" + ) + consequence = "the cache may allocate each block's decode output fresh" + else: + derived_minimum = sum(sorted(largest.values(), reverse=True)[:2]) + minimum_is = "the two largest decoded groups" + consequence = "every block's decode output will be allocated fresh" limit = cache_limit_for( constants, largest, @@ -183,12 +252,12 @@ def plan_call_for( ) if limit < derived_minimum: log.warning( - "cache_limit %d is below the derived minimum %d (the two largest decoded groups): " - "every block's decode output will be allocated fresh", + "cache_limit %d is below the derived minimum %d (%s): %s", limit, derived_minimum, + minimum_is, + consequence, ) - allowance = activation_allowance(constants, height=height, width=width, text_tokens=text_tokens) common: dict[str, Any] = { "sizes": sizes, "largest": largest, diff --git a/src/mlx_dfloat/mflux/families.py b/src/mlx_dfloat/mflux/families.py index 58ded2b..87308bd 100644 --- a/src/mlx_dfloat/mflux/families.py +++ b/src/mlx_dfloat/mflux/families.py @@ -31,6 +31,12 @@ class ModelEntry: uses_negative_prompt: bool # A distilled model runs at this guidance only (the CLI refuses any other); None: any guidance. fixed_guidance: float | None = None + # Classifier-free guidance runs only with a negative prompt as well as a guidance above 1.0 (mflux's rule for + # Qwen-Image 2.1); the CLI warns when --guidance asks for it without --negative-prompt. + cfg_needs_negative: bool = False + # 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 @dataclass(frozen=True, slots=True, kw_only=True) @@ -79,6 +85,18 @@ def _flux2_names() -> NameMap: return klein_name_map() +def _qwen21_class() -> type: + from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 + + return DFloatQwenImage21 + + +def _qwen21_names() -> NameMap: + from mlx_dfloat.mflux.qwen21.names import qwen21_name_map + + return qwen21_name_map() + + FAMILIES: dict[str, FamilySpec] = { "flux1": FamilySpec( name="flux1", @@ -101,6 +119,12 @@ def _flux2_names() -> NameMap: load_model_class=_flux2_class, load_name_map=_flux2_names, ), + "qwen21": FamilySpec( + name="qwen21", + refused_flags={}, + load_model_class=_qwen21_class, + load_name_map=_qwen21_names, + ), } @@ -218,6 +242,25 @@ def _klein(name: str, label: str, steps: int, revision: str, *, base: bool) -> M "45a202a0bd19ec01ce8db2d5236890586cbac6a6", base=False, ), + # mflux 0.20.0: 40 steps (cli/defaults/defaults.py:52), guidance 1.0 (models/qwen21/cli/qwen21_generate.py:60), + # scheduler "linear" (cli/parser/parsers.py:180). A second transformer call per step runs only above guidance + # 1.0 with a non-empty negative prompt (models/qwen21/variants/txt2img/qwen_image_21.py:89). + ModelEntry( + name="qwen-image-2.1", + family="qwen21", + label="Qwen-Image-2.1", + df11_repo="mingyi456/Qwen-Image-2.1-DF11-ComfyUI", + df11_revision="1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + base_repo="Qwen/Qwen-Image-2.1", + default_steps=40, + default_guidance=1.0, + default_scheduler="linear", + cfg_two_calls=True, + uses_negative_prompt=True, + cfg_needs_negative=True, + # The snapshot the identity check and the MEASURED row ran on. + base_revision="d26bb61231c349cf6b7896fa83353113880e1ba3", + ), ) } @@ -245,7 +288,8 @@ def DFloatModel(name: str, /, **kwargs: Any) -> Any: # noqa: N802 """A ready DFloat11 model for a registered name: the family's model class built with ``name`` and ``kwargs``. Imports mflux. ``DFloatModel("z-image-turbo")`` returns a ``DFloatZImage``; ``DFloatModel("schnell")`` a - ``DFloatFlux1``; ``DFloatModel("flux2-klein-4b")`` a ``DFloatFlux2Klein``. + ``DFloatFlux1``; ``DFloatModel("flux2-klein-4b")`` a ``DFloatFlux2Klein``; ``DFloatModel("qwen-image-2.1")`` a + ``DFloatQwenImage21``. Args: name: A registered model name, such as ``"schnell"`` or ``"z-image-turbo"``. diff --git a/src/mlx_dfloat/mflux/flux1/memory.py b/src/mlx_dfloat/mflux/flux1/memory.py index 80041c5..1f4702f 100644 --- a/src/mlx_dfloat/mflux/flux1/memory.py +++ b/src/mlx_dfloat/mflux/flux1/memory.py @@ -169,7 +169,8 @@ def fit_for( ``height`` and ``width`` scale the VAE transient (above 1024²) and, with ``text_tokens``, the denoise activation term. ``vae_with_set=False`` plans the VAE phase after the compressed set has - been dropped (what a call does when the resident variant would not fit). + been dropped (what a call does when the resident variant would not fit); the extras stay, as the + transformer's own BF16 parameters loaded once at construction. """ phases = flux_phases( compressed_bytes=sizes.compressed, @@ -188,5 +189,4 @@ def fit_for( ) if not vae_with_set: phases["vae"].pop("compressed") - phases["vae"].pop("extras") return fit_estimate(phases, budget_bytes=budget) diff --git a/src/mlx_dfloat/mflux/flux1/model.py b/src/mlx_dfloat/mflux/flux1/model.py index 6a734ec..39063bd 100644 --- a/src/mlx_dfloat/mflux/flux1/model.py +++ b/src/mlx_dfloat/mflux/flux1/model.py @@ -2,8 +2,9 @@ The class subclasses mflux's ``Flux1`` so its loop, scheduler, callbacks, VAE decode and image metadata run unchanged; only construction and the prelude of ``generate_image`` differ. The text -encoders and the compressed transformer are never resident together (see ``lifecycle``). The -Python API installs no memory caps and no watchdog; the ``mlx-dfloat`` command does both. +encoders and the compressed transformer are never resident together (see ``lifecycle``). Each +``generate_image`` call runs under the memory caps the ``mlx-dfloat`` command installs (unless a wired +limit is already in force) and restores MLX's limits afterwards; only the command adds a watchdog. """ import logging @@ -28,7 +29,7 @@ from mlx_dfloat.integrate.memory import budget_bytes, largest_decoded_bytes from mlx_dfloat.integrate.names import NameMap, Shapes from mlx_dfloat.integrate.providers import DF11Provider -from mlx_dfloat.mflux import require_mflux +from mlx_dfloat.mflux import _pipeline, require_mflux from mlx_dfloat.mflux.flux1 import init as base_init from mlx_dfloat.mflux.flux1.init import BaseComponents, ResolvedRepo from mlx_dfloat.mflux.flux1.lifecycle import Lifecycle @@ -423,6 +424,7 @@ def encode(self, *prompts: str) -> None: """Encode prompts now, so several generations pay the encoder reload once (drops a resident set first).""" self._lifecycle.ensure_embeddings(*prompts) + @_pipeline.with_call_caps def generate_image( self, seed: int, diff --git a/src/mlx_dfloat/mflux/flux1/names.py b/src/mlx_dfloat/mflux/flux1/names.py index 908c8c4..3bdbe68 100644 --- a/src/mlx_dfloat/mflux/flux1/names.py +++ b/src/mlx_dfloat/mflux/flux1/names.py @@ -63,7 +63,8 @@ def check_flux_groups(ckpt: Any) -> tuple[int, int]: seen: dict[str, set[int]] = {DOUBLE_PREFIX: set(), SINGLE_PREFIX: set()} for name, group in ckpt.groups.items(): kind, _dot, idx = name.partition(".") - if kind not in seen or not (idx.isascii() and idx.isdigit()): + # A zero-padded index ("01") is refused: the seam asks for the group by its canonical name. + if kind not in seen or not (idx.isascii() and idx.isdigit() and idx == str(int(idx))): raise DFloatFormatError(f"{name}: not a FLUX block group") n = len(group.matrix_names) if n != MATRICES_PER_KIND[kind]: diff --git a/src/mlx_dfloat/mflux/flux2/model.py b/src/mlx_dfloat/mflux/flux2/model.py index f61864a..2ad24ba 100644 --- a/src/mlx_dfloat/mflux/flux2/model.py +++ b/src/mlx_dfloat/mflux/flux2/model.py @@ -2,8 +2,9 @@ The class subclasses mflux's ``Flux2Klein`` so its loop, scheduler, 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``). The Python API installs no memory -caps and no watchdog; the ``mlx-dfloat`` command does both. +transformer are never resident together (see ``mlx_dfloat.mflux.lifecycle``). Each ``generate_image`` call runs +under the memory caps the ``mlx-dfloat`` command installs (unless a wired limit is already in force) and restores +MLX's limits afterwards; only the command adds a watchdog. """ import logging @@ -396,6 +397,7 @@ def _encode_prompt_pair( negative_embeds, negative_ids = self._embeddings[str(negative_prompt)] return embeds, text_ids, negative_embeds, negative_ids + @_pipeline.with_call_caps def generate_image( self, seed: int, diff --git a/src/mlx_dfloat/mflux/flux2/names.py b/src/mlx_dfloat/mflux/flux2/names.py index fccc410..2c04bd1 100644 --- a/src/mlx_dfloat/mflux/flux2/names.py +++ b/src/mlx_dfloat/mflux/flux2/names.py @@ -78,7 +78,8 @@ def check_klein_groups(ckpt: Any) -> dict[str, int]: ) continue kind, _dot, idx = name.partition(".") - if kind not in seen or not (idx.isascii() and idx.isdigit()): + # A zero-padded index ("01") is refused: the seam asks for the group by its canonical name. + if kind not in seen or not (idx.isascii() and idx.isdigit() and idx == str(int(idx))): raise DFloatFormatError( f"{name}: not a FLUX.2 Klein block group (blocks: {', '.join(RUN_ORDER)}; non-block groups " f"accepted: {', '.join(NONBLOCK_GROUPS)})" diff --git a/src/mlx_dfloat/mflux/generate.py b/src/mlx_dfloat/mflux/generate.py index 36c1670..c6d478f 100644 --- a/src/mlx_dfloat/mflux/generate.py +++ b/src/mlx_dfloat/mflux/generate.py @@ -119,7 +119,7 @@ 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)", + "FLUX.2 Klein base 50, distilled 4, qwen-image-2.1 40)", ) p.add_argument("--height", type=int, default=1024) p.add_argument("--width", type=int, default=1024) @@ -129,19 +129,20 @@ def _add_arguments(p: argparse.ArgumentParser) -> None: default=None, help="guidance (default per model, as mflux: FLUX.1 dev 3.5, ignored by schnell and z-image-turbo; " "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)", + "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)", ) p.add_argument( "--scheduler", default=None, - help="scheduler (default per model: FLUX.1 linear; FLUX.2 Klein always runs " + help="scheduler (default per model: FLUX.1 and Qwen-Image 2.1 linear; FLUX.2 Klein always runs " "flow_match_euler_discrete and refuses the flag)", ) p.add_argument( "--negative-prompt", default=None, - help="used by z-image with --guidance above 1; accepted and ignored for the other models, " - "as mflux does", + help="used by z-image and qwen-image-2.1 with --guidance above 1; accepted and ignored for the " + "other models, as mflux does", ) p.add_argument( "--output", @@ -398,9 +399,21 @@ def run( ) print(f"warning: --negative-prompt is ignored: {why}", file=sys.stderr) elif args.negative_prompt and (_guidance(args) or 0.0) <= 1.0: + 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 " - "guidance 1.0 (the default is 0). Pass --guidance above 1.0 to enable it.", + f"guidance 1.0 (the default is {default}). Pass --guidance above 1.0 to enable it.", + file=sys.stderr, + ) + if ( + entry(args.model).cfg_needs_negative + and (_guidance(args) or 0.0) > 1.0 + and not args.negative_prompt + ): + print( + "warning: --guidance above 1.0 runs no classifier-free guidance for " + f"{entry(args.model).label} without --negative-prompt (as mflux); pass --negative-prompt " + "to enable it", file=sys.stderr, ) output = Path(args.output) diff --git a/src/mlx_dfloat/mflux/qwen21/__init__.py b/src/mlx_dfloat/mflux/qwen21/__init__.py new file mode 100644 index 0000000..3b263c4 --- /dev/null +++ b/src/mlx_dfloat/mflux/qwen21/__init__.py @@ -0,0 +1 @@ +"""Qwen-Image 2.1 on mflux from a DFloat11 transformer.""" diff --git a/src/mlx_dfloat/mflux/qwen21/init.py b/src/mlx_dfloat/mflux/qwen21/init.py new file mode 100644 index 0000000..dbe15ec --- /dev/null +++ b/src/mlx_dfloat/mflux/qwen21/init.py @@ -0,0 +1,176 @@ +"""The Qwen-Image 2.1 base repository without its transformer: text encoder, VAE and tokenizer, resolved and loaded. + +mflux's ``Qwen21Initializer.init`` wants every component under one path, so it would download and size the +transformer; 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/``. +""" + +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 DFloatFormatError +from mlx_dfloat.mflux import _hub, require_mflux +from mlx_dfloat.mflux._hub import ResolvedRepo as ResolvedRepo +from mlx_dfloat.mflux._hub import resolve as resolve + +# The ComfyUI DF11 repository holds one weight file and no config; the config pattern is harmless there. +DF11_PATTERNS: tuple[str, ...] = ("*.safetensors", "config.json") +ENCODER_PATTERNS: tuple[str, ...] = ("text_encoder/*.safetensors", "text_encoder/*.json") +VAE_PATTERNS: tuple[str, ...] = ("vae/*.safetensors", "vae/*.json") +# The tokenizer's own files in processor/ (transformers' AutoTokenizer for Qwen2), named one by one: the directory +# also holds the image and video preprocessor configs, which mflux never builds, and anything a later revision adds. +TOKENIZER_PATTERNS: tuple[str, ...] = ( + "processor/tokenizer.json", + "processor/tokenizer_config.json", + "processor/vocab.json", + "processor/merges.txt", + "processor/added_tokens.json", + "processor/special_tokens_map.json", + "processor/chat_template.jinja", +) +BASE_PATTERNS: tuple[str, ...] = ENCODER_PATTERNS + VAE_PATTERNS + TOKENIZER_PATTERNS + +_WHAT = "Qwen-Image 2.1 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 ``Qwen3VLRotaryEmbedding.__init__`` (qwen3_vl_rope.py:21). The +# VAE has none: every one of its parameters is a target of mflux's VAE mapping, and its latent mean and standard +# deviation are NumPy class attributes, not parameters (qwen21_vae.py:12-32). +COMPUTED: dict[str, frozenset[str]] = { + "text_encoder": frozenset({"rotary_emb.inv_freq"}), + "vae": frozenset(), +} +_SHOWN = 3 # missing names quoted in the error + + +@dataclass(slots=True, kw_only=True) +class Qwen21Components: + """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 Qwen-Image 2.1 components.""" + require_mflux() + from mflux.models.qwen21.weights.qwen21_weight_definition import Qwen21WeightDefinition + + return _hub.subset_definition(Qwen21WeightDefinition, names, patterns) + + +def encoder_definition() -> type: + """The text encoder component of mflux's Qwen-Image 2.1 definition, with its download patterns.""" + return _definition(("text_encoder",), ENCODER_PATTERNS) + + +def vae_definition() -> type: + """The VAE component of mflux's Qwen-Image 2.1 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 ``Qwen21TextEncoder`` 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 mapped; the vision tower and ``lm_head`` in the same shards are never loaded. + 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.qwen21.model.qwen21_text_encoder.qwen21_text_encoder import ( + Qwen21TextEncoder, + ) + + encoder = Qwen21TextEncoder(**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 ``Qwen21VAE`` module 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.qwen21.model.qwen21_vae.qwen21_vae import Qwen21VAE + + vae = Qwen21VAE() + 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 Qwen-Image 2.1 tokenizer (one, keyed ``"qwen21"``, read from ``processor/``). + + Raises: + DFloatFormatError: The tokenizer files are missing or incomplete. + """ + require_mflux() + from mflux.models.common.tokenizer.tokenizer_loader import TokenizerLoader + from mflux.models.qwen21.weights.qwen21_weight_definition import Qwen21WeightDefinition + + try: + return dict( + TokenizerLoader.load_all( + definitions=Qwen21WeightDefinition.get_tokenizers(), model_path=str(root) + ) + ) + except (OSError, RuntimeError) as exc: # OSError: transformers' error for a partial processor/ + raise DFloatFormatError(f"{root}: no usable Qwen-Image 2.1 tokenizer: {exc}") from exc + + +def load_base(root: Path) -> Qwen21Components: + """Text encoder, VAE and tokenizer of a base repository (weights lazy, every parameter from the base).""" + return Qwen21Components( + vae=load_vae(root), + text_encoder=load_text_encoder(root), + tokenizers=load_tokenizers(root), + ) diff --git a/src/mlx_dfloat/mflux/qwen21/memory.py b/src/mlx_dfloat/mflux/qwen21/memory.py new file mode 100644 index 0000000..7604007 --- /dev/null +++ b/src/mlx_dfloat/mflux/qwen21/memory.py @@ -0,0 +1,158 @@ +"""Qwen-Image 2.1 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``). Qwen-Image 2.1 differs in two inputs: +its text encoder's shards also hold a vision tower and an output head that mflux never loads (see +``encoder_bytes_used``), and its prompt length varies per call (up to the tokenizer's 2048 tokens), so the text +tokens are counted per prompt (``prompt_tokens``) rather than read from the model config. + +The transformer has one block kind, so the cache limit holds one decoded block next to the activations (per-block +evaluation frees one block-sized buffer at a time), and the family carries its own activation allowance. The VAE +transient is mflux's VAE decode, which runs in float32 (mflux 0.20.0 ``qwen21_vae.py:46-48``: the float32 latent +mean and standard deviation promote the latents): re-measure it when an mflux upgrade changes that dtype. No Qwen-Image 2.1 per-block overhead scenario is benched. +""" + +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.qwen21.names import NONBLOCK_GROUPS + +# The only text-encoder tensors mflux maps (its Qwen-Image 2.1 weight mapping names model.language_model.* only). +ENCODER_PREFIX = "model.language_model." +# The tokenizer's max_length (mflux's weight definition; tokenization truncates at it, system prefix included). +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 +# (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. +# Re-measure before changing any. Each run's phase peaks: bench/results/calibration/qwen-image-2.1-1024.json. +# +# REFERENCE_TOKENS: the calibration point, 1024² (4096 image tokens) + the run's text tokens, 33 (the prompt's; the +# negative prompt's 9 are fewer, and the call is sized by its longest prompt). +REFERENCE_TOKENS = 4096 + 33 +# OVERHEAD: the footprint MLX's counters do not see after the build phase (footprint - MLX active - MLX cache): +# 606_274_808 - 146_825_736 - 3_974. +OVERHEAD = 459_445_098 # 0.428 GiB +# 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 +# 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 +# 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. +VAE_TRANSIENT = 11_119_436_651 # 10.356 GiB +# 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, +# 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 +# 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 +# no measurable difference in step time or denoise footprint. +ALLOWANCE = 2_295_554_018 # 2.138 GiB +ALLOWANCE_TOKENS = 4096 + 33 +# ENCODE_ACTIVATION: what the prompt encode held beyond the language model's tensors by MLX's measure: the encode's +# MLX peak minus what was active at the phase's start (15_350_204_544 - 146_825_736 = 15_203_378_808) minus the language +# model (15_136_811_008). It sizes the encode phase's activations in place of the denoise allowance. By footprint the +# encode peaked 195_114_382 B above the phase this predicts (non-MLX memory the build-time overhead does not cover). +# One 33-token prompt: the term is not scaled with the prompt length. +ENCODE_ACTIVATION = 66_567_800 +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, +) +# What a prompt encode may hold beyond the estimate's encode phase before the model warns (the bound is the language +# model + ENCODE_ACTIVATION + this slack). The run's encode excess by footprint (peak - language model - overhead: +# 15_857_938_288 - 15_136_811_008 - 459_445_098 = 261_682_182, 249.6 MiB) rounded up to 256 MiB. By MLX's measure the +# 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). +ENCODE_ACTIVATION_TOKENS = 33 + + +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 language model + ``ENCODE_ACTIVATION``) 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 more of the encoder than its language model, or a prompt longer than the + one ``ENCODE_ACTIVATION`` 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 more of the encoder than its " + f"language model, 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 Qwen-Image 2.1 prompt encode makes resident, read from the safetensors headers. + + Only the language model's tensors (``model.language_model.*``) are loaded, and every one of its layers runs (the + prompt embedding is the final hidden state). The vision tower and the output head in the same shards are never + loaded, so they do not count. + + Raises: + DFloatFormatError: An encoder file's header cannot be read. + """ + return sum( + info.nbytes + for path in sorted((base_root / "text_encoder").glob("*.safetensors")) + for name, info in read_header(path).items() + if name.startswith(ENCODER_PREFIX) + ) + + +def prompt_tokens(tokenizer: Any, prompt: str) -> int: + """The text tokens mflux feeds the transformer for ``prompt`` (imports mflux). + + As mflux encodes it: a blank prompt becomes ``" "``, the templated prompt is tokenized (at most + ``TEXT_MAX_LENGTH`` tokens, the system prefix included) and the system prefix is dropped. + """ + require_mflux() + from mflux.models.qwen21.model.qwen21_text_encoder.qwen21_prompt_encoder import ( + Qwen21PromptEncoder, + ) + + text = prompt if prompt and prompt.strip() else " " + total = min(int(tokenizer.tokenize(text).input_ids.shape[1]), TEXT_MAX_LENGTH) + return total - int(Qwen21PromptEncoder._system_prefix_length(tokenizer)) + + +def sizes_for(ckpt: DF11Checkpoint, base_root: Path) -> FamilySizes: + """The family sizes of a Qwen-Image 2.1 checkpoint and base, the encoder counted by its language model.""" + sizes = family_sizes_for(ckpt, base_root, nonblock_groups=NONBLOCK_GROUPS) + return replace(sizes, encoders=encoder_bytes_used(base_root)) diff --git a/src/mlx_dfloat/mflux/qwen21/model.py b/src/mlx_dfloat/mflux/qwen21/model.py new file mode 100644 index 0000000..346bf1f --- /dev/null +++ b/src/mlx_dfloat/mflux/qwen21/model.py @@ -0,0 +1,547 @@ +"""``DFloatQwenImage21``: mflux's Qwen-Image 2.1 pipeline over a DFloat11 transformer, decoded one block at a time. + +The class subclasses mflux's ``QwenImage21`` so its loop, scheduler, 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``): prompts are encoded first into +mflux's own prompt cache, which its loop reads without touching the encoder. Each ``generate_image`` call runs under +the memory caps the ``mlx-dfloat`` command installs (unless a wired limit is already in force) and restores MLX's +limits afterwards; only the command adds a watchdog. +""" + +import logging +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._phases import FamilySizes +from mlx_dfloat.mflux.lifecycle import Lifecycle +from mlx_dfloat.mflux.qwen21 import init as qinit +from mlx_dfloat.mflux.qwen21 import memory as qmem +from mlx_dfloat.mflux.qwen21.init import Qwen21Components, ResolvedRepo +from mlx_dfloat.mflux.qwen21.names import NONBLOCK_GROUPS, qwen21_name_map +from mlx_dfloat.mflux.qwen21.transformer import Qwen21Build, build_transformer, is_eager + +require_mflux() +from mflux.models.qwen21.variants.txt2img.qwen_image_21 import QwenImage21 # noqa: E402 + +log = logging.getLogger("mlx_dfloat.mflux.qwen21") + +# (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 == "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 +# accepted, and the report labels it unmeasured. +MEASURED_POLICIES: frozenset[str] = frozenset({"per-block"}) +RETAINED_SLACK_BYTES = 2 * 1024**3 +# mflux 0.20.0 ``QwenImage21.generate_image``: its guidance and scheduler defaults. +DEFAULT_GUIDANCE = 1.0 +DEFAULT_SCHEDULER = "linear" + + +def _clear_frames(exc: BaseException) -> None: + """Release the locals of the frames ``exc`` and the exception it was raised while handling passed through.""" + traceback.clear_frames(exc.__traceback__) + if exc.__context__ is not None and exc.__context__.__traceback__ is not None: + traceback.clear_frames(exc.__context__.__traceback__) + + +class _EncoderNotResident: + """Takes the text encoder's place while it is dropped: mflux calls the encoder only for a prompt not cached.""" + + 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 + 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)" + ) + + +class DFloatQwenImage21(QwenImage21): # type: ignore[misc] # mflux ships no type information + """Qwen-Image 2.1 generation from a DFloat11 transformer through mflux. + + The transformer stays compressed and each of its 32 blocks is decoded on the GPU as it runs; the shared + modulation matrix is decoded once when the set loads. The forward pass runs uncompiled (mflux compiles it; 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 = "qwen-image-2.1", + *, + 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. ``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 Qwen-Image 2.1 DF11 checkpoint (a config-less file of + another layout 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. + 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 (a name map ambiguity, an uncovered extra). + """ + _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) + df11_default, base_default = MODELS[model] + # The default checkpoint lives on one person's Hub account: pinned to the commit that was verified here. + # A user's own checkpoint is taken as given. + pin = None if df11_path else registry.MODELS[model].df11_revision + df11 = qinit.resolve(df11_path or df11_default, patterns=qinit.DF11_PATTERNS, revision=pin) + ckpt = open_checkpoint(df11.root) + # The default base is pinned to the snapshot the recorded runs used; a user's own base is taken as given. + base_pin = None if base_path else registry.MODELS[model].base_revision + base = qinit.resolve( + base_path or base_default, patterns=qinit.BASE_PATTERNS, revision=base_pin + ) + # mflux builds the published transformer with its constructor defaults (Qwen21Initializer._init_models). + build = build_transformer(ckpt) + components = qinit.load_base(base.root) + self._assemble( + model=model, + model_config=model_config, + ckpt=ckpt, + df11=df11, + base=base, + components=components, + build=build, + sizes=qmem.sizes_for(ckpt, base.root), + eval_policy=eval_policy, + cache_limit=cache_limit, + fit_check=fit_check, + budget_bytes=budget_bytes, + ) + + @classmethod + def _from_parts(cls, **parts: Any) -> "DFloatQwenImage21": + """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: Qwen21Components, + build: Qwen21Build, + 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.qwen21.qwen21_initializer import Qwen21Initializer + + _pipeline.check_eval_policy(eval_policy, POLICIES) + nn.Module.__init__(self) # type: ignore[attr-defined] # not QwenImage21.__init__: mflux's initializer loads the whole base + Qwen21Initializer._init_config( + self, model_config + ) # prompt_cache, model_config, callbacks, tiling_config + self.vae = components.vae + self.text_encoder = components.text_encoder + self.tokenizers = components.tokenizers + self.transformer = build.transformer + self.bits = None + self._model = model + self._constants = qmem.CONSTANTS + self._ckpt, self._df11, self._base, self._sizes = ckpt, df11, base, sizes + self._shapes, self._nonblock = build.shapes, build.nonblock + self._decode = decode + self._names = qwen21_name_map() if name_map is None else name_map + self._cfg_calls = 1 + self._text_tokens: int | 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 pairs go into mflux's own prompt cache, keyed by the prompt as the caller wrote it: mflux's loop looks + # them up there before it would touch the (dropped) encoder. + 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.prompt_cache, + 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 pairs, the VAE, slack.""" + pairs = sum(int(a.nbytes) for arrays in self.prompt_cache.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 + pairs + vae + RETAINED_SLACK_BYTES + + def _load_encoders(self) -> None: + self.text_encoder = qinit.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 and mask, as mflux's ``Qwen21PromptEncoder.encode_prompt`` computes them. + + It is called with a cache of its own, so its normalisation (a blank prompt encoded as ``" "``) stays mflux's; + the lifecycle stores the pair under the prompt as given. + """ + from mflux.models.qwen21.model.qwen21_text_encoder.qwen21_prompt_encoder import ( + Qwen21PromptEncoder, + ) + + embeds, mask = Qwen21PromptEncoder.encode_prompt( + prompt=prompt, + prompt_cache={}, + tokenizer=self.tokenizers["qwen21"], + text_encoder=self.text_encoder, + ) + return embeds, mask + + def _load_set(self) -> None: + 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) + 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} + 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 ----------------------------------------------------------------------- + + def encode(self, *prompts: str) -> None: + """Encode prompts now, so several generations pay the encoder reload once (drops a resident set first).""" + self._lifecycle.ensure_embeddings(*dict.fromkeys(prompts)) + + 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 at this size (rounded down to multiples of 16, as mflux does). + + ``text_tokens`` is the call's longest prompt in tokens (``memory.prompt_tokens``). 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 + + def cfg_prompts( + self, prompt: str, *, negative_prompt: str | None, guidance: float | None + ) -> tuple[str, ...]: + """The prompts a call encodes: the prompt, and the negative prompt when classifier-free guidance runs. + + mflux 0.20's rule (``QwenImage21.generate_image``): the negative branch runs, and the transformer is called + twice per step, only for a guidance above 1.0 together with a non-empty negative prompt. + """ + if guidance is None or guidance <= 1.0 or not negative_prompt: + return (prompt,) + return (prompt, negative_prompt) + + @_pipeline.with_call_caps + def generate_image( + self, + seed: int, + prompt: str, + num_inference_steps: int = 40, + height: int = 1024, + width: int = 1024, + guidance: float | None = DEFAULT_GUIDANCE, + image_path: Path | str | None = None, + image_strength: float | None = None, + scheduler: str | None = DEFAULT_SCHEDULER, + negative_prompt: str | None = None, + ) -> Any: + """Mflux's ``generate_image`` behind a prelude: refusals, cache limit, fit check, prompt encoding, set load. + + ``guidance`` or ``scheduler`` given as None take mflux's defaults (1.0, ``"linear"``). As in mflux, + classifier-free guidance runs only for a guidance above 1.0 with a non-empty ``negative_prompt``. + + Raises: + DFloatUnsupportedError: ``image_path`` / ``image_strength`` (img2img). + 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 transformer's forward pass. + """ + if image_path is not None or image_strength is not None: + _pipeline.refuse("image_path/image_strength", "img2img") + guidance = DEFAULT_GUIDANCE if guidance is None else guidance + scheduler = DEFAULT_SCHEDULER if scheduler is None else scheduler + prompts = self.cfg_prompts(prompt, negative_prompt=negative_prompt, guidance=guidance) + tokenizer = self.tokenizers["qwen21"] + text_tokens = max(qmem.prompt_tokens(tokenizer, p) for p in 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, + image_path=None, + image_strength=None, + scheduler=scheduler, + negative_prompt=negative_prompt, + ) + 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 for the one it was raised while handling. + _clear_frames(exc) + 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 (language model + encode term) + slack. + + The estimate counts only the encoder's language-model tensors (``memory.encoder_bytes_used``); an mflux + change that also loads the vision tower or the output head would make it under-predict the encode phase, and + so may a prompt much longer than the 33 tokens the encode term was measured on (the warning names the call's + prompt length; ``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 = qmem.encode_peak_warning( + record["mlx_peak"] - record["active_at_start"], + qmem.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`` says whether the transformer's forward pass is the plain method (``"uncompiled"``) or a compiled + one (``"compiled"``, which the next call refuses), checked when the report is built. + ``cfg_calls_per_step`` and ``text_tokens`` are the last call's (2 calls when classifier-free guidance ran; + the longest prompt's tokens). ``sizes["encoders"]`` is the text encoder's language-model bytes, not its file + size. ``df11["config_source"]`` says where the checkpoint's config came from (a config-less file names the + pinned layout it matched). ``eval_policy_status`` is ``"unmeasured"`` for depth2, which no Qwen-Image 2.1 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": "qwen21", + "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, + }, + "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": "uncompiled" if is_eager(self.transformer) else "compiled", + "nonblock_groups": sorted(NONBLOCK_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/qwen21/names.py b/src/mlx_dfloat/mflux/qwen21/names.py new file mode 100644 index 0000000..969a054 --- /dev/null +++ b/src/mlx_dfloat/mflux/qwen21/names.py @@ -0,0 +1,160 @@ +"""Qwen-Image 2.1 naming: the block kind, the decoded sub-paths, the map derived from mflux, the group check. + +The block kind and its run order: mflux 0.20.0 ``Qwen21Transformer`` runs one list, ``transformer_blocks``. The +decoded sub-paths are those of the ComfyUI single-file DF11 export after its fused ``img_mlp.gate_up`` matrix is cut +into ``img_mlp.gate_layer`` and ``img_mlp.proj`` (the reader does the cut), so the map never sees ``gate_up``. mflux's +``Qwen21WeightMapping`` lists the block targets by concrete index rather than a ``{block}`` template; +``templated_targets`` folds them into one template per sub-path before the shared derivation runs. +""" + +import re +import types +from collections.abc import Iterable, Sequence +from typing import Any + +from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError +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 = "transformer_blocks" +RUN_ORDER: tuple[str, ...] = (KIND,) +MATRIX_SUBS: dict[str, tuple[str, ...]] = { + KIND: ( + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "img_mlp.gate_layer", + "img_mlp.proj", + "img_mlp.out", + ) +} +NONBLOCK_GROUPS: dict[str, tuple[str, ...]] = {"modulation.1": ("modulation.1.weight",)} +DEFAULT_LAYERS = 32 # mflux 0.20.0 ``Qwen21Transformer.__init__(num_layers=32)`` + + +def _sources(target: Any) -> list[str]: + src = target.from_pattern + return list(src) if isinstance(src, list | tuple) else [src] + + +def templated_targets(targets: Iterable[Any], *, kinds: Sequence[str], n_blocks: int) -> list[Any]: + """Fold concrete-index block targets into one ``{block}``-templated target per sub-path. + + A target whose target and single source both name ``..`` (the same ``i``) is folded with the + other blocks' targets of the same source sub-path. Every block ``0 .. n_blocks - 1`` must have one, each with the + same target sub-path and no transform. Other targets pass through unchanged, in their order. + + Raises: + DFloatIntegrationError: A block target with several sources, a transform, another index or kind on one side, + or another rename than block 0's; a sub-path missing for a block, given twice, or given beyond + ``n_blocks``. + """ + pattern = re.compile(rf"^({'|'.join(re.escape(k) for k in kinds)})\.([0-9]+)\.(.+)$") + out: list[Any] = [] + folded: dict[tuple[str, str], dict[int, str]] = {} + slot: dict[tuple[str, str], int] = {} + for target in targets: + sources = _sources(target) + to = pattern.match(target.to_pattern) + frm = [pattern.match(s) for s in sources] + if to is None and not any(frm): + out.append(target) + continue + if len(sources) != 1 or to is None or frm[0] is None: + raise DFloatIntegrationError( + f"{target.to_pattern} <- {sources}: a block target must be one block source in its own block" + ) + source = frm[0] + kind, index, sub = source.group(1), int(source.group(2)), source.group(3) + if (to.group(1), int(to.group(2))) != (kind, index): + raise DFloatIntegrationError( + f"{sources[0]}: mflux loads it into {target.to_pattern}, another block" + ) + if target.transform is not None: + raise DFloatIntegrationError(f"{kind}.{sub}: a transform on block {index}'s target") + if index >= n_blocks: + raise DFloatIntegrationError( + f"{kind}.{sub}: a target for block {index}, beyond the {n_blocks} blocks" + ) + per = folded.setdefault((kind, sub), {}) + if index in per: + raise DFloatIntegrationError(f"{kind}.{sub}: two targets for block {index}") + per[index] = to.group(3) + if (kind, sub) not in slot: + slot[(kind, sub)] = len(out) + out.append(None) + for (kind, sub), per in folded.items(): + for index in range(n_blocks): + if index not in per: + raise DFloatIntegrationError(f"{kind}.{sub}: no target for block {index}") + if per[index] != per[0]: + raise DFloatIntegrationError( + f"{kind}.{sub}: block {index} loads it into {per[index]}, block 0 into {per[0]}" + ) + out[slot[(kind, sub)]] = types.SimpleNamespace( + to_pattern=f"{kind}.{{block}}.{per[0]}", + from_pattern=[f"{kind}.{{block}}.{sub}"], + transform=None, + ) + return out + + +def qwen21_name_map() -> StaticNameMap: + """The map read from mflux's ``Qwen21WeightMapping`` (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.qwen21.weights.qwen21_weight_mapping import Qwen21WeightMapping + + return name_map_from_mapping( + templated_targets( + Qwen21WeightMapping.get_transformer_mapping(), + kinds=RUN_ORDER, + n_blocks=Qwen21WeightMapping.NUM_TRANSFORMER_BLOCKS, + ), + matrix_subs=MATRIX_SUBS, + ) + + +def check_qwen21_groups(ckpt: Any) -> dict[str, int]: + """The block count: block groups contiguous from 0, each holding the seven matrices, plus ``modulation.1``. + + ``modulation.1`` is optional: a checkpoint that stores it as a BF16 extra loads it as a plain weight. 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 seven, or a ``modulation.1`` group with other matrices. + """ + expected = MATRIX_SUBS[KIND] + 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 {group.matrix_names}, expected {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 Qwen-Image 2.1 group (blocks: {KIND}.; non-block groups accepted: " + f"{', '.join(NONBLOCK_GROUPS)})" + ) + if tuple(group.matrix_names) != tuple(f"{name}.{sub}.weight" for sub in expected): + raise DFloatFormatError( + f"{name}: holds {len(group.matrix_names)} matrices {list(group.matrix_names)}; a Qwen-Image 2.1 " + f"block holds {', '.join(expected)} (the fused img_mlp.gate_up cut into its halves)" + ) + 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)} diff --git a/src/mlx_dfloat/mflux/qwen21/transformer.py b/src/mlx_dfloat/mflux/qwen21/transformer.py new file mode 100644 index 0000000..cc67bf0 --- /dev/null +++ b/src/mlx_dfloat/mflux/qwen21/transformer.py @@ -0,0 +1,214 @@ +"""mflux's Qwen-Image 2.1 transformer with the class-swap seam and its forward pass kept uncompiled. + +mflux 0.20.0 ``Qwen21Transformer.__call__`` compiles its ``_forward`` on the first call, on every chip, and keeps it +in ``_step_fn``. The block seam evaluates at block boundaries, which cannot happen inside ``mx.compile``, so +``EagerForward`` sets ``_step_fn`` to the plain bound ``_forward`` first and refuses any other value. + +The build installs placeholders on the seven block matrices and on ``modulation.1`` (a non-block group, decoded once +when the set loads), loads every other parameter from the checkpoint's extras (or a BF16 base's) and leaves the +parameters mflux computes at construction as they are. +""" + +import types +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from functools import cache +from pathlib import Path +from typing import Any + +import mlx.core as mx +from mlx.utils import tree_flatten + +from mlx_dfloat._safetensors import TensorInfo +from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError +from mlx_dfloat.format import DF11Checkpoint +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 NameMap, Shapes +from mlx_dfloat.integrate.placeholders import PLACEHOLDER, install_placeholders +from mlx_dfloat.integrate.resident import NonBlockShapes, 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.qwen21.names import ( + DEFAULT_LAYERS, + KIND, + NONBLOCK_GROUPS, + RUN_ORDER, + check_qwen21_groups, + qwen21_name_map, +) + +MAX_BUILD_ACTIVE_BYTES = 2 * 1024**3 + +# Parameters mflux computes when it builds the transformer; no checkpoint holds them, so the extras coverage leaves +# them out: the RoPE tables (mflux 0.20.0 qwen21_rope.py:21-26) and the timestep frequencies +# (qwen21_time_text_embed.py:16). +COMPUTED_PARAMS: frozenset[str] = frozenset( + { + "pos_embed.cos_tables.0", + "pos_embed.cos_tables.1", + "pos_embed.cos_tables.2", + "pos_embed.sin_tables.0", + "pos_embed.sin_tables.1", + "pos_embed.sin_tables.2", + "time_text_embed.time_proj.freqs", + } +) + + +def is_eager(transformer: Any) -> bool: + """Whether the transformer's ``_step_fn`` is unset or its own plain, uncompiled ``_forward`` (imports mflux). + + ``mx.compile`` returns ``mlx.gc_func``, a ``FunctionType`` subclass, so only an exact type check tells a compiled + function from a bound method. + """ + require_mflux() + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + fn = transformer._step_fn + return fn is None or ( + type(fn) is types.MethodType + and fn.__func__ is Qwen21Transformer._forward + and fn.__self__ is transformer + ) + + +class EagerForward: + """Compose in front of ``Qwen21Transformer``: its forward pass runs as the plain method, never compiled.""" + + _step_fn: Any + _forward: Any + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Set ``_step_fn`` to the plain ``_forward`` when unset, then run mflux's own call. + + Raises: + DFloatIntegrationError: ``_step_fn`` is set to anything else (a compiled forward pass). + """ + if self._step_fn is None: + self._step_fn = self._forward + elif not is_eager(self): + fn = self._step_fn + raise DFloatIntegrationError( + f"Qwen21Transformer._step_fn is {type(fn).__module__}.{type(fn).__qualname__}: mflux compiled the " + "forward pass, and the block seam cannot evaluate inside mx.compile" + ) + return super().__call__(*args, **kwargs) # type: ignore[misc] + + +@cache +def seam_transformer_class() -> type: + """``StepSeam`` and ``EagerForward`` composed in front of mflux's ``Qwen21Transformer`` (imports mflux). + + Raises: + DFloatDependencyError: The optional ``mflux`` extra is not installed. + """ + require_mflux() + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + return type("SeamQwen21Transformer", (StepSeam, EagerForward, Qwen21Transformer), {}) + + +def block_lists(transformer: Any) -> list[tuple[str, Sequence[Any]]]: + """The one block list, the one ``_forward`` iterates.""" + return [(kind, getattr(transformer, kind)) for kind in RUN_ORDER] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Qwen21Build: + """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, + *, + 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, +) -> Qwen21Build: + """Construct the seamed transformer: block and non-block matrices as placeholders, every other parameter loaded. + + ``transformer_kwargs`` are the transformer's constructor arguments (none for the published model); the + checkpoint's block count must equal the depth they build (mflux's default 32 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`` ``modulation.1`` loads from them as a + plain weight instead of a placeholder (the BF16 side of an identity check). A checkpoint that stores + ``modulation.1`` as an extra loads it as a plain weight either way. The parameters mflux computes at construction + (``COMPUTED_PARAMS``) are left out of the coverage check and keep their values. The active memory the build added + is asserted. + + Raises: + DFloatFormatError: The groups are not Qwen-Image 2.1'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). + """ + names = qwen21_name_map() if name_map is None else name_map + counts = check_qwen21_groups(ckpt) + kwargs = dict(transformer_kwargs or {}) + want = int(kwargs.get("num_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]}") + make = seam_transformer_class() if transformer_class is None else transformer_class + before = int(mx.get_active_memory()) + tf = make(**{**kwargs, "num_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 Qwen21Build(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, by their decoded names (``modulation.1`` stays). + + The block groups' ``matrix_names`` are the decoded ones (``img_mlp.gate_layer``, ``img_mlp.proj``), which is how + the base names them, so a fused stored matrix never leaves its halves behind as extras. + """ + 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} + + +__all__ = [ + "COMPUTED_PARAMS", + "MAX_BUILD_ACTIVE_BYTES", + "EagerForward", + "Qwen21Build", + "base_extras", + "base_transformer_files_index", + "block_lists", + "build_transformer", + "is_eager", + "seam_transformer_class", +] diff --git a/src/mlx_dfloat/mflux/zimage/memory.py b/src/mlx_dfloat/mflux/zimage/memory.py index 7246d6e..7ed1a8e 100644 --- a/src/mlx_dfloat/mflux/zimage/memory.py +++ b/src/mlx_dfloat/mflux/zimage/memory.py @@ -171,8 +171,9 @@ def fit_for( ) -> FitEstimate: """The phase estimate for one generate call against ``budget`` (a prediction, labelled as such by the caller). - ``vae_with_set=False`` plans the VAE phase after the compressed set has been dropped (compressed, - extras and the decoded non-block weight all leave with it). + ``vae_with_set=False`` plans the VAE phase after the compressed set has been dropped: the compressed + groups and the decoded non-block weight leave with it; the extras stay (the transformer's own BF16 + parameters, loaded once at construction). """ phases = zimage_phases( compressed_bytes=sizes.compressed, @@ -191,6 +192,6 @@ def fit_for( ), ) if not vae_with_set: - for term in ("compressed", "extras", "nonblock"): + for term in ("compressed", "nonblock"): phases["vae"].pop(term) return fit_estimate(phases, budget_bytes=budget) diff --git a/src/mlx_dfloat/mflux/zimage/model.py b/src/mlx_dfloat/mflux/zimage/model.py index 6c6b657..228cd13 100644 --- a/src/mlx_dfloat/mflux/zimage/model.py +++ b/src/mlx_dfloat/mflux/zimage/model.py @@ -2,8 +2,9 @@ The class subclasses mflux's ``ZImage`` so its loop, scheduler, 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``). The Python API installs no memory -caps and no watchdog; the ``mlx-dfloat`` command does both. +transformer are never resident together (see ``mlx_dfloat.mflux.lifecycle``). Each ``generate_image`` call runs +under the memory caps the ``mlx-dfloat`` command installs (unless a wired limit is already in force) and restores +MLX's limits afterwards; only the command adds a watchdog. """ import logging @@ -33,8 +34,8 @@ 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 import require_mflux from mlx_dfloat.mflux._compile import predict_mode, uncompiled from mlx_dfloat.mflux.lifecycle import Lifecycle from mlx_dfloat.mflux.zimage import init as zinit @@ -461,6 +462,7 @@ def _encode_prompts( text = self._embeddings[wanted[0]][0] return text, (self._embeddings[wanted[1]][0] if len(wanted) > 1 else None) + @_pipeline.with_call_caps def generate_image( self, seed: int, diff --git a/src/mlx_dfloat/mflux/zimage/names.py b/src/mlx_dfloat/mflux/zimage/names.py index 56f98b1..a4fbe04 100644 --- a/src/mlx_dfloat/mflux/zimage/names.py +++ b/src/mlx_dfloat/mflux/zimage/names.py @@ -62,7 +62,8 @@ def check_zimage_groups(ckpt: Any) -> dict[str, int]: ) continue kind, _dot, idx = name.partition(".") - if kind not in seen or not (idx.isascii() and idx.isdigit()): + # A zero-padded index ("01") is refused: the seam asks for the group by its canonical name. + if kind not in seen or not (idx.isascii() and idx.isdigit() and idx == str(int(idx))): raise DFloatFormatError(f"{name}: not a Z-Image block group") want = {f"{name}.{sub}.weight" for sub in MATRIX_SUBS[kind]} if set(group.matrix_names) != want or len(group.matrix_names) != len(want): diff --git a/tests/_df11_fixtures.py b/tests/_df11_fixtures.py index 68436be..fc8ab7f 100644 --- a/tests/_df11_fixtures.py +++ b/tests/_df11_fixtures.py @@ -100,6 +100,62 @@ def write_checkpoint( return root +def pin_layout( + file, + *, + pattern_dict, + row_splits, + groups, + extras, + key="tiny", + probe_tensor=None, + probe_offset=0, + probe_length=8, +): + """A ``SynthesizedLayout`` pinned to a written config-less ``file``, as the package pins a published one. + + The header digest is the file's own; one spot check covers ``probe_length`` bytes of ``probe_tensor`` (default: + the first group's ``sign_mantissa``) at ``probe_offset``. ``file_sha256`` is a placeholder: a plain file name is + never a Hugging Face blob name, so the whole-file pin is not consulted. + """ + import hashlib + + from mlx_dfloat._layouts import ContentProbe, SynthesizedLayout + from mlx_dfloat._safetensors import parse_header, read_header_bytes + + raw, start, size = read_header_bytes(file) + infos = parse_header(raw, data_start=start, file_size=size, source=file.name) + tensor = probe_tensor or next(n for n in sorted(infos) if n.endswith(".sign_mantissa")) + at = infos[tensor].offset + probe_offset + data = file.read_bytes()[at : at + probe_length] + return SynthesizedLayout( + key=key, + label=key, + repo_id="t/t", + revision="0" * 40, + file_name=file.name, + header_sha256=hashlib.sha256(raw).hexdigest(), + file_sha256="0" * 64, + groups=groups, + extras=extras, + raw_config={ + "version": "0.5.0", + "threads_per_block": [512], + "bytes_per_thread": 8, + "pattern_dict": {k: list(v) for k, v in pattern_dict.items()}, + }, + row_splits=dict(row_splits), + probes=( + ContentProbe( + tensor=tensor, + offset=probe_offset, + length=probe_length, + sha256=hashlib.sha256(data).hexdigest(), + ), + ), + ) + + def write_bf16_original(root, tensors, *, with_index=True, dtype=mx.bfloat16): """Write 'original' tensors (uint16 bits) as up to two shards plus an HF-style index.""" root.mkdir(parents=True, exist_ok=True) diff --git a/tests/_mlx_limits.py b/tests/_mlx_limits.py new file mode 100644 index 0000000..a3daf3e --- /dev/null +++ b/tests/_mlx_limits.py @@ -0,0 +1,35 @@ +"""MLX's memory limits as a Python-API user starts with them, for the per-call caps tests.""" + +from collections.abc import Iterator +from contextlib import contextmanager + +import mlx.core as mx + +from mlx_dfloat._memory_caps import compute_safe_caps_gb +from mlx_dfloat.bench.capped import current_limits + +GIB = 1024**3 + +__all__ = ["command_caps", "current_limits", "mlx_without_wired_cap"] + + +def command_caps() -> tuple[int, int]: + """The (wired, memory) caps in bytes the ``mlx-dfloat`` command installs here (``install_memory_caps``).""" + wired_gb, memory_gb = compute_safe_caps_gb() + return (wired_gb * GIB, memory_gb * GIB) + + +@contextmanager +def mlx_without_wired_cap() -> Iterator[dict[str, int]]: + """MLX's default wired limit (0, no cap) for the block, as a Python-API user's process starts. + + The suite installs the caps at import (``tests/conftest.py``); this lifts the wired one and puts every limit back + afterwards. Yields the limits in force on entry, the wired one then 0. + """ + before = current_limits() + mx.set_wired_limit(0) + try: + yield {**before, "wired": 0} + finally: + mx.set_memory_limit(before["memory"]) + mx.set_wired_limit(before["wired"]) diff --git a/tests/_qwen21_tiny.py b/tests/_qwen21_tiny.py new file mode 100644 index 0000000..638d364 --- /dev/null +++ b/tests/_qwen21_tiny.py @@ -0,0 +1,284 @@ +"""Tiny Qwen-Image 2.1 transformer helpers for the mflux lane. Every mflux import happens inside a function.""" + +import zlib +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__ = [ + "PREFIX_IDS", + "STORED_SUBS", + "TINY", + "TINY_ROW_SPLITS", + "TINY_SEED", + "StubTextEncoder", + "StubTokenizer", + "StubVAE", + "fake_model", + "stub_components", + "tiny_config", + "tiny_groups", + "tiny_inputs", + "tiny_seamed", + "write_tiny_checkpoint", +] + +# Head dim 16 = the sum of the three RoPE axes (4 + 6 + 6, each even: mflux 0.20.0 qwen21_rope.py:23-26 builds one +# table per axis at half its width); inner dim 2 x 16 = 32. in_channels stays the default 64: Qwen21LatentCreator +# always makes 64-channel latents (latent_creator/qwen21_latent_creator.py:13-14), and the model-level tests run +# mflux's own loop. +TINY = { + "num_layers": 2, + "attention_head_dim": 16, + "num_attention_heads": 2, + "context_in_dim": 32, + "mlp_ratio": 3, + "axes_dims_rope": (4, 6, 6), +} + + +def tiny_config(steps=2, size=64): + """mflux's ``Config`` for Qwen-Image 2.1 at ``size`` pixels square (4 x 4 = 16 latent tokens at 64).""" + from mflux.models.common.config.config import Config + from mflux.models.common.config.model_config import ModelConfig + + return Config( + model_config=ModelConfig.qwen_image_21(), + num_inference_steps=steps, + height=size, + width=size, + guidance=1.0, + scheduler="linear", + ) + + +def tiny_inputs(seed=0): + """The keyword inputs of ``Qwen21Transformer.__call__``: 16 image tokens of 64 channels, 8 prompt tokens of 32. + + The mask is all ones, so mflux takes its padding-free path (``qwen21_transformer.py:117-123``). + """ + keys = mx.random.split(mx.random.key(seed), 2) + return { + "t": 0, + "config": tiny_config(), + "hidden_states": mx.random.normal((1, 16, 64), key=keys[0]).astype(mx.bfloat16), + "encoder_hidden_states": mx.random.normal((1, 8, 32), key=keys[1]).astype(mx.bfloat16), + "encoder_hidden_states_mask": mx.ones((1, 8), dtype=mx.int32), + } + + +def tiny_seamed(n_layers=2): + """A seamed mflux ``Qwen21Transformer`` 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.qwen21.names import qwen21_name_map + from mlx_dfloat.mflux.qwen21.transformer import block_lists, seam_transformer_class + + tf = seam_transformer_class()(**{**TINY, "num_layers": n_layers}) + shapes = install_placeholders(block_lists(tf), qwen21_name_map()) + seam_blocks(block_lists(tf), tf.seam_cell) + return tf, shapes + + +def tiny_groups(shapes, rng): + """One real DF11 group per block of ``shapes``; Qwen'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) + + +# The ComfyUI export's stored order (S0): gate_layer and proj stacked by rows as img_mlp.gate_up, gate first. +STORED_SUBS = ( + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "img_mlp.gate_up", + "img_mlp.out", +) +TINY_ROW_SPLITS = {"img_mlp.gate_up": ("img_mlp.gate_layer", "img_mlp.proj")} + + +def write_tiny_checkpoint( + root, rng, *, n_layers=2, modulation_as_group=True, inner_override=None, random_extras=None +): + """A config-less single-file DF11 checkpoint for mflux's ``Qwen21Transformer`` at ``TINY`` size, as ComfyUI stores it. + + One group per block in the stored order (``STORED_SUBS``: ``img_mlp.gate_up`` = ``concat(gate_layer, proj)`` by + rows), ``modulation.1`` as a one-matrix group (or, with ``modulation_as_group=False``, a BF16 extra), and every + other parameter mflux does not compute as a BF16 extra under its checkpoint name, each filled with a distinct + constant ``0x3F80 + i`` so a mis-routed extra shows. ``inner_override`` builds the probe that many heads wide. + ``random_extras`` (a NumPy generator) fills the extras with random BF16 instead: constant matrices make the forward + pass blind to its inputs (a LayerNorm's output sums to zero across channels, so a constant ``proj_out`` maps every + input to its bias), which an end-to-end comparison cannot use. + + Returns ``(matrices, constants, layout)``: ``{group: {sub-path: uint16 matrix}}`` (the seven decoded sub-paths, + or ``"weight"`` for ``modulation.1``), ``{mflux parameter: constant}``, and the ``SynthesizedLayout`` pinned to + the written file, for ``open_checkpoint(root, layouts=(layout,))``. With ``random_extras`` the constants' + values are the uint16 arrays written. + """ + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + from mlx_dfloat.integrate.placeholders import install_placeholders + from mlx_dfloat.mflux.qwen21.names import qwen21_name_map + from mlx_dfloat.mflux.qwen21.transformer import COMPUTED_PARAMS, block_lists + + dims = {**TINY, "num_layers": n_layers} + if inner_override is not None: + dims["num_attention_heads"] = inner_override + probe = Qwen21Transformer(**dims) + modulation_shape = tuple(probe.modulation.layers[1].weight.shape) + shapes = install_placeholders(block_lists(probe), qwen21_name_map()) + matrices, groups = {}, {} + for block, per in shapes.items(): + mats = {attr: random_bf16(rng, shape) for attr, shape in per.items()} + matrices[block] = mats + gate_up = np.concatenate([mats["img_mlp.gate_layer"], mats["img_mlp.proj"]]) + stored = {**mats, "img_mlp.gate_up": gate_up} + groups[block] = [stored[sub] for sub in STORED_SUBS] + taken = {f"{b}.{a}.weight" for b, per in shapes.items() for a in per} | set(COMPUTED_PARAMS) + patterns = {r"transformer_blocks\.\d+": STORED_SUBS} + if modulation_as_group: + matrices["modulation.1"] = {"weight": random_bf16(rng, modulation_shape)} + groups["modulation.1"] = [matrices["modulation.1"]["weight"]] + patterns[r"modulation\.1"] = () + taken.add("modulation.layers.1.weight") + constants, extras = {}, {} + for i, (param, array) in enumerate(sorted(tree_flatten(probe.parameters()))): + if param in taken: + continue + name = "modulation.1.weight" if param == "modulation.layers.1.weight" else 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=TINY_ROW_SPLITS, + groups=len(groups), + extras=len(extras), + key="tiny-qwen", + ) + return matrices, constants, layout + + +TINY_SEED = 5 # ``fake_model`` writes its checkpoint with this seed: a test regenerates the source matrices from it +PREFIX_IDS = 3 # the stub tokenizer's system prefix: mflux drops this many ids (qwen21_prompt_encoder.py:28-31) + + +class StubTokenizer: + """mflux's tokenizer protocol as ``Qwen21PromptEncoder`` uses it (mflux 0.20.0 qwen21_prompt_encoder.py:26-38). + + ``tokenize(prompt)`` gives ids and an attention mask of shape (1, L), L = ``lengths.get(prompt, 8)``; the wrapped + tokenizer encodes the system prefix to ``PREFIX_IDS`` ids. The ids depend on the prompt (its CRC-32 and the + position, in 1..251 so they stay exact in bfloat16), so the prompts the tests use get different ids. + """ + + def __init__(self, lengths=None): + self.lengths = dict(lengths or {}) + self.tokenizer = lambda text, add_special_tokens=True: {"input_ids": [1, 2, 3]} + + def tokenize(self, prompt, **_kwargs): + n = self.lengths.get(prompt, 8) + seed = zlib.crc32(prompt.encode()) + ids = [(seed + 7 * i) % 251 + 1 for i in range(n)] + return SimpleNamespace( + input_ids=mx.array([ids], dtype=mx.int32), + attention_mask=mx.ones((1, n), dtype=mx.int32), + ) + + +class StubTextEncoder(nn.Module): + """Stands in for the Qwen3-VL text stack: ``__call__(input_ids, attention_mask)`` -> (1, L, 32) hidden states, + each position the weight times its id. + + 32 is ``TINY``'s ``context_in_dim``; mflux calls it with keyword arguments (qwen21_prompt_encoder.py:27). + """ + + def __init__(self): + super().__init__() + self.weight = mx.full((1, 32), 0.25, dtype=mx.bfloat16) + self.calls = 0 + + def __call__(self, input_ids, attention_mask=None): + self.calls += 1 + # Each position's hidden state is its id x the weight: the output depends on the prompt's ids. + return self.weight[None] * input_ids[..., None].astype(mx.bfloat16) + + +class StubVAE(nn.Module): + """One weight (so the retained bound can size it) and the decode ``VAEUtil.decode`` calls (vae_util.py:66). + + mflux unpacks the latents to (1, 64, 1, h / 16, w / 16) first (qwen21_latent_creator.py:26-29); the stub returns + a zero image of 16 pixels per latent cell. + """ + + def __init__(self): + super().__init__() + self.weight = mx.zeros((4,), dtype=mx.bfloat16) + + def decode(self, latents): + return mx.zeros((1, 3, 16 * latents.shape[-2], 16 * latents.shape[-1])) + + +def stub_components(tokenizer=None): + """Base components as ``Qwen21Components`` with the stub encoder, VAE and tokenizer.""" + from mlx_dfloat.mflux.qwen21.init import Qwen21Components + + return Qwen21Components( + vae=StubVAE(), + text_encoder=StubTextEncoder(), + tokenizers={"qwen21": tokenizer or StubTokenizer()}, + ) + + +def fake_model(tmp_path, monkeypatch, *, tokenizer=None, **overrides): + """A ``DFloatQwenImage21`` over the tiny real transformer (2 blocks, config-less checkpoint) and stub base parts. + + No weights, no network. As ``tests/_flux2_tiny.py:fake_model``: the budget is fixed at 23 GiB, the encoder reload + builds a fresh stub, every decode runs on the NumPy reference. + """ + 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.qwen21 import init as qinit + from mlx_dfloat.mflux.qwen21 import model as model_module + from mlx_dfloat.mflux.qwen21.transformer import build_transformer + + _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) + (tmp_path / "base").mkdir(exist_ok=True) + monkeypatch.setattr(qinit, "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": "qwen-image-2.1", + "model_config": ModelConfig.qwen_image_21(), + "ckpt": ckpt, + "df11": ResolvedRepo(root=tmp_path / "df11", repo_id="mingyi456/tiny", revision="a" * 40), + "base": ResolvedRepo(root=tmp_path / "base", repo_id="Qwen/tiny", revision="b" * 40), + "components": stub_components(tokenizer), + "build": build, + "sizes": FamilySizes(compressed=1_000, extras=100, nonblock=10, encoders=1_000, vae=10), + "decode": partial(decode_group, backend="reference"), + } + parts.update(overrides) + return model_module.DFloatQwenImage21._from_parts(**parts) diff --git a/tests/test_ci_config.py b/tests/test_ci_config.py index e4d8449..48125a6 100644 --- a/tests/test_ci_config.py +++ b/tests/test_ci_config.py @@ -87,7 +87,7 @@ 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] - assert len(uv_runs) == 6 + assert len(uv_runs) == 7 assert all("uv run --extra mflux --group dev " in run for run in uv_runs) pytest_runs = [ shlex.split(run) @@ -170,6 +170,31 @@ def test_the_mflux_lane_measures_and_gates_the_flux2_model_module_on_its_own(): assert "src/mlx_dfloat/mflux/flux2/model.py" in pyproject["tool"]["coverage"]["run"]["omit"] +def test_the_mflux_lane_measures_and_gates_the_qwen21_model_module_on_its_own(): + # Bug caught: qwen21/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 Qwen's init.py and transformer.py not + # reported in the lane beside its model (the lane's mflux-only tests cover lines the main job's report cannot). + 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.qwen21.model" in args + assert "--cov=mlx_dfloat.mflux.qwen21.init" in args + assert "--cov=mlx_dfloat.mflux.qwen21.transformer" in args + report = "uv run --extra mflux --group dev coverage report --include='*/mflux/qwen21/model.py' --fail-under=80" + assert report in runs + assert lane < runs.index(report) < lane + 5 + pyproject = tomllib.loads((CI.parents[2] / "pyproject.toml").read_text()) + assert "src/mlx_dfloat/mflux/qwen21/model.py" in pyproject["tool"]["coverage"]["run"]["omit"] + + def _config() -> dict: return yaml.safe_load(CI.read_text()) diff --git a/tests/test_cli.py b/tests/test_cli.py index bd101b4..82d2688 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -811,6 +811,7 @@ def test_finish_writes_the_report_with_the_home_directory_as_a_tilde(tmp_path): ("flux2-klein-4b", 4), ("flux2-klein-base-9b", 50), ("flux2-klein-9b", 4), + ("qwen-image-2.1", 40), ], ) def test_default_steps_follow_mflux_per_model(model, steps): @@ -831,6 +832,7 @@ def test_the_model_choices_are_the_registry_names(): "flux2-klein-4b", "flux2-klein-base-9b", "flux2-klein-9b", + "qwen-image-2.1", ): assert gen.build_parser().parse_args(["--model", name, "--prompt", "p"]).model == name with pytest.raises(SystemExit): @@ -1173,3 +1175,99 @@ def factory(**kw): assert code == 0 assert order == ["clock", "build", "clock"] assert json.loads((tmp_path / "r.json").read_text())["elapsed_seconds"] == 72.5 + + +# --- Qwen-Image 2.1 --------------------------------------------------------------------------------------------- + + +def test_qwen_gets_mfluxs_defaults(tmp_path, capsys): + # Bug caught: Qwen-Image 2.1 on FLUX's 3.5 guidance (a guidance mflux never passes; with a negative prompt it would + # run CFG at 3.5) or on Z-Image's None. mflux 0.20.0: 40 steps (cli/defaults/defaults.py:52), guidance 1.0 + # (qwen21_generate.py:60), scheduler "linear" (cli/parser/parsers.py:180). + call = _generate_kwargs("qwen-image-2.1", tmp_path) + assert (call["num_inference_steps"], call["guidance"], call["scheduler"]) == (40, 1.0, "linear") + assert call["negative_prompt"] is None + assert "warning" not in capsys.readouterr().err + + +def test_qwen_passes_the_negative_prompt_through(tmp_path, capsys): + # Bug caught: the negative prompt dropped for Qwen-Image 2.1 (CFG never runs: mflux needs both a guidance above + # 1.0 and a negative prompt, qwen_image_21.py:89), or the "ignored" warning printed when it is used. + call = _generate_kwargs( + "qwen-image-2.1", tmp_path, "--guidance", "4", "--negative-prompt", "blurry" + ) + assert (call["guidance"], call["negative_prompt"]) == (4.0, "blurry") + assert "warning" not in capsys.readouterr().err + + +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 + # 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") + assert (call["guidance"], call["negative_prompt"]) == (4.0, None) + err = capsys.readouterr().err + assert ( + "warning: --guidance above 1.0 runs no classifier-free guidance for Qwen-Image-2.1 without " + "--negative-prompt (as mflux); pass --negative-prompt to enable it" + ) in err + + +@pytest.mark.parametrize("guidance", ["1.0", "1"]) +def test_qwen_guidance_at_one_without_a_negative_prompt_does_not_warn(guidance, tmp_path, capsys): + # Bug caught: the CFG warning on `>= 1.0` (guidance 1.0 is mflux's default and runs no CFG by design). + _generate_kwargs("qwen-image-2.1", tmp_path, "--guidance", guidance) + assert "warning" not in capsys.readouterr().err + + +def test_the_cfg_warning_is_qwen_only(tmp_path, capsys): + # Bug caught: the warning shown for a Z-Image base (mflux runs its CFG with an empty negative, so --guidance 4 + # alone does run it) or a base Klein (mflux's blank negative). + _generate_kwargs("z-image", tmp_path, "--guidance", "4") + _generate_kwargs("flux2-klein-base-4b", tmp_path, "--guidance", "4") + assert "classifier-free guidance" not in capsys.readouterr().err + + +def test_qwen_negative_prompt_at_default_guidance_warns_with_its_default(tmp_path, capsys): + # Bug caught: the "no effect" warning quoting Z-Image's default 0 for Qwen-Image 2.1, whose default is 1.0 + # (qwen21_generate.py:60). + _generate_kwargs("qwen-image-2.1", tmp_path, "--negative-prompt", "blurry") + err = capsys.readouterr().err + assert "warning: --negative-prompt has no effect" in err + assert "(the default is 1.0)" in err + assert "(the default is 0)" not in err + + +def test_zimage_negative_prompt_warning_text_is_unchanged(tmp_path, capsys): + # Bug caught: the Z-Image warning changed by the per-model default (its default is None, mflux's own rule: 0). + _generate_kwargs("z-image", tmp_path, "--negative-prompt", "blurry") + assert ( + "warning: --negative-prompt has no effect: classifier-free guidance runs only above " + "guidance 1.0 (the default is 0). Pass --guidance above 1.0 to enable it." + ) in capsys.readouterr().err + + +def test_qwen_takes_a_scheduler(tmp_path): + # Bug caught: Klein's scheduler refusal applied to Qwen-Image 2.1 (mflux's Qwen command takes --scheduler, + # cli/parser/parsers.py:180). + assert ( + _generate_kwargs("qwen-image-2.1", tmp_path, "--scheduler", "linear")["scheduler"] + == "linear" + ) + + +@pytest.mark.parametrize( + "flag", [["--quantize", "8"], ["--lora-paths", "a"], ["--image-path", "x.png"], ["-q", "4"]] +) +def test_qwen_refuses_the_common_flags_before_anything_loads(flag, tmp_path, capsys): + # Bug caught: a new family whose refusal list replaced the common one (quantize or LoRA reaching + # DFloatQwenImage21, which refuses them only after the caps and the watchdog). + log = [] + code, watchdogs = _run( + ["--model", "qwen-image-2.1", "--prompt", "p", "--output", str(tmp_path / "o.png"), *flag], + log, + tmp_path, + factory=lambda **kw: pytest.fail("built"), + ) + assert (code, watchdogs, log) == (2, [], []) + assert flag[0] in capsys.readouterr().err diff --git a/tests/test_flux1_memory.py b/tests/test_flux1_memory.py index de9182e..c5738a0 100644 --- a/tests/test_flux1_memory.py +++ b/tests/test_flux1_memory.py @@ -370,11 +370,12 @@ def test_fit_for_can_plan_the_vae_phase_with_the_set_dropped(): text_tokens=256, vae_with_set=False, ) - assert dropped.phases["vae"] == SCHNELL_SIZES.vae + VAE_TRANSIENT_BYTES + OVERHEAD_BYTES + # The extras stay on the transformer when the set goes (_unload_set detaches the provider only). assert ( - resident.phases["vae"] - dropped.phases["vae"] - == SCHNELL_SIZES.compressed + SCHNELL_SIZES.extras + dropped.phases["vae"] + == SCHNELL_SIZES.extras + SCHNELL_SIZES.vae + VAE_TRANSIENT_BYTES + OVERHEAD_BYTES ) + assert resident.phases["vae"] - dropped.phases["vae"] == SCHNELL_SIZES.compressed assert dropped.phases["denoise"] == resident.phases["denoise"] assert not resident.fits assert dropped.fits diff --git a/tests/test_flux1_model.py b/tests/test_flux1_model.py index dd50929..fc549f1 100644 --- a/tests/test_flux1_model.py +++ b/tests/test_flux1_model.py @@ -348,11 +348,13 @@ def _patch_upstream_generate(monkeypatch, model, *, raise_after_loop=None, raise """Replace Flux1.generate_image with a probe: records the cache limit in force, fires our after-loop subscriber the way mflux's GenerationContext would, returns a sentinel.""" from mflux.models.flux.variants.txt2img.flux import Flux1 + 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: @@ -660,3 +662,27 @@ def fake(self, **kwargs): model.generate_image(seed=1, prompt="p", height=256, width=256) assert seen == ["encode", "set_load", "denoise", "vae"] assert model.open_phase is None + + +@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 was 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, model, 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 diff --git a/tests/test_flux1_names.py b/tests/test_flux1_names.py index ae19e3e..5297348 100644 --- a/tests/test_flux1_names.py +++ b/tests/test_flux1_names.py @@ -66,6 +66,14 @@ def test_flux_groups_refuses_holes_missing_kinds_and_foreign_groups(names, reaso check_flux_groups(_ckpt(names)) +def test_flux_groups_refuses_a_zero_padded_block_index(): + # Bug caught: transformer_blocks.01 counted as block 1 (int("01") == 1), then the seam asks the provider for + # transformer_blocks.1, a group the checkpoint does not have: a failure at the first block after the set load. + names = ["transformer_blocks.0", "transformer_blocks.01", "single_transformer_blocks.0"] + with pytest.raises(DFloatFormatError, match=r"transformer_blocks\.01: not a FLUX block"): + check_flux_groups(_ckpt(names)) + + def test_flux_groups_refuses_a_group_with_the_wrong_matrix_count(): # Bug caught: a group with the wrong pattern_dict subs count (a different FLUX variant, or a # corrupted/truncated config) treated as an ordinary block instead of refused. diff --git a/tests/test_flux2_model.py b/tests/test_flux2_model.py index fdeac56..7b61512 100644 --- a/tests/test_flux2_model.py +++ b/tests/test_flux2_model.py @@ -461,11 +461,13 @@ def _patch_upstream_generate(monkeypatch, *, raise_after_loop=None, raise_before """Replace Flux2Klein.generate_image with a probe: records the cache limit in force and the kwargs, fires our after-loop subscriber the way mflux's GenerationContext would, returns a sentinel.""" from mflux.models.flux2.variants.txt2img.flux2_klein import Flux2Klein + 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: @@ -845,3 +847,27 @@ def test_save_model_is_refused_and_freeze_skips_a_dropped_encoder(tmp_path, monk model.freeze() with pytest.raises(DFloatUnsupportedError, match="save"): model.save_model(str(tmp_path)) + + +@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 was 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 diff --git a/tests/test_flux2_names.py b/tests/test_flux2_names.py index 1301d07..dbd45db 100644 --- a/tests/test_flux2_names.py +++ b/tests/test_flux2_names.py @@ -158,6 +158,15 @@ def test_check_accepts_a_checkpoint_without_nonblock_groups(tmp_path): assert check_klein_groups(_ckpt(tmp_path, nonblock=())) == {DOUBLE: 2, SINGLE: 3} +def test_check_refuses_a_zero_padded_block_index(tmp_path): + # Bug caught: single_transformer_blocks.01 counted as block 1 (int("01") == 1), then the seam asks the provider + # for single_transformer_blocks.1, a group the checkpoint does not have. + with pytest.raises( + DFloatFormatError, match=r"single_transformer_blocks\.01: not a FLUX\.2 Klein block" + ): + check_klein_groups(_ckpt(tmp_path, singles=("0", "01", "2"))) + + def test_check_refuses_a_hole_in_the_single_blocks(tmp_path): # Bug caught: a missing single block run as placeholders (zero-size matmul at that depth). with pytest.raises(DFloatFormatError, match="not contiguous"): diff --git a/tests/test_mflux_families.py b/tests/test_mflux_families.py index 9719d5f..b28e7c4 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() == "9" + assert out.stdout.strip() == "10" def test_resolving_an_adapter_module_and_importing_the_package_loads_no_mlx(): @@ -27,6 +27,18 @@ def test_resolving_an_adapter_module_and_importing_the_package_loads_no_mlx(): assert out.stdout.strip() == "False" +def test_resolving_the_qwen21_coverage_specs_loads_no_mlx(): + # Bug caught: an eager import in mlx_dfloat/mflux/qwen21/__init__.py. The mflux lane's + # --cov=mlx_dfloat.mflux.qwen21.model / .init resolve those specs (importing the qwen21 package) when coverage + # starts; MLX loaded there registers its native types twice and the lane aborts (exit 134), as it did for FLUX.1. + code = ( + "import importlib.util, sys; importlib.util.find_spec('mlx_dfloat.mflux.qwen21.model'); " + "importlib.util.find_spec('mlx_dfloat.mflux.qwen21.init'); 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 @@ -71,6 +83,10 @@ def test_an_unknown_name_is_refused_with_the_known_ones(): ("flux2-klein-4b", "flux2", 4, 1.0, "flow_match_euler_discrete", False, False), ("flux2-klein-base-9b", "flux2", 50, 1.0, "flow_match_euler_discrete", True, False), ("flux2-klein-9b", "flux2", 4, 1.0, "flow_match_euler_discrete", False, False), + # mflux 0.20.0 cli/defaults/defaults.py:52 (40 steps); qwen21_generate.py:60 (guidance 1.0); + # cli/parser/parsers.py:180 (scheduler "linear"); qwen_image_21.py:89 (a second call per step above 1.0 with + # a negative prompt, which mflux takes from --negative-prompt, parsers.py:177). + ("qwen-image-2.1", "qwen21", 40, 1.0, "linear", True, True), ], ) def test_the_entries_carry_the_mflux_defaults( @@ -164,6 +180,37 @@ def test_only_the_distilled_klein_models_fix_their_guidance(): "flux2-klein-4b": 1.0, "flux2-klein-base-9b": None, "flux2-klein-9b": 1.0, + "qwen-image-2.1": None, + } + + +def test_only_qwen_image_needs_a_negative_prompt_for_cfg(): + # Bug caught: the "no CFG without --negative-prompt" warning shown for Z-Image (whose base model does run CFG with + # an empty negative) or missing for Qwen-Image 2.1 (mflux runs CFG only with both, qwen_image_21.py:89). + assert {n for n, e in families.MODELS.items() if e.cfg_needs_negative} == {"qwen-image-2.1"} + + +def test_the_qwen_entry_names_its_published_repos_and_pinned_checkpoint(): + # Bug caught: the default checkpoint (one person's Hub account, a ComfyUI single file) following whatever lands on + # main instead of the revision whose header and bytes were verified, or a swapped repository. + e = families.entry("qwen-image-2.1") + assert (e.label, e.df11_repo, e.df11_revision, e.base_repo) == ( + "Qwen-Image-2.1", + "mingyi456/Qwen-Image-2.1-DF11-ComfyUI", + "1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "Qwen/Qwen-Image-2.1", + ) + + +def test_the_qwen_base_is_pinned_to_the_revision_the_measured_runs_used_and_no_other_base_is(): + # Bug caught: the Qwen base (text encoder, VAE, tokenizer) following whatever lands on main instead of the snapshot + # the identity check and the MEASURED row ran on (d26bb61, bench/results/tiers/qwen-image-2.1-1024.json), or a pin + # set on a family whose base no recorded run names. + assert ( + families.entry("qwen-image-2.1").base_revision == "d26bb61231c349cf6b7896fa83353113880e1ba3" + ) + assert {n for n, e in families.MODELS.items() if e.base_revision is not None} == { + "qwen-image-2.1" } @@ -172,6 +219,7 @@ def test_family_of_returns_the_spec_of_the_models_family(): assert families.family_of("schnell") is families.FAMILIES["flux1"] assert families.family_of("z-image") is families.FAMILIES["zimage"] assert families.family_of("flux2-klein-9b") is families.FAMILIES["flux2"] + assert families.family_of("qwen-image-2.1") is families.FAMILIES["qwen21"] @pytest.mark.mflux @@ -214,6 +262,16 @@ def test_the_flux2_model_table_is_derived_from_the_registry(): } +@pytest.mark.mflux +def test_the_qwen21_model_table_is_derived_from_the_registry(): + # Bug caught: the Qwen class's own MODELS drifting from the registry the CLI resolves through. + from mlx_dfloat.mflux.qwen21.model import MODELS as QWEN + + assert QWEN == { + "qwen-image-2.1": ("mingyi456/Qwen-Image-2.1-DF11-ComfyUI", "Qwen/Qwen-Image-2.1") + } + + @pytest.mark.mflux def test_the_family_loaders_return_the_real_classes_and_maps(): # Bug caught: a loader lambda pointing at the other family's class or map. @@ -234,3 +292,8 @@ def test_the_family_loaders_return_the_real_classes_and_maps(): "transformer_blocks", "single_transformer_blocks", ) + from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 + + assert families.FAMILIES["qwen21"].load_model_class() is DFloatQwenImage21 + # mflux 0.20.0 Qwen21Transformer._forward (qwen21_transformer.py:96-100): one block list. + assert families.FAMILIES["qwen21"].load_name_map().kinds == ("transformer_blocks",) diff --git a/tests/test_mflux_phases.py b/tests/test_mflux_phases.py index 841fd2e..96691ce 100644 --- a/tests/test_mflux_phases.py +++ b/tests/test_mflux_phases.py @@ -92,13 +92,46 @@ def _fit(*, policy="per-block", vae_with_set=True): def test_phases_count_each_term_once_and_drop_the_set_from_the_vae_phase_when_asked(): - # Bug caught: the decoded non-block bytes left out of denoise or counted twice; vae_with_set=False keeping them. + # Bug caught: the decoded non-block bytes left out of denoise or counted twice; vae_with_set=False keeping the + # compressed set or the non-block bytes, or dropping the extras (they stay on the transformer when the set goes). # encode = 8_000 + 700 + 50 = 8_750; denoise = 5_000 + 10 + 300 + 400 + 900 + 0 + 50 = 6_660; - # vae = 5_000 + 10 + 300 + 160 + 1_000 + 50 = 6_520; without the set: 160 + 1_000 + 50 = 1_210. + # vae = 5_000 + 10 + 300 + 160 + 1_000 + 50 = 6_520; without the set: 10 + 160 + 1_000 + 50 = 1_220. fit = _fit() assert fit.phases == {"encode": 8_750, "denoise": 6_660, "vae": 6_520} assert (fit.peak_phase, fit.peak_bytes, fit.budget_bytes) == ("encode", 8_750, 10_000) - assert _fit(vae_with_set=False).phases["vae"] == 1_210 + assert _fit(vae_with_set=False).phases["vae"] == 1_220 + + +@pytest.mark.parametrize(("term", "encode"), [(30, 8_080), (0, 8_050)]) +def test_a_measured_encode_term_replaces_the_allowance_in_the_encode_phase_only(term, encode): + # Bug caught: the family's measured encode term ignored (Qwen's encode sized by the denoise allowance, +1.08 GiB), + # added on top of the allowance, a 0 term read as "unset" (`or` for `is None`), or the term leaking into the + # denoise phase. encode = 8_000 + term + 50; denoise 6_660 and vae 6_520 as without the term. + c = _constants(overhead_bytes=50, vae_transient_bytes=1_000, encode_activation_bytes=term) + fit = fit_for( + c, + sizes=SIZES, + largest=SMALL_LARGEST, + policy="per-block", + cache_limit=900, + allowance=700, + budget=10_000, + height=1024, + width=1024, + text_tokens=512, + ) + assert fit.phases == {"encode": encode, "denoise": 6_660, "vae": 6_520} + + +def test_a_one_kind_familys_cache_limit_is_one_decoded_group_plus_the_allowance(): + # Bug caught: a one-kind family budgeted two decoded buffers (per-block evaluation frees one block-sized buffer + # at a time), or depth2's look-ahead buffer missing. Allowance at 1024^2 + 512 tokens 1_588_235_294; + # 436_207_616 + 1_588_235_294 = 2_024_442_910; depth2 adds 436_207_616. + c = _constants() + one = {"transformer_blocks": 436_207_616} + common = {"height": 1024, "width": 1024, "text_tokens": 512} + assert cache_limit_for(c, one, policy="per-block", **common) == 2_024_442_910 + assert cache_limit_for(c, one, policy="depth2", **common) == 2_460_650_526 def test_depth2_holds_two_of_the_largest_decoded_groups_in_flight(): @@ -142,6 +175,23 @@ def test_family_phases_name_every_term(): "overhead": 10, }, } + with_term = family_phases( + compressed_bytes=1, + extras_bytes=2, + nonblock_bytes=3, + largest={"a": 4}, + policy="per-block", + cache_limit=5, + allowance=6, + encoders_bytes=7, + vae_bytes=8, + vae_transient_bytes=9, + overhead_bytes=10, + denoise_activation_bytes=11, + encode_activation_bytes=12, + ) + assert with_term["encode"] == {"encoders": 7, "activations": 12, "overhead": 10} + assert with_term["denoise"] == phases["denoise"] @pytest.mark.parametrize( diff --git a/tests/test_mflux_pipeline.py b/tests/test_mflux_pipeline.py index 3624e7e..1c6e252 100644 --- a/tests/test_mflux_pipeline.py +++ b/tests/test_mflux_pipeline.py @@ -152,6 +152,125 @@ def test_a_cache_limit_below_the_two_largest_groups_warns(caplog): with caplog.at_level("WARNING", logger="tests.pipeline"): _plan(cache_limit_override=two - 1) assert "below the derived minimum" in caplog.text + assert "(the two largest decoded groups)" in caplog.text + + +def test_a_one_kind_familys_derived_minimum_includes_the_activation_allowance(caplog): + # Bug caught: a one-kind family's minimum taken as its one decoded group (any override of a block or more passes + # silently, though the activations then evict the decoded buffer and every block is allocated fresh). One kind of + # 436_207_616 B; allowance at 1024^2 + 512 tokens 1_588_235_294: minimum 2_024_442_910. + one = {"transformer_blocks": 436_207_616} + with caplog.at_level("WARNING", logger="tests.pipeline"): + _plan(largest=one, cache_limit_override=2_024_442_910) + assert "below the derived minimum" not in caplog.text + with caplog.at_level("WARNING", logger="tests.pipeline"): + _plan(largest=one, cache_limit_override=2_024_442_909) + assert ( + "below the derived minimum 2024442910 (the largest decoded group and the activation allowance)" + in caplog.text + ) + + +def test_a_one_kind_familys_depth2_minimum_holds_the_look_ahead_group_too(caplog): + # Bug caught: depth2's look-ahead buffer left out of a one-kind family's minimum (an override with room for one + # decoded group passes silently though two are in flight). One kind of 436_207_616 B and the allowance at + # 1024^2 + 512 tokens 1_588_235_294: per-block minimum 2_024_442_910, depth2 2_460_650_526. + one = {"transformer_blocks": 436_207_616} + with caplog.at_level("WARNING", logger="tests.pipeline"): + _plan(largest=one, policy="depth2", cache_limit_override=2_460_650_526) + assert "below the derived minimum" not in caplog.text + with caplog.at_level("WARNING", logger="tests.pipeline"): + _plan(largest=one, policy="depth2", cache_limit_override=2_460_650_525) + assert ( + "below the derived minimum 2460650526 (the largest decoded group twice, for depth2's look-ahead, and the " + "activation allowance)" in caplog.text + ) + + +def test_a_one_kind_warning_states_the_fresh_allocation_as_a_prediction(caplog): + # Bug caught: the warning asserting as fact that every block is allocated fresh: Qwen-Image 2.1's A/B of real + # steps at the derived limit and below it measured no difference in step time, so the warning may only predict it. + with caplog.at_level("WARNING", logger="tests.pipeline"): + _plan(largest={"transformer_blocks": 436_207_616}, cache_limit_override=1) + assert "may allocate each block's decode output fresh" in caplog.text + assert "will be allocated fresh" not in caplog.text + + +# --- the per-call memory caps ------------------------------------------------------------------------------------- + + +class FakeLimits: + """MLX's limit setters as mlx 0.32.2 has them (each returns the previous value) and its device report.""" + + def __init__(self, *, wired, memory, recommended=26_800_603_136): + self.wired, self.memory, self.recommended = wired, memory, recommended + + def set_wired_limit(self, n): + previous, self.wired = self.wired, n + return previous + + def set_memory_limit(self, n): + previous, self.memory = self.memory, n + return previous + + def device_info(self): + return {"max_recommended_working_set_size": self.recommended} + + +# MLX's default memory limit on this Mac: 1.5 x the recommended working set. +DEFAULT_MEMORY = 40_200_904_704 + + +def test_a_call_without_caps_runs_under_the_commands_caps_and_restores_mlxs_defaults(): + # Bug caught: the Python API running a call at MLX's default wired limit 0 (a Klein VAE decode's MLX peak went + # 6.89 -> 9.49 GiB there, and every VAE term was measured under the caps), or the caps left installed after the + # call. This Mac's 26_800_603_136 B working set floors to 24 GiB: wired 20 GiB, memory 22 GiB (the command's caps). + limits = FakeLimits(wired=0, memory=DEFAULT_MEMORY) + with _pipeline.call_caps(limits) as installed: + assert (limits.wired, limits.memory) == (21_474_836_480, 23_622_320_128) + assert installed == (21_474_836_480, 23_622_320_128) + assert (limits.wired, limits.memory) == (0, DEFAULT_MEMORY) + + +def test_a_call_that_raises_still_restores_mlxs_defaults(): + # Bug caught: the restore skipped on an exception (an interrupted generation leaves the process capped). + limits = FakeLimits(wired=0, memory=DEFAULT_MEMORY) + with pytest.raises(KeyboardInterrupt), _pipeline.call_caps(limits): + raise KeyboardInterrupt + assert (limits.wired, limits.memory) == (0, DEFAULT_MEMORY) + + +def test_caps_already_in_force_are_left_as_they_are(): + # Bug caught: a call replacing the caps the command (or a CAPPED tier: wired 8 GiB, memory 10 GiB at 16 GB) + # installed with the host's, or restoring them to MLX's defaults afterwards. + limits = FakeLimits(wired=8 * GIB, memory=10 * GIB) + with _pipeline.call_caps(limits) as installed: + assert (limits.wired, limits.memory) == (8 * GIB, 10 * GIB) + assert installed == (0, 0) + assert (limits.wired, limits.memory) == (8 * GIB, 10 * GIB) + + +def test_a_device_that_reports_no_working_set_gets_no_caps(): + # Bug caught: caps of 0 bytes installed (a memory limit of 0) on a device that reports nothing, as + # install_memory_caps itself installs nothing there. + limits = FakeLimits(wired=0, memory=DEFAULT_MEMORY, recommended=0) + with _pipeline.call_caps(limits) as installed: + assert (limits.wired, limits.memory) == (0, DEFAULT_MEMORY) + assert installed == (0, 0) + assert (limits.wired, limits.memory) == (0, DEFAULT_MEMORY) + + +def test_the_command_path_keeps_its_caps_through_a_call(): + # Bug caught: the command's own caps (installed before the model is built, as tests/conftest.py installs them + # here) changed by a call: real MLX, the limits read before, inside and after. + from tests._mlx_limits import current_limits + + before = current_limits() + assert before["wired"] > 0 + with _pipeline.call_caps(): + inside = current_limits() + assert inside == before + assert current_limits() == before # --- the phase tracker -------------------------------------------------------------------------------------------- diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 1e00d2d..6594ee0 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -81,3 +81,13 @@ def test_dfloat_flux2_klein_is_reachable_lazily(): from mlx_dfloat.mflux.flux2.model import DFloatFlux2Klein assert adapters.DFloatFlux2Klein is DFloatFlux2Klein + + +@pytest.mark.mflux +def test_dfloat_qwen_image_21_is_reachable_lazily(): + # Bug caught: DFloatQwenImage21 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.qwen21.model import DFloatQwenImage21 + + assert adapters.DFloatQwenImage21 is DFloatQwenImage21 diff --git a/tests/test_qwen21_build.py b/tests/test_qwen21_build.py new file mode 100644 index 0000000..6fa7a1f --- /dev/null +++ b/tests/test_qwen21_build.py @@ -0,0 +1,244 @@ +"""The Qwen-Image 2.1 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._df11_fixtures import write_bf16_original +from tests._qwen21_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.qwen21.transformer import base_extras, build_transformer + +# The seven block matrices and their shapes at TINY size (inner 2 x 16 = 32, mlp_ratio 3), from mflux 0.20.0: +# Qwen21Attention's four Linear(32, 32) (qwen21_attention.py:13-16) and Qwen21SwiGLUFeedForward's +# proj/gate_layer Linear(32, 96) and out Linear(96, 32) (qwen21_feed_forward.py:8-10); the weights are (out, in). +SEVEN = { + "attn.to_q": (32, 32), + "attn.to_k": (32, 32), + "attn.to_v": (32, 32), + "attn.to_out.0": (32, 32), + "img_mlp.gate_layer": (96, 32), + "img_mlp.proj": (96, 32), + "img_mlp.out": (32, 96), +} +# The shared modulation: nn.Linear(inner, 4 * inner) (qwen21_transformer.py:37), so (128, 32). +MODULATION = {"modulation.layers.1.weight": (128, 32)} + + +def test_base_extras_drop_the_expanded_block_matrices_and_keep_modulation(): + # Bug caught: the stored gate_up names compared (nothing dropped, so the BF16 side would load its block matrices + # as extras), or modulation.1 dropped (the BF16 side's coverage would refuse it). + groups = { + "transformer_blocks.0": SimpleNamespace( + matrix_names=( + "transformer_blocks.0.img_mlp.gate_layer.weight", + "transformer_blocks.0.img_mlp.proj.weight", + ) + ), + "modulation.1": SimpleNamespace(matrix_names=("modulation.1.weight",)), + } + index = { + name: (Path(name), None) + for name in ( + "transformer_blocks.0.img_mlp.gate_layer.weight", + "transformer_blocks.0.img_mlp.proj.weight", + "modulation.1.weight", + "img_in.weight", + ) + } + assert sorted(base_extras(index, SimpleNamespace(groups=groups))) == [ + "img_in.weight", + "modulation.1.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: the modulation rename lost (an extra-coverage refusal), a compressed matrix left live (memory), + # or an extra loaded onto the wrong parameter (a constant in the wrong place). + ckpt, _matrices, constants = _open(tmp_path) + assert "modulation.1" in ckpt.groups + build = build_transformer(ckpt, transformer_kwargs=TINY) + params = _params(build) + assert build.nonblock == MODULATION + assert params["modulation.layers.1.weight"].size == 0 + assert build.shapes == {f"transformer_blocks.{i}": SEVEN for i in range(2)} + 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). + assert sorted(constants) == [ + "img_in.weight", + "norm_out.linear.weight", + "proj_out.weight", + "time_text_embed.timestep_embedder.linear_1.weight", + "time_text_embed.timestep_embedder.linear_2.weight", + "transformer_blocks.0.attn.norm_k.weight", + "transformer_blocks.0.attn.norm_q.weight", + "transformer_blocks.1.attn.norm_k.weight", + "transformer_blocks.1.attn.norm_q.weight", + "txt_in.in_layer.weight", + "txt_in.out_layer.weight", + "txt_in.text_norm.weight", + ] + for name, value in constants.items(): + assert np.all(np.array(params[name].view(mx.uint16)) == value), name + assert build.counts == {"transformer_blocks": 2} + + +@pytest.mark.mflux +def test_the_computed_tables_are_left_as_mflux_built_them(tmp_path): + # Bug caught: the coverage exemption replaced by loading something into the computed tables (a zero or stale + # RoPE table would wreck the positions while every check passes). + from mflux.models.qwen21.model.qwen21_transformer.qwen21_rope import Qwen21Rope + from mflux.models.qwen21.model.qwen21_transformer.qwen21_time_text_embed import ( + Qwen21Timesteps, + ) + + ckpt, _m, _c = _open(tmp_path) + params = _params(build_transformer(ckpt, transformer_kwargs=TINY)) + rope = Qwen21Rope(axes_dim=[4, 6, 6]) + for i in range(3): + for table, fresh in (("cos", rope.cos_tables[i]), ("sin", rope.sin_tables[i])): + got = params[f"pos_embed.{table}_tables.{i}"] + assert mx.array_equal(got.view(mx.uint32), fresh.view(mx.uint32)).item(), (table, i) + freqs = params["time_text_embed.time_proj.freqs"] + assert mx.array_equal(freqs.view(mx.uint32), Qwen21Timesteps().freqs.view(mx.uint32)).item() + + +@pytest.mark.mflux +def test_modulation_stored_as_an_extra_loads_as_a_plain_weight(tmp_path): + # Bug caught: a checkpoint without the non-block group refused, or its modulation left on a placeholder (a + # zero-size matmul before the first block). + ckpt, _m, constants = _open(tmp_path, modulation_as_group=False) + assert "modulation.1.weight" in ckpt.extras + build = build_transformer(ckpt, transformer_kwargs=TINY) + weight = _params(build)["modulation.layers.1.weight"] + assert build.nonblock == {} + assert weight.shape == (128, 32) + assert np.all(np.array(weight.view(mx.uint16)) == constants["modulation.layers.1.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 (a shape error, or 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\b"): + build_transformer(ckpt, transformer_kwargs=TINY) + + +@pytest.mark.mflux +def test_the_count_check_defaults_to_mflux_s_depth(tmp_path): + # Bug caught: kwargs without num_layers checked against something other than mflux's default 32 + # (qwen21_transformer.py:21): the real model passes no overrides. + ckpt, _m, _c = _open(tmp_path) + no_depth = {k: v for k, v in TINY.items() if k != "num_layers"} + with pytest.raises(DFloatFormatError, match=r"checkpoint has 2 blocks; this model builds 32\b"): + 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 == {"transformer_blocks": 1} + assert len(build.transformer.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_a_checkpoint_of_another_width_is_refused_by_the_extras_shape_check(tmp_path): + # Bug caught: the same block count at another inner width (a hand-made checkpoint) accepted. + ckpt, _m, _c = _open(tmp_path, inner_override=4) + with pytest.raises(DFloatFormatError, match="shape"): + build_transformer(ckpt, transformer_kwargs=TINY) + + +@pytest.mark.mflux +def test_the_decoded_halves_land_on_gate_layer_and_proj(tmp_path): + # Bug caught: the halves swapped anywhere between the layout, the reader and the map (the model would run + # silu(proj) * gate), or a decoded matrix cut at the wrong row. + from mlx_dfloat.integrate.coverage import load_resident_set + from mlx_dfloat.integrate.providers import DF11Provider + from mlx_dfloat.mflux.qwen21.names import qwen21_name_map + + ckpt, matrices, _c = _open(tmp_path) + build = build_transformer(ckpt, transformer_kwargs=TINY) + blocks = [g for g in ckpt.groups if g != "modulation.1"] + provider = DF11Provider( + load_resident_set(ckpt, blocks), + {g: ckpt.groups[g].matrix_names for g in blocks}, + qwen21_name_map(), + decode=partial(decode_group, backend="reference"), + ) + got = provider.weights_for("transformer_blocks.1", build.shapes["transformer_blocks.1"]) + provider.verify() + for attr in SEVEN: + want = matrices["transformer_blocks.1"][attr] + assert np.array_equal(np.array(got[attr].view(mx.uint16)), want), attr + + +@pytest.mark.mflux +def test_the_bf16_side_builds_from_a_sharded_base_with_modulation_as_a_plain_weight(tmp_path): + # Bug caught (the BF16 side of an identity check): the sharded base's index not followed, the block matrices + # loaded as extras, or modulation.1 left on a placeholder instead of loading as a plain weight. + from mlx_dfloat.mflux.qwen21.transformer import base_transformer_files_index + + ckpt, matrices, _constants = _open(tmp_path) + originals = {} + for name, (_path, info) in ckpt.extras.items(): + originals[name] = np.full( + info.shape, 0x4000, dtype=np.uint16 + ) # 2.0: distinct from the DF11 extras + originals["modulation.1.weight"] = matrices["modulation.1"]["weight"] + for block in ("transformer_blocks.0", "transformer_blocks.1"): + for attr, mat in matrices[block].items(): + originals[f"{block}.{attr}.weight"] = mat + root = write_bf16_original(tmp_path / "base", originals) + index = base_transformer_files_index(root) + assert {p.name for p, _i in index.values()} == { + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors", + } + extras = base_extras(index, ckpt) + assert not any(".img_mlp." in n or ".attn.to_" in n for n in extras) + build = build_transformer( + ckpt, transformer_kwargs=TINY, extras=extras, nonblock_from_extras=True + ) + params = _params(build) + assert build.nonblock == {} + modulation = np.array(params["modulation.layers.1.weight"].view(mx.uint16)) + assert np.array_equal(modulation, matrices["modulation.1"]["weight"]) + assert np.all(np.array(params["img_in.weight"].view(mx.uint16)) == 0x4000) + assert params["transformer_blocks.0.img_mlp.proj.weight"].size == 0 + + +@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.qwen21 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) diff --git a/tests/test_qwen21_init.py b/tests/test_qwen21_init.py new file mode 100644 index 0000000..f1b658b --- /dev/null +++ b/tests/test_qwen21_init.py @@ -0,0 +1,344 @@ +"""Qwen-Image 2.1 base components (text encoder, VAE, tokenizer) without the transformer.""" + +import os +import re +from fnmatch import fnmatch +from pathlib import Path + +import mlx.core as mx +import numpy as np +import pytest +from mlx.utils import tree_flatten, tree_unflatten +from tests._df11_fixtures import write_bf16_original + +from mlx_dfloat.errors import DFloatFormatError +from mlx_dfloat.mflux import _hub +from mlx_dfloat.mflux.qwen21.init import BASE_PATTERNS, DF11_PATTERNS + +# Qwen/Qwen-Image-2.1 @ d26bb61231c349cf6b7896fa83353113880e1ba3, every file on the Hub (listed 2026-10-08). +BASE_FILES = ( + ".gitattributes", + "LICENSE", + "README.md", + "assets/qr.png", + "model_index.json", + "processor/added_tokens.json", + "processor/chat_template.jinja", + "processor/merges.txt", + "processor/preprocessor_config.json", + "processor/special_tokens_map.json", + "processor/tokenizer.json", + "processor/tokenizer_config.json", + "processor/video_preprocessor_config.json", + "processor/vocab.json", + "scheduler/scheduler_config.json", + "text_encoder/config.json", + "text_encoder/generation_config.json", + "text_encoder/model-00001-of-00004.safetensors", + "text_encoder/model-00002-of-00004.safetensors", + "text_encoder/model-00003-of-00004.safetensors", + "text_encoder/model-00004-of-00004.safetensors", + "text_encoder/model.safetensors.index.json", + "transformer/config.json", + "transformer/diffusion_pytorch_model-00001-of-00002.safetensors", + "transformer/diffusion_pytorch_model-00002-of-00002.safetensors", + "transformer/diffusion_pytorch_model.safetensors.index.json", + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", +) + + +# The tokenizer's files in processor/ (what transformers' AutoTokenizer reads for a Qwen2 tokenizer); the two +# preprocessor configs there belong to the image and video processors, which mflux never builds. +TOKENIZER_FILES = ( + "processor/added_tokens.json", + "processor/chat_template.jinja", + "processor/merges.txt", + "processor/special_tokens_map.json", + "processor/tokenizer.json", + "processor/tokenizer_config.json", + "processor/vocab.json", +) + + +def test_base_patterns_fetch_the_encoder_vae_and_tokenizer_and_never_the_transformer(): + # Bug caught: a pattern such as "*.safetensors" or "transformer/*" pulling the 14.2 GB BF16 transformer on first + # use, or one that misses a tokenizer file mflux's tokenizer loader reads (its hf_subdir is "processor"). + hit = sorted(f for f in BASE_FILES if any(fnmatch(f, p) for p in BASE_PATTERNS)) + assert hit == sorted( + [f for f in BASE_FILES if f.split("/", 1)[0] in ("text_encoder", "vae")] + + list(TOKENIZER_FILES) + ) + assert not [f for f in hit if f.startswith("transformer/")] + + +def test_the_processor_pattern_fetches_tokenizer_files_only(): + # Bug caught: "processor/*" fetching whatever the repository adds there on a future revision (code, pickles, large + # binaries), when the tokenizer needs seven named files. The listing is the Hub's plus files a later revision + # could add. + listing = ( + *BASE_FILES, + "processor/tokenization_qwen.py", + "processor/processor.pkl", + "processor/extra/tokenizer.json", + ) + hit = sorted( + f + for f in listing + if f.startswith("processor/") and any(fnmatch(f, p) for p in BASE_PATTERNS) + ) + assert hit == sorted(TOKENIZER_FILES) + + +def test_df11_patterns_fetch_the_one_comfyui_file(): + # Bug caught: the config-less repository resolved without its only weight file (a pattern requiring a config). + # mingyi456/Qwen-Image-2.1-DF11-ComfyUI @ 1b22a3a1 holds these three files (Hub, 2026-10-08). + files = (".gitattributes", "README.md", "qwen_image_2.1_bf16-DF11.safetensors") + assert [f for f in files if any(fnmatch(f, p) for p in DF11_PATTERNS)] == [ + "qwen_image_2.1_bf16-DF11.safetensors" + ] + + +# --- mflux lane ----------------------------------------------------------------------------------------------- + + +@pytest.mark.mflux +def test_the_definitions_name_only_their_component_and_the_processor_tokenizer(): + # Bug caught: a subset keeping the transformer component (mflux's loader would read transformer/), or the + # tokenizer looked up anywhere but processor/ (mflux 0.20.0 qwen21_weight_definition.py:38-50). + from mlx_dfloat.mflux.qwen21.init import encoder_definition, vae_definition + + assert [c.name for c in encoder_definition().get_components()] == ["text_encoder"] + 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) == ( + "qwen21", + "processor", + 2048, + ) + + +@pytest.mark.mflux +def test_the_computed_exemptions_are_exactly_what_mfluxs_mapping_does_not_load(): + # Bug caught: a computed buffer missing from COMPUTED (every real base refused), or a mapped weight listed there + # (a missing tensor never noticed). mflux 0.20.0: the VAE maps every parameter (qwen21_weight_mapping.py:91 on), + # the text encoder all but rotary_emb.inv_freq (qwen3_vl_rope.py:21). A two-layer encoder: the mapping names 36. + from mflux.models.qwen21.model.qwen21_text_encoder.qwen21_text_encoder import Qwen21TextEncoder + from mflux.models.qwen21.model.qwen21_vae.qwen21_vae import Qwen21VAE + from mflux.models.qwen21.weights.qwen21_weight_mapping import Qwen21WeightMapping + + from mlx_dfloat.mflux.qwen21.init import COMPUTED + + def unmapped(module, mapping): + targets = {t.to_pattern for t in mapping} + return {n for n, _ in tree_flatten(module.parameters())} - targets + + encoder = Qwen21TextEncoder(num_hidden_layers=2) + assert unmapped(encoder, Qwen21WeightMapping.get_text_encoder_mapping()) == { + "rotary_emb.inv_freq" + } + assert unmapped(Qwen21VAE(), Qwen21WeightMapping.get_vae_mapping()) == set() + assert sorted(COMPUTED["text_encoder"]) == ["rotary_emb.inv_freq"] + assert sorted(COMPUTED["vae"]) == [] + + +@pytest.mark.mflux +def test_a_base_with_another_encoder_width_is_refused_before_any_encode(monkeypatch, tmp_path): + # Bug caught: Qwen-Image 1's base (a Qwen2.5-VL encoder 3584 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.qwen21.init import load_text_encoder + + def wrong_width(root, definition, models, *, what): + models["text_encoder"].embed_tokens.weight = mx.zeros((151936, 3584), dtype=mx.bfloat16) + + monkeypatch.setattr(_hub, "load_into", wrong_width) + with pytest.raises(DFloatFormatError, match=r"embed_tokens\.weight.*\(151936, 3584\).*4096"): + load_text_encoder(tmp_path) + + +# A two-layer Qwen3-VL text stack (tiny widths; mflux 0.20.0 Qwen21TextEncoder'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, +} + + +def _write_tiny_encoder(root, *, drop=()): + """The tiny encoder under the base's names: ``model.language_model.`` + mflux's path, plus a vision tower tensor + and an untied ``lm_head`` as the real shards carry (mflux maps neither), in two shards with an index. + + ``rotary_emb.inv_freq`` is left out: no checkpoint holds it (computed in ``Qwen3VLRotaryEmbedding.__init__``). + """ + from mflux.models.qwen21.model.qwen21_text_encoder.qwen21_text_encoder import Qwen21TextEncoder + + bits = np.uint16(0x3F00) # bf16 0.5 + params = dict(tree_flatten(Qwen21TextEncoder(**TINY_ENCODER).parameters())) + tensors = { + f"model.language_model.{name}": np.full(p.shape, bits, np.uint16) + for name, p in params.items() + if name != "rotary_emb.inv_freq" + } + tensors["model.visual.blocks.0.attn.qkv.weight"] = np.zeros((4, 4), np.uint16) + tensors["lm_head.weight"] = np.zeros((100, 64), np.uint16) + write_bf16_original(root / "text_encoder", {n: t for n, t in tensors.items() if n not in drop}) + + +@pytest.mark.mflux +def test_an_encoder_tensor_missing_from_the_base_is_refused_by_name(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.qwen21.init import load_text_encoder + + _write_tiny_encoder(tmp_path, drop={"model.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 or the + # lm_head in the same shards breaking the load, or the file's weights not landing (still the float32 init). + from mlx_dfloat.mflux.qwen21.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 + + +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_vae_tensor_missing_from_the_base_is_refused_by_name(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.qwen21.init import load_vae + + missing = "decoder.conv_out.conv.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(monkeypatch, tmp_path): + # Bug caught: a complete VAE refused (an exemption list naming something the module does not have, or a check + # comparing on-disk names instead of the parameters the applier replaced). + from mlx_dfloat.mflux.qwen21.init import load_vae + + monkeypatch.setattr(_hub, "load_into", _replace_all_but("vae", None)) + vae = load_vae(tmp_path) + assert vae.decoder.conv_out.conv.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.qwen21.init import load_text_encoder, load_vae + + with pytest.raises(DFloatFormatError, match=r"Qwen-Image 2\.1 base"): + load_text_encoder(tmp_path) + with pytest.raises(DFloatFormatError, match=r"Qwen-Image 2\.1 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.qwen21.init import load_tokenizers + + with pytest.raises(DFloatFormatError, match=r"no usable Qwen-Image 2\.1 tokenizer"): + load_tokenizers(tmp_path) + + +@pytest.mark.mflux +def test_a_partial_processor_directory_is_a_format_error(monkeypatch, tmp_path): + # Bug caught (fable M7): transformers' plain OSError for a processor/ 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.qwen21.init import load_tokenizers + + def partial(**_kwargs): + raise OSError("Can't load tokenizer for 'processor'") + + monkeypatch.setattr(TokenizerLoader, "load_all", staticmethod(partial)) + with pytest.raises( + DFloatFormatError, match=r"no usable Qwen-Image 2\.1 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.qwen21 import init as qinit + + monkeypatch.setattr(qinit, "load_vae", lambda root: ("vae", root)) + monkeypatch.setattr(qinit, "load_text_encoder", lambda root: ("encoder", root)) + monkeypatch.setattr(qinit, "load_tokenizers", lambda root: {"qwen21": root}) + parts = qinit.load_base(tmp_path) + assert (parts.vae, parts.text_encoder, parts.tokenizers) == ( + ("vae", tmp_path), + ("encoder", tmp_path), + {"qwen21": tmp_path}, + ) + + +@pytest.mark.mflux +@pytest.mark.network +def test_the_real_processor_tokenizer_loads(): + # Bug caught: the processor subdir not where mflux's loader looks, or the template not applied (the prefix mflux + # drops would then be missing from the ids). Needs the base's processor/ files (a few MB). + base = os.environ.get("MLX_DFLOAT_QWEN21_BASE") + if not base: + pytest.skip("MLX_DFLOAT_QWEN21_BASE is not set (a Qwen-Image 2.1 snapshot with processor/)") + from mlx_dfloat.mflux.qwen21.init import load_tokenizers + + tokenizer = load_tokenizers(Path(base))["qwen21"] + ids = tokenizer.tokenize("a").input_ids + assert ids.shape[0] == 1 + assert ids.shape[1] > 1 + + +@pytest.mark.mflux +@pytest.mark.network +def test_the_tokenizer_loads_from_the_files_the_base_patterns_fetch(tmp_path): + # Bug caught: a narrowed processor pattern that leaves out a file mflux's TokenizerLoader (transformers' + # AutoTokenizer) needs, so a fresh download loads no tokenizer. Copies only the matched processor/ files of a real + # snapshot and loads from the copy. Needs the base's processor/ files (a few MB). + import shutil + + base = os.environ.get("MLX_DFLOAT_QWEN21_BASE") + if not base: + pytest.skip("MLX_DFLOAT_QWEN21_BASE is not set (a Qwen-Image 2.1 snapshot with processor/)") + from mlx_dfloat.mflux.qwen21.init import load_tokenizers + from mlx_dfloat.mflux.qwen21.memory import prompt_tokens + + (tmp_path / "processor").mkdir() + for path in sorted(Path(base, "processor").iterdir()): + name = f"processor/{path.name}" + if any(fnmatch(name, p) for p in BASE_PATTERNS): + shutil.copyfile(path, tmp_path / name) + copied = sorted(f"processor/{p.name}" for p in (tmp_path / "processor").iterdir()) + assert copied == sorted(TOKENIZER_FILES) + tokenizer = load_tokenizers(tmp_path)["qwen21"] + # The calibration's count: the negative prompt " " is 9 text tokens with the full processor/ (test_qwen21_memory). + assert prompt_tokens(tokenizer, " ") == 9 diff --git a/tests/test_qwen21_memory.py b/tests/test_qwen21_memory.py new file mode 100644 index 0000000..2d84cb7 --- /dev/null +++ b/tests/test_qwen21_memory.py @@ -0,0 +1,447 @@ +"""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 +(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, +mlx 0.32.2, mflux 0.20.0). Footprint peaks: encode 15_857_938_288, set load 10_347_487_632, denoise 13_526_704_528, +vae 22_139_643_024 (the call's peak). Sizes: compressed 9_587_103_379, extras 137_388_032, decoded non-block +(``modulation.1``) 134_217_728, the encoder's language model 15_136_811_008 (shards 17_534_339_488), VAE +1_350_989_512; largest decoded block 436_207_616; cache limit 1_859_346_402 (the run's; the planner now derives +2_731_761_634, see ``CACHE_LIMIT``). MLX's encode peak 15_350_204_544 with 146_825_736 active after the build: the +encode held 66_567_800 over the language model. The committed MEASURED record +``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 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. + +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 +tests below read them from there. +""" + +import json +import logging +import os +from pathlib import Path +from types import SimpleNamespace + +import mlx.core as mx +import numpy as np +import pytest +from tests._df11_fixtures import random_bf16, write_checkpoint + +from mlx_dfloat.bench.capped import tier_limits +from mlx_dfloat.errors import DFloatResourceError +from mlx_dfloat.format import open_checkpoint +from mlx_dfloat.integrate import memory as imem +from mlx_dfloat.mflux._phases import FamilySizes, activation_allowance, fit_for +from mlx_dfloat.mflux._pipeline import plan_call_for +from mlx_dfloat.mflux.qwen21 import memory as qmem + +GIB = 1024**3 +SIZES = FamilySizes( + compressed=9_587_103_379, + extras=137_388_032, + nonblock=134_217_728, + encoders=15_136_811_008, + vae=1_350_989_512, +) +LARGEST = {"transformer_blocks": 436_207_616} # one block kind; 32 groups of 218_103_808 elements +# 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. +CALIBRATION_CACHE_LIMIT = 1_859_346_402 +TEXT_TOKENS = 33 # the de-risk 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. +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. +RECOMMENDED_32 = 26_800_603_136 +RAM_32 = 34_359_738_368 + + +def _bf16(*shape): + return mx.zeros(shape, dtype=mx.bfloat16) + + +def _write_encoder(root): + """A Qwen3-VL-shaped text encoder in two shards, the way the base stores it (``model.language_model.*``, the + vision tower ``model.visual.*`` and an untied ``lm_head``). + + bf16 bytes: embed_tokens 10 x 4 = 80, up_proj 4 x 4 = 32, the vision tensor 8 x 4 = 64, lm_head 10 x 4 = 80. + """ + enc = root / "text_encoder" + enc.mkdir(parents=True) + mx.save_safetensors( + str(enc / "model-00001-of-00002.safetensors"), + { + "model.language_model.embed_tokens.weight": _bf16(10, 4), + "model.visual.blocks.0.attn.qkv.weight": _bf16(8, 4), + }, + ) + mx.save_safetensors( + str(enc / "model-00002-of-00002.safetensors"), + { + "model.language_model.layers.0.mlp.up_proj.weight": _bf16(4, 4), + "lm_head.weight": _bf16(10, 4), + }, + ) + (enc / "model.safetensors.index.json").write_text("{}") + + +def test_encoder_bytes_count_the_language_model_only(tmp_path): + # Bug caught: the encoder's file size counted (as for FLUX.1), or the vision tower or lm_head counted: mflux maps + # model.language_model.* only (qwen21_weight_mapping.py:56-88), so the other tensors are never loaded. 80 + 32. + _write_encoder(tmp_path) + assert qmem.encoder_bytes_used(tmp_path) == 112 + + +def test_sizes_for_replaces_the_encoder_file_size_with_the_language_model_bytes(tmp_path): + # Bug caught: the sizes still carrying the shards' file size (vision tower and lm_head included: the encode phase + # over-predicted), or the DF11 / non-block / VAE terms lost when the encoder term is replaced. + rng = np.random.default_rng(0) + write_checkpoint( + tmp_path / "df11", + groups={ + "transformer_blocks.0": [random_bf16(rng, (8, 4))], + "modulation.1": [random_bf16(rng, (16, 4))], + }, + patterns={r"transformer_blocks\.\d+": ("attn.to_q",), r"modulation\.1": ()}, + extras={"img_in.weight": np.zeros((3, 5), dtype=np.uint16)}, + single_file=True, + ) + _write_encoder(tmp_path / "base") + (tmp_path / "base" / "vae").mkdir() + (tmp_path / "base" / "vae" / "diffusion_pytorch_model.safetensors").write_bytes(b"\0" * 100) + + sizes = qmem.sizes_for(open_checkpoint(tmp_path / "df11"), tmp_path / "base") + total = (tmp_path / "df11" / "model.safetensors").stat().st_size + assert sizes.encoders == 112 + assert sizes.extras == 30 # 3 x 5 BF16 + assert sizes.compressed == total - 30 + assert sizes.nonblock == 128 # modulation.1 decoded: 16 x 4 BF16 + assert sizes.vae == 100 + + +class _StubTokenizer: + """``tokenize(p)`` gives ``len(p) + 5`` ids (a 5-token system prefix around the prompt, no truncation); the + wrapped tokenizer encodes the system prefix to 5 ids, as ``Qwen21PromptEncoder._system_prefix_length`` asks it.""" + + def __init__(self): + self.tokenizer = lambda text, add_special_tokens=True: {"input_ids": [0] * 5} + + def tokenize(self, prompt): + return SimpleNamespace(input_ids=mx.zeros((1, len(prompt) + 5), dtype=mx.int32)) + + +@pytest.mark.mflux +@pytest.mark.parametrize( + ("prompt", "tokens"), + [ + ("abc", 3), + ("", 1), # mflux encodes a blank prompt as " " (qwen21_prompt_encoder.py:24-25) + (" ", 1), + ("x" * 2042, 2042), # 2047 ids in all: under the cap + ("x" * 2043, 2043), # 2048: on it + ( + "x" * 2044, + 2043, + ), # 2049 would be cut to 2048 (max_length 2048, truncation, tokenizer.py:98-103) + ], + ids=["abc", "empty", "blank", "under-cap", "on-cap", "over-cap"], +) +def test_prompt_tokens_drop_the_system_prefix_count_a_blank_as_a_space_and_cap_at_2048( + prompt, tokens +): + # Bug caught: the system prefix counted (mflux drops it: qwen21_prompt_encoder.py:28-31), "" counted as 0 tokens + # (the allowance computed for a prompt mflux does not encode), or a long prompt counted past the tokenizer's + # 2048-token cut (`>` / `>=` at the cap, or no cap: the activation term over-predicted). + assert qmem.prompt_tokens(_StubTokenizer(), prompt) == tokens + + +@pytest.mark.mflux +def test_the_token_cap_is_mfluxs_tokenizer_max_length(): + # Bug caught: an mflux bump changing the tokenizer's max_length (qwen21_weight_definition.py:45) while the cap + # here keeps 2048. + from mflux.models.qwen21.weights.qwen21_weight_definition import Qwen21WeightDefinition + + (definition,) = Qwen21WeightDefinition.get_tokenizers() + assert definition.max_length == 2048 + assert qmem.TEXT_MAX_LENGTH == 2048 + + +# The de-risk 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" +) + + +@pytest.mark.mflux +@pytest.mark.network +@pytest.mark.parametrize( + ("prompt", "tokens"), + [(" ", 9), ("", 9), (DERISK_PROMPT, 33)], + ids=["space", "empty", "derisk-prompt"], +) +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 " " + # (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: + pytest.skip("MLX_DFLOAT_QWEN21_BASE is not set (a Qwen-Image 2.1 snapshot with processor/)") + from mlx_dfloat.mflux.qwen21.init import load_tokenizers + + tokenizer = load_tokenizers(Path(base))["qwen21"] + assert qmem.prompt_tokens(tokenizer, prompt) == tokens + + +def _fit(*, text_tokens=TEXT_TOKENS, budget=24_653_119_488): + c = qmem.CONSTANTS + return fit_for( + c, + sizes=SIZES, + largest=LARGEST, + policy="per-block", + cache_limit=CACHE_LIMIT, + allowance=activation_allowance(c, height=1024, width=1024, text_tokens=text_tokens), + budget=budget, + height=1024, + width=1024, + text_tokens=text_tokens, + ) + + +def _plan(budget, *, fit_check=True, cache_limit_override=None): + return plan_call_for( + constants=qmem.CONSTANTS, + sizes=SIZES, + largest=LARGEST, + policy="per-block", + cache_limit_override=cache_limit_override, + fit_check=fit_check, + budget=budget, + height=1024, + width=1024, + text_tokens=TEXT_TOKENS, + log=logging.getLogger("test"), + ) + + +# The constants are calibrated on the de-risk 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. + + +def test_the_calibration_record_holds_the_runs_the_constants_name(): + # Bug caught: the committed record and the constants drifting apart (a re-derived term whose run is not recorded, + # or a record edit that moves a peak): the three source peaks are the literals the constants were derived from. + assert CALIBRATION["constants_from"] == { + "encode": "derisk-1", + "denoise": "ab-new-limit-1", + "vae": "ab-old-limit-2", + } + assert MEASURED == {"encode": 15_857_938_288, "denoise": 13_939_106_072, "vae": 22_788_580_400} + assert CALIBRATION_RUNS["ab-new-limit-1"]["cache_limit"] == CACHE_LIMIT + + +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 + # sample dropped from the record. + assert len(VAE_SAMPLES) == 11 + fit = _fit() + assert fit.phases["vae"] == max(VAE_SAMPLES.values()) + assert all(peak <= fit.phases["vae"] for peak in VAE_SAMPLES.values()) + + +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); + # 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. + fit = _fit() + assert fit.peak_phase == "vae" + assert fit.phases["vae"] == MEASURED["vae"] + assert fit.phases["encode"] == 15_662_823_906 + assert fit.phases["denoise"] == MEASURED["denoise"] + + +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 + # 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 + # 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) + assert "below the derived minimum" not in caplog.text + with caplog.at_level("WARNING", logger="test"): + _plan(24_653_119_488, cache_limit_override=CALIBRATION_CACHE_LIMIT) + assert "below the derived minimum 2731761634" in caplog.text + + +def test_a_32_gb_mac_keeps_the_set_through_the_vae_decode(monkeypatch): + # The shared rule (_pipeline.plan_call_for) drops the set before the VAE decode iff the VAE phase with the set is + # over the budget, with no margin beyond the budget's own: the recommended working set minus a 2 GiB reserve + # (integrate.memory.budget_bytes), 26_800_603_136 - 2_147_483_648 = 24_653_119_488 B here. The measured VAE peak, + # 22_788_580_400 B (the estimate's by calibration), is 1_864_539_088 B (1.74 GiB) under it: the set stays resident, + # and a 32 GB Mac pays no set reload per call. Bug caught: the decision inverted, a margin added on top of the + # reserve that flips it (more than 1.74 GiB), or the budget read without the reserve. + monkeypatch.setattr( + imem.mx, "device_info", lambda: {"max_recommended_working_set_size": RECOMMENDED_32} + ) + budget = imem.budget_bytes() + assert budget == 24_653_119_488 + plan = _plan(budget) + assert plan.drop_set_before_vae is False + assert plan.estimate.peak_phase == "vae" + assert plan.estimate.peak_bytes == MEASURED["vae"] + # The boundary: one byte less of budget than the VAE phase drops the set (`>`, not `>=`, at the budget). + assert _plan(MEASURED["vae"]).drop_set_before_vae is False + assert _plan(MEASURED["vae"] - 1).drop_set_before_vae is True + + +@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 + # (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 = + # 15_662_823_906 B, 14.59 GiB) is then the peak and over both fit budgets. The fit budget is what a real Mac of + # that size applies (generate passes limits.fit_budget_bytes): the tier's recommended working set (2/3 of its RAM) + # minus the fit rule's 2 GiB reserve. 24 GB: 17_179_869_184 - 2_147_483_648 = 15_032_385_536 B (14.00 GiB), over by + # 630_438_370 B; 16 GB: 11_453_246_122 - 2_147_483_648 = 9_305_762_474 B (8.67 GiB). The denoise phase + # (13_939_106_072 B, 12.98 GiB) is under the 24 GB fit budget. Bug caught: the encode phase sized by less than the + # language model (a call that cannot hold the encoder planned to fit), or the drop not applied (the refusal would + # name the VAE phase). + budget = tier_limits( + tier, host_ram_bytes=RAM_32, host_recommended_bytes=RECOMMENDED_32 + ).fit_budget_bytes + assert budget == {16: 9_305_762_474, 24: 15_032_385_536}[tier] + with pytest.raises(DFloatResourceError, match="in the encode phase"): + _plan(budget) + forced = _plan(budget, fit_check=False) + assert forced.drop_set_before_vae is True + assert forced.estimate.peak_phase == "encode" + assert forced.estimate.peak_bytes == 15_662_823_906 + assert forced.estimate.phases["vae"] == 13_067_259_293 + + +# 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 +# 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 + + +def _encode_bound(text_tokens): + return qmem.encode_peak_bound(_fit(text_tokens=text_tokens).phases["encode"]) + + +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 + # 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 + + +@pytest.mark.parametrize("text_tokens", [9, TEXT_TOKENS, qmem.TEXT_MAX_LENGTH]) +def test_an_encode_that_also_loads_the_vision_tower_and_lm_head_is_over_the_bound_by_2_gb_at_any_prompt( + text_tokens, +): + # Bug caught: the encode bound carrying the denoise allowance again (it grows with the prompt: at 2048 tokens the + # vision-loaded encode was over it by only 78_013_766 B), or a slack wide enough that an mflux change loading the + # vision tower and lm_head (+2_397_528_480 B) goes unnoticed. The excess is the same at every prompt length: + # 2_397_528_480 - 268_435_456. + assert ENCODE_HELD + NOT_LOADED - _encode_bound(text_tokens) == 2_129_093_024 + + +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 (33 tokens; a 2048-token negative prompt + # may hold more activations, which the term does not scale for). + bound = _encode_bound(TEXT_TOKENS) + assert qmem.encode_peak_warning(bound, bound, text_tokens=2048) is None + message = qmem.encode_peak_warning(bound + 1, bound, text_tokens=2048) + assert message is not None + assert "for a 2048-token prompt" in message + assert "measured on one 33-token prompt" in message + 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 +# 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 +# checked against the code as it stands. Both peak in the VAE decode with the set resident, whose estimate at 1024² is +# the calibration value itself (the VAE transient is floored there): they test that the step count and guidance do not +# grow the 1024² peak, not any scaling term. The check is one-sided: the VAE phase varies by about 1 GiB run to run, +# and a run under the estimate is not the bug. +STEPS_AND_CFG_RECORDS = [ + pytest.param("bench/results/tiers/qwen-image-2.1-1024.json", id="measured-40-steps"), + pytest.param( + "bench/results/identity/qwen-image-2.1-1024/df11-result.json", + id="identity-df11-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", STEPS_AND_CFG_RECORDS) +def test_a_committed_1024_record_peaks_within_the_current_estimate_in_the_vae_phase(path): + # Bug caught: a 1024² peak that grows with the step count or with classifier-free guidance (a second transformer + # call per step holding its activations into the VAE decode) while the estimate stays at the one-step value: a + # recorded run more than 0.25 GiB over the current prediction, or the estimate peaking in another phase. Also red + # when the VAE term under-predicts these runs: 1 GiB less of VAE transient puts the identity run 0.08 GiB over. + report, height, width, watched = _record_run(path) + c = qmem.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 == "vae" + assert watched <= fit.peak_bytes + GIB // 4 diff --git a/tests/test_qwen21_model.py b/tests/test_qwen21_model.py new file mode 100644 index 0000000..e3f99a1 --- /dev/null +++ b/tests/test_qwen21_model.py @@ -0,0 +1,1039 @@ +"""Tests for `DFloatQwenImage21`: construction, refusals, the lifecycle and the set, the prelude 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 + +import mlx.core as mx +import mlx.nn as nn +import numpy as np +import pytest +from tests._qwen21_tiny import ( + PREFIX_IDS, + TINY_SEED, + StubTextEncoder, + fake_model, + write_tiny_checkpoint, +) + +from mlx_dfloat.errors import DFloatIntegrationError, DFloatUnsupportedError +from mlx_dfloat.mflux.qwen21 import init as qinit + +pytestmark = pytest.mark.mflux + +MODULATION = "modulation.layers.1.weight" # mflux's path for the checkpoint's modulation.1 (qwen21_weight_mapping.py:18) + + +def _modulation_source(tmp_path): + """The modulation matrix ``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 matrices["modulation.1"]["weight"] + + +def _modulation(model): + return model.transformer.modulation.layers[1].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 QwenImage21.generate_image our assembly never + # sets (mflux's own initializer is not run). + from mflux.models.qwen21.variants.txt2img.qwen_image_21 import QwenImage21 + + names = set(re.findall(r"self\.(\w+)", inspect.getsource(QwenImage21.generate_image))) + assert names >= {"vae", "transformer", "text_encoder", "prompt_cache", "tokenizers", "bits"} + model = fake_model(tmp_path, monkeypatch) + assert [n for n in sorted(names) if not hasattr(model, n)] == [] + assert isinstance(model, QwenImage21) + assert (model.bits, model.prompt_cache, model.tiling_config) == (None, {}, None) + + +@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"), + ], +) +def test_construction_refusals_come_before_any_resolution(monkeypatch, kwargs, match): + # Bug caught: a quantize / LoRA argument or the "none" policy silently honoured (mflux's applier would quantize the + # decoded weights), or resolved (downloading 9.7 GB) before the refusal. + from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 + + monkeypatch.setattr( + qinit, "resolve", lambda *a, **k: pytest.fail("resolved before the refusal") + ) + with pytest.raises(DFloatUnsupportedError, match=match): + DFloatQwenImage21(**{"model": "qwen-image-2.1", **kwargs}) + + +def test_an_unknown_name_is_refused_listing_the_names_this_path_runs(monkeypatch): + # Bug caught: Qwen-Image 1 ("qwen-image", another mflux class and another checkpoint layout) accepted by this + # class, or the refusal not telling the user which names it runs. + from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 + + monkeypatch.setattr(qinit, "resolve", lambda *a, **k: pytest.fail("resolved")) + with pytest.raises(DFloatUnsupportedError, match=r"'qwen-image'.*qwen-image-2\.1"): + DFloatQwenImage21("qwen-image") + + +def test_the_set_load_installs_modulation_and_attaches_with_verify_in_call(tmp_path, monkeypatch): + # Bug caught: modulation.1 decoded but not installed (a zero-size matmul before the first block), or the set + # attached without verify_in_call (a CFG call's second transformer call would be refused). + model = fake_model(tmp_path, monkeypatch) + model._lifecycle.ensure_set() + weight = _modulation(model) + assert weight.shape == (128, 32) + assert np.array_equal(np.array(weight.view(mx.uint16)), _modulation_source(tmp_path)) + assert model.transformer.seam_cell.state.verify_in_call is True + + +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 + # forces, or the encoder reload reading another base than the model was built from. + model = fake_model(tmp_path, monkeypatch) + loads = [] + monkeypatch.setattr( + qinit, + "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 + assert _modulation(model).size == 0 + model._lifecycle.ensure_set() + assert counters.set_loads == 2 + again = np.array(_modulation(model).view(mx.uint16)) + assert again.shape == (128, 32) + assert np.array_equal(again, _modulation_source(tmp_path)) + + +def test_encoded_pairs_live_in_mflux_prompt_cache_under_the_raw_prompt(tmp_path, monkeypatch): + # Bug caught: the empty prompt stored under mflux's normalised " " (mflux's own lookup of "" would miss and call + # the dropped encoder: qwen21_prompt_encoder.py:22-25), or a cache of our own that mflux never reads. + model = fake_model(tmp_path, monkeypatch) + model.encode("a", "") + assert set(model.prompt_cache) == {"a", ""} + for embeds, mask in model.prompt_cache.values(): + # 8 stub ids minus the 3-id system prefix mflux drops; TINY's context_in_dim 32. + assert embeds.shape == (1, 8 - PREFIX_IDS, 32) + assert mask.shape == (1, 8 - PREFIX_IDS) + + +def _bits(array): + return np.array(array.view(mx.uint16)) if array.dtype == mx.bfloat16 else np.array(array) + + +def test_a_cfg_call_caches_each_prompt_with_its_own_embeddings_as_mflux_encodes_it( + tmp_path, monkeypatch +): + # Bug caught: the negative prompt's embeddings stored under the prompt (or the reverse), one prompt encoded for + # both, or a fixed text encoded instead of the caller's. The stub encoder's output depends on the stub + # tokenizer's ids, which depend on the prompt, so each pair must equal mflux's own encode of that prompt. + from mflux.models.qwen21.model.qwen21_text_encoder.qwen21_prompt_encoder import ( + Qwen21PromptEncoder, + ) + from tests._qwen21_tiny import StubTokenizer + + model = fake_model(tmp_path, monkeypatch) + _patch_upstream_generate(monkeypatch) + model.generate_image(seed=1, prompt="p", height=64, width=64, guidance=4.0, negative_prompt="n") + assert not np.array_equal(_bits(model.prompt_cache["p"][0]), _bits(model.prompt_cache["n"][0])) + for prompt in ("p", "n"): + embeds, mask = Qwen21PromptEncoder.encode_prompt( + prompt=prompt, + prompt_cache={}, + tokenizer=StubTokenizer(), + text_encoder=StubTextEncoder(), + ) + got_embeds, got_mask = model.prompt_cache[prompt] + assert np.array_equal(_bits(got_embeds), _bits(embeds)), prompt + assert np.array_equal(_bits(got_mask), _bits(mask)), prompt + + +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(input_ids=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 modulation 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 _modulation(model).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_modulation_group_at_set_load_releases_the_set(tmp_path, monkeypatch): + # Bug caught: a status failure while the set loads (modulation.1 decoded 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("modulation.1") + model = fake_model(tmp_path, monkeypatch, decode=decode) + model.encode("p") + gc.collect() + mx.clear_cache() + before = int(mx.get_active_memory()) + with pytest.raises(DFloatFormatError, match=r"modulation\.1") 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_the_retained_bound_counts_the_cached_pairs_and_the_vae(tmp_path, monkeypatch): + # Bug caught: the bound reading a cache of our own (empty: every drop after an encode would refuse) instead of + # mflux's prompt_cache, where the pairs live. + model = fake_model(tmp_path, monkeypatch) + empty = model._retained_bound() + model.encode("a") + # One pair: embeds (1, 5, 32) bf16 = 320 B, mask (1, 5) int32 = 20 B. + assert model._retained_bound() - empty == 320 + 20 + + +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)) + + +# --- the generate prelude, the call 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 QwenImage21.generate_image with a probe: records the cache limit in force and the kwargs, fires our + after-loop subscriber the way mflux's GenerationContext would, returns a sentinel.""" + from mflux.models.qwen21.variants.txt2img.qwen_image_21 import QwenImage21 + 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(QwenImage21, "generate_image", fake) + return seen + + +def _record_encodes(monkeypatch, model): + encodes = [] + real = model._lifecycle._encode + monkeypatch.setattr(model._lifecycle, "_encode", lambda p: (encodes.append(p), real(p))[1]) + return encodes + + +@pytest.mark.parametrize( + ("guidance", "negative", "expected"), + [ + (1.0, "n", ("p",)), + (4.0, None, ("p",)), + (4.0, "", ("p",)), + (4.0, " ", ("p", " ")), + (1.0001, "n", ("p", "n")), + (None, "n", ("p",)), + ], +) +def test_cfg_prompts_follow_mflux(tmp_path, monkeypatch, guidance, negative, expected): + # Bug caught: `>=` for `>` (guidance 1.0 encoding the negative: a wasted encoder reload), an empty negative running + # CFG (mflux disables it: `config.guidance > 1.0 and bool(negative_prompt)`, qwen_image_21.py:89), or a blank-but- + # non-empty negative dropped (bool(" ") is True: mflux encodes it). + model = fake_model(tmp_path, monkeypatch) + assert model.cfg_prompts("p", negative_prompt=negative, guidance=guidance) == expected + + +def test_the_generate_defaults_are_mfluxs(tmp_path, monkeypatch): + # Bug caught: our signature drifting from QwenImage21.generate_image's (qwen_image_21.py:42-54: 40 steps, 1024^2, + # guidance 1.0, scheduler "linear", no negative prompt), so the Python API ran another recipe than mflux's. + from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 + + defaults = { + n: p.default + for n, p in inspect.signature(DFloatQwenImage21.generate_image).parameters.items() + if p.default is not inspect.Parameter.empty + } + assert defaults == { + "num_inference_steps": 40, + "height": 1024, + "width": 1024, + "guidance": 1.0, + "image_path": None, + "image_strength": None, + "scheduler": "linear", + "negative_prompt": None, + } + + +def test_a_missing_guidance_and_scheduler_become_mflux_defaults(tmp_path, monkeypatch): + # Bug caught: the CLI's None (no --guidance, no --scheduler) reaching mflux's Config, which builds the scheduler by + # name, instead of mflux's defaults (guidance 1.0, "linear"). + model = fake_model(tmp_path, monkeypatch) + seen = _patch_upstream_generate(monkeypatch) + model.generate_image(seed=1, prompt="p", height=64, width=64, guidance=None, scheduler=None) + assert seen["kwargs"]["guidance"] == 1.0 + assert seen["kwargs"]["scheduler"] == "linear" + assert model._cfg_calls == 1 + + +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 + # 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) + encodes = _record_encodes(monkeypatch, model) + model.generate_image( + seed=1, + prompt="", + num_inference_steps=1, + height=64, + width=64, + guidance=4.0, + negative_prompt="", + ) + assert (encodes, model._cfg_calls) == ([""], 1) + model.generate_image( + seed=1, + prompt="x", + num_inference_steps=1, + height=64, + width=64, + guidance=4.0, + negative_prompt="x", + ) + assert (encodes, model._cfg_calls) == (["", "x"], 2) + assert set(model.prompt_cache) == {"", "x"} + # 1 step x 1 call x 2 blocks, then 1 step x 2 calls x 2 blocks: the second call ran the negative branch. + assert model.report()["decode_launches"] == 2 + 4 + + +def test_text_tokens_follow_the_longest_prompt_of_the_call(tmp_path, monkeypatch): + # Bug caught: the plan sized by the positive prompt only (a longer negative prompt's activations unplanned), or the + # system prefix counted. Stub lengths 7 and 12 ids, a 3-id prefix: 9 text tokens. + from tests._qwen21_tiny import StubTokenizer + + 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=4.0, negative_prompt="n") + assert planned == [12 - PREFIX_IDS] + assert model.report()["text_tokens"] == 12 - PREFIX_IDS + + +@pytest.mark.parametrize("kwargs", [{"image_path": "in.png"}, {"image_strength": 0.5}]) +def test_img2img_is_refused_before_anything_loads(tmp_path, monkeypatch, kwargs): + # Bug caught: encoding the prompt (an encoder load) or loading the set before refusing img2img. + 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_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 pairs exist, img2img + # reaching QwenImage21.generate_image, or the negative prompt not passed through (mflux reads it for the CFG rule + # and the image metadata). + from mflux.models.qwen21.variants.txt2img.qwen_image_21 import QwenImage21 + + 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 = QwenImage21.generate_image + + def recording(self, **kwargs): + events.append("upstream") + return upstream(self, **kwargs) + + monkeypatch.setattr(QwenImage21, "generate_image", recording) + image = model.generate_image( + seed=7, + prompt="p", + num_inference_steps=1, + height=256, + width=256, + guidance=4.0, + 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"]) == (None, None) + assert (kwargs["guidance"], kwargs["negative_prompt"], kwargs["num_inference_steps"]) == ( + 4.0, + "n", + 1, + ) + assert model._cfg_calls == 2 + + +def test_plan_call_reads_the_familys_constants_sizes_and_groups(tmp_path, monkeypatch): + # Bug caught: plan_call not wired to the shared planner with this model's own largest group and the call's text + # tokens (the cache limit would miss the allowance or the decoded group), or the plan not kept for the VAE guard. + from mlx_dfloat.errors import DFloatResourceError + + model = fake_model(tmp_path, monkeypatch) + plan = model.plan_call(height=1030, width=1030, text_tokens=5) # rounded to 1024^2 first + # One block kind: the tiny block group, 13_312 elements (26_624 B), next to Qwen's allowance at 4096 image + 5 + # text tokens: 2_295_554_018 x 4101 / 4129 = 2_279_987_170 (truncated). + assert plan.cache_limit == 26_624 + 2_279_987_170 + assert model._plan is plan + with pytest.raises(DFloatResourceError, match="1024x1024"): + model.plan_call(height=1040, width=1040, text_tokens=5) + + +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 + # dropped on a roomy Mac: a reload per call), or a transformer-sized cache limit left for the VAE decode. + # Tight: VAE phase 12e9 + 11_119_436_651 + 459_445_098 = 23_578_881_749 > 16e9; denoise ~15.2e9 fits. + from mlx_dfloat.mflux._phases import FamilySizes + + 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) + 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) + 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 + + +@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; + # checked before the after-loop guard ran (inside the loop) and after it (during the decode). + 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 _modulation(model).size == 0 + assert model.open_phase is None + + +def test_keyboard_interrupt_in_the_loop_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, qwen_image_21.py:125-129) + # 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": 4.0, + "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 + + +def _innermost_locals(tb): + while tb.tb_next is not None: + tb = tb.tb_next + return tb.tb_frame.f_locals + + +@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, through __context__) 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.qwen21.variants.txt2img.qwen_image_21 import QwenImage21 + + 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: + raise RuntimeError("outer") from None # __context__ is kept even with `from None` + + monkeypatch.setattr(QwenImage21, "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.__context__, ValueError) + assert "big" not in _innermost_locals(exc.__context__.__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: the phase must follow the call through + # encode, set load, the loop and the VAE decode, and be None again once the call returns. + from mflux.models.qwen21.variants.txt2img.qwen_image_21 import QwenImage21 + + 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(QwenImage21, "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 test_a_failed_set_load_in_a_call_closes_the_phase(tmp_path, monkeypatch): + # Bug caught: the tracker left at "set_load" after the set load raised (the watchdog's abort context would name a + # phase that is not running). + model = fake_model(tmp_path, monkeypatch) + _patch_upstream_generate(monkeypatch) + + def broken_attach(*args, **kwargs): + raise RuntimeError("attach") + + monkeypatch.setattr(model.transformer, "attach", broken_attach) + with pytest.raises(RuntimeError, match="attach"): + model.generate_image(seed=1, prompt="p", height=256, width=256) + assert model.open_phase is None + + +def _fake_encode_peak(monkeypatch, model, above_start): + """Make the tracker record an encode phase whose MLX peak is ``above_start`` bytes over the phase's start.""" + real_end = model._tracker.end + + def end(name): + real_end(name) + if name == "encode": + record = model._tracker.peaks["encode"] + record["mlx_peak"] = record["active_at_start"] + above_start + + monkeypatch.setattr(model._tracker, "end", end) + + +# fake_model's encoders (1_000) + Qwen's encode term (66_567_800, in place of the allowance) + the 256 MiB slack +# (268_435_456). +ENCODE_WARN_ABOVE = 335_004_256 + + +def test_an_encode_peak_over_the_estimate_warns_that_mflux_may_load_more_of_the_encoder( + tmp_path, monkeypatch +): + # Bug caught: an mflux change that loads the vision tower or the lm_head with the language model going unnoticed: + # the encode phase grows by ~2.4 GB while the fit estimate, sized by the language model, keeps passing calls. + model = fake_model(tmp_path, monkeypatch) + _patch_upstream_generate(monkeypatch) + _fake_encode_peak(monkeypatch, model, ENCODE_WARN_ABOVE + 1) + # The stub tokenizer gives 8 ids, 3 of them the system prefix mflux drops: a 5-token prompt. + with pytest.warns( + UserWarning, match=r"for a 5-token prompt.*more of the encoder than its language model" + ): + model.generate_image(seed=1, prompt="p", height=256, width=256) + + +def test_an_encode_peak_at_the_estimate_does_not_warn(tmp_path, monkeypatch): + # Bug caught: the slack or the encode term left out of the bound, or `>=` for `>` at the bound. Pytest turns a + # warning into an error here. + model = fake_model(tmp_path, monkeypatch) + _patch_upstream_generate(monkeypatch) + _fake_encode_peak(monkeypatch, model, ENCODE_WARN_ABOVE) + assert model.generate_image(seed=1, prompt="p", height=256, width=256) == "image" + + +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.qwen21 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_carries_the_family_keys_and_the_sizes_before_any_call(tmp_path, monkeypatch): + # Bug caught: the family, sizes or the checkpoint's config source missing from the report (the CLI and the table + # read them; a config-less checkpoint's provenance is the layout that matched it), or a Klein-only key. + model = fake_model(tmp_path, monkeypatch) + report = model.report() + assert report["family"] == "qwen21" + assert report["model"] == "qwen-image-2.1" + assert "size" not in report + assert report["sizes"] == { + "compressed": 1_000, + "extras": 100, + "nonblock": 10, + "encoders": 1_000, + "vae": 10, + } # fake_model's own FamilySizes + assert report["predict"] == "uncompiled" + assert report["nonblock_groups"] == ["modulation.1"] + assert (report["cfg_calls_per_step"], report["text_tokens"], report["fit"]) == (1, None, None) + assert report["df11"]["repo_id"] == "mingyi456/tiny" + assert "layout tiny-qwen" in report["df11"]["config_source"] + assert report["base"]["revision"] == "b" * 40 + assert report["eval_policy"] == "per-block" + + +@pytest.mark.parametrize( + ("policy", "status"), [("per-block", "measured"), ("depth2", "unmeasured")] +) +def test_the_report_labels_a_depth2_plan_unmeasured(tmp_path, monkeypatch, policy, status): + # Bug caught: a depth2 run reported like the measured per-block path (no Qwen-Image 2.1 run used depth2: its + # look-ahead buffer and step time are unmeasured), or depth2 refused instead of labelled. + report = fake_model(tmp_path, monkeypatch, eval_policy=policy).report() + assert (report["eval_policy"], report["eval_policy_status"]) == (policy, status) + + +def test_the_report_after_a_cfg_call_names_two_calls_per_step_and_the_estimate( + tmp_path, monkeypatch +): + # Bug caught: cfg_calls_per_step not following the call (the bench table would mislabel a CFG run), or the + # estimate missing or labelled as a measurement. + model = fake_model(tmp_path, monkeypatch) + model.generate_image( + seed=1, + prompt="a", + num_inference_steps=1, + height=64, + width=64, + guidance=4.0, + negative_prompt="n", + ) + report = model.report() + assert report["cfg_calls_per_step"] == 2 + assert report["predict"] == "uncompiled" + assert report["fit"]["label"] == "predicted" + assert report["text_tokens"] == 8 - PREFIX_IDS + + +def test_a_compiled_forward_is_refused_in_the_loop_and_the_report_says_compiled( + tmp_path, monkeypatch +): + # Bug caught: the eager guard not reached through mflux's loop (the seam would evaluate inside mx.compile), or + # report() still printing "uncompiled" after mflux compiled the forward pass. + model = fake_model(tmp_path, monkeypatch) + assert model.report()["predict"] == "uncompiled" + model.transformer._step_fn = mx.compile(model.transformer._forward) + assert model.report()["predict"] == "compiled" + with pytest.raises(DFloatIntegrationError, match="compiled"): + model.generate_image(seed=1, prompt="a", num_inference_steps=1, height=64, width=64) + assert model.open_phase is None + + +def test_the_default_checkpoint_resolves_at_its_pin_and_a_user_checkpoint_unpinned(monkeypatch): + # Bug caught: the pin not reaching the resolver (the default, a single file on one person's Hub account, following + # whatever lands on main instead of the revision whose header and bytes were verified), or a pin applied to a + # user's own --df11. + from mlx_dfloat import _metal_decode + from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 + + class StopError(Exception): + pass + + calls = [] + + def fake_resolve(spec, *, patterns, revision=None): + calls.append((spec, revision)) + raise StopError + + monkeypatch.setattr(_metal_decode, "ensure_canary", lambda **kwargs: None) + monkeypatch.setattr(qinit, "resolve", fake_resolve) + for kwargs in ({}, {"df11_path": "me/fork"}): + with pytest.raises(StopError): + DFloatQwenImage21("qwen-image-2.1", **kwargs) + assert calls == [ + ("mingyi456/Qwen-Image-2.1-DF11-ComfyUI", "1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc"), + ("me/fork", None), + ] + + +def test_the_default_base_resolves_at_its_pin_and_a_user_base_unpinned(tmp_path, monkeypatch): + # Bug caught: the base's pin not reaching the resolver (the text encoder, VAE and tokenizer following whatever + # lands on Qwen/Qwen-Image-2.1's main instead of the snapshot the recorded runs used), or a pin applied to a user's + # own --base. + from mlx_dfloat import _metal_decode + from mlx_dfloat.mflux._hub import ResolvedRepo + from mlx_dfloat.mflux.qwen21 import model as model_module + from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 + + class StopError(Exception): + pass + + calls = [] + + def fake_resolve(spec, *, patterns, revision=None): + if patterns == qinit.DF11_PATTERNS: + return ResolvedRepo(root=tmp_path, repo_id=spec, revision=revision) + calls.append((spec, revision)) + raise StopError + + monkeypatch.setattr(_metal_decode, "ensure_canary", lambda **kwargs: None) + monkeypatch.setattr(model_module, "open_checkpoint", lambda root: None) + monkeypatch.setattr(qinit, "resolve", fake_resolve) + for kwargs in ({}, {"base_path": "me/base"}): + with pytest.raises(StopError): + DFloatQwenImage21("qwen-image-2.1", **kwargs) + assert calls == [ + ("Qwen/Qwen-Image-2.1", "d26bb61231c349cf6b7896fa83353113880e1ba3"), + ("me/base", None), + ] + + +# --- the constructor, end to end past the resolver ------------------------------------------------------------ + + +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 + # 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 + # tower and lm_head included) instead of the language model's tensor bytes. + from tests._df11_fixtures import write_bf16_original + from tests._qwen21_tiny import TINY, stub_components + + from mlx_dfloat import _layouts, _metal_decode + from mlx_dfloat.mflux.qwen21 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" + # The language model: embed_tokens (10, 4) 80 B, one layer's up_proj (4, 4) 32 B, final norm (4,) 8 B. Never + # loaded: a vision tensor (8, 4) 64 B and an untied lm_head (10, 4) 80 B. 80 + 32 + 8 = 120. + write_bf16_original( + base / "text_encoder", + { + "model.language_model.embed_tokens.weight": np.zeros((10, 4), np.uint16), + "model.language_model.layers.0.mlp.up_proj.weight": np.zeros((4, 4), np.uint16), + "model.language_model.norm.weight": np.zeros((4,), np.uint16), + "model.visual.blocks.0.attn.qkv.weight": np.zeros((8, 4), np.uint16), + "lm_head.weight": np.zeros((10, 4), np.uint16), + }, + ) + (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( + qinit, "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) + + monkeypatch.setattr(model_module, "build_transformer", tiny_build) + model = model_module.DFloatQwenImage21( + "qwen-image-2.1", df11_path=str(tmp_path / "df11"), base_path=str(base) + ) + assert built == [{}] + assert loaded == [base] + report = model.report() + assert report["sizes"]["encoders"] == 120 + assert report["sizes"]["vae"] == 100 + assert report["sizes"]["nonblock"] == 128 * 32 * 2 # modulation.1 decoded, (128, 32) BF16 + assert report["df11"]["config_source"].startswith( + "header and spot checks match layout tiny-qwen" + ) + assert (report["df11"]["root"], report["df11"]["repo_id"]) == (str(tmp_path / "df11"), None) + assert report["base"]["root"] == str(base) + + +@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 the VAE term, kept with the set on a 32 GB Mac with 1.74 GiB to spare, was 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 + + +class _RecordingVAE(nn.Module): + """The stub VAE's decode, keeping the latents mflux hands it (the denoised result, 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, 16 * latents.shape[-2], 16 * latents.shape[-1])) + + +class _VaryingTextEncoder(nn.Module): + """A text stack whose hidden states differ in direction from token to token (the stub's are one vector scaled by + the id, which the transformer's RMSNorm on the text stream maps to the same vector for every prompt).""" + + def __init__(self): + super().__init__() + self.freq = mx.arange(1, 33, dtype=mx.float32) / 7 + + def __call__(self, input_ids, attention_mask=None): + return mx.sin(input_ids[..., None].astype(mx.float32) * self.freq).astype(mx.bfloat16) + + +def _stock_qwen(tokenizer, vae, matrices, extras): + """Stock mflux ``QwenImage21`` over a plain tiny ``Qwen21Transformer`` holding the BF16 source of a tiny checkpoint + (``write_tiny_checkpoint``'s matrices and extras), uncompiled like our side.""" + from mflux.models.common.config.model_config import ModelConfig + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + from mflux.models.qwen21.qwen21_initializer import Qwen21Initializer + from mflux.models.qwen21.variants.txt2img.qwen_image_21 import QwenImage21 + from mlx.utils import tree_flatten, tree_unflatten + from tests._qwen21_tiny import TINY + + transformer = Qwen21Transformer(**TINY) + shapes = {name: p.shape for name, p in tree_flatten(transformer.parameters())} + weights = [ + (f"{group}.{sub}.weight", mx.array(matrix).view(mx.bfloat16)) + for group, subs in matrices.items() + if group != "modulation.1" + for sub, matrix in subs.items() + ] + weights.append( + ( + "modulation.layers.1.weight", + mx.array(matrices["modulation.1"]["weight"]).view(mx.bfloat16), + ) + ) + weights += [(name, mx.array(bits).view(mx.bfloat16)) for name, bits in extras.items()] + assert all(tuple(shapes[name]) == bits.shape for name, bits in extras.items()) + transformer.update(tree_unflatten(weights)) + transformer._step_fn = ( + transformer._forward + ) # mflux compiles it otherwise; our side runs it plain + stock = QwenImage21.__new__(QwenImage21) + nn.Module.__init__(stock) + Qwen21Initializer._init_config(stock, ModelConfig.qwen_image_21()) + stock.vae = vae + stock.text_encoder = _VaryingTextEncoder() + stock.tokenizers = {"qwen21": tokenizer} + stock.transformer = transformer + stock.bits = None + return stock + + +@pytest.mark.parametrize( + ("guidance", "negative"), [(1.0, None), (4.0, "n")], ids=["guidance-1", "cfg"] +) +def test_a_tiny_generation_gives_stock_mfluxs_latents_bit_for_bit( + tmp_path, monkeypatch, guidance, negative +): + # Bug caught: anything between the checkpoint and mflux's loop changing the result: a decoded matrix on the wrong + # parameter, a gate_up half swapped, an extra mis-routed, or the CFG branch fed the other prompt's embeddings or + # mask (the negative prompt is 12 ids here, the prompt 8, so a swap changes the latents or fails). The reference is + # stock mflux over the same BF16 weights, uncompiled as ours; equality is bit for bit (an integer view of the + # latents). The extras are random: constant ones make the forward pass blind to its inputs. + from tests._qwen21_tiny import TINY, StubTokenizer, stub_components + + from mlx_dfloat.format import open_checkpoint + from mlx_dfloat.mflux.qwen21.transformer import build_transformer + + 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 = _RecordingVAE(), _RecordingVAE() + components = stub_components(tokenizer) + components.vae, components.text_encoder = ours_vae, _VaryingTextEncoder() + model = fake_model( + tmp_path, + monkeypatch, + ckpt=ckpt, + build=build_transformer(ckpt, transformer_kwargs=TINY), + components=components, + ) + stock = _stock_qwen(tokenizer, stock_vae, matrices, extras) + kwargs = { + "seed": 3, + "prompt": "p", + "num_inference_steps": 2, + "height": 64, + "width": 64, + "guidance": guidance, + "negative_prompt": negative, + } + model.generate_image(**kwargs) + stock.generate_image(**kwargs, scheduler="linear") + (ours,), (theirs,) = ours_vae.seen, stock_vae.seen + assert model.report()["cfg_calls_per_step"] == (2 if negative else 1) + assert bool(mx.all(mx.isfinite(theirs))) + assert ours.dtype == theirs.dtype + assert ours.shape == theirs.shape + view = mx.uint32 if ours.dtype == mx.float32 else mx.uint16 + assert np.array_equal(np.array(ours.view(view)), np.array(theirs.view(view))) diff --git a/tests/test_qwen21_names.py b/tests/test_qwen21_names.py new file mode 100644 index 0000000..b70f577 --- /dev/null +++ b/tests/test_qwen21_names.py @@ -0,0 +1,349 @@ +"""Qwen-Image 2.1 naming: the folded mapping and the group check (offline); the real mapping (mflux lane).""" + +from types import SimpleNamespace + +import numpy as np +import pytest +from tests._df11_fixtures import random_bf16, write_checkpoint + +from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError +from mlx_dfloat.format import open_checkpoint +from mlx_dfloat.mflux._mapping import name_map_from_mapping +from mlx_dfloat.mflux.qwen21.names import ( + KIND, + MATRIX_SUBS, + RUN_ORDER, + check_qwen21_groups, + templated_targets, +) + +# The seven block matrices in the order a block decodes into after the gate_up split (S0: the stored order 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", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "img_mlp.gate_layer", + "img_mlp.proj", + "img_mlp.out", +) + + +def _t(to, frm, transform=None): + return SimpleNamespace(to_pattern=to, from_pattern=[frm], transform=transform) + + +def _mapping(n_blocks=3, *, odd_block=None): + # A copy of Qwen21WeightMapping.get_transformer_mapping() (mflux 0.20.0 + # models/qwen21/weights/qwen21_weight_mapping.py:14-53) for n_blocks blocks, in its order: nine non-block targets, + # one rename (modulation.1 -> modulation.layers.1, :18), then per block by concrete index (:32-52). + # test_the_real_mapping_matches_the_copy compares it with mflux's. + targets = [_t("img_in.weight", "img_in.weight"), _t("proj_out.weight", "proj_out.weight")] + targets.append(_t("modulation.layers.1.weight", "modulation.1.weight")) + targets += [ + _t(n, n) + for n in ( + "norm_out.linear.weight", + "txt_in.text_norm.weight", + "txt_in.in_layer.weight", + "txt_in.out_layer.weight", + "time_text_embed.timestep_embedder.linear_1.weight", + "time_text_embed.timestep_embedder.linear_2.weight", + ) + ] + for b in range(n_blocks): + p = f"transformer_blocks.{b}." + targets += [ + _t(f"{p}attn.{s}.weight", f"{p}attn.{s}.weight") + for s in ("to_q", "to_k", "to_v", "norm_q", "norm_k") + ] + to_out = f"{p}attn.to_out.0.weight" + targets.append(_t(f"{p}attn.out_proj.weight" if b == odd_block else to_out, to_out)) + targets += [ + _t(f"{p}img_mlp.{s}.weight", f"{p}img_mlp.{s}.weight") + for s in ("proj", "out", "gate_layer") + ] + return targets + + +def _folded(targets=None): + targets = _mapping() if targets is None else targets + return name_map_from_mapping( + templated_targets(targets, kinds=RUN_ORDER, n_blocks=3), matrix_subs=MATRIX_SUBS + ) + + +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). + m = _folded() + assert m.attrs_of("transformer_blocks") == SEVEN + assert m.place("transformer_blocks.2.img_mlp.proj.weight").attr == "img_mlp.proj" + assert m.place("transformer_blocks.0.attn.to_out.0.weight").attr == "attn.to_out.0" + assert ( + m.param_name("transformer_blocks.1.attn.norm_k.weight") + == "transformer_blocks.1.attn.norm_k.weight" + ) + + +def test_modulation_is_renamed_into_the_sequential(): + # Bug caught: the one non-block rename lost (modulation.1.weight has no parameter of that name in mflux). + assert _folded().param_name("modulation.1.weight") == "modulation.layers.1.weight" + + +def test_non_block_targets_pass_through_unchanged(): + # Bug caught: the fold swallowing (or templating) a non-block target, so img_in or the rename disappears. + out = templated_targets(_mapping(), kinds=RUN_ORDER, n_blocks=3) + flat = {(t.to_pattern, tuple(t.from_pattern)) for t in out} + assert ("img_in.weight", ("img_in.weight",)) in flat + assert ("modulation.layers.1.weight", ("modulation.1.weight",)) in flat + assert ( + "transformer_blocks.{block}.img_mlp.gate_layer.weight", + ("transformer_blocks.{block}.img_mlp.gate_layer.weight",), + ) in flat + assert len(out) == 9 + 9 # nine non-block targets, nine block subs (seven matrices, two norms) + + +def test_a_block_whose_rename_differs_is_refused(): + # Bug caught: block 0's rename applied to every block when block 1 maps elsewhere. + with pytest.raises(DFloatIntegrationError, match=r"attn\.to_out\.0.*block 1\b"): + templated_targets(_mapping(odd_block=1), kinds=RUN_ORDER, n_blocks=3) + + +def test_a_missing_block_target_is_refused(): + # Bug caught: a mapping that stops at block 30 folded as if it covered 32. + with pytest.raises(DFloatIntegrationError, match=r"no target for block 2\b"): + templated_targets(_mapping(n_blocks=2), kinds=RUN_ORDER, n_blocks=3) + + +def test_a_block_target_beyond_the_count_is_refused(): + # Bug caught: a mapping for 33 blocks folded against 32, the extra block's target silently dropped. + with pytest.raises(DFloatIntegrationError, match=r"block 3\b"): + templated_targets(_mapping(n_blocks=4), kinds=RUN_ORDER, n_blocks=3) + + +def test_a_block_target_with_a_transform_is_refused(): + # Bug caught: a transform on one block's target dropped by the fold (the BF16 bytes would load untransformed). + targets = _mapping() + i = next( + i + for i, t in enumerate(targets) + if t.to_pattern == "transformer_blocks.2.attn.norm_q.weight" + ) + targets[i] = _t(targets[i].to_pattern, targets[i].from_pattern[0], transform=lambda a: a) + with pytest.raises(DFloatIntegrationError, match=r"attn\.norm_q.*block 2\b"): + templated_targets(targets, kinds=RUN_ORDER, n_blocks=3) + + +def test_a_block_source_mapped_to_another_block_index_is_refused(): + # Bug caught: block 1's source loaded into block 2's parameter folded as if it were the identity. + targets = _mapping() + i = next( + i for i, t in enumerate(targets) if t.to_pattern == "transformer_blocks.1.attn.to_q.weight" + ) + targets[i] = _t( + "transformer_blocks.2.attn.to_q.weight", "transformer_blocks.1.attn.to_q.weight" + ) + with pytest.raises(DFloatIntegrationError, match=r"transformer_blocks\.1\.attn\.to_q"): + templated_targets(targets, kinds=RUN_ORDER, n_blocks=3) + + +# --- the group check, on written tiny checkpoints ------------------------------------------------------------------- + +_BLOCK = r"transformer_blocks\.\d+" + + +def _write(root, *, blocks=(0, 1), subs=SEVEN, modulation=False, extra_groups=()): + rng = np.random.default_rng(3) + groups = {f"transformer_blocks.{b}": [random_bf16(rng, (2, 4)) for _ in subs] for b in blocks} + patterns = {_BLOCK: subs} + if modulation: + groups["modulation.1"] = [random_bf16(rng, (8, 4))] + patterns[r"modulation\.1"] = () + for name in extra_groups: + groups[name] = [random_bf16(rng, (2, 4))] + patterns[name.replace(".", r"\.")] = () + write_checkpoint(root, groups=groups, patterns=patterns) + return open_checkpoint(root) + + +@pytest.mark.parametrize("modulation", [False, True]) +def test_the_block_count_is_returned_with_and_without_the_modulation_group(tmp_path, modulation): + # Bug caught: the optional non-block group counted as a block (or refused), or the count off by one. + assert check_qwen21_groups(_write(tmp_path, modulation=modulation)) == {"transformer_blocks": 2} + + +def test_a_hole_in_the_block_indices_is_refused(tmp_path): + # Bug caught: blocks 0 and 2 accepted, block 1 left on placeholders (a zero-size matmul at run time). + with pytest.raises(DFloatFormatError, match="not contiguous"): + check_qwen21_groups(_write(tmp_path, blocks=(0, 2))) + + +def test_a_zero_padded_block_index_is_refused(tmp_path): + # Bug caught (fable M3): transformer_blocks.01 counted as block 1 (int("01") == 1), then the seam asks the + # provider for transformer_blocks.1, a group the checkpoint does not have: a failure at the first block after the + # encoder and set loads. + with pytest.raises( + DFloatFormatError, match=r"transformer_blocks\.01: not a Qwen-Image 2\.1 group" + ): + check_qwen21_groups(_write(tmp_path, blocks=("0", "01"))) + + +def test_a_checkpoint_without_blocks_is_refused(tmp_path): + # Bug caught: a modulation-only checkpoint accepted as zero blocks. + rng = np.random.default_rng(4) + write_checkpoint( + tmp_path, + groups={"modulation.1": [random_bf16(rng, (8, 4))]}, + patterns={r"modulation\.1": ()}, + ) + with pytest.raises(DFloatFormatError, match="no transformer_blocks groups"): + check_qwen21_groups(open_checkpoint(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 + # 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", + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "img_mlp.net.0.proj", + "img_mlp.net.2", + "txt_mod.1", + "attn.add_q_proj", + "attn.add_k_proj", + "attn.add_v_proj", + "attn.to_add_out", + "txt_mlp.net.0.proj", + "txt_mlp.net.2", + ) + with pytest.raises(DFloatFormatError, match=r"transformer_blocks\.0.*img_mlp\.gate_layer"): + check_qwen21_groups(_write(tmp_path, subs=qwen_image_1)) + + +def test_a_stray_group_is_refused_naming_the_accepted_non_block_group(tmp_path): + # Bug caught: an unknown group (another model's) silently ignored, its matrices never placed. + with pytest.raises(DFloatFormatError, match=r"txt_mod\.1.*modulation\.1"): + check_qwen21_groups(_write(tmp_path, extra_groups=("txt_mod.1",))) + + +def test_a_modulation_group_with_other_matrices_is_refused(tmp_path): + # Bug caught: a modulation.1 group of several matrices accepted, the extra ones never installed. + rng = np.random.default_rng(5) + groups = { + "transformer_blocks.0": [random_bf16(rng, (2, 4)) for _ in SEVEN], + "modulation.1": [random_bf16(rng, (8, 4)), random_bf16(rng, (8, 4))], + } + write_checkpoint( + tmp_path, + groups=groups, + patterns={_BLOCK: SEVEN, r"modulation\.1": ("a", "b")}, + ) + with pytest.raises(DFloatFormatError, match=r"modulation\.1"): + check_qwen21_groups(open_checkpoint(tmp_path)) + + +def _shape_of(targets): + return [(t.to_pattern, tuple(t.from_pattern), t.transform) for t in targets] + + +@pytest.mark.mflux +def test_the_real_mapping_matches_the_copy(): + # Bug caught: the offline tests running on a copy that no longer is mflux's mapping (a target renamed, added, + # reordered or given a transform upstream), so they pass while the real fold would differ. mflux's 32 blocks are + # cut to the copy's three; everything else must be equal, in order. + from mflux.models.qwen21.weights.qwen21_weight_mapping import Qwen21WeightMapping + + real = Qwen21WeightMapping.get_transformer_mapping() + kept = [ + t + for t in real + if not t.to_pattern.startswith("transformer_blocks.") or int(t.to_pattern.split(".")[1]) < 3 + ] + assert _shape_of(kept) == _shape_of(_mapping(3)) + assert len(real) == 9 + 32 * 9 + + +@pytest.mark.mflux +def test_the_default_layer_count_is_mfluxs_constructor_default(): + # Bug caught: an mflux bump changing Qwen21Transformer's default depth (what the published model is built with) + # while the count check keeps 32: every real checkpoint refused, or one of another depth accepted. + import inspect + + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + from mlx_dfloat.mflux.qwen21.names import DEFAULT_LAYERS + + default = inspect.signature(Qwen21Transformer.__init__).parameters["num_layers"].default + assert default == 32 + assert DEFAULT_LAYERS == 32 + + +@pytest.mark.mflux +def test_the_real_mapping_folds_and_every_matrix_lands_on_a_linear(): + # Bug caught: a decoded matrix with no nn.Linear to land on in a real tiny Qwen21Transformer, or the real mapping + # not folding into the seven block matrices and the modulation rename. + import mlx.nn as nn + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + from mlx_dfloat.integrate.placeholders import get_attr_path + from mlx_dfloat.mflux.qwen21.names import qwen21_name_map + + m = qwen21_name_map() + assert m.attrs_of(KIND) == SEVEN + assert m.param_name("modulation.1.weight") == "modulation.layers.1.weight" + tf = Qwen21Transformer( + num_layers=1, + attention_head_dim=16, + num_attention_heads=2, + context_in_dim=32, + axes_dims_rope=(4, 6, 6), + ) + for attr in m.attrs_of(KIND): + assert isinstance(get_attr_path(tf.transformer_blocks[0], attr), nn.Linear), attr + modulation = get_attr_path(tf, m.param_name("modulation.1.weight").removesuffix(".weight")) + assert isinstance(modulation, nn.Linear) + + +QWEN21_MODULES = ( + "mlx_dfloat.mflux.qwen21", + "mlx_dfloat.mflux.qwen21.names", + "mlx_dfloat.mflux.qwen21.transformer", +) + + +@pytest.mark.parametrize("module", QWEN21_MODULES) +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" + + +@pytest.mark.parametrize( + ("extra", "message"), + [ + # Bug caught: a non-block parameter loaded from a block's tensor passed through the fold untouched (the + # shared derivation would place the block source by the table, not where mflux loads it). + ( + _t("img_in.weight", "transformer_blocks.0.attn.to_q.weight"), + "one block source in its own block", + ), + # Bug caught: a second target for block 1's source folded over the first (one of them silently lost). + ( + _t("transformer_blocks.1.attn.to_q.weight", "transformer_blocks.1.attn.to_q.weight"), + r"two targets for block 1\b", + ), + ], +) +def test_a_block_source_that_does_not_fold_cleanly_is_refused(extra, message): + with pytest.raises(DFloatIntegrationError, match=message): + templated_targets([*_mapping(), extra], kinds=RUN_ORDER, n_blocks=3) diff --git a/tests/test_qwen21_real_blocks.py b/tests/test_qwen21_real_blocks.py new file mode 100644 index 0000000..5d35ac9 --- /dev/null +++ b/tests/test_qwen21_real_blocks.py @@ -0,0 +1,221 @@ +"""The real Qwen-Image 2.1 DF11 file: its layout, Metal decode against the BF16 base, and a one-block build run. + +Env-gated: every test skips unless the snapshot directories are named by ``MLX_DFLOAT_QWEN21_DF11`` (the +``mingyi456/Qwen-Image-2.1-DF11-ComfyUI`` snapshot holding the single config-less file) and, for the base comparisons, +``MLX_DFLOAT_QWEN21_BASE`` (the ``Qwen/Qwen-Image-2.1`` snapshot with its ``transformer/`` index and shards). Expected +counts are literals from the file's own header (33 groups: ``modulation.1`` and ``transformer_blocks.0..31``; 72 BF16 +extras) and from mflux 0.20.0 (``Qwen21Transformer`` builds 32 blocks, qwen21_transformer.py:17-28); the expected bits +come from the base's tensors, never from the code under test. +""" + +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.integrate.coverage import load_resident_set +from mlx_dfloat.integrate.providers import DF11Provider, StreamingBF16Provider, read_bf16 +from mlx_dfloat.integrate.resident import decode_nonblock, install_nonblock +from mlx_dfloat.mflux.qwen21.names import KIND, check_qwen21_groups + +DF11_ENV, BASE_ENV = "MLX_DFLOAT_QWEN21_DF11", "MLX_DFLOAT_QWEN21_BASE" +MODULATION = "modulation.1" +# Every block's expanded split positions: the stored (16777216, 33554432, 50331648, 67108864, 167772160) with the +# gate/up seam inserted at 67108864 + 12288 x 4096 = 117440512 (the fused img_mlp.gate_up cut into its halves). +BLOCK_SPLITS = (16777216, 33554432, 50331648, 67108864, 117440512, 167772160) +SEVEN = ( + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "img_mlp.gate_layer", + "img_mlp.proj", + "img_mlp.out", +) + + +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)) + + +def _base_index(base: Path) -> dict[str, Any]: + from mlx_dfloat.mflux.qwen21.transformer import base_transformer_files_index + + return base_transformer_files_index(base / "transformer") + + +@pytest.mark.slow +def test_the_real_file_opens_through_its_layout() -> None: + # Bug caught: the published file not matched by its pinned layout (a header or spot-check drift: the file would + # be refused, or worse, read with another order), a group count other than the header's 33, a block holding + # other than the seven expanded matrices (the fused gate_up left whole), or an extra count other than 72 (two + # RMSNorm weights per block plus eight non-block tensors; a stray tensor the coverage check would refuse). + ckpt = open_checkpoint(_env(DF11_ENV)) + assert check_qwen21_groups(ckpt) == {KIND: 32} + assert len(ckpt.groups) == 33 + assert MODULATION in ckpt.groups + assert tuple(ckpt.groups["transformer_blocks.31"].matrix_names) == tuple( + f"transformer_blocks.31.{sub}.weight" for sub in SEVEN + ) + assert len(ckpt.extras) == 72 + assert sorted({info.dtype for _path, info in ckpt.extras.values()}) == ["BF16"] + assert ckpt.config_source.startswith( + "header and spot checks match layout qwen-image-2.1-comfyui " + "(mingyi456/Qwen-Image-2.1-DF11-ComfyUI@1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc)" + ) + + +@pytest.mark.slow +@pytest.mark.metal +def test_four_real_groups_decode_on_metal_bit_exact_against_the_base() -> None: + # Bug caught: the exponent bits of this file differing from the base (only sign/mantissa bytes were compared + # before), the gate_up seam cut at the wrong row or its halves swapped (img_mlp.gate_layer and img_mlp.proj are + # each compared with the base tensor of their own name), a stored matrix order off in a middle or the last block, + # or modulation.1 (16384 x 4096, a non-block group) decoded to the wrong layout. + ckpt = open_checkpoint(_env(DF11_ENV)) + index = _base_index(_env(BASE_ENV)) + for name in ( + "transformer_blocks.0", + "transformer_blocks.15", + "transformer_blocks.31", + MODULATION, + ): + group = ckpt.groups[name] + loaded = load_group_mx(group) + if name != MODULATION: + assert loaded.split_positions == BLOCK_SPLITS, name + result = decode_group(loaded, backend="metal") + check(result, name=name) + parts = split_matrices(result.bits, loaded.split_positions) + assert len(parts) == len(group.matrix_names), name + for matrix, part in zip(group.matrix_names, parts, strict=True): + path, info = index[matrix] + want = np.array(read_bf16(path, info).view(mx.uint16)).ravel() + assert np.array_equal(np.array(part.view(mx.uint16)).ravel(), want), matrix + del loaded, result, parts + gc.collect() + mx.clear_cache() + + +def _config(size: int) -> Any: + """mflux's ``Config`` for Qwen-Image 2.1 at ``size`` pixels square (the scheduler's sigma shift reads it).""" + from mflux.models.common.config.config import Config + from mflux.models.common.config.model_config import ModelConfig + + return Config( + model_config=ModelConfig.qwen_image_21(), + num_inference_steps=40, + height=size, + width=size, + guidance=1.0, + scheduler="linear", + ) + + +def _inputs(*, size: int, text_tokens: int, context_dim: int, seed: int = 1) -> dict[str, Any]: + """One call's keyword inputs: (size/16)² image tokens of 64 channels, ``text_tokens`` of ``context_dim``, t=0.""" + tokens = (size // 16) ** 2 + keys = mx.random.split(mx.random.key(seed), 2) + return { + "t": 0, + "config": _config(size), + "hidden_states": mx.random.normal((1, tokens, 64), key=keys[0]).astype(mx.bfloat16), + "encoder_hidden_states": mx.random.normal( + (1, text_tokens, context_dim), key=keys[1] + ).astype(mx.bfloat16), + "encoder_hidden_states_mask": mx.ones((1, text_tokens), dtype=mx.int32), + } + + +def _run(transformer: Any, inputs: Mapping[str, Any]) -> mx.array: + from mlx_dfloat.mflux.qwen21.transformer import is_eager + + out = transformer(**inputs) + mx.eval(out) + assert is_eager(transformer) # the uncompiled forward ran, as on the model's path + return out + + +def df11_one_block( + ckpt: DF11Checkpoint, inputs: Mapping[str, Any], **build_kwargs: Any +) -> mx.array: + """Build one block from the checkpoint, decode and install ``modulation.1``, attach the DF11 decode, run once.""" + from mlx_dfloat.mflux.qwen21.names import qwen21_name_map + from mlx_dfloat.mflux.qwen21.transformer import build_transformer + + names = qwen21_name_map() + build = build_transformer(ckpt, n_layers=1, **build_kwargs) + blocks = list(build.shapes) + assert blocks == ["transformer_blocks.0"] + groups = load_resident_set(ckpt, [*blocks, MODULATION]) + matrix_names = {n: ckpt.groups[n].matrix_names for n in groups} + nonblock = decode_nonblock( + {MODULATION: groups[MODULATION]}, + {MODULATION: matrix_names[MODULATION]}, + build.nonblock, + names, + ) + assert sorted(nonblock) == ["modulation.layers.1.weight"] + install_nonblock(build.transformer, nonblock) + 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, + index: Mapping[str, Any], + inputs: Mapping[str, Any], + **build_kwargs: Any, +) -> mx.array: + """The same one-block build with the base's BF16 extras, ``modulation.1`` a plain weight, blocks streamed.""" + from mlx_dfloat.mflux.qwen21.names import qwen21_name_map + from mlx_dfloat.mflux.qwen21.transformer import base_extras, build_transformer + + names = qwen21_name_map() + build = build_transformer( + ckpt, n_layers=1, extras=base_extras(index, ckpt), nonblock_from_extras=True, **build_kwargs + ) + assert build.nonblock == {} + 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 real extra name the coverage or the renames do not handle (the build refuses it on the real + # files; a rename applied the same way to both builds is not caught here, the two sides share it), modulation.1 + # decoded to a wrong layout or left on its placeholder (a zero-size matmul before the first block), the forward + # pass compiled on one side, or any block-seam difference between the DF11 and the streamed BF16 paths. 1024²: + # 4096 image tokens, 32 text tokens of mflux's 4096-wide context (qwen21_transformer.py:17-28). + ckpt = open_checkpoint(_env(DF11_ENV)) + index = _base_index(_env(BASE_ENV)) + inputs = _inputs(size=1024, text_tokens=32, context_dim=4096) + out_a = df11_one_block(ckpt, inputs) + out_b = bf16_one_block(ckpt, index, inputs) + assert out_a.shape == out_b.shape == (1, 4096, 64) + assert out_a.dtype == out_b.dtype + assert np.array_equal(_bits(out_a), _bits(out_b)) diff --git a/tests/test_qwen21_tiny.py b/tests/test_qwen21_tiny.py new file mode 100644 index 0000000..0c97856 --- /dev/null +++ b/tests/test_qwen21_tiny.py @@ -0,0 +1,69 @@ +"""One real tiny Qwen-Image 2.1 generation through mflux's own loop: our prelude, the uncompiled forward pass, the +shared 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._qwen21_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( + ("guidance", "negative", "calls"), + [(4.0, "n", 2), (1.0, "n", 1), (4.0, "", 1)], + ids=["cfg", "guidance-1", "empty-negative"], +) +def test_one_tiny_generation_decodes_once_per_block_per_transformer_call( + tmp_path, monkeypatch, guidance, negative, calls +): + # Bug caught: anything between our prelude and mflux's loop not composing (the eager forward pass, the pairs in + # mflux's prompt cache, the VAE guard), or the negative branch run at guidance 1.0 or with an empty negative + # (mflux's rule: guidance > 1.0 and bool(negative_prompt), qwen_image_21.py:89). Tiny transformer: 2 blocks; 2 + # steps: launches == 2 steps x calls x 2 blocks == 8 with CFG, 4 without. + from mflux.utils.apple_silicon import AppleSiliconUtil + + # mflux compiles this transformer's forward on every chip; a Max chip shows no chip check is relied on. + monkeypatch.setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: False)) + model = fake_model(tmp_path, monkeypatch) + 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.astype(mx.float32)) + + 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, + negative_prompt=negative, + ) + assert np.asarray(image.image).shape == (64, 64, 3) + # The denoised latents: 4 x 4 = 16 tokens of 64 channels at 64^2 (qwen21_latent_creator.py:11-15), finite. + assert seen["latents"].shape == (1, 16, 64) + 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" + # The tiny block group (4 x 32 x 32 + 2 x 96 x 32 + 32 x 96 = 13_312 elements, 26_624 B) next to the allowance's + # floor (500_000_000 B: 16 image tokens + 5 text tokens is far below the reference). + assert report["cache_limit_in_force"] == 26_624 + 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_qwen21_transformer.py b/tests/test_qwen21_transformer.py new file mode 100644 index 0000000..13aad3b --- /dev/null +++ b/tests/test_qwen21_transformer.py @@ -0,0 +1,299 @@ +"""The Qwen-Image 2.1 transformer: the class-swap seam and the eager-forward guard, on a real tiny mflux transformer.""" + +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_unflatten +from tests._qwen21_tiny import TINY, tiny_groups, tiny_inputs, tiny_seamed + +from mlx_dfloat.decode import decode_group +from mlx_dfloat.errors import DFloatDependencyError, DFloatIntegrationError +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.qwen21.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 = ["transformer_blocks.0", "transformer_blocks.1"] + + +def test_block_lists_is_the_one_list_the_forward_runs(): + # Bug caught: another attribute than the list _forward iterates (qwen21_transformer.py:96-100), so no block + # would run through the seam. + tf = SimpleNamespace(transformer_blocks=[1, 2]) + assert block_lists(tf) == [("transformer_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). The class factory is cached, so the cache is emptied around the call. + real = importlib.util.find_spec + monkeypatch.setattr( + importlib.util, + "find_spec", + lambda name, *a, **k: None if name == "mflux" else real(name, *a, **k), + ) + seam_transformer_class.cache_clear() + try: + with pytest.raises(DFloatDependencyError, match=r"mlx-dfloat\[mflux\]"): + seam_transformer_class() + finally: + seam_transformer_class.cache_clear() + + +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 _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 _counting_provider(shapes, calls): + from mlx_dfloat.mflux.qwen21.names import qwen21_name_map + + groups, names, _src = tiny_groups(shapes, np.random.default_rng(9)) + + def count(group): + calls.append(group.name) + return decode_group(group, backend="reference") + + return DF11Provider(groups, names, qwen21_name_map(), decode=count) + + +@pytest.mark.mflux +def test_the_stock_transformer_compiles_its_forward_on_the_first_call(): + # Bug caught (the control): mflux no longer compiling _forward in __call__ (qwen21_transformer.py:69-72), which + # would leave the guard and its note unjustified. mflux compiles there on every chip, no chip check. + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + tf = Qwen21Transformer(**TINY) + assert tf._step_fn is None + mx.eval(tf(**tiny_inputs())) + assert type(tf._step_fn) is COMPILED + + +@pytest.mark.mflux +def test_the_seamed_transformer_runs_the_plain_forward(): + # Bug caught: mflux's compile reached through the seamed class (the per-block eval would then raise inside it). + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + from mlx_dfloat.mflux.qwen21.transformer import is_eager + + tf, shapes = tiny_seamed() + tf.attach(ResidentProvider(_constant_weights(shapes)), shapes) + mx.eval(tf(**tiny_inputs())) + assert type(tf._step_fn) is types.MethodType + assert tf._step_fn.__func__ is Qwen21Transformer._forward + assert tf._step_fn.__self__ is tf + assert is_eager(tf) + + +@pytest.mark.mflux +@pytest.mark.parametrize("step_fn", ["compiled", "another transformer's forward"]) +def test_a_step_fn_other_than_this_transformers_plain_forward_is_refused_before_any_block(step_fn): + # Bug caught: the guard checking only for None (a compiled _step_fn left from mflux, or set by a later mflux + # elsewhere, runs the blocks inside mx.compile), or accepting any bound _forward (another transformer's weights). + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + from mlx_dfloat.mflux.qwen21.transformer import is_eager + + tf, shapes = tiny_seamed() + calls = [] + tf.attach(_counting_provider(shapes, calls), shapes, verify_in_call=True) + tf._step_fn = ( + mx.compile(tf._forward) if step_fn == "compiled" else Qwen21Transformer(**TINY)._forward + ) + assert not is_eager(tf) + with pytest.raises(DFloatIntegrationError, match="compiled the forward pass"): + tf(**tiny_inputs()) + assert calls == [] + + +@pytest.mark.mflux +def test_the_compiled_forward_cannot_run_the_seam(): + # Bug caught (the reason for the guard, as a control): if mflux's compiled _forward could run per-block evals, + # the guard and its uncompiled cost would be unjustified. The arguments are built as __call__ builds them + # (qwen21_transformer.py:61-68). + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + tf, shapes = tiny_seamed() + zeros = { + b: {a: mx.zeros(s, dtype=mx.bfloat16) for a, s in per.items()} for b, per in shapes.items() + } + tf.attach(ResidentProvider(zeros), shapes) + inputs = tiny_inputs() + timestep = Qwen21Transformer._compute_timestep(inputs["t"], inputs["config"]) + rows = mx.concatenate([timestep, mx.zeros((1,), dtype=timestep.dtype)]) + cos, sin, mask = tf._geometry( + text_len=8, + latent_height=4, + latent_width=4, + encoder_hidden_states_mask=inputs["encoder_hidden_states_mask"], + ) + compiled = mx.compile(tf._forward) + with pytest.raises(ValueError, match=r"\[eval\] Attempting to eval"): + mx.eval( + compiled(inputs["hidden_states"], inputs["encoder_hidden_states"], rows, cos, sin, mask) + ) + + +@pytest.mark.mflux +def test_the_swap_keeps_parameter_paths_and_the_computed_tables(): + # Bug caught: the class swap breaking on Qwen21TransformerBlock, parameters moving under another path + # (placeholders, load_weights and coverage would miss them), or a computed table renamed upstream (the extras + # coverage exemption would go stale silently). + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer_block import ( + Qwen21TransformerBlock, + ) + + from mlx_dfloat.integrate.blockseam import seam_blocks + from mlx_dfloat.mflux.qwen21.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) == ( + "transformer_blocks.0", + "transformer_blocks.1", + ) + after = sorted(k for k, _v in tree_flatten(tf.parameters())) + assert after == before + assert isinstance(tf.transformer_blocks[1], Qwen21TransformerBlock) + # The seven names mflux 0.20.0 computes at construction (qwen21_rope.py:21-26, qwen21_time_text_embed.py:16). + assert set(COMPUTED_PARAMS) == { + "pos_embed.cos_tables.0", + "pos_embed.cos_tables.1", + "pos_embed.cos_tables.2", + "pos_embed.sin_tables.0", + "pos_embed.sin_tables.1", + "pos_embed.sin_tables.2", + "time_text_embed.time_proj.freqs", + } + assert set(COMPUTED_PARAMS) <= set(after) + + +def _eager_reference(tf, shapes, weights): + """Stock mflux holding tf's non-block parameters and ``weights`` on its blocks, running its plain forward. + + The seamed side runs uncompiled, so the reference runs mflux's own ``_forward`` uncompiled too. + """ + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + plain = Qwen21Transformer(**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.transformer_blocks[idx], attr).weight = w + plain._step_fn = plain._forward + return plain + + +@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, or a matrix landing on another attribute of its own block + # (to_q and to_k swapped: same shape, so only per-matrix values show it), either way 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 = _random_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(**inputs) + mx.eval(out) + tf.verify_step() + want = _eager_reference(tf, shapes, weights)(**inputs) + mx.eval(want) + + assert len(evals) == 2 + # Each block returns one array: the joint [8 text | 16 image] sequence at inner dim 32. + assert [tuple(e.shape) for e in evals] == [(1, 24, 32), (1, 24, 32)] + assert out.shape == (1, 16, 64) # the image tokens, projected back to 64 channels + assert bool(mx.isfinite(out).all().item()) + assert out.dtype == want.dtype + 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.transformer_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(**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 +def test_two_calls_per_step_decode_once_per_block_per_call(): + # Bug caught: the second (negative-prompt) call of a CFG step (qwen_image_21.py:104-118) refused for the first + # call's unchecked status words, or a block decoded twice in one call. 2 calls x 2 blocks = 4 decodes. + tf, shapes = tiny_seamed() + calls = [] + tf.attach(_counting_provider(shapes, calls), shapes, verify_in_call=True) + positive = tiny_inputs(seed=0) + negative = {**positive, "encoder_hidden_states": tiny_inputs(seed=1)["encoder_hidden_states"]} + pos = tf(**positive) + neg = tf(**negative) + mx.eval(pos, neg) + assert calls == BLOCKS * 2 + assert pos.shape == neg.shape == (1, 16, 64) + assert not mx.array_equal(pos, neg).item() diff --git a/tests/test_verify_checkpoint.py b/tests/test_verify_checkpoint.py index e80e650..c82631e 100644 --- a/tests/test_verify_checkpoint.py +++ b/tests/test_verify_checkpoint.py @@ -1,4 +1,3 @@ -import hashlib import json import threading @@ -7,11 +6,9 @@ import pytest import scripts.verify_checkpoint as vc from scripts.verify_checkpoint import main, run_key, verify -from tests._df11_fixtures import random_bf16, write_bf16_original, write_checkpoint +from tests._df11_fixtures import pin_layout, random_bf16, write_bf16_original, write_checkpoint from mlx_dfloat import _layouts -from mlx_dfloat._layouts import ContentProbe, SynthesizedLayout -from mlx_dfloat._safetensors import parse_header, read_header_bytes KEY = { "df11_revision": "a" * 40, @@ -827,35 +824,12 @@ def _config_less_fused(tmp_path, monkeypatch): extras={"norm.weight": NORM}, write_config=False, ) - file = df11 / "model.safetensors" - raw, start, size = read_header_bytes(file) - sm = parse_header(raw, data_start=start, file_size=size, source="t")["blocks.0.sign_mantissa"] - head = file.read_bytes()[sm.offset : sm.offset + 8] - layout = SynthesizedLayout( - key="tiny", - label="tiny", - repo_id="t/t", - revision="0" * 40, - file_name="model.safetensors", - header_sha256=hashlib.sha256(raw).hexdigest(), - file_sha256="0" * 64, # a plain file name here, so the whole-file pin is not consulted + layout = pin_layout( + df11 / "model.safetensors", + pattern_dict={r"blocks\.\d+": ("q", "gate_up")}, + row_splits={"gate_up": ("gate", "up")}, groups=1, extras=1, - raw_config={ - "version": "0.5.0", - "threads_per_block": [512], - "bytes_per_thread": 8, - "pattern_dict": {r"blocks\.\d+": ["q", "gate_up"]}, - }, - row_splits={"gate_up": ("gate", "up")}, - probes=( - ContentProbe( - tensor="blocks.0.sign_mantissa", - offset=0, - length=8, - sha256=hashlib.sha256(head).hexdigest(), - ), - ), ) monkeypatch.setattr(_layouts, "KNOWN_LAYOUTS", (layout,)) bf16 = tmp_path / "bf16" diff --git a/tests/test_verify_image.py b/tests/test_verify_image.py index 30357b4..bfb1fa3 100644 --- a/tests/test_verify_image.py +++ b/tests/test_verify_image.py @@ -564,3 +564,189 @@ def test_a_negative_prompt_is_refused_for_klein(tmp_path, capsys): ) assert info.value.code == 2 assert "--negative-prompt" in capsys.readouterr().err + + +# --- Qwen-Image 2.1 ------------------------------------------------------------------------------------------- + +QWEN = "qwen-image-2.1" + + +def _qwen_base( + tmp_path, + shards=( + "diffusion_pytorch_model-00001-of-00002.safetensors", + "diffusion_pytorch_model-00002-of-00002.safetensors", + ), + *, + present=None, +): + """A base snapshot's transformer/ as diffusers publishes it: an index naming its shards, the shards present.""" + root = tmp_path / "qbase" + transformer = root / "transformer" + transformer.mkdir(parents=True, exist_ok=True) + weight_map = {f"transformer_blocks.{i}.attn.to_q.weight": s for i, s in enumerate(shards)} + (transformer / "diffusion_pytorch_model.safetensors.index.json").write_text( + json.dumps({"metadata": {"total_size": 1}, "weight_map": weight_map}) + ) + for s in shards if present is None else present: + (transformer / s).write_bytes(b"\0") + return root + + +def _parse_qwen(tmp_path, base, *extra, side=("--side", "df11")): + return vi.parse_args( + [ + *side, + "--model", + QWEN, + "--df11", + str(tmp_path), + "--base", + str(base), + "--out", + str(tmp_path / "out"), + *extra, + ] + ) + + +def test_runner_for_dispatches_qwen_to_its_sides(): + # Bug caught: Qwen-Image 2.1 routed to the FLUX.1 sides (a wrong-model identity verdict), or no qwen21 row (a + # KeyError after the df11 side already ran for minutes). + assert vi.runner_for(QWEN, "df11") is vi.run_df11_qwen21 + assert vi.runner_for(QWEN, "bf16") is vi.run_bf16_qwen21 + + +def test_qwen_base_problem_is_none_for_an_index_with_its_shards(tmp_path): + # Bug caught: a complete transformer/ refused (the identity check could never run). + assert vi.qwen21_base_problem(QWEN, _qwen_base(tmp_path)) is None + + +@pytest.mark.parametrize("layout", ["no_dir", "empty_dir", "no_index"]) +def test_qwen_base_problem_names_the_missing_transformer(tmp_path, layout): + # Bug caught: the bf16 side launched on a base without its BF16 transformer (the encoder and VAE are enough for + # generate, and the transformer is the last, 14 GB download), failing after the df11 side's minutes. + base = tmp_path / "qbase" + base.mkdir() + if layout != "no_dir": + (base / "transformer").mkdir() + if layout == "no_index": + (base / "transformer" / "diffusion_pytorch_model-00001-of-00002.safetensors").write_bytes( + b"\0" + ) + problem = vi.qwen21_base_problem(QWEN, base) + assert problem is not None + assert f"--base {base}: the bf16 side streams the BF16 transformer from transformer/" in problem + assert "download Qwen/Qwen-Image-2.1's transformer/* first" in problem + + +def test_qwen_base_problem_names_a_shard_the_index_lists_but_the_snapshot_lacks(tmp_path): + # Bug caught: only the index checked, so a download stopped between the two shards passes the check and the bf16 + # side fails at the first block of the second shard. + second = "diffusion_pytorch_model-00002-of-00002.safetensors" + base = _qwen_base(tmp_path, present=("diffusion_pytorch_model-00001-of-00002.safetensors",)) + problem = vi.qwen21_base_problem(QWEN, base) + assert problem is not None + assert second in problem + + +def test_qwen_base_problem_reports_an_unreadable_index_instead_of_raising(tmp_path): + # Bug caught: a truncated index JSON raising a traceback out of parse_args instead of a usage message. + base = _qwen_base(tmp_path) + (base / "transformer" / "diffusion_pytorch_model.safetensors.index.json").write_text("{") + problem = vi.qwen21_base_problem(QWEN, base) + assert problem is not None + assert "diffusion_pytorch_model.safetensors.index.json" in problem + + +def test_qwen_base_problem_refuses_an_index_that_names_no_shards(tmp_path): + # Bug caught: an empty weight map passing as "every named shard present" (nothing to stream; the bf16 side fails + # at the first block). + base = _qwen_base(tmp_path) + (base / "transformer" / "diffusion_pytorch_model.safetensors.index.json").write_text( + json.dumps({"weight_map": {}}) + ) + assert vi.qwen21_base_problem(QWEN, base) is not None + + +@pytest.mark.parametrize("shard", ["../outside.safetensors", "sub/inner.safetensors"]) +def test_qwen_base_problem_refuses_a_shard_name_that_is_not_a_plain_file_name(tmp_path, shard): + # Bug caught: an index naming a path outside transformer/ (or below it) accepted because the file happens to exist + # there, so the bf16 side streams another file's tensors; the streaming reader's index check refuses the same names + # (Path(name).name must equal name). + base = _qwen_base(tmp_path) + target = base / "transformer" / shard + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(b"\0") + index = base / "transformer" / "diffusion_pytorch_model.safetensors.index.json" + index.write_text(json.dumps({"weight_map": {"a.weight": shard}})) + problem = vi.qwen21_base_problem(QWEN, base) + assert problem is not None + assert f"not a plain file name: {shard!r}" in problem + + +def test_qwen_base_problem_is_none_for_the_other_families(tmp_path): + # Bug caught: the transformer/ check applied to FLUX.1 or Z-Image bases (their own checks differ). + assert vi.qwen21_base_problem("schnell", tmp_path) is None + assert vi.qwen21_base_problem("z-image", tmp_path) is None + + +@pytest.mark.parametrize("side", [("--side", "bf16"), ("--orchestrate",)]) +def test_a_qwen_run_that_streams_the_bf16_side_is_refused_without_its_transformer( + tmp_path, capsys, side +): + # Bug caught: the check not wired into parse_args for the runs that stream the BF16 transformer (the bf16 side, + # and --orchestrate, whose bf16 child starts after the df11 side's minutes). + base = tmp_path / "qbase" + (base / "transformer").mkdir(parents=True) + with pytest.raises(SystemExit) as info: + _parse_qwen(tmp_path, base, side=side) + assert info.value.code == 2 + assert "transformer/" in capsys.readouterr().err + + +@pytest.mark.parametrize("side", [("--side", "df11"), ("--side", "compare")]) +def test_a_qwen_df11_side_or_compare_needs_no_bf16_transformer(tmp_path, side): + # Bug caught: the df11 side refused before the BF16 transformer is downloaded (it needs only the encoder, the VAE + # and the DF11 file; the identity's df11 side runs first, before the 14 GB transformer download). + base = tmp_path / "qbase" + base.mkdir() + assert _parse_qwen(tmp_path, base, side=side).model == QWEN + + +def test_qwen_negative_prompt_is_accepted_and_keyed_with_mfluxs_guidance(tmp_path, monkeypatch): + # Bug caught: --negative-prompt refused for Qwen as for FLUX.2 Klein (Qwen's CFG needs it), dropped from the key + # (a stored side reused after it changed) or from the children; or the guidance left None instead of mflux's 1.0 + # (qwen21_generate.py:60), which the bf16 side's Config cannot take. + monkeypatch.setattr(vi, "source_hash", lambda: "src") + args = _parse_qwen( + tmp_path, _qwen_base(tmp_path), "--negative-prompt", " ", side=("--orchestrate",) + ) + assert vi.run_key(args)["negative_prompt"] == " " + assert args.guidance == 1.0 + command = vi.child_command(args, "bf16") + assert command[command.index("--negative-prompt") + 1] == " " + assert command[command.index("--guidance") + 1] == "1.0" + + +def test_qwen_embeds_hold_the_positive_pair_and_the_negative_one_when_cfg_ran(): + # Bug caught: the negative pair saved under the positive names (the bf16 side would run CFG against the prompt + # itself), a mask not saved (the bf16 side's attention would see padding), or a negative saved when CFG did not + # run (the bf16 side would then run two calls per step while the df11 side ran one). + pos = (mx.array([[[1.0]]]), mx.array([[1]], dtype=mx.int32)) + neg = (mx.array([[[2.0]]]), mx.array([[0]], dtype=mx.int32)) + cache = {"p": pos, " ": neg} + one = vi.qwen21_embeds(cache, ("p",)) + assert sorted(one) == ["prompt_embeds", "prompt_mask"] + assert one["prompt_embeds"] is pos[0] + assert one["prompt_mask"] is pos[1] + two = vi.qwen21_embeds(cache, ("p", " ")) + assert sorted(two) == [ + "negative_prompt_embeds", + "negative_prompt_mask", + "prompt_embeds", + "prompt_mask", + ] + assert two["negative_prompt_embeds"] is neg[0] + assert two["negative_prompt_mask"] is neg[1] + assert two["prompt_embeds"] is pos[0] diff --git a/tests/test_verify_remote.py b/tests/test_verify_remote.py index d93de7b..3c9b5b6 100644 --- a/tests/test_verify_remote.py +++ b/tests/test_verify_remote.py @@ -257,6 +257,31 @@ def test_hf_range_source_reads_a_seeked_range(): assert 0 < length < src.size(shard) +def test_hf_range_source_reads_with_huggingface_hubs_seekable_default_only(monkeypatch): + # Bug caught: a read-ahead knob reaching huggingface_hub (block_size=0 gives a streaming file whose seek() raises, + # so every read past offset 0 fails) when no caller needs one. The fake file system mirrors that: a 0 block size + # opens a stream that cannot seek; the default opens a seekable file. + import io + + from scripts import verify_remote_group + + data = bytes(range(32)) + + class Stream(io.BytesIO): + def seek(self, *args): + raise ValueError("Cannot seek streaming HF file") + + class FakeFs: + def open(self, path, mode, **kwargs): + return Stream(data) if kwargs.get("block_size") == 0 else io.BytesIO(data) + + monkeypatch.setattr(verify_remote_group, "HfFileSystem", FakeFs) + with pytest.raises(TypeError): + HfRangeSource("owner/repo", "0" * 40, block_size=0) + src = HfRangeSource("owner/repo", "0" * 40) + assert src.read("model.safetensors", 5, 4) == bytes([5, 6, 7, 8]) + + def test_hf_range_source_path_guard_rejects_escapes_and_subdirs(): # Network-free: constructing HfRangeSource and calling its path guard never touches the Hub. src = HfRangeSource("owner/repo", "0" * 40) @@ -719,3 +744,187 @@ def test_main_refuses_the_cast_flag_with_structural_only(tmp_path, capsys): main(argv, source_factory=_local_factory) assert info.value.code == 2 assert "--cast-fp32-to-bf16" in capsys.readouterr().err + + +# --- config-less repositories, read through a pinned layout ---------------------------------------------------------- + +_FUSED_PATTERN = {r"blocks\.\d+": ("q", "gate_up")} +_FUSED_SPLITS = {"gate_up": ("gate", "up")} +_NORM = np.full((4,), 0x3F80, dtype=np.uint16) + + +def _config_less_fused(tmp_path, *, layout_groups=1, layout_extras=1): + """A config-less single file storing q (4x4) and concat(gate, up) (12x4), its originals, its pinned layout.""" + from tests._df11_fixtures import pin_layout + + rng = np.random.default_rng(17) + q, g, u = (random_bf16(rng, s) for s in [(4, 4), (6, 4), (6, 4)]) + df11 = write_checkpoint( + tmp_path / "df11", + groups={"blocks.0": [q, np.concatenate([g, u])]}, + pattern=r"blocks\.\d+", + sub_paths=("q", "gate_up"), + single_file=True, + extras={"norm.weight": _NORM}, + write_config=False, + ) + originals = {"blocks.0.q.weight": q, "blocks.0.gate.weight": g, "blocks.0.up.weight": u} + bf16 = write_bf16_original(tmp_path / "bf16", originals) + layout = pin_layout( + df11 / "model.safetensors", + pattern_dict=_FUSED_PATTERN, + row_splits=_FUSED_SPLITS, + groups=layout_groups, + extras=layout_extras, + ) + return df11, bf16, layout + + +def test_a_config_less_repository_is_verified_through_its_layout(tmp_path): + # Bug caught: the remote verifier stops at the missing config.json, or compares the stored gate_up whole (no + # original of that name: exit 2), or the halves against each other's originals (exit 1). + df11, bf16, layout = _config_less_fused(tmp_path) + out = tmp_path / "out" / "RESULT.json" + code = main(_argv(df11, bf16, out), source_factory=_local_factory, layouts=(layout,)) + result = json.loads(out.read_text()) + assert code == 0, result + assert result["config_source"] == "layout tiny" + got = [(m["name"], m["n"], m["equal"]) for m in result["groups"][0]["matrices"]] + assert got == [ # 4x4, then the 6x4 halves of the stored 12x4 matrix + ("blocks.0.q.weight", 16, True), + ("blocks.0.gate.weight", 24, True), + ("blocks.0.up.weight", 24, True), + ] + assert (result["compared"], result["mismatched"]) == (3, 0) + + +def test_a_config_json_repository_records_where_its_config_came_from(tmp_path): + # Bug caught: a repository with a config.json sent through the layout path (or the field left out, so a result + # cannot say which config its matrix names came from). + df11, bf16, _ = _pair(tmp_path, single_file=True) + out = tmp_path / "out" / "RESULT.json" + assert main(_argv(df11, bf16, out), source_factory=_local_factory, layouts=()) == 0 + assert json.loads(out.read_text())["config_source"] == "config.json" + + +def _flip_probed_byte(df11, layout): + probe = layout.probes[0] + file = df11 / "model.safetensors" + from mlx_dfloat._safetensors import read_header + + at = read_header(file)[probe.tensor].offset + probe.offset + data = bytearray(file.read_bytes()) + data[at] ^= 0x01 + file.write_bytes(bytes(data)) + + +@pytest.mark.parametrize( + ("case", "message"), + [ + # Bug caught: an unknown config-less file verified with a guessed matrix order. + ("no layout", "matches no known layout"), + # Bug caught: the spot checks skipped on the remote path, so a re-saved file with the same header and other + # bytes is read with the pinned order. + ("probe byte flipped", "spot check"), + # Bug caught: the layout's group and extra counts not checked, so a file holding more or fewer groups passes. + ("counts differ", "expects 2 groups and 1 extras"), + ], +) +def test_a_config_less_repository_is_refused_before_any_decode(tmp_path, capsys, case, message): + df11, bf16, layout = _config_less_fused( + tmp_path, layout_groups=2 if case == "counts differ" else 1 + ) + if case == "probe byte flipped": + _flip_probed_byte(df11, layout) + out = tmp_path / "out" / "RESULT.json" + layouts = () if case == "no layout" else (layout,) + assert main(_argv(df11, bf16, out), source_factory=_local_factory, layouts=layouts) == 2 + assert message in capsys.readouterr().err + assert not out.exists() + + +def test_a_diffusers_config_without_a_dfloat11_block_is_an_error_not_the_layout_path(tmp_path): + # Bug caught: a config.json without dfloat11_config (a diffusers transformer config) treated as "no config" and + # the file read through a layout that happens to match. + df11, bf16, layout = _config_less_fused(tmp_path) + (df11 / "config.json").write_text(json.dumps({"_class_name": "QwenImageTransformer2DModel"})) + out = tmp_path / "out" / "RESULT.json" + assert main(_argv(df11, bf16, out), source_factory=_local_factory, layouts=(layout,)) == 2 + assert not out.exists() + + +def test_a_config_less_repository_must_be_one_file(tmp_path): + # Bug caught: the first of several safetensors files identified and the rest ignored (their groups never seen). + from scripts.verify_remote_group import config_from_source + + df11, _bf16, layout = _config_less_fused(tmp_path) + (df11 / "other.safetensors").write_bytes((df11 / "model.safetensors").read_bytes()) + with pytest.raises(VerifyError, match="2 safetensors files"): + config_from_source(LocalRangeSource(df11), layouts=(layout,)) + + +def test_a_stored_count_that_disagrees_with_the_pattern_is_an_error_not_a_crash(tmp_path): + # Bug caught: a pattern listing three stored matrices for a group that stores two reaching the row-split cut and + # the zip (a confusing crash message, or the wrong segment cut) instead of an error record that says so. + from mlx_dfloat.format import parse_df11_config, with_row_splits + + df11, bf16, _layout = _config_less_fused(tmp_path) + raw = { + "version": "0.5.0", + "threads_per_block": [512], + "bytes_per_thread": 8, + "pattern_dict": {r"blocks\.\d+": ["q", "gate_up", "extra"]}, + } + config = with_row_splits(parse_df11_config(raw, source="t"), _FUSED_SPLITS, source="t") + d, b = LocalRangeSource(df11), LocalRangeSource(bf16) + record = verify_group( + d, b, "blocks.0", config=config, dindex=df11_index(d), bindex=bf16_index(b) + ) + assert record["status"] == "error" + assert "stores 2 matrices" in record["error"] + assert record["matrices"] == [] + + +@pytest.mark.network +def test_the_published_qwen_image_21_file_is_identified_by_range_reads(): + # Bug caught: the remote path reading the header or the spot checks at the wrong offsets (or not at all), so the + # one published config-less file is refused, or accepted without its content checked. Reads the header (28,408 + # bytes) and nine 4096-byte spot checks, each with huggingface_hub's default read-ahead. + from scripts.verify_remote_group import config_from_source + + from mlx_dfloat._layouts import KNOWN_LAYOUTS, QWEN_IMAGE_21_COMFYUI + + layout = QWEN_IMAGE_21_COMFYUI + src = HfRangeSource(layout.repo_id, layout.revision) + config, source = config_from_source(src, layouts=KNOWN_LAYOUTS) + assert source == "layout qwen-image-2.1-comfyui" + assert dict(config.row_splits) == {"img_mlp.gate_up": ("img_mlp.gate_layer", "img_mlp.proj")} + assert config.pattern_dict[r"transformer_blocks\.\d+"][4] == "img_mlp.gate_up" + + +@pytest.mark.parametrize( + ("probes", "message"), + [ + # Bug caught: a layout with no spot check accepted on its header alone (the content never checked). + ((), "pins no spot check"), + # Bug caught: a spot check of a tensor the file lacks skipped instead of refused. + ("missing", "which the file lacks"), + # Bug caught: a spot check past its tensor's end read from the next tensor's bytes (or short-read) silently. + ("outside", "lies outside its"), + ], +) +def test_a_layout_whose_spot_checks_cannot_run_is_refused(tmp_path, probes, message): + import dataclasses + + from scripts.verify_remote_group import config_from_source + + df11, _bf16, layout = _config_less_fused(tmp_path) + probe = layout.probes[0] + if probes == "missing": + probes = (dataclasses.replace(probe, tensor="blocks.9.sign_mantissa"),) + elif probes == "outside": + probes = (dataclasses.replace(probe, offset=10**6),) + with pytest.raises(VerifyError, match=message): + config_from_source( + LocalRangeSource(df11), layouts=(dataclasses.replace(layout, probes=probes),) + ) diff --git a/tests/test_zimage_memory.py b/tests/test_zimage_memory.py index 98047a1..c836d1b 100644 --- a/tests/test_zimage_memory.py +++ b/tests/test_zimage_memory.py @@ -154,8 +154,8 @@ def test_depth2_doubles_the_in_flight_decoded_buffer(): def test_fit_for_sums_the_phases_and_drops_the_set_from_the_vae_phase_when_asked(): - # Bug caught: vae_with_set=False keeping the compressed set / extras / non-block bytes in the VAE phase, or the - # constants not the measured ones. Hand sums (denoise at 1024², 512 tokens): + # Bug caught: vae_with_set=False keeping the compressed set / non-block bytes in the VAE phase or dropping the + # extras (they stay on the transformer when the set goes), or the constants not the measured ones. Hand sums (denoise at 1024², 512 tokens): # 8_000_000_000 + 6_000_000 + 20_000_000 + 360_000_000 + 2_300_000_000 + 1_523_200_242 + 510_030_936 # = 12_719_231_178 ; vae with set: 8_026_000_000 + 160_000_000 + 5_042_526_170 + 510_030_936 = 13_738_557_106. sizes = ZImageSizes( @@ -182,8 +182,8 @@ def test_fit_for_sums_the_phases_and_drops_the_set_from_the_vae_phase_when_asked assert kept.phases["encode"] == 8_000_000_000 + 1_500_000_000 + 510_030_936 assert kept.peak_phase == "vae" dropped = fit_for(vae_with_set=False, **common) - # 160_000_000 + 5_042_526_170 + 510_030_936 - assert dropped.phases["vae"] == 5_712_557_106 + # 6_000_000 + 160_000_000 + 5_042_526_170 + 510_030_936 + assert dropped.phases["vae"] == 5_718_557_106 assert dropped.peak_phase == "denoise" assert dropped.phases["denoise"] == kept.phases["denoise"] diff --git a/tests/test_zimage_model.py b/tests/test_zimage_model.py index 001a17e..dc62f3a 100644 --- a/tests/test_zimage_model.py +++ b/tests/test_zimage_model.py @@ -232,11 +232,13 @@ def _patch_upstream_generate(monkeypatch, model, *, raise_after_loop=None, raise """Replace ZImage.generate_image with a probe: records the cache limit in force, fires our after-loop subscriber the way mflux's GenerationContext would, returns a sentinel.""" from mflux.models.z_image.variants.z_image import ZImage + 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: @@ -731,3 +733,27 @@ def broken_attach(*args, **kwargs): assert model._provider is None assert model.transformer.cap_embedder[1].weight.size == 0 assert not model._lifecycle.set_resident + + +@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 was 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, model, 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 diff --git a/tests/test_zimage_names.py b/tests/test_zimage_names.py index a3bb9ab..b5bcf5d 100644 --- a/tests/test_zimage_names.py +++ b/tests/test_zimage_names.py @@ -56,6 +56,16 @@ def test_the_non_block_group_is_optional(tmp_path): assert check_zimage_groups(_ckpt(tmp_path, spec))["layers"] == 2 +def test_a_zero_padded_block_index_is_refused(tmp_path): + # Bug caught: layers.01 counted as block 1 (int("01") == 1), then the seam asks the provider for layers.1, a + # group the checkpoint does not have: a failure at the first block after the set load. + spec = {**GOOD} + del spec["layers.1"] + spec["layers.01"] = 8 + with pytest.raises(DFloatFormatError, match=r"layers\.01: not a Z-Image block group"): + check_zimage_groups(_ckpt(tmp_path, spec)) + + def test_a_hole_in_an_index_sequence_is_refused(tmp_path): # Bug caught: a missing block (layers.1) accepted as if the download were complete. spec = {**GOOD}