diff --git a/CHANGELOG.md b/CHANGELOG.md
index dccc21119..bfc03fbaa 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
- FLUX.3 klein rewrites each prompt with a chat model of your choice before
encoding it, set as the entry's `upsampler` or by `--upsampler`.
+- FLUX.2 klein, Krea 2, Z-Image and Qwen-Image 2.1 take an `upsampler` too,
+ and `upsample_prompt` sets your own system prompt for any family.
- `gmlx pull`, `image generate`, `serve` and the ComfyUI pack run FLUX.3 klein
and FLUX.3 klein Base GGUFs, with edits from up to 10 reference images.
- `gmlx pull`, `image generate`, `serve` and the ComfyUI pack run FLUX.2 klein
@@ -40,9 +42,13 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
prefills faster. `GMLX_ARENA_SPLIT_MAX_SHARE` tunes the choice.
- `gmlx pull` of an image model reuses a VAE GGUF on disk packed for its
family from another source, and names it, where it downloaded the VAE again.
+- `presence_penalty` and `frequency_penalty` count every generated token, as
+ vLLM does, where they counted the last 20 tokens, the prompt's included.
### Fixed
+- `gmlx distill gen` stopped with Ctrl-C could leave out a reply that had
+ already arrived.
- A chat template with a `{% generation %}` block was refused by
`--chat-template` as invalid Jinja.
- `--thinking off` was silently ignored on a model with no thinking switch
diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md
index dc0b1d39a..5bcb89515 100644
--- a/THIRD_PARTY_NOTICES.md
+++ b/THIRD_PARTY_NOTICES.md
@@ -22,7 +22,9 @@ in [`licenses/`](licenses/).
| [mflux](https://github.com/filipstrand/mflux) | MIT | The Z-Image transformer at commit `e066d98`, ported to `gmlx/image/families/z_image/model.py` with the fused qkv kept and modules named after the GGUF tensors, and its attention, timestep and RoPE helpers in `gmlx/image/ops/attention.py` and `gmlx/image/ops/embeds.py`. |
| [diffusers](https://github.com/huggingface/diffusers) | Apache-2.0 | Code at commit `e0abab8`, ported to MLX and modified for gmlx: the Z-Image forward of `transformer_z_image.py` in `gmlx/image/families/z_image/model.py`, its timestep embedding, RoPE tables and rotary product in `gmlx/image/ops/embeds.py` and `gmlx/image/ops/attention.py`, the RMSNorm rounding of `normalization.py` in `gmlx/image/ops/norms.py`, and the FLUX.1 autoencoder of `autoencoders/vae.py`, `resnet.py`, `unet_2d_blocks.py`, `downsampling.py`, `upsampling.py` and `attention_processor.py` in `gmlx/image/vae/flux1_ae.py`, the quant convs and BatchNorm statistics of `autoencoders/autoencoder_kl_flux2.py` in `gmlx/image/vae/flux2_ae.py`, the Qwen-Image 2.1 autoencoder of `autoencoders/autoencoder_kl_qwenimage21.py`, for one frame, in `gmlx/image/vae/qwen_image21_vae.py`, and the UniPC bh2 sampler of `schedulers/scheduling_unipc_multistep.py` in `gmlx/image/unipc.py`. Also the Qwen-Image 2.1 forward, token layout and RoPE of `transformer_qwenimage21.py` in `gmlx/image/families/qwen_image21/model.py` and `layout.py`, and the dynamic shift of `pipelines/qwenimage21/pipeline_qwenimage21.py` and `schedulers/scheduling_flow_match_euler_discrete.py` in `gmlx/image/schedule.py`, and the reference image sizes, the prompt template with image slots and the image branch of the encoder tap of that pipeline in `gmlx/image/families/qwen_image21/edit.py`. For Krea 2, at the same commit: the forward of `transformer_krea2.py` in `gmlx/image/families/krea2/model.py`, the Qwen-Image autoencoder of `autoencoders/autoencoder_kl_qwenimage.py`, for one frame, in `gmlx/image/vae/qwen_image_vae.py`, the prompt template, encoder layers and constants of `pipelines/krea2/pipeline_krea2.py` and `modular_pipelines/krea2/` in `gmlx/image/families/krea2/family.py`, and the schedule of that pipeline in `krea2_sigmas` of `gmlx/image/schedule.py`. |
| [flux-3-image-klein](https://github.com/black-forest-labs/flux-3-image-klein) | Apache-2.0 | BFL's FLUX.3 klein reference at commit `45e57c9`, ported to MLX and modified for gmlx: the DiT constants and tensor layout of `src/flux3/dit.py`, the image RoPE scale and schedule alpha of `src/flux3/pipeline.py` and the encoder layers, pad multiple, length, pad token and position and chat template of `src/flux3/text_encoder.py` in `gmlx/image/families/flux3_klein/family.py`, the DiT forward of `src/flux3/dit.py` with the position ids of `pipeline.py`, `autoencoder.py` and `text_encoder.py` in `gmlx/image/families/flux3_klein/model.py`, the schedule of `src/flux3/sampling.py` in `flux3_sigmas` of `gmlx/image/schedule.py`, the latent patchify and BatchNorm of `src/flux3/autoencoder.py` in `gmlx/image/vae/flux2_ae.py`, and the reference canvas, reference preparation and default size of `src/flux3/pipeline.py` with the reference choice and output crop of `src/flux3/generate.py` in `gmlx/image/families/flux3_klein/edit.py`; and at commit `bfc8889`, the upsampler's system prompts `src/flux3/pup_prompts_external_vlm.json`, shipped unmodified in `gmlx/image/families/flux3_klein/`, with the chat turn and sampling of `src/flux3/pup.py` in `gmlx/image/families/flux3_klein/family.py`. |
-| [flux2](https://github.com/black-forest-labs/flux2) | Apache-2.0 | BFL's FLUX.2 reference at commit `50fe516`, ported to MLX and modified for gmlx: the DiT constants of `Klein4BParams`, `Klein9BParams` and `Flux2Params` and the tensor names and shapes of `Flux2` in `src/flux2/model.py`, the sampling defaults of `FLUX2_MODEL_INFO` in `src/flux2/util.py`, the prompt rules of `Qwen3Embedder` in `src/flux2/text_encoder.py`, the rounding of `denoise` and `denoise_cfg` and the reference image rules of `encode_image_refs` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein/family.py`, the constants of `Klein9BParams`, the 9B rows of `FLUX2_MODEL_INFO` and the Qwen3-8B encoder of `load_qwen3_embedder` in `src/flux2/text_encoder.py` in `gmlx/image/families/flux2_klein_9b/family.py`, the `flux.2-klein-9b-kv` row of `FLUX2_MODEL_INFO` in `gmlx/image/families/flux2_klein_9b_kv/family.py` and the step of `denoise_cached` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein_9b_kv/step.py`, the DiT forward of `src/flux2/model.py`, with `forward_kv_extract`, `forward_kv_cached` and `causal_attn_fn`, and the token ids of `prc_img`, `prc_txt` and `encode_image_refs` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein/model.py`, and the schedule of `compute_empirical_mu` and `get_schedule` in `src/flux2/sampling.py` in `flux2_mu` and `flux2_sigmas` of `gmlx/image/schedule.py`, and the reference image preparation of `encode_image_refs` and `default_prep` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein/edit.py`. |
+| [flux2](https://github.com/black-forest-labs/flux2) | Apache-2.0 | BFL's FLUX.2 reference at commit `50fe516`, ported to MLX and modified for gmlx: the DiT constants of `Klein4BParams`, `Klein9BParams` and `Flux2Params` and the tensor names and shapes of `Flux2` in `src/flux2/model.py`, the sampling defaults of `FLUX2_MODEL_INFO` in `src/flux2/util.py`, the prompt rules of `Qwen3Embedder` in `src/flux2/text_encoder.py`, the rounding of `denoise` and `denoise_cfg` and the reference image rules of `encode_image_refs` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein/family.py`, the constants of `Klein9BParams`, the 9B rows of `FLUX2_MODEL_INFO` and the Qwen3-8B encoder of `load_qwen3_embedder` in `src/flux2/text_encoder.py` in `gmlx/image/families/flux2_klein_9b/family.py`, the `flux.2-klein-9b-kv` row of `FLUX2_MODEL_INFO` in `gmlx/image/families/flux2_klein_9b_kv/family.py` and the step of `denoise_cached` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein_9b_kv/step.py`, the DiT forward of `src/flux2/model.py`, with `forward_kv_extract`, `forward_kv_cached` and `causal_attn_fn`, and the token ids of `prc_img`, `prc_txt` and `encode_image_refs` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein/model.py`, and the schedule of `compute_empirical_mu` and `get_schedule` in `src/flux2/sampling.py` in `flux2_mu` and `flux2_sigmas` of `gmlx/image/schedule.py`, and the reference image preparation of `encode_image_refs` and `default_prep` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein/edit.py`; and the upsampler's text to image system message `SYSTEM_MESSAGE_UPSAMPLING_T2I` of `src/flux2/system_messages.py`, shipped unmodified as `gmlx/image/families/flux2_klein/system_message_upsampling_t2i.txt`, with the chat turn of `scripts/cli.py` in `gmlx/image/families/flux2_klein/family.py`. |
+| [Z-Image-Turbo Space](https://huggingface.co/spaces/Tongyi-MAI/Z-Image-Turbo) | apache-2.0, from the Space's card; `pe.py` carries no license header | At revision `768cb50`, the prompt enhancer's system prompt `prompt_template` of `pe.py`, shipped as its string in `gmlx/image/families/z_image/pe_prompt_template.txt`, with its use in `app.py` in `gmlx/image/families/z_image/family.py`. |
+| [krea-2](https://github.com/krea-ai/krea-2) | Apache-2.0 | At commit `db3984f`, the prompt expansion system prompt `docs/expansion.txt`, shipped unmodified in `gmlx/image/families/krea2/`, with its use of `docs/prompting.md` in `gmlx/image/families/krea2/family.py`. The Krea 2 weights are under their own license and do not ship with gmlx. |
| [ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF) | Apache-2.0 | The model detection keys of `convert.py` at commit `6ea2651`, which `gmlx/image/fingerprint.py` uses to refuse other diffusion model layouts by name. |
| [misaki](https://github.com/hexgrad/misaki) | Apache-2.0 | Whole package vendored at `gmlx/_vendor/misaki` (commit `fba1236`): Kokoro's grapheme-to-phoneme front-end, used by the TTS service until upstream's next PyPI release. |
| [musl](https://musl.libc.org) | MIT | Statically linked into the guest entry binary `gmlx/container/guest/gmlx-entry` as its C library ([`licenses/musl-COPYRIGHT`](licenses/musl-COPYRIGHT)). |
diff --git a/docs/api.md b/docs/api.md
index f6d94437c..a2b75a7a1 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -135,7 +135,11 @@ and their `*_context_size` companions, `enable_thinking`,
`thinking_budget`, and the OpenAI `reasoning` and `reasoning_effort`
controls. The three chat routes ignore `n`, `user`, `parallel_tool_calls`
and `metadata`: they return one choice, and the template decides how many
-tool calls to make. The other fields differ by route:
+tool calls to make. `presence_penalty` and `frequency_penalty` count every
+token generated so far and never the prompt, as vLLM and SGLang do.
+`presence_context_size` or `frequency_context_size` set to N counts only the
+last N generated tokens, and 0 or no value counts all of them. The other
+fields differ by route:
| Parameter | `/v1/chat/completions` | `/v1/responses` | `/v1/messages` | Notes |
|-----------|------------------------|-----------------|----------------|-------|
diff --git a/docs/cli.md b/docs/cli.md
index caf618150..827eb9080 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -1068,8 +1068,9 @@ and `crop`, the left, top, right and bottom of the image in it.
| `--prompt TEXT` | Required, or `--prompt-file` | Generate from this prompt. |
| `--prompt-file FILE` | None | Read the prompt from this file. |
| `--negative-prompt TEXT` | Empty | Guide away from this prompt, or from the empty prompt without it. It is unused when `--guidance` is 1. |
-| `--upsampler GGUF` | The entry's `upsampler` | On FLUX.3 klein, rewrite the prompt with this chat model, a GGUF or a `models` id, as [Prompt upsampling](images.md#prompt-upsampling) describes. |
+| `--upsampler GGUF` | The entry's `upsampler` | Rewrite the prompt with this chat model, a GGUF or a `models` id, as [Prompt upsampling](images.md#prompt-upsampling) describes. |
| `--no-upsample` | Off | Run the prompt as typed, without the upsampler. |
+| `--upsample-prompt FILE` | The entry's `upsample_prompt` | The upsampler takes the system prompt in this text file in place of the family's. Qwen-Image 2.1 ships none and needs one. |
| `--size WxH` | The init image's, else a FLUX.2 klein or FLUX.3 klein reference's, else the entry's `defaults.size`, else the variant's | Set the width and height, each a multiple of 32 on Qwen-Image 2.1 and of 16 on the others. Another size exits 2 and names the multiple. |
| `--steps N` | The variant's | Sample for N steps. |
| `--seed N` | Random, and printed | Seed the first image. |
diff --git a/docs/config.md b/docs/config.md
index c88bbf22a..a1e2dc56b 100644
--- a/docs/config.md
+++ b/docs/config.md
@@ -271,7 +271,8 @@ entries.
| `family` | detected | The image family, such as `z-image`, `qwen-image-2.1` or `flux-3-klein`. A file of another family is refused at load. |
| `variant` | detected | The DiT's variant, such as `turbo`, for a file whose weights match no known variant. A Krea 2 file of unknown weights needs `variant: turbo` to load. |
| `adapters` | none | LoRA files applied to every request, each a path or `{path, strength}` with strength 1 by default |
-| `upsampler` | none | The `models` id of the chat model that rewrites each FLUX.3 klein prompt, as [Prompt upsampling](images.md#prompt-upsampling) describes |
+| `upsampler` | none | The `models` id of the chat model that rewrites each prompt, as [Prompt upsampling](images.md#prompt-upsampling) describes |
+| `upsample_prompt` | the family's | A text file with the system prompt for the `upsampler`, found in `model_dirs` when relative. Qwen-Image 2.1 needs one, which `gmlx pull qwen-image` fetches |
| `upsample` | `true` | Rewrite prompts with the `upsampler`. `false` runs them as typed, unless a request sends `prompt_upsampling: true`. |
| `defaults` | the family's | Request defaults: `size` as `WxH`, `steps`, `guidance`, `sampler` (`euler` or `uni_pc_bh2`) and `negative_prompt` |
| `encoder_release` | `auto` | Release the text encoder after each encode: `always`, `never`, or `auto`, which keeps it while the whole model fits in 80 percent of the working budget |
diff --git a/docs/image-api.md b/docs/image-api.md
index a016b1e2e..800f1558f 100644
--- a/docs/image-api.md
+++ b/docs/image-api.md
@@ -159,7 +159,8 @@ for event in stream:
| One of several FLUX.3 klein references is under 16 pixels on a side | 400 naming the image |
| A FLUX.2 klein reference is under 64 pixels on a side, or one side is more than 8 times the other | 400 naming the image |
| `prompt_upsampling: true` on an entry without an upsampler | 400 naming the `upsampler` setting |
-| The upsampler writes an empty caption | 500 `server_error` |
+| The upsampler writes an empty caption, or a reply its family's parser finds no caption in | 500 `server_error` |
+| The entry's `upsample_prompt` file is missing, or a Qwen-Image 2.1 entry with an upsampler has none | 500 `server_error` saying where to get the file |
| The upsampler's context has no room for its system prompt and caption | 500 `server_error` naming the upsampler |
| More requests waiting than [`image.max_queue`](config.md#image-requests) | 429 with `Retry-After` |
| The client disconnects | The job stops at its next step, and the log marks the request `cancelled` |
diff --git a/docs/images.md b/docs/images.md
index dca64a7e7..5d882d73c 100644
--- a/docs/images.md
+++ b/docs/images.md
@@ -166,17 +166,26 @@ refused by name, as is a LoRA of another model. ComfyUI takes LoRAs as
## Prompt upsampling
-FLUX.3 klein and FLUX.3 klein Base were trained on long, detailed
-captions, not on the short prompts people type. Prompt upsampling turns
-your prompt into such a caption before the text encoder reads it: a chat
-model rewrites it with the system prompt that Black Forest Labs publishes
-for this, and the image follows the caption. A FLUX.3 klein prompt sent as
-typed gives weaker images.
-
-Qwen3.5-9B writes captions close to those of BFL's own upsampler, and
-its Q6_K GGUF is a good choice. Any chat model of your configuration can
-write them. Add it under `models`, and name it as the `upsampler` of the
-image model's entry:
+Prompt upsampling turns your prompt into a long, detailed caption before
+the text encoder reads it: a chat model of your choice rewrites it with a
+system prompt written for that image model, and the image follows the
+caption. FLUX.3 klein and FLUX.3 klein Base were trained on such captions,
+not on the short prompts people type, and a prompt sent as typed gives
+weaker images. On the other families the rewrite is optional, and helps
+most with short prompts and with text in the image.
+
+| Family | System prompt |
+|---|---|
+| FLUX.3 klein, FLUX.3 klein Base | Ships with gmlx, from Black Forest Labs |
+| FLUX.2 klein, all sizes and variants | Ships with gmlx, from Black Forest Labs |
+| Krea 2 | Ships with gmlx, from Krea |
+| Z-Image, Z-Image-Turbo | Ships with gmlx, from the Z-Image-Turbo demo of Tongyi-MAI |
+| Qwen-Image 2.1, Qwen-Image 2.1 Turbo | Your own file, as [Qwen-Image 2.1](#qwen-image-21-rewriter-prompt) describes |
+
+Qwen3.5-9B writes captions close to those of BFL's own FLUX.3 upsampler,
+and its Q6_K GGUF is a good choice for every family. Any chat model of your
+configuration can write them. Add it under `models`, and name it as the
+`upsampler` of the image model's entry:
```yaml
models:
@@ -207,17 +216,63 @@ With an upsampler set, each prompt is rewritten:
The chat model writes the same caption for the same prompt. The PNG keeps
the caption beside your prompt. A prompt with reference images or an
-init image runs as typed, and the negative prompt is never rewritten. In ComfyUI the text
-encode node does not see the references of an edit, so the server also
-encodes the prompt as typed, and the pack samples an edit and a negative
-prompt on that encoding.
+init image runs as typed, and the negative prompt is never rewritten. In
+ComfyUI the text encode node does not see the references of an edit, so
+the server also encodes the prompt as typed, and the pack samples an edit
+and a negative prompt on that encoding.
To run a prompt as typed, pass `--no-upsample`, send
`"prompt_upsampling": false` in the request, or set `upsample: false` on
the entry, which a request with `"prompt_upsampling": true` turns back on.
Without an upsampler, FLUX.3 klein runs each prompt as typed and prints a
warning that names the `upsampler` setting and Qwen3.5-9B, unless the
-entry has `upsample: false`.
+entry has `upsample: false`. The other families run it as typed without a
+warning.
+
+### Your own system prompt
+
+`upsample_prompt` on an entry names a text file whose system prompt the
+upsampler takes in place of the family's, and `--upsample-prompt FILE`
+does the same for one `gmlx image generate` run. A relative path is found
+in `model_dirs`, as a model path is. The server reads the file for each
+request, so an edit to it takes effect on the next prompt, and captions
+written with the old file are not reused.
+
+```yaml
+image_models:
+ krea-2-turbo:
+ path: Krea-2-Turbo-Q6_K.gguf
+ encoder: Qwen3-VL-4B-Instruct-Q8_0.gguf
+ vae: qwen-image-vae-BF16.gguf
+ upsampler: qwen3.5-9b
+ upsample_prompt: prompts/krea-2-short.txt
+```
+
+The reply is read without its reasoning, as the family's own system prompt
+asks for it: the caption as written on FLUX.3 klein and FLUX.2 klein, the
+last paragraph on Krea 2, `rewritten_prompt` of a JSON object on
+Qwen-Image 2.1, or the reply as written when it has no `{`, and the
+caption as written, or `revised_prompt` of a JSON block, on Z-Image. A
+Qwen-Image reply with a `{` but no `rewritten_prompt` in it, such as one
+cut off at the token limit, is an error. On Z-Image, a `{prompt}` in your
+file is replaced by the typed prompt. Your file should ask for the same
+form.
+
+### Qwen-Image 2.1 rewriter prompt
+
+Qwen publishes the system prompt of its text to image rewriter as
+`system_prompt.txt` in the Hugging Face repository
+`Qwen/Qwen-Image-2.1-PE-T2I`, under the Qwen Research License, which
+allows non-commercial use only. gmlx does not ship it.
+`gmlx pull qwen-image` downloads it next to the model's other files, says
+its license, and names it as the entry's `upsample_prompt`, so setting an
+`upsampler` is all that is left to do. A file of that name already there
+that is not Qwen's, such as your own prompt, stays as it is, and the pull
+says to move or delete it. A new entry then goes without it, and an entry
+that already names the file keeps naming it, as the warning says. Without
+the file, a Qwen-Image request with an upsampler set is refused with a
+message that says where to get it. The rewrite keeps the image size you ask for; the aspect ratio in
+Qwen's reply is not used.
## ComfyUI
@@ -331,6 +386,13 @@ content before it is shared, and with a disclosure that the image is made
with AI where the law requires one. Commercial use of them needs a license
from Black Forest Labs.
+The system prompts gmlx ships for [prompt upsampling](#prompt-upsampling)
+are under Apache 2.0, from Black Forest Labs' flux-3-image-klein and flux2
+repositories, Krea's krea-2 repository and Tongyi-MAI's Z-Image-Turbo
+Space. Qwen's rewriter prompt, which
+`gmlx pull qwen-image` downloads, is under the Qwen Research License, for
+non-commercial use only.
+
gmlx has no content filter of its own, so a deployment that a license
asks to filter adds its own.
diff --git a/docs/internals/testing.md b/docs/internals/testing.md
index d002af185..e8122d507 100644
--- a/docs/internals/testing.md
+++ b/docs/internals/testing.md
@@ -41,7 +41,7 @@ CI runs the CPU tier on a small macOS runner, not on a developer's Mac. Every
test sees that runner by default, a Mac VM of 7.5 GB with a virtual GPU, so a
test that would fail in CI fails on your machine first.
-A test that needs more memory declares the Mac it assumes:
+A unit test that needs more memory declares the Mac it assumes:
```python
pytestmark = pytest.mark.host_gb(16) # a 16 GB Mac, as most users run
@@ -53,7 +53,7 @@ as the governor rig does, needs no machine at all.
| CI differs in | So a test |
|---------------|-----------|
-| About 7.5 GB of memory and a virtual GPU | Never reads this Mac's memory, chip or kernel counters. It declares a machine with `host_gb`. |
+| About 7.5 GB of memory and a virtual GPU | Never reads this Mac's memory, chip or kernel counters unless it runs on model files from this Mac. A unit test that needs more declares `host_gb`. |
| Python 3.11 to 3.14, tests spread over parallel workers | Passes on each version, alone or beside any other test, in any order. |
| A slower, shared machine | Bounds a time by its failure case with seconds to spare, never by a fraction of a second. |
| Only the `all` extras, so no OpenCV | That needs an optional package skips without it, with `pytest.importorskip`. |
@@ -66,6 +66,21 @@ pip install -U -e ".[all]" pytest httpx pytest-xdist
KQUANT_FORCE_CPU=1 pytest -n auto
```
+A test on model files from this Mac never runs in CI, so it sees this Mac,
+live, as its models load and free. That holds for a test marked
+`integration`, a test that fetches a gate fixture or a fixture that
+depends on one, and a test that calls a gate helper. A gate fixture, such
+as `gguf_dir`, carries `@host_box.gate` under `@pytest.fixture`. A gate
+helper, such as `image_dir()`, calls `host_box.real_files()` once the files
+it needs are there. A new gate does one or the other. Above
+`@pytest.fixture`, `@host_box.gate` marks nothing, and the fixture's tests
+stay on the runner.
+
+A `host_gb` on the module is for its unit rows, and a test on model files
+in that module still sees this Mac. A `host_gb` on such a test or its class
+fails the test. Keep rows that read no files out of a module marked
+`integration`, or they see this Mac too.
+
`GMLX_TEST_HOST_RAM_GB` changes the size of the default runner.
`GMLX_TEST_REAL_HOST=1` turns the simulation off, so a run sees the real
host. A test that needs the real GPU checks it with
diff --git a/docs/python.md b/docs/python.md
index 353e9d195..9457a3413 100644
--- a/docs/python.md
+++ b/docs/python.md
@@ -85,7 +85,7 @@ Sampling:
| `xtc_probability` / `xtc_threshold` | `0.0` | XTC sampling, on when the probability is above 0. |
| `repetition_penalty` | `0.0` | Repetition penalty over the last tokens. `0.0` turns it off. |
| `repetition_context_size` | `20` | How many recent tokens the repetition penalty looks at. |
-| `presence_penalty` / `frequency_penalty` | `0.0` | OpenAI-style penalties. |
+| `presence_penalty` / `frequency_penalty` | `0.0` | OpenAI-style penalties over every generated token, never the prompt. |
| `logit_bias` | `None` | A `{token_id: bias}` dict added to the logits. |
| `stop` | `None` | Strings that end generation. The match is trimmed. |
diff --git a/gmlx/_vendor/manifest.json b/gmlx/_vendor/manifest.json
index 032e79880..256ef5364 100644
--- a/gmlx/_vendor/manifest.json
+++ b/gmlx/_vendor/manifest.json
@@ -145,8 +145,8 @@
"mlx_lm/quant/dynamic_quant.py": {"sha256": "e4f49e5cc185466e50489e8ac1e5e5edc1abe46a4c0390d7751207a6cf8d8821", "status": "codemod", "upstream": "c1031bd9b2046a93fe3ffaa991001055a7b591f29b549cc1ed5959ad0bc87020"},
"mlx_lm/quant/gptq.py": {"sha256": "1adf426e4d998c6e83a6ad60961ab5bddc015f1818c2d068369788d3d070114c", "status": "codemod", "upstream": "8ba42877f45e86262146c6c962691c19a819478b561f55f04371d28ae3a74c9e"},
"mlx_lm/quant/utils.py": {"status": "verbatim", "upstream": "fbae54a7e39b9ae999bedfebf833e865e6912bada29f4c3fe53383b9d8655e58"},
- "mlx_lm/sample_utils.py": {"status": "verbatim", "upstream": "c0ce439f8dbf0d4e6d0f37f728a324f3f72878e6a9123df41be520804c596d67"},
- "mlx_lm/server.py": {"sha256": "d901a5a5bfaebf5cc03a92732642c9ca5e70fa7ed0b3e2373fc6de6b7ab6c90d", "status": "edited", "upstream": "cdfcb4ac848636f9927851a0ec7a951584526530cb7832ba58049e4a9144db8b"},
+ "mlx_lm/sample_utils.py": {"sha256": "5a378f35926d2f0f6b191370dc9fd6659b022ccdb9f08aae80f926587db09536", "status": "edited", "upstream": "c0ce439f8dbf0d4e6d0f37f728a324f3f72878e6a9123df41be520804c596d67"},
+ "mlx_lm/server.py": {"sha256": "f8ee9966dc1d0d2219f45a44e2b17628889fbc141453866c4dd8a2b53a96dcb2", "status": "edited", "upstream": "cdfcb4ac848636f9927851a0ec7a951584526530cb7832ba58049e4a9144db8b"},
"mlx_lm/share.py": {"sha256": "a216a734543aff85839a30fdcf7636b56c4ef6c710c340d2fd6d644e6f4b0811", "status": "codemod", "upstream": "3c25e46d4b413d67cf5bde546f47d09fbcea9ccc446e878721543af43cf91c19"},
"mlx_lm/tokenizer_utils.py": {"sha256": "1843480722194ebf1afe7350b1e99245f202ea586c9ea8d20f7d43844760a4da", "status": "codemod", "upstream": "25784bb03c922d0d7832ce6c66a6cd4eb3a4820b6c5a8e583dedb63a018fb56a"},
"mlx_lm/tool_parsers/__init__.py": {"status": "verbatim", "upstream": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"},
@@ -193,12 +193,12 @@
"mlx_vlm/fp8.py": {"status": "verbatim", "upstream": "91395897d4b9f9987db67391d4034714135faf2d6cdd171066de90f958dc2069"},
"mlx_vlm/generate/__init__.py": {"sha256": "593e7201c7bd578375efaa7d20367da5347dde74727c1692c6e47a74b54d5e61", "status": "codemod", "upstream": "e4f2794a17f4096a02c5786a98c32113ace2d0a3cc3c63d5a8f16ae873c8b478"},
"mlx_vlm/generate/__main__.py": {"status": "verbatim", "upstream": "3129adff95e0f38b87ab34cddfc27082c7892bcaec267ff5d40f160fdf55b44a"},
- "mlx_vlm/generate/ar.py": {"sha256": "8ee00e8429c361c44a39a5f71535302a7bacc1a84ecc5be254845d7a46725735", "status": "edited", "upstream": "23106ad929b88c47cb67dd2f010519f954dbf3aa61b63f20eeca2ebae0053e7e"},
+ "mlx_vlm/generate/ar.py": {"sha256": "d2dc36f632e29cc1d52a0b6881b0103bb99516193733b434aca31da5f589ba86", "status": "edited", "upstream": "23106ad929b88c47cb67dd2f010519f954dbf3aa61b63f20eeca2ebae0053e7e"},
"mlx_vlm/generate/audio.py": {"status": "verbatim", "upstream": "33dfde340761855e8171446e365f235144e4f845e5dde9d48dbd6495cef05433"},
"mlx_vlm/generate/cli.py": {"status": "verbatim", "upstream": "86e9348e69596623ffe36b76111fa252b55cfac75633b048f7afba051086aa9a"},
"mlx_vlm/generate/common.py": {"sha256": "32975a1bd06d4c34f067af39ac39d5413bf9142259537ba4a7da04480af6ec35", "status": "edited", "upstream": "c69e7e38a09990404d299b0a8d4be55c8220e60652e67a2456e8177633813ba0"},
"mlx_vlm/generate/diffusion.py": {"sha256": "2a7f6626fe01872fc7f44db07178bfb12b3e35a692979c44e58cca1d53e41948", "status": "codemod", "upstream": "b91dbb29d8486e9adf20bc408fdf0956d6eac68f5563ce775b4afcf931d3b4fd"},
- "mlx_vlm/generate/dispatch.py": {"sha256": "3af79e0809fd4bb707ea87aa0e1fbe406a27f11d6d717d9dc08e3486ef514c9e", "status": "edited", "upstream": "e49e508bb0a8cc8eda24b53430a60450697a52b1d3d5b63b2dd14d3dd413f0f3"},
+ "mlx_vlm/generate/dispatch.py": {"sha256": "67c73473a12db7e55f665e8ea9f5ed605baecc7fdfec823d0cd7f3c99d80b834", "status": "edited", "upstream": "e49e508bb0a8cc8eda24b53430a60450697a52b1d3d5b63b2dd14d3dd413f0f3"},
"mlx_vlm/generate/edit_image.py": {"sha256": "aaca736da50024ef0b2ebbbc798ee23ff203fd0195df822c196c6ae3b9419ec9", "status": "codemod", "upstream": "22a36e39ec7b8e0bb4f5a25c9e1c991d2e04422052c88dced8b1d631c53d2442"},
"mlx_vlm/generate/image.py": {"sha256": "4cec6e6be4209219adcf9297d4996756d6edde7d5f3a065ac3548880b1256971", "status": "codemod", "upstream": "c322621f307e8f6556841f36510352272bbd9aed8ebf2159a1861c57173dec47"},
"mlx_vlm/generate/types.py": {"status": "verbatim", "upstream": "dff487807bedaa3549c39dea02986ada31a45e9e4c27ab447a207e3fff009f6c"},
@@ -1268,7 +1268,7 @@
"mlx_vlm/models/qwen3_omni_moe/language.py": {"status": "verbatim", "upstream": "13a815b3f4cd73c37e27c4c4b6739606e4ed6b61b0d21ba79f106025373e0243"},
"mlx_vlm/models/qwen3_omni_moe/omni_utils.py": {"sha256": "410cc3d5c7a481cbc5d5b0b3ad1173886d43ceb43254c473d104e3d8ea5449b8", "status": "codemod", "upstream": "fadadb95e0ae12c9e34a77bef9f9b689e5d3bef02eebc2604ff1d0d3c9229506"},
"mlx_vlm/models/qwen3_omni_moe/processing_qwen3_omni_moe.py": {"status": "verbatim", "upstream": "c90949b1ce6554317113e837bd628029c83050836ab568c5365fc66c5096ad2f"},
- "mlx_vlm/models/qwen3_omni_moe/qwen3_omni_moe.py": {"sha256": "8ab2cc8b0c49e979bc2aa29dc8b7aa99dd9e52335277c96d6e9d34408b683ae4", "status": "codemod", "upstream": "c218452e2e05762f365fb3aff8d7300e9c5054b9f7803da914cb81e2cf745389"},
+ "mlx_vlm/models/qwen3_omni_moe/qwen3_omni_moe.py": {"sha256": "0d880c79712d081206b58ff9d151857138589b78724ce8424119624a2bba7424", "status": "edited", "upstream": "c218452e2e05762f365fb3aff8d7300e9c5054b9f7803da914cb81e2cf745389"},
"mlx_vlm/models/qwen3_omni_moe/talker.py": {"sha256": "aef77782d150ae4f4a1368c2868ab41a50bf72dc56219a2d933f7ba01d9ecf2a", "status": "codemod", "upstream": "a66e4fda3d154b0042f81bc3ca3a4ab094b3b8a617ec9accf7fa32f82ddc09d7"},
"mlx_vlm/models/qwen3_omni_moe/thinker.py": {"sha256": "049ea8134c9f80b87c1d699b35891efc494552f59f3681171e948f5aa457197f", "status": "codemod", "upstream": "f741396e0ce1a43e85b856f797887500d7e742735e130d02b59c5deec1c31336"},
"mlx_vlm/models/qwen3_omni_moe/vision.py": {"status": "verbatim", "upstream": "7c89063f511c18699d98ffc924be1be2dbc650ff1b274e384c272859aa359ec6"},
@@ -1511,7 +1511,7 @@
"mlx_vlm/quantization/one_bit.py": {"status": "verbatim", "upstream": "6771b9003e4529658c6f051f208265cceb8a25fe7f7ee27369d2bed4098e46c7"},
"mlx_vlm/reranker.py": {"status": "verbatim", "upstream": "fe7fac6dabbd25ede3bdd7e6f0bd909331ccab5d2ba9c679dad4f9ca97e98690"},
"mlx_vlm/reranker_loader.py": {"status": "verbatim", "upstream": "09357d98a5efa64268e7e4191312a2893ec57d94ffa3ffa10c0f3d8e2a475f32"},
- "mlx_vlm/sample_utils.py": {"status": "verbatim", "upstream": "b3057b6dcaefe5b0a7c50cb96b70baad4334a2f88f70adc04fe7ec2060f7851c"},
+ "mlx_vlm/sample_utils.py": {"sha256": "5f2ba6b6e3374f9df0534155ab2fc0c718e3a861ecf834e1d1869afe0c32365d", "status": "edited", "upstream": "b3057b6dcaefe5b0a7c50cb96b70baad4334a2f88f70adc04fe7ec2060f7851c"},
"mlx_vlm/server/__init__.py": {"status": "verbatim", "upstream": "dc1ae2a78cf2fd68e6135c296c8cd4373f29b29341a56970401147d39c46f05a"},
"mlx_vlm/server/__main__.py": {"status": "verbatim", "upstream": "e418cdbb27adf0063e3cec28179ac6b7bdb6ac743bb49d157f450551fcf38be2"},
"mlx_vlm/server/anthropic.py": {"sha256": "663596363dafd12770c178ebcfafac5142bff3af8746bebf4233af150ffd6dc4", "status": "codemod", "upstream": "23b29d5341dc1fb807362dac7c0f9491e0bd72c1063d5824ae397c63254ee6f7"},
@@ -1519,16 +1519,16 @@
"mlx_vlm/server/audio.py": {"sha256": "2b02bc59073b3b231e18d9f3ee3374353cf555e961f04958b73cd2f1ab618bc7", "status": "codemod", "upstream": "53cef68307be67cb850d39465012cea9d7003c10b5408116fca37f9cd6b895b6"},
"mlx_vlm/server/cli.py": {"sha256": "7cd75c552d891bd7dfddb4731aab55a8d638a70686f162abcf274c6c2bfa910e", "status": "codemod", "upstream": "15d49f41caea3acba6238e4ac8e0839279f4d0a0c9c620ac56370aab1262b597"},
"mlx_vlm/server/embeddings.py": {"status": "verbatim", "upstream": "a2af958265cae9cb2c82348baf38422a0d242909010dc8876980d591e7444875"},
- "mlx_vlm/server/generation.py": {"sha256": "96fbb3ede65c3610e974bc3ae91c36d39250e99d349a9eb694552d979876f84b", "status": "edited", "upstream": "6c0b663d7adafbc686c6049649953fd78bea1572d2a6280532fc063dda511b77"},
+ "mlx_vlm/server/generation.py": {"sha256": "61d6b2389bf49830a4a778445a95546f3ab9cdf88fab75f9c7da0812275a271d", "status": "edited", "upstream": "6c0b663d7adafbc686c6049649953fd78bea1572d2a6280532fc063dda511b77"},
"mlx_vlm/server/model_discovery.py": {"status": "verbatim", "upstream": "0d6fb031c579b7a1b4c63adb94f926e5913bd588234271942f52e4b9e21e72d2"},
"mlx_vlm/server/openai.py": {"sha256": "00a590ad84b2cdaaa7bf058a633117b1fc004c3b1b6e61b1336868f26c2a181a", "status": "codemod", "upstream": "fd5dafaba30c5779390657110aa16c1fa2409335853d58bf4ceec1fa701fe6c7"},
"mlx_vlm/server/realtime.py": {"sha256": "942671328088936df3145eb72ba7e6650d0171c22ca17321d79ad113cd0482b9", "status": "codemod", "upstream": "815998c835bd348649779ce480152c48cf8ff770ca53503f1a1c31649668b416"},
- "mlx_vlm/server/request_normalization.py": {"sha256": "38c69ac6dd59534f5b6ad93adf15be145bf299c1aae9af0c9d0216338c006c59", "status": "edited", "upstream": "37eb0727a61eb8a647cf010af20098f388688336d0e9df9f217c2409fc3a00e5"},
+ "mlx_vlm/server/request_normalization.py": {"sha256": "3bd9bed19b50e094ade4899439f74401717010b150b82df2a6f7812fe01460f6", "status": "edited", "upstream": "37eb0727a61eb8a647cf010af20098f388688336d0e9df9f217c2409fc3a00e5"},
"mlx_vlm/server/reranking.py": {"status": "verbatim", "upstream": "36b03fb077741261696795a3b32bfcae64339d0d50d091ee39f2df4b97343179"},
"mlx_vlm/server/responses_state.py": {"sha256": "0a195e66a156624c7fc8710c54fb3a782f6a14f7bfa2cd75b59e099e060eebf4", "status": "edited", "upstream": "138184cb972c43c6f48fc104642608add2605180c1682d3cc1dfb0bd408f3657"},
"mlx_vlm/server/runtime.py": {"status": "verbatim", "upstream": "de80f06d985f98270d7f5d16d2c4e41f1f8ecd69c15c9a819ec8e891ff1f83d2"},
"mlx_vlm/server/runtime_config.py": {"sha256": "ae05cb5e399e51700340a7ed31e24f7385f5373325e37fe15425ceee5e570b8a", "status": "codemod", "upstream": "cccded87d36ee60bb302be4898b2a674335448e9eb6838584ca797095e1b4c07"},
- "mlx_vlm/server/schemas.py": {"status": "verbatim", "upstream": "80c863e32dd543ca878876e3b183559adfa1a27848d44599ae0852efe74fea68"},
+ "mlx_vlm/server/schemas.py": {"sha256": "e4b49e3359bdb6d9a6f2884427470b19221f5c8f22f67baadba8ee566b03ecfe", "status": "edited", "upstream": "80c863e32dd543ca878876e3b183559adfa1a27848d44599ae0852efe74fea68"},
"mlx_vlm/speculative/__init__.py": {"status": "verbatim", "upstream": "1f98ff6655dc4629b8d8c5001a32d5a95d232e156722be5b1adac848ff6f0559"},
"mlx_vlm/speculative/cache_state.py": {"status": "verbatim", "upstream": "39d35ef0aae0c9298f2fcd3ff6b91b106c6b6c1c3ab6becae216a8fafb2c5300"},
"mlx_vlm/speculative/common.py": {"status": "verbatim", "upstream": "e3d3c0294a6d6fc915a32e96460370bfe57baef19539c5a290d4254525717cf1"},
diff --git a/gmlx/_vendor/mlx_lm/sample_utils.py b/gmlx/_vendor/mlx_lm/sample_utils.py
index 05a45fc60..961f4a20e 100644
--- a/gmlx/_vendor/mlx_lm/sample_utils.py
+++ b/gmlx/_vendor/mlx_lm/sample_utils.py
@@ -6,6 +6,9 @@
import mlx.core as mx
+# gmlx: presence and frequency penalties count the generated tokens only.
+from gmlx.gen.penalties import OutputPenalty
+
def make_sampler(
temp: float = 0.0,
@@ -74,10 +77,11 @@ def make_logits_processors(
repetition_penalty: Optional[float] = None,
repetition_context_size: Optional[int] = 20,
presence_penalty: Optional[float] = None,
- presence_context_size: Optional[int] = 20,
+ presence_context_size: Optional[int] = None, # gmlx: whole output
frequency_penalty: Optional[float] = None,
- frequency_context_size: Optional[int] = 20,
+ frequency_context_size: Optional[int] = None, # gmlx: whole output
):
+ # gmlx: the docstring gives the generated-token rule of the context sizes.
"""
Make logits processors for use with ``generate_step``.
@@ -88,13 +92,15 @@ def make_logits_processors(
consider for repetition penalty. Default: ``20``.
presence_penalty (float, optional): An additive penalty to reduce
repeating tokens.
- presence_context_size (int, optional): The number of tokens to consider
- for the presence penalty. Default: ``20``.
+ presence_context_size (int, optional): The number of last generated
+ tokens to consider for the presence penalty. ``None`` or ``0``
+ considers every generated token. Default: ``None``.
frequency_penalty (float, optional): An additive penalty to reduce
repeating tokens. The tokens are penalized proportionally to their
frequency.
- frequency_context_size (int, optional): The number of tokens to consider
- for the frequency penalty. Default: ``20``.
+ frequency_context_size (int, optional): The number of last generated
+ tokens to consider for the frequency penalty. ``None`` or ``0``
+ considers every generated token. Default: ``None``.
logit_bias (dictionary, optional): Additive logit bias.
Returns:
@@ -312,56 +318,50 @@ def repetition_penalty_processor(tokens, logits):
return repetition_penalty_processor
-def make_presence_penalty(penalty: float, context_size: int = 20):
+def make_presence_penalty(penalty: float, context_size: Optional[int] = None):
+ # gmlx: the docstring gives the generated-token rule.
"""
Make a presence penalty processor.
Corresponds to the OpenAI option with the same name. Namely, subtracts
- ``penalty`` from a logit if the token has occured at least once in the
- ``context_size`` previous tokens.
+ ``penalty`` from a logit if the token has occurred at least once in the
+ ``context_size`` last generated tokens, or in any generated token
+ when ``context_size`` is ``None`` or ``0``. Prompt tokens are never
+ penalized.
Args:
penalty (float): The presence penalty to be applied.
- context_size (int): The number of previous tokens to use.
- Default: ``20``.
+ context_size (int, optional): The number of last generated
+ tokens to use, or ``None`` for all of them. Default: ``None``.
Returns:
Callable[[mx.array, List[int]], mx.array]
"""
-
- def presence_penalty_processor(tokens, logits):
- if len(tokens) > 0:
- tokens = tokens[-context_size:]
- logits[:, tokens] -= penalty
- return logits
-
- return presence_penalty_processor
+ # gmlx: the generated tokens only, the whole output by default.
+ return OutputPenalty(penalty, context_size)
-def make_frequency_penalty(penalty: float, context_size: int = 20):
+def make_frequency_penalty(penalty: float, context_size: Optional[int] = None):
+ # gmlx: the docstring gives the generated-token rule.
"""
Make a frequency penalty processor.
Corresponds to the OpenAI option with the same name. Namely, subtracts
- ``penalty`` from a logit for every time that the token has occured in the
- ``context_size`` previous tokens.
+ ``penalty`` from a logit for every time that the token has occurred in the
+ ``context_size`` last generated tokens, or in any generated token
+ when ``context_size`` is ``None`` or ``0``. Prompt tokens are never
+ penalized.
The difference with the presence penalty is that the more often a token
occurs the more it will be penalized.
Args:
penalty (float): The frequency penalty to be applied.
- context_size (int): The number of previous tokens to use.
- Default: ``20``.
+ context_size (int, optional): The number of last generated
+ tokens to use, or ``None`` for all of them. Default: ``None``.
Returns:
Callable[[mx.array, List[int]], mx.array]
"""
-
- def frequency_penalty_processor(tokens, logits):
- if len(tokens) > 0:
- tokens = tokens[-context_size:]
- logits = logits.at[:, tokens].subtract(penalty)
- return logits
-
- return frequency_penalty_processor
+ # gmlx: the generated tokens only, the whole output by default.
+ return OutputPenalty(penalty, context_size, frequency=True)
diff --git a/gmlx/_vendor/mlx_lm/server.py b/gmlx/_vendor/mlx_lm/server.py
index edd68f2d0..6d6a6cd5f 100644
--- a/gmlx/_vendor/mlx_lm/server.py
+++ b/gmlx/_vendor/mlx_lm/server.py
@@ -1183,9 +1183,10 @@ def do_POST(self):
self.repetition_penalty = self.body.get("repetition_penalty", 0.0)
self.repetition_context_size = self.body.get("repetition_context_size", 20)
self.presence_penalty = self.body.get("presence_penalty", 0.0)
- self.presence_context_size = self.body.get("presence_context_size", 20)
+ # gmlx: None counts the whole output.
+ self.presence_context_size = self.body.get("presence_context_size")
self.frequency_penalty = self.body.get("frequency_penalty", 0.0)
- self.frequency_context_size = self.body.get("frequency_context_size", 20)
+ self.frequency_context_size = self.body.get("frequency_context_size")
self.xtc_probability = self.body.get("xtc_probability", 0.0)
self.xtc_threshold = self.body.get("xtc_threshold", 0.0)
self.logit_bias = self.body.get("logit_bias", None)
@@ -1241,9 +1242,11 @@ def validate_model_parameters(self):
self._validate("repetition_penalty", (float, int), min_val=0)
self._validate("repetition_context_size", int, min_val=0)
self._validate("presence_penalty", (float, int))
- self._validate("presence_context_size", int, min_val=0)
+ # gmlx: None counts the whole output.
+ self._validate("presence_context_size", int, min_val=0, optional=True)
self._validate("frequency_penalty", (float, int))
- self._validate("frequency_context_size", int, min_val=0)
+ # gmlx: None counts the whole output.
+ self._validate("frequency_context_size", int, min_val=0, optional=True)
self._validate("logprobs", bool)
self._validate("top_logprobs", int, min_val=0, max_val=11, whitelist=[-1])
self._validate("xtc_probability", float, min_val=0, max_val=1)
diff --git a/gmlx/_vendor/mlx_vlm/generate/ar.py b/gmlx/_vendor/mlx_vlm/generate/ar.py
index 6cb0b3629..ba205160f 100644
--- a/gmlx/_vendor/mlx_vlm/generate/ar.py
+++ b/gmlx/_vendor/mlx_vlm/generate/ar.py
@@ -189,9 +189,9 @@ def generate_step(
repetition_penalty: Optional[float] = None,
repetition_context_size: Optional[int] = DEFAULT_REPETITION_CONTEXT_SIZE,
presence_penalty: Optional[float] = None,
- presence_context_size: Optional[int] = DEFAULT_REPETITION_CONTEXT_SIZE,
+ presence_context_size: Optional[int] = None, # gmlx: whole output
frequency_penalty: Optional[float] = None,
- frequency_context_size: Optional[int] = DEFAULT_REPETITION_CONTEXT_SIZE,
+ frequency_context_size: Optional[int] = None, # gmlx: whole output
top_p: float = DEFAULT_TOP_P,
min_p: float = DEFAULT_MIN_P,
top_k: int = DEFAULT_TOP_K,
@@ -222,6 +222,7 @@ def generate_step(
verbose: bool = False,
**kwargs,
) -> Generator[Tuple[mx.array, mx.array], None, None]:
+ # gmlx: the docstring gives the generated-token rule of the context sizes.
"""
A generator producing token ids based on the given prompt from the model.
@@ -238,12 +239,12 @@ def generate_step(
consider for repetition penalty.
presence_penalty (float, optional): Additive penalty for tokens that
already appeared in recent generated context.
- presence_context_size (int, optional): The number of tokens to
- consider for presence penalty.
+ presence_context_size (int, optional): The number of last generated
+ tokens to consider for presence penalty, or None for all of them.
frequency_penalty (float, optional): Additive penalty scaled by token
frequency in recent generated context.
- frequency_context_size (int, optional): The number of tokens to
- consider for frequency penalty.
+ frequency_context_size (int, optional): The number of last generated
+ tokens to consider for frequency penalty, or None for all of them.
top_p (float, optional): Nucleus sampling, higher means model considers
more less likely words.
min_p (float, optional): Minimum probability threshold relative to the
diff --git a/gmlx/_vendor/mlx_vlm/generate/dispatch.py b/gmlx/_vendor/mlx_vlm/generate/dispatch.py
index 9b289b7b0..7e2971eeb 100644
--- a/gmlx/_vendor/mlx_vlm/generate/dispatch.py
+++ b/gmlx/_vendor/mlx_vlm/generate/dispatch.py
@@ -436,8 +436,9 @@ def parse_arguments():
parser.add_argument(
"--presence-context-size",
type=int,
- default=DEFAULT_REPETITION_CONTEXT_SIZE,
- help="Number of recent generated tokens used for presence penalty.",
+ default=None, # gmlx: whole output
+ help="Number of recent generated tokens used for presence penalty, "
+ "default all of them.",
)
parser.add_argument(
"--frequency-penalty",
@@ -448,8 +449,9 @@ def parse_arguments():
parser.add_argument(
"--frequency-context-size",
type=int,
- default=DEFAULT_REPETITION_CONTEXT_SIZE,
- help="Number of recent generated tokens used for frequency penalty.",
+ default=None, # gmlx: whole output
+ help="Number of recent generated tokens used for frequency penalty, "
+ "default all of them.",
)
parser.add_argument("--chat", action="store_true", help="Chat in multi-turn style.")
parser.add_argument(
diff --git a/gmlx/_vendor/mlx_vlm/models/qwen3_omni_moe/qwen3_omni_moe.py b/gmlx/_vendor/mlx_vlm/models/qwen3_omni_moe/qwen3_omni_moe.py
index 38692dc95..50279262a 100644
--- a/gmlx/_vendor/mlx_vlm/models/qwen3_omni_moe/qwen3_omni_moe.py
+++ b/gmlx/_vendor/mlx_vlm/models/qwen3_omni_moe/qwen3_omni_moe.py
@@ -186,9 +186,9 @@ def _generate_thinker_with_hidden_states(
repetition_penalty: Optional[float] = None,
repetition_context_size: Optional[int] = 20,
presence_penalty: Optional[float] = None,
- presence_context_size: Optional[int] = 20,
+ presence_context_size: Optional[int] = None, # gmlx: whole output
frequency_penalty: Optional[float] = None,
- frequency_context_size: Optional[int] = 20,
+ frequency_context_size: Optional[int] = None, # gmlx: whole output
**kwargs,
):
"""Generate Thinker tokens while retaining the Talker accept-layer hidden.
diff --git a/gmlx/_vendor/mlx_vlm/sample_utils.py b/gmlx/_vendor/mlx_vlm/sample_utils.py
index 0fb3161df..7a2563c88 100644
--- a/gmlx/_vendor/mlx_vlm/sample_utils.py
+++ b/gmlx/_vendor/mlx_vlm/sample_utils.py
@@ -6,6 +6,9 @@
import mlx.core as mx
+# gmlx: presence and frequency penalties count the generated tokens only.
+from gmlx.gen.penalties import OutputPenalty
+
def make_sampler(
temp: float = 0.0,
@@ -94,10 +97,11 @@ def make_logits_processors(
repetition_penalty: Optional[float] = None,
repetition_context_size: Optional[int] = 20,
presence_penalty: Optional[float] = None,
- presence_context_size: Optional[int] = 20,
+ presence_context_size: Optional[int] = None, # gmlx: whole output
frequency_penalty: Optional[float] = None,
- frequency_context_size: Optional[int] = 20,
+ frequency_context_size: Optional[int] = None, # gmlx: whole output
):
+ # gmlx: the docstring gives the generated-token rule of the context sizes.
"""
Make logits processors for use with ``generate_step``.
@@ -108,13 +112,15 @@ def make_logits_processors(
consider for repetition penalty. Default: ``20``.
presence_penalty (float, optional): An additive penalty to reduce
repeating tokens.
- presence_context_size (int, optional): The number of tokens to consider
- for the presence penalty. Default: ``20``.
+ presence_context_size (int, optional): The number of last generated
+ tokens to consider for the presence penalty. ``None`` or ``0``
+ considers every generated token. Default: ``None``.
frequency_penalty (float, optional): An additive penalty to reduce
repeating tokens. The tokens are penalized proportionally to their
frequency.
- frequency_context_size (int, optional): The number of tokens to consider
- for the frequency penalty. Default: ``20``.
+ frequency_context_size (int, optional): The number of last generated
+ tokens to consider for the frequency penalty. ``None`` or ``0``
+ considers every generated token. Default: ``None``.
logit_bias (dictionary, optional): Additive logit bias.
Returns:
@@ -420,59 +426,53 @@ def repetition_penalty_processor(tokens, logits):
return repetition_penalty_processor
-def make_presence_penalty(penalty: float, context_size: int = 20):
+def make_presence_penalty(penalty: float, context_size: Optional[int] = None):
+ # gmlx: the docstring gives the generated-token rule.
"""
Make a presence penalty processor.
Corresponds to the OpenAI option with the same name. Namely, subtracts
``penalty`` from a logit if the token has occurred at least once in the
- ``context_size`` previous tokens.
+ ``context_size`` last generated tokens, or in any generated token
+ when ``context_size`` is ``None`` or ``0``. Prompt tokens are never
+ penalized.
Args:
penalty (float): The presence penalty to be applied.
- context_size (int): The number of previous tokens to use.
- Default: ``20``.
+ context_size (int, optional): The number of last generated
+ tokens to use, or ``None`` for all of them. Default: ``None``.
Returns:
Callable[[mx.array, List[int]], mx.array]
"""
-
- def presence_penalty_processor(tokens, logits):
- if len(tokens) > 0:
- tokens = tokens[-context_size:]
- logits[:, tokens] -= penalty
- return logits
-
- return presence_penalty_processor
+ # gmlx: the generated tokens only, the whole output by default.
+ return OutputPenalty(penalty, context_size)
-def make_frequency_penalty(penalty: float, context_size: int = 20):
+def make_frequency_penalty(penalty: float, context_size: Optional[int] = None):
+ # gmlx: the docstring gives the generated-token rule.
"""
Make a frequency penalty processor.
Corresponds to the OpenAI option with the same name. Namely, subtracts
``penalty`` from a logit for every time that the token has occurred in the
- ``context_size`` previous tokens.
+ ``context_size`` last generated tokens, or in any generated token
+ when ``context_size`` is ``None`` or ``0``. Prompt tokens are never
+ penalized.
The difference with the presence penalty is that the more often a token
occurs the more it will be penalized.
Args:
penalty (float): The frequency penalty to be applied.
- context_size (int): The number of previous tokens to use.
- Default: ``20``.
+ context_size (int, optional): The number of last generated
+ tokens to use, or ``None`` for all of them. Default: ``None``.
Returns:
Callable[[mx.array, List[int]], mx.array]
"""
-
- def frequency_penalty_processor(tokens, logits):
- if len(tokens) > 0:
- tokens = tokens[-context_size:]
- logits = logits.at[:, tokens].subtract(penalty)
- return logits
-
- return frequency_penalty_processor
+ # gmlx: the generated tokens only, the whole output by default.
+ return OutputPenalty(penalty, context_size, frequency=True)
def top_p_sampling(logits: mx.array, top_p: float, temperature: float) -> mx.array:
diff --git a/gmlx/_vendor/mlx_vlm/server/generation.py b/gmlx/_vendor/mlx_vlm/server/generation.py
index c1c8ed687..f8ec1609d 100644
--- a/gmlx/_vendor/mlx_vlm/server/generation.py
+++ b/gmlx/_vendor/mlx_vlm/server/generation.py
@@ -621,9 +621,9 @@ class GenerationArguments:
repetition_penalty: Optional[float] = None
repetition_context_size: Optional[int] = DEFAULT_REPETITION_CONTEXT_SIZE
presence_penalty: Optional[float] = None
- presence_context_size: Optional[int] = DEFAULT_REPETITION_CONTEXT_SIZE
+ presence_context_size: Optional[int] = None # gmlx: whole output
frequency_penalty: Optional[float] = None
- frequency_context_size: Optional[int] = DEFAULT_REPETITION_CONTEXT_SIZE
+ frequency_context_size: Optional[int] = None # gmlx: whole output
max_denoising_steps: Optional[int] = None
block_length: Optional[int] = None
num_to_transfer: Optional[int] = None
diff --git a/gmlx/_vendor/mlx_vlm/server/request_normalization.py b/gmlx/_vendor/mlx_vlm/server/request_normalization.py
index c2f7a5ff7..991c2afeb 100644
--- a/gmlx/_vendor/mlx_vlm/server/request_normalization.py
+++ b/gmlx/_vendor/mlx_vlm/server/request_normalization.py
@@ -215,17 +215,10 @@ def _build_gen_args(
DEFAULT_REPETITION_CONTEXT_SIZE,
),
presence_penalty=getattr(request, "presence_penalty", None),
- presence_context_size=_request_field_or_default(
- request,
- "presence_context_size",
- DEFAULT_REPETITION_CONTEXT_SIZE,
- ),
+ # gmlx: None counts the whole output.
+ presence_context_size=getattr(request, "presence_context_size", None),
frequency_penalty=getattr(request, "frequency_penalty", None),
- frequency_context_size=_request_field_or_default(
- request,
- "frequency_context_size",
- DEFAULT_REPETITION_CONTEXT_SIZE,
- ),
+ frequency_context_size=getattr(request, "frequency_context_size", None),
max_denoising_steps=_request_field_or_default(
request, "max_denoising_steps", None
),
diff --git a/gmlx/_vendor/mlx_vlm/server/schemas.py b/gmlx/_vendor/mlx_vlm/server/schemas.py
index f4cdd135c..8078d5b5b 100644
--- a/gmlx/_vendor/mlx_vlm/server/schemas.py
+++ b/gmlx/_vendor/mlx_vlm/server/schemas.py
@@ -373,12 +373,14 @@ class OpenAIRequest(FlexibleBaseModel):
None, description="Repetition penalty context size."
)
presence_penalty: Optional[float] = Field(None, description="Presence penalty.")
+ # gmlx: None or 0 counts every generated token, N the last N.
presence_context_size: Optional[int] = Field(
- None, description="Presence penalty context size."
+ None, ge=0, description="Presence penalty context size."
)
frequency_penalty: Optional[float] = Field(None, description="Frequency penalty.")
+ # gmlx: None or 0 counts every generated token, N the last N.
frequency_context_size: Optional[int] = Field(
- None, description="Frequency penalty context size."
+ None, ge=0, description="Frequency penalty context size."
)
logit_bias: Optional[Any] = Field(None, description="Logit bias dict.")
enable_thinking: Optional[bool] = Field(
@@ -775,12 +777,14 @@ class VLMRequest(FlexibleBaseModel):
None, description="Repetition penalty context size."
)
presence_penalty: Optional[float] = Field(None, description="Presence penalty.")
+ # gmlx: None or 0 counts every generated token, N the last N.
presence_context_size: Optional[int] = Field(
- None, description="Presence penalty context size."
+ None, ge=0, description="Presence penalty context size."
)
frequency_penalty: Optional[float] = Field(None, description="Frequency penalty.")
+ # gmlx: None or 0 counts every generated token, N the last N.
frequency_context_size: Optional[int] = Field(
- None, description="Frequency penalty context size."
+ None, ge=0, description="Frequency penalty context size."
)
logit_bias: Optional[Any] = Field(None, description="Logit bias dict.")
enable_thinking: Optional[bool] = Field(
@@ -970,12 +974,14 @@ class AnthropicRequest(FlexibleBaseModel):
None, description="Repetition penalty context size."
)
presence_penalty: Optional[float] = Field(None, description="Presence penalty.")
+ # gmlx: None or 0 counts every generated token, N the last N.
presence_context_size: Optional[int] = Field(
- None, description="Presence penalty context size."
+ None, ge=0, description="Presence penalty context size."
)
frequency_penalty: Optional[float] = Field(None, description="Frequency penalty.")
+ # gmlx: None or 0 counts every generated token, N the last N.
frequency_context_size: Optional[int] = Field(
- None, description="Frequency penalty context size."
+ None, ge=0, description="Frequency penalty context size."
)
logit_bias: Optional[Any] = Field(None, description="Logit bias dict.")
enable_thinking: Optional[bool] = None
diff --git a/gmlx/commands/image.py b/gmlx/commands/image.py
index ba2dab41c..ee19f927b 100644
--- a/gmlx/commands/image.py
+++ b/gmlx/commands/image.py
@@ -169,12 +169,16 @@ def _generate_parser(prog: str) -> argparse.ArgumentParser:
up = p.add_mutually_exclusive_group()
up.add_argument("--upsampler", default=None, metavar="GGUF",
help="The chat model that rewrites the prompt before it "
- "is encoded, on a family that takes prompt "
- "upsampling (FLUX.3 klein): a GGUF, or a models: id "
- "of the config (default: the entry's upsampler). "
- "A prompt with --ref or --init-image runs as typed.")
+ "is encoded: a GGUF, or a models: id of the config "
+ "(default: the entry's upsampler). A prompt with "
+ "--ref or --init-image runs as typed.")
up.add_argument("--no-upsample", action="store_true",
help="Run the prompt as typed, without the upsampler.")
+ p.add_argument("--upsample-prompt", default=None, metavar="FILE",
+ help="A text file whose system prompt the upsampler takes "
+ "in place of the family's (default: the entry's "
+ "upsample_prompt). Qwen-Image 2.1 ships none and "
+ "needs one.")
p.add_argument("--size", type=_size, default=None, metavar="WxH",
help="Image width and height, multiples of 16, or of 32 "
"on Qwen-Image 2.1 (default: the init image's, else "
@@ -468,6 +472,8 @@ def _from_config(a) -> None:
a.no_upsample = True
elif rm.upsampler:
a.upsampler = rm.upsampler
+ if getattr(a, "upsample_prompt", "") is None and rm.upsample_prompt:
+ a.upsample_prompt = rm.upsample_prompt
def _upsampler_path(a) -> str:
@@ -510,27 +516,37 @@ def _upsample(a, family, prompt: str, refs, log) -> str | None:
"upsampler is not loaded")
return None
if a.upsampler is None:
- log(f"warning: {family.title} expects its prompts upsampled; set "
- f"--upsampler to a {family.upsample.suggested} GGUF or its "
- "models: id, or upsampler on the image_models entry, or pass "
- "--no-upsample")
+ if family.upsample.expected:
+ log(f"warning: {family.title} expects its prompts upsampled; set "
+ f"--upsampler to a {family.upsample.suggested} GGUF or its "
+ "models: id, or upsampler on the image_models entry, or pass "
+ "--no-upsample")
return None
+ from gmlx.image.upsample import UpsampleError, system_prompt
+
+ try:
+ system = system_prompt(family, a.upsample_prompt)
+ except UpsampleError as e:
+ raise _Usage(str(e)) from None
a.upsampler = _upsampler_path(a)
- return run_upsampler(a.upsampler, family, prompt, log, quiet=a.json)
+ return run_upsampler(a.upsampler, family, prompt, log, quiet=a.json,
+ system=system)
def run_upsampler(path: str, family, prompt: str, log,
- quiet: bool = False) -> str:
+ quiet: bool = False, system: str | None = None) -> str:
"""Load the chat model at ``path``, write the caption of ``prompt``,
- and release the model before returning. ``quiet`` keeps the load's
- lines off stdout."""
+ and release the model before returning, under ``system``, by default
+ the family's shipped system prompt. ``quiet`` keeps the load's lines
+ off stdout."""
import gc
import mlx.core as mx
import gmlx.load.loadlog as loadlog
from gmlx.gen.generation import generate
- from gmlx.image.upsample import UpsampleError, system_prompt, upsample
+ from gmlx.image.upsample import (UpsampleError, reply_lead,
+ system_prompt, upsample)
from gmlx.load.loader import load_model
from gmlx.serve.capacity import working_set_bytes
from gmlx.serve.memory import resolve_cache_limit
@@ -538,8 +554,10 @@ def run_upsampler(path: str, family, prompt: str, log,
from .tool_preflight import check_or_exit
+ if system is None:
+ system = system_prompt(family)
# About 3 characters a token, which overcounts English.
- ctx = (len(system_prompt(family)) + len(prompt)) // 3 + (
+ ctx = (len(system) + len(prompt)) // 3 + (
family.upsample.max_tokens)
def stdout():
return (contextlib.redirect_stdout(sys.stderr) if quiet
@@ -564,13 +582,21 @@ def chat(req):
kwargs = map_thinking_controls(
{}, req.thinking, None, template,
warn=lambda msg: log(f"[upsample] {msg}"))
- return generate(model, tok, req.user, system_prompt=req.system,
- max_tokens=req.max_tokens, temp=req.temperature,
- presence_penalty=req.presence_penalty,
- template_kwargs=kwargs, verbose=False)
+ reply = generate(model, tok, req.user, system_prompt=req.system,
+ max_tokens=req.max_tokens, temp=req.temperature,
+ presence_penalty=req.presence_penalty,
+ template_kwargs=kwargs, verbose=False)
+ # A template that opens the thinking block or the message
+ # header in the prompt leaves its marker out of the reply, which
+ # the parser reads it by.
+ rendered = tok.apply_chat_template(
+ list(req.messages), tokenize=False, add_generation_prompt=True,
+ **kwargs)
+ return reply_lead(rendered, tokenizer=tok) + reply
try:
- caption = upsample(family, prompt, chat, name=f"--upsampler {name}")
+ caption = upsample(family, prompt, chat, name=f"--upsampler {name}",
+ system=system)
except UpsampleError as e:
raise _Usage(str(e)) from None
finally:
diff --git a/gmlx/commands/image_pull.py b/gmlx/commands/image_pull.py
index dc16bfb76..61c0ee6be 100644
--- a/gmlx/commands/image_pull.py
+++ b/gmlx/commands/image_pull.py
@@ -8,8 +8,13 @@
encoder. The rest downloads after a summary of files, sizes and free disk.
The VAE comes as safetensors from its source repository, is checked
against the recipe's sha256, packed to GGUF and checked again, and the
-source is deleted unless ``--keep-source``. The entry is then added under
-``image_models:`` in the server config.
+source is deleted unless ``--keep-source``. For a family that ships no
+upsampler prompt, the pull also fetches the one its recipe pins, checks
+its sha256 before it lands, and the entry names it as
+``upsample_prompt``. A fetch that fails, or a file of another digest
+already at its path, which stays as it is, leaves a new entry without it,
+and an entry that already names the file keeps it. The entry is then
+added under ``image_models:`` in the server config.
"""
from __future__ import annotations
@@ -23,7 +28,7 @@
PARTS = ("dit", "encoder", "vae")
_ROLE = {"dit": "DiT", "encoder": "text encoder", "mmproj": "vision tower",
- "vae": "VAE"}
+ "vae": "VAE", "prompt": "upsampler prompt"}
@dataclass
@@ -41,6 +46,7 @@ class Part:
chat_id: str | None = None # the configured chat model it is
wanted: bool = True # inside --only
note: str | None = None # a line the plan prints under it
+ optional: bool = False # a failed fetch leaves the pull whole
@property
def download(self) -> bool:
@@ -57,6 +63,7 @@ class Plan:
default: bool = True # the quants of the memory class
blocked: list = field(default_factory=list) # parts whose path is taken
size: str | None = None # the defaults.size the entry gets
+ notes: list = field(default_factory=list) # lines the plan ends with
def part(self, role: str) -> Part | None:
return next((p for p in self.parts if p.role == role), None)
@@ -257,8 +264,30 @@ def of(role):
size=None if size is None else size + cfg_size,
need=None if need is None else need + cfg_size), dest)
- for p in plan.parts:
- role = "encoder" if p.role == "mmproj" else p.role
+ # The upsampler prompt, in the folder of its repository as the other
+ # parts are. A file there of another digest is the user's: it stays as
+ # it is, and the entry does not name it.
+ if r.prompt_repo:
+ d = _dest(r.prompt_repo, base_dir, nest)
+ dest = os.path.join(d, r.prompt_file)
+ if not os.path.lexists(dest):
+ size = _repo_sizes(r.prompt_repo, r.prompt_revision, sizes).get(
+ r.prompt_file)
+ plan.parts.append(Part(
+ "prompt", r.prompt_file, dest, repo=r.prompt_repo,
+ files=(r.prompt_file,), root=d, revision=r.prompt_revision,
+ size=size, need=_remaining(dest, size), optional=True))
+ elif os.path.isfile(dest) and \
+ catalog.sha256_file(dest) == r.prompt_sha256:
+ plan.parts.append(Part("prompt", r.prompt_file, dest,
+ size=os.path.getsize(dest), need=0,
+ optional=True))
+ else:
+ plan.blocked.append(Part("prompt", r.prompt_file, dest,
+ repo=r.prompt_repo, optional=True))
+
+ for p in plan.parts + plan.blocked:
+ role = {"mmproj": "encoder", "prompt": "dit"}.get(p.role, p.role)
p.wanted = role in only
return plan
@@ -270,7 +299,7 @@ def _gb(n: int | None, digits: int = 1) -> str:
def _combined(plan: Plan) -> int | None:
- sizes = [p.size for p in plan.parts]
+ sizes = [p.size for p in plan.parts if p.role != "prompt"]
if any(s is None for s in sizes):
return None
return sum(sizes) # type: ignore[arg-type]
@@ -337,10 +366,17 @@ def _print_plan(plan: Plan, *, ram: int | None, free: int | None,
if plan.size:
print(f" the entry gets defaults.size {plan.size}, the largest image "
f"this DiT makes on a {plan.memory_class} GB Mac")
+ for n in plan.notes:
+ print(f" {n}")
note = f", {r.license_note}" if r.license_note else ""
vae = f"; VAE {r.vae_license}" if r.vae_license else ""
print(f" license: {r.title} {r.license}{note} ({r.license_url}); "
f"encoder {r.encoder_license}{vae}")
+ prompt = plan.part("prompt")
+ if prompt is not None and prompt.download and prompt.wanted:
+ print(f" license: the upsampler prompt {r.prompt_file} of "
+ f"{r.prompt_repo}, {r.prompt_license} "
+ f"({r.prompt_license_url})")
_PACK_ITEMSIZE = {"BF16": 2, "F16": 2, "F32": 4}
@@ -448,7 +484,7 @@ def _cast(plan: Plan) -> int:
def _need(plan: Plan) -> int | None:
total = 0
for p in plan.parts:
- if not (p.download and p.wanted):
+ if not (p.download and p.wanted) or (p.optional and p.need is None):
continue
if p.need is None:
return None
@@ -506,6 +542,67 @@ def _pack(plan: Plan, part: Part, got: list, *, keep_source: bool) -> str:
return part.path
+def _fetch_prompt(plan: Plan, part: Part) -> str:
+ """Download the upsampler prompt into a folder of its own beside its
+ path, check its sha256 there, and move it into place, so a copy of
+ another digest never lands on the path."""
+ import shutil
+ import tempfile
+
+ from gmlx.commands import manage
+
+ r = plan.recipe
+ os.makedirs(os.path.dirname(part.path), exist_ok=True)
+ tmp = tempfile.mkdtemp(prefix=".gmlx-pull-", dir=os.path.dirname(
+ part.path))
+ try:
+ print(f"downloading {part.repo}/{r.prompt_file}", file=sys.stderr)
+ got = manage._hf_download(part.repo, r.prompt_file, part.revision,
+ tmp)
+ sha = catalog.sha256_file(got)
+ if sha != r.prompt_sha256:
+ raise remote.RemoteError(
+ f"the download has sha256 {sha}, and {r.prompt_file} of "
+ f"{r.prompt_repo} at {r.prompt_revision[:12]} has "
+ f"{r.prompt_sha256}")
+ if os.path.lexists(part.path):
+ raise remote.RemoteError(f"{part.path} appeared during the pull")
+ os.replace(got, part.path)
+ finally:
+ shutil.rmtree(tmp, ignore_errors=True)
+ return part.path
+
+
+def _naming(cfg, path: str) -> list:
+ """The image entries whose upsample_prompt is the file at ``path``,
+ on disk or not: a relative one under any of the model_dirs."""
+ if cfg is None:
+ return []
+
+ def expand(p):
+ return os.path.expanduser(os.path.expandvars(p))
+
+ ap = os.path.abspath(path)
+ out = []
+ for iid, m in cfg.image_models.items():
+ if not m.upsample_prompt:
+ continue
+ p = expand(m.upsample_prompt)
+ where = [p] if os.path.isabs(p) else [
+ os.path.join(expand(d), p) for d in cfg.model_dirs]
+ if ap in map(os.path.abspath, where):
+ out.append(iid)
+ return out
+
+
+def _unnamed(cfg, path: str) -> str:
+ """What a missing or unread upsampler prompt leaves of the entry."""
+ ids = _naming(cfg, path)
+ if ids:
+ return f"{', '.join(ids)} still names {path}"
+ return "the entry does not name it"
+
+
def _entry_id(r: catalog.Recipe, plan: Plan, cfg) -> str:
taken = set(cfg.models) | set(cfg.aliases) | set(cfg.assistants)
if r.name not in taken and r.name not in cfg.image_models:
@@ -532,6 +629,9 @@ def _entry(plan: Plan, cfg) -> dict:
if mm is not None and not enc.chat_id:
out["encoder_mmproj"] = _rel(mm.path, dirs)
out["vae"] = _rel(plan.part("vae").path, dirs)
+ prompt = plan.part("prompt")
+ if prompt is not None and not prompt.download:
+ out["upsample_prompt"] = _rel(prompt.path, dirs)
if plan.size:
out["defaults"] = {"size": plan.size}
return out
@@ -562,6 +662,8 @@ def register(plan: Plan, cfg_path: str) -> str | None:
roots = [os.path.abspath(os.path.expanduser(os.path.expandvars(d)))
for d in cfg.model_dirs]
for p in plan.parts:
+ if p.optional and p.download:
+ continue
ap = os.path.abspath(p.path)
if p.chat_id is None and not any(
ap == r or ap.startswith(r + os.sep) for r in roots):
@@ -571,7 +673,7 @@ def register(plan: Plan, cfg_path: str) -> str | None:
old = _registered(plan, cfg)
if old is not None:
print(f"{plan.recipe.name} is already configured as {old}")
- return None
+ return _name_prompt(plan, cfg, cfg_path, old)
iid = _entry_id(plan.recipe, plan, cfg)
entry = _entry(plan, cfg)
set_default = not cfg.defaults.image_model
@@ -601,6 +703,26 @@ def mutate(doc):
return iid
+def _name_prompt(plan: Plan, cfg, cfg_path: str, iid: str) -> str | None:
+ """Name the fetched upsampler prompt on the entry ``iid`` that has
+ none. Returns ``iid`` when it did."""
+ from gmlx.config import edit_config_yaml
+ from gmlx.load.discovery import _rel
+
+ prompt = plan.part("prompt")
+ if prompt is None or prompt.download or \
+ cfg.image_models[iid].upsample_prompt:
+ return None
+ rel = _rel(prompt.path, cfg.model_dirs)
+
+ def mutate(doc):
+ doc["image_models"][iid]["upsample_prompt"] = rel
+
+ edit_config_yaml(cfg_path, mutate)
+ print(f"set upsample_prompt of {iid} to {rel} in {cfg_path}")
+ return iid
+
+
def _config_path(a) -> str | None:
import gmlx.config as cfgmod
@@ -698,6 +820,7 @@ def pull_recipe(a, r: catalog.Recipe) -> int:
"license_note": r.license_note or None,
"encoder_license": r.encoder_license,
"vae_license": r.vae_license or None,
+ "prompt_license": r.prompt_license or None,
"parts": [{"role": p.role, "name": p.name, "path": p.path,
"size_bytes": p.size, "download": p.download
and p.wanted, "on_disk": not p.download,
@@ -709,7 +832,12 @@ def pull_recipe(a, r: catalog.Recipe) -> int:
# The plan goes to stdout and the download lines to stderr: a pipe or a
# file gets the plan and its license line first.
sys.stdout.flush()
- blocked = [p for p in plan.blocked if p.wanted]
+ for p in plan.blocked:
+ if p.optional and p.wanted:
+ print(f"warning: {p.path} is not the {_ROLE[p.role]} of {p.repo} "
+ f"at {r.prompt_revision[:8]}, so {_unnamed(cfg, p.path)}; "
+ "move or delete it, and pull again", file=sys.stderr)
+ blocked = [p for p in plan.blocked if p.wanted and not p.optional]
if blocked:
for p in blocked:
print(f"error: {p.path} is not the {r.title} {_ROLE[p.role]}, "
@@ -730,7 +858,10 @@ def pull_recipe(a, r: catalog.Recipe) -> int:
if not (part.download and part.wanted):
continue
try:
- got = _fetch(plan, part)
+ if part.role == "prompt":
+ _fetch_prompt(plan, part)
+ else:
+ got = _fetch(plan, part)
if part.role == "vae":
_pack(plan, part, got, keep_source=a.keep_source)
part.repo = None
@@ -738,11 +869,15 @@ def pull_recipe(a, r: catalog.Recipe) -> int:
if part.size is None or part.role != "vae":
part.size = os.path.getsize(part.path)
except (remote.RemoteError, OSError, ValueError) as e:
+ if part.optional:
+ print(f"warning: the {_ROLE[part.role]} was not fetched, so "
+ f"{_unnamed(cfg, part.path)}: {e}", file=sys.stderr)
+ continue
print(f"error: {part.role} failed: {e}", file=sys.stderr)
return 1
print(f" {part.role}: {part.path}", flush=True)
- missing = [p.role for p in plan.parts if p.download]
+ missing = [p.role for p in plan.parts if p.download and not p.optional]
if missing:
print(f"not registered: {', '.join(missing)} not on disk yet; "
f"`gmlx pull {r.name}` fetches the rest", file=sys.stderr)
diff --git a/gmlx/config.py b/gmlx/config.py
index fa578ab70..9e5b915a2 100644
--- a/gmlx/config.py
+++ b/gmlx/config.py
@@ -149,7 +149,8 @@ def default_apc_disk_path() -> str:
_IMAGE_MODEL_KEYS = frozenset({"path", "family", "variant", "encoder",
"encoder_mmproj", "vae", "adapters",
"defaults", "encoder_release", "pin",
- "ttl_s", "upsampler", "upsample"})
+ "ttl_s", "upsampler", "upsample",
+ "upsample_prompt"})
_IMAGE_DEFAULTS_KEYS = frozenset({"size", "steps", "guidance", "sampler",
"negative_prompt"})
_IMAGE_ADAPTER_KEYS = frozenset({"path", "strength"})
@@ -617,7 +618,8 @@ class ImageModelCfg:
every request, and ``defaults`` the request defaults the entry sets
over its family's. ``upsampler`` is the chat model id that rewrites
the prompt of a family that takes prompt upsampling, by default while
- ``upsample`` is on."""
+ ``upsample`` is on, and ``upsample_prompt`` a text file whose system
+ prompt it takes in place of the family's."""
id: str
path: str
encoder: str
@@ -632,6 +634,7 @@ class ImageModelCfg:
variant: str | None = None # None => from the weights
upsampler: str | None = None
upsample: bool = True
+ upsample_prompt: str | None = None
@dataclass(frozen=True)
@@ -1219,7 +1222,10 @@ class ResolvedImageModel:
the encoder GGUF, or None when ``encoder_model`` names the chat model
the entry shares; ``encoder_mmproj`` is the vision tower that comes
with either. ``upsampler`` is the chat model id of the entry's prompt
- upsampling, which runs by default while ``upsample`` is on."""
+ upsampling, which runs by default while ``upsample`` is on, and
+ ``upsample_prompt`` the file of its system prompt, resolved where it
+ exists and as written where it does not, so a missing file fails the
+ rewrite and not the entry."""
id: str
path: str
vae: str
@@ -1235,6 +1241,7 @@ class ResolvedImageModel:
variant: str | None = None
upsampler: str | None = None
upsample: bool = True
+ upsample_prompt: str | None = None
def load_signature(self) -> tuple:
"""What makes two loads of the entry the same: every file, the
@@ -2086,7 +2093,19 @@ def resolve_image_model(image_id: str, cfg: ServerCfg) -> ResolvedImageModel:
for a, st in m.adapters),
defaults=dict(m.defaults), encoder_release=m.encoder_release,
pin=bool(m.pin), ttl_s=ttl_s, variant=m.variant,
- upsampler=m.upsampler, upsample=m.upsample)
+ upsampler=m.upsampler, upsample=m.upsample,
+ upsample_prompt=_prompt_path(m.upsample_prompt, dirs))
+
+
+def _prompt_path(p: str | None, dirs: list[str]) -> str | None:
+ """An entry's ``upsample_prompt``: the file under ``model_dirs`` as a
+ model path resolves, or the expanded path when there is none."""
+ if p is None:
+ return None
+ try:
+ return resolve_path(p, dirs)
+ except ConfigError:
+ return os.path.expanduser(os.path.expandvars(p))
def env_for(resolved: ResolvedModel) -> dict[str, str]:
@@ -2733,7 +2752,8 @@ def _parse_image_model(image_id: str, raw) -> ImageModelCfg:
if not isinstance(raw.get(key), str) or not raw[key]:
raise ConfigError(f"{where} needs `{key}` (keys present: "
f"{sorted(raw)})")
- for key in ("family", "variant", "encoder_mmproj", "upsampler"):
+ for key in ("family", "variant", "encoder_mmproj", "upsampler",
+ "upsample_prompt"):
if raw.get(key) is not None and not isinstance(raw[key], str):
raise ConfigError(f"{where}.{key}: expected a string, got "
f"{raw[key]!r}")
@@ -2757,6 +2777,7 @@ def _parse_image_model(image_id: str, raw) -> ImageModelCfg:
raw.get("defaults")),
encoder_release=release, pin=bool(raw.get("pin", False)),
upsampler=raw.get("upsampler") or None, upsample=upsample,
+ upsample_prompt=raw.get("upsample_prompt") or None,
# As in `models`: an explicit null never unloads.
ttl_s=(0.0 if "ttl_s" in raw and raw["ttl_s"] is None
else _coerce_num("ttl_s", raw.get("ttl_s"), float,
diff --git a/gmlx/distill/gen.py b/gmlx/distill/gen.py
index 31331ae04..1275a7585 100644
--- a/gmlx/distill/gen.py
+++ b/gmlx/distill/gen.py
@@ -846,6 +846,7 @@ def work(i, r):
ex = ThreadPoolExecutor(max_workers=opts.concurrency)
queue = iter(todo)
pending: set = set()
+ order: dict = {}
with open(out, "a", encoding="utf-8") as ofh:
try:
# twice the concurrency is in flight at a time, refilled as
@@ -856,11 +857,15 @@ def work(i, r):
item = next(queue, None)
if item is None:
break
- pending.add(ex.submit(work, *item))
+ fut = ex.submit(work, *item)
+ order[fut] = item[0]
+ pending.add(fut)
if not pending:
break
finished, pending = wait(pending, return_when=FIRST_COMPLETED)
- for fut in finished:
+ # replies that land in one wake are written in prompt
+ # order, so an interrupt keeps every earlier reply
+ for fut in sorted(finished, key=order.pop):
try:
i, r, c = fut.result()
except Exception as e: # noqa: BLE001 - one bad request never ends the run
diff --git a/gmlx/gen/generation.py b/gmlx/gen/generation.py
index 67bb2fe67..14865feef 100644
--- a/gmlx/gen/generation.py
+++ b/gmlx/gen/generation.py
@@ -753,8 +753,14 @@ def _generate_over(
suppressed_eos,
)
+ from .penalties import OutputPenalty
+
if over_logits_processors is None:
over_logits_processors = logits_processors
+ # Phase 2's context starts at its bridge, so its penalties count afresh.
+ over_logits_processors = [
+ p.clone() if isinstance(p, OutputPenalty) else p
+ for p in over_logits_processors or ()]
real_eos = set(getattr(tokenizer, "eos_token_ids", None) or [])
cache = make_prompt_cache(model, max_kv_size=max_kv_size)
diff --git a/gmlx/gen/penalties.py b/gmlx/gen/penalties.py
new file mode 100644
index 000000000..61a4e93e6
--- /dev/null
+++ b/gmlx/gen/penalties.py
@@ -0,0 +1,153 @@
+"""Presence and frequency penalties over the generated tokens.
+
+The OpenAI ``presence_penalty`` subtracts the penalty once from every token
+the response holds, and ``frequency_penalty`` subtracts it once per
+occurrence. vLLM and SGLang count every token generated so far and never
+the prompt, and model cards tune their values for that. The copied mlx-lm
+and mlx-vlm samplers return :class:`OutputPenalty` for both.
+
+A processor takes the decode loop's ``(tokens, logits)``: ``tokens`` is
+the row's history, which ends with the last sampled token. ``context_size``
+None or 0 counts the whole output in a count per vocab entry. A positive
+``context_size`` counts only the last N generated tokens.
+
+A caller that knows what it committed says so. ``anchor(n)`` marks the
+history before index n as accounted for, and ``commit(ids)`` counts the
+ids generated after it. A call then counts the history past the anchor for
+that call only, as the speculative verify rounds need for their drafts.
+``resumed(ids)`` gives a fresh processor for a row replayed as its prompt
+plus ``ids``.
+
+A caller that only passes a history gets a rule from history lengths. The
+history at the first call is context, since that call scores the first
+generated token. A longer history adds the tokens past the last one. A
+history no longer than the last one is a rollback of drafts: the history
+before its last token is unchanged, and the tokens counted after that
+point are taken back out.
+
+Counts are indexed by token id and sized to the widest logits seen, so a
+drafter with a narrower vocab than its target shares the processor.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+
+class OutputPenalty:
+ """A presence penalty, or a frequency penalty with ``frequency=True``,
+ over one row's generated tokens. ``clone`` gives a fresh processor for
+ another row."""
+
+ def __init__(self, penalty: float, context_size: int | None = None, *,
+ frequency: bool = False):
+ if context_size is not None and context_size < 0:
+ raise ValueError(
+ f"context_size must be None or at least 0, got {context_size}")
+ self.penalty = penalty
+ self.context_size = context_size or None
+ self.frequency = frequency
+ self.reset()
+
+ def reset(self) -> None:
+ self._anchor = None # history length accounted for, when told
+ self._base = None # otherwise the history length at the first call
+ self._seen = 0 # and the history length at the last call
+ self._hist = None # and the last call's history
+ self._counts = None # [width] int32 counts of the generated tokens
+ self._width = 0
+ self._tail = [] # window: last generated ids outside the history
+
+ def clone(self) -> OutputPenalty:
+ return OutputPenalty(self.penalty, self.context_size,
+ frequency=self.frequency)
+
+ def resumed(self, generated) -> OutputPenalty:
+ """A fresh processor for a row replayed as its prompt followed by
+ ``generated``. Its first call scores the token after the replay,
+ so that history is all context, and ``generated`` is counted."""
+ out = self.clone()
+ out._take([int(t) for t in generated])
+ return out
+
+ def anchor(self, n: int) -> None:
+ """The history before index ``n`` is accounted for. Calls count the
+ history past it for that call only, until ``commit``."""
+ self._anchor = n
+
+ def commit(self, ids) -> None:
+ """Count ``ids``, generated right after the anchor, and move the
+ anchor past them."""
+ if self._anchor is None:
+ raise RuntimeError("commit before anchor")
+ ids = [int(t) for t in ids]
+ self._take(ids)
+ self._anchor += len(ids)
+
+ def _take(self, ids) -> None:
+ if not ids:
+ return
+ if self.context_size is not None:
+ self._tail = (self._tail + ids)[-self.context_size:]
+ return
+ self._fit(max(ids) + 1)
+ self._counts = self._counts.at[mx.array(ids, dtype=mx.int32)].add(1)
+
+ def _fit(self, width: int) -> None:
+ if width <= self._width:
+ return
+ grow = mx.zeros((width - self._width,), dtype=mx.int32)
+ self._counts = (grow if self._counts is None
+ else mx.concatenate([self._counts, grow]))
+ self._width = width
+
+ def __call__(self, tokens, logits: mx.array) -> mx.array:
+ if not isinstance(tokens, mx.array):
+ tokens = mx.array(tokens, dtype=mx.int32)
+ tokens = tokens.reshape(-1)
+ n = tokens.shape[0]
+ self._fit(logits.shape[-1])
+ if self._anchor is not None:
+ start = self._anchor
+ else:
+ if self._base is None:
+ self._base = n
+ start = self._base
+ if self.context_size is not None:
+ return self._apply(logits, self._window(tokens, n, start))
+ if self._anchor is None:
+ self._roll(tokens, n)
+ return self._apply(logits, self._counts)
+ counts = self._counts
+ if n > start:
+ counts = counts.at[tokens[start:n]].add(1)
+ return self._apply(logits, counts)
+
+ def _roll(self, tokens, n) -> None:
+ lo = max(self._base, self._seen if n > self._seen else n - 1)
+ counts = self._counts
+ if lo < self._seen:
+ counts = counts.at[self._hist[lo:self._seen]].add(-1)
+ if lo < n:
+ counts = counts.at[tokens[lo:n]].add(1)
+ self._counts, self._seen, self._hist = counts, n, tokens
+
+ def _window(self, tokens, n, start):
+ size = self.context_size
+ t = max(n - start, 0)
+ parts = []
+ if t < size and self._tail:
+ parts.append(mx.array(self._tail[-(size - t):], dtype=mx.int32))
+ if t:
+ parts.append(tokens[max(start, n - size):n].astype(mx.int32))
+ if not parts:
+ return None
+ recent = parts[0] if len(parts) == 1 else mx.concatenate(parts)
+ return mx.zeros((self._width,), dtype=mx.int32).at[recent].add(1)
+
+ def _apply(self, logits, counts):
+ if counts is None:
+ return logits
+ counts = counts[:logits.shape[-1]]
+ hits = counts if self.frequency else counts > 0
+ return logits - hits.astype(logits.dtype) * self.penalty
diff --git a/gmlx/image/catalog.py b/gmlx/image/catalog.py
index c0918b17e..169746f92 100644
--- a/gmlx/image/catalog.py
+++ b/gmlx/image/catalog.py
@@ -1,7 +1,8 @@
"""The image models gmlx can pull and the image files a model folder holds.
``RECIPES`` names each pullable image model: its DiT repository and quants,
-its text encoder (and vision tower), its VAE source and the licenses.
+its text encoder (and vision tower), its VAE source, the upsampler's
+system prompt of a family that ships none, and the licenses.
``DEFAULT_QUANTS`` gives the quants a recipe pulls on each memory class,
and ``NO_IMAGE`` the classes where none of them makes an image.
``classify`` and ``scan`` find the image files on disk by their weights:
@@ -51,6 +52,14 @@ class Recipe:
license_note: str = "" # what the license allows, if limited
vae_license: str = "" # the VAE's license, when it has its
# own
+ # The upsampler's system prompt of a family that ships none, which the
+ # pull fetches when it can and the entry names as upsample_prompt.
+ prompt_repo: str = ""
+ prompt_revision: str = ""
+ prompt_file: str = ""
+ prompt_sha256: str = ""
+ prompt_license: str = ""
+ prompt_license_url: str = ""
def dit_name(self, quant: str) -> str:
return self.dit_file.format(quant=quant)
@@ -131,6 +140,18 @@ def encoder_name(self, quant: str) -> str:
"that it is made with AI where the law requires one"),
encoder_license="apache-2.0", encoder_bytes=4_280_406_816)
+# Qwen's text to image rewriter prompt, which gmlx does not ship: a Qwen
+# Research License file of 10045 bytes in the repository of its rewriter
+# model.
+_QWEN_PROMPT = dict(
+ prompt_repo="Qwen/Qwen-Image-2.1-PE-T2I",
+ prompt_revision="f3ed7985c788ad75b3ab7223e0c4c51e2a43545b",
+ prompt_file="system_prompt.txt",
+ prompt_sha256="a77c9a06c59b120741141d9514b95682bb8761d02bec49ca61def7b2b3d9fb99",
+ prompt_license="the Qwen Research License, non-commercial use only",
+ prompt_license_url="https://huggingface.co/Qwen/Qwen-Image-2.1-PE-T2I/"
+ "blob/f3ed7985c788ad75b3ab7223e0c4c51e2a43545b/LICENSE")
+
RECIPES = {r.name: r for r in (
Recipe(
name="z-image-turbo", title="Z-Image-Turbo", family="z_image",
@@ -168,7 +189,7 @@ def encoder_name(self, quant: str) -> str:
license_note="non-commercial use only",
encoder_license="apache-2.0",
encoder_bytes=8_709_520_224 + 1_162_569_280, vae_bytes=1_350_985_376,
- aliases=("qwen-image-2.1",)),
+ aliases=("qwen-image-2.1",), **_QWEN_PROMPT),
Recipe(
name="qwen-image-turbo", title="Qwen-Image 2.1 Turbo",
family="qwen_image21", variant="turbo",
@@ -189,7 +210,7 @@ def encoder_name(self, quant: str) -> str:
license_note="non-commercial use only",
encoder_license="apache-2.0",
encoder_bytes=8_709_520_224 + 1_162_569_280, vae_bytes=1_350_985_376,
- aliases=("qwen-image-2.1-turbo",)),
+ aliases=("qwen-image-2.1-turbo",), **_QWEN_PROMPT),
Recipe(
name="krea-2-turbo", title="Krea 2 Turbo", family="krea2",
variant="turbo",
diff --git a/gmlx/image/families/flux2_klein/family.py b/gmlx/image/families/flux2_klein/family.py
index a9dce5c21..4a11ee399 100644
--- a/gmlx/image/families/flux2_klein/family.py
+++ b/gmlx/image/families/flux2_klein/family.py
@@ -6,8 +6,10 @@
# and shapes of Flux2 in src/flux2/model.py, the prompt rules of
# Qwen3Embedder in src/flux2/text_encoder.py, the rounding of denoise
# and denoise_cfg and the reference image rules of encode_image_refs and
-# cap_min_pixels in src/flux2/sampling.py, and modified for gmlx
-# (Apache-2.0; see licenses/flux2-LICENSE).
+# cap_min_pixels in src/flux2/sampling.py, the text to image system
+# message of src/flux2/system_messages.py and its chat turn in
+# scripts/cli.py, and modified for gmlx (Apache-2.0; see
+# licenses/flux2-LICENSE).
"""FLUX.2 klein fingerprint, constants, config synthesis and the hooks of
its registry row.
@@ -28,7 +30,8 @@
from gmlx.image.fingerprint import HeaderView, Match, Refusal
from gmlx.image.registry import (ConstantsRow, EditParts, Family,
- LoopRules, SamplingRow, VariantRow)
+ LoopRules, SamplingRow, UpsampleParts,
+ VariantRow)
TITLE = "FLUX.2 klein"
# What a refusal of another FLUX model names as the ones gmlx image runs.
@@ -387,6 +390,29 @@ def synth_config(view: HeaderView) -> dict:
)
+# The optional upsampler of scripts/cli.py for a prompt without
+# references, on a chat model of the user's choice: the text to image
+# system message SYSTEM_MESSAGE_UPSAMPLING_T2I of src/flux2/system_messages.py,
+# shipped as its string, with the typed prompt as the user turn. BFL
+# samples at the defaults of the model it calls; gmlx decodes greedily,
+# so the same prompt writes the same caption, with at most the 1024
+# tokens of BFL's FLUX.2 klein Space and thinking off, as the message
+# asks for the revised prompt and nothing else. The 9B rows share it.
+UPSAMPLE = UpsampleParts(
+ prompts="system_message_upsampling_t2i.txt",
+ sha256="0ad9a42c565716b0a8eb82a442dd4f6cc82d583eb3474f82ccb610c1a29233c7",
+ key=None,
+ max_tokens=1024,
+ temperature=0.0,
+ presence_penalty=0.0,
+ thinking="off",
+ repo="github.com/black-forest-labs/flux2",
+ revision="50fe5162777813d869182b139e83b10743caef15",
+ package=__package__,
+ suggested="Qwen3.5-9B",
+)
+
+
def _prepare_all(images):
from .edit import prepare
@@ -427,4 +453,5 @@ def _size(refs):
loop=LOOP,
edit=EditParts(max_refs=REFERENCES.values["max_refs"],
prepare_all=_prepare_all, size=_size),
+ upsample=UPSAMPLE,
)
diff --git a/gmlx/image/families/flux2_klein/system_message_upsampling_t2i.txt b/gmlx/image/families/flux2_klein/system_message_upsampling_t2i.txt
new file mode 100644
index 000000000..42587ff13
--- /dev/null
+++ b/gmlx/image/families/flux2_klein/system_message_upsampling_t2i.txt
@@ -0,0 +1,8 @@
+You are an expert prompt engineer for FLUX.2 by Black Forest Labs. Rewrite user prompts to be more descriptive while strictly preserving their core subject and intent.
+
+Guidelines:
+1. Structure: Keep structured inputs structured (enhance within fields). Convert natural language to detailed paragraphs.
+2. Details: Add concrete visual specifics - form, scale, textures, materials, lighting (quality, direction, color), shadows, spatial relationships, and environmental context.
+3. Text in Images: Put ALL text in quotation marks, matching the prompt's language. Always provide explicit quoted text for objects that would contain text in reality (signs, labels, screens, etc.) - without it, the model generates gibberish.
+
+Output only the revised prompt and nothing else.
\ No newline at end of file
diff --git a/gmlx/image/families/flux2_klein_9b/family.py b/gmlx/image/families/flux2_klein_9b/family.py
index 182c77730..4ceccf82f 100644
--- a/gmlx/image/families/flux2_klein_9b/family.py
+++ b/gmlx/image/families/flux2_klein_9b/family.py
@@ -153,4 +153,5 @@ def lora_map(config: dict):
schedule=klein.schedule,
loop=klein.LOOP,
edit=klein.FAMILY.edit,
+ upsample=klein.UPSAMPLE,
)
diff --git a/gmlx/image/families/flux2_klein_9b_kv/family.py b/gmlx/image/families/flux2_klein_9b_kv/family.py
index fa388cbf5..c2cad69d7 100644
--- a/gmlx/image/families/flux2_klein_9b_kv/family.py
+++ b/gmlx/image/families/flux2_klein_9b_kv/family.py
@@ -115,4 +115,5 @@ def lora_map(config: dict):
schedule=klein.schedule,
loop=klein.LOOP,
edit=klein.FAMILY.edit,
+ upsample=klein.UPSAMPLE,
)
diff --git a/gmlx/image/families/flux3_klein/family.py b/gmlx/image/families/flux3_klein/family.py
index 22658aedf..a0cdb2fcf 100644
--- a/gmlx/image/families/flux3_klein/family.py
+++ b/gmlx/image/families/flux3_klein/family.py
@@ -342,20 +342,23 @@ def schedule(steps: int, *, tokens: int | None, variant: str | None = None):
# The upsampler of src/flux3/pup.py for a prompt without references, run
# on a chat model of the user's choice: the t2i prompt that BFL publishes
# for other models, shipped unmodified, with the typed prompt as the user
-# turn, greedy decoding of at most 3584 tokens and thinking off. pup.py
-# sets a presence penalty of 1.5 for BFL's own fine-tune and its own
-# prompts; a stock chat model runs this prompt without one.
+# turn, greedy decoding of at most 3584 tokens and thinking off. BFL
+# gives no sampling settings for this prompt, so the presence penalty is
+# the 1.5 that pup.py sets for its own fine-tune on text to image.
UPSAMPLE = UpsampleParts(
prompts="pup_prompts_external_vlm.json",
sha256="361bb6e081d50e54a3dee9abf26327436a9386c83465f881c7848989232028be",
key="t2i",
max_tokens=3584,
temperature=0.0,
- presence_penalty=0.0,
+ presence_penalty=1.5,
thinking="off",
repo="github.com/black-forest-labs/flux-3-image-klein",
revision="bfc8889a4497f846416f253a3972a1941ac698b9",
+ package=__package__,
suggested="Qwen3.5-9B",
+ # BFL's README says not to run the model without it.
+ expected=True,
)
diff --git a/gmlx/image/families/krea2/expansion.txt b/gmlx/image/families/krea2/expansion.txt
new file mode 100644
index 000000000..2868d1fcb
--- /dev/null
+++ b/gmlx/image/families/krea2/expansion.txt
@@ -0,0 +1,19 @@
+You are an expert prompt engineer for text-to-image models. Your task is to expand the user's prompt into a highly effective image-generation prompt.
+
+Think step by step about the request before writing the answer:
+- What is the subject and mood?
+- What visual styles, mediums, and lighting options would fit? Consider two or three alternatives and pick the one that best serves the caption.
+- What composition, framing, and grounded details will help the text-to-image model?
+
+Then output a single expanded prompt paragraph.
+
+Follow these rules strictly:
+1. **Faithfulness First:** Preserve all original subjects, actions, colors, and spatial relationships. Do not add new objects, props, characters, or animals unless the user clearly implies them.
+2. **Practical T2I Structure:** Write a prompt that a text-to-image model can parse cleanly. Group subjects with their own attributes and actions. Use grounded phrasing for poses, interactions, and spatial layout.
+3. **Style Planning Stays Internal:** Use your internal reasoning to choose style, medium, framing, and lighting. Do not emit planning tags or wrappers in the visible answer body.
+4. **Text Rendering:** If the user requests visible text, quotes, labels, or typography, specify the exact text clearly and wrap requested words in quotes.
+5. **Avoid Over-Specification:** Do not invent highly specific clothing, colors, materials, or scene details unless the input supports them.
+6. **Structure:** Write one cohesive paragraph after the thinking block. No bullets, JSON, or markdown.
+7. **Respect Existing Detail:** If the user's prompt is already detailed, lightly polish and finalize rather than heavily expanding — preserve their phrasing and direction.
+8. **Respect the Human Form:** Treat depictions of people with dignity. Assume clothing covers genitals and intimate anatomy.
+9. **Preserve User Medium:** When the user explicitly requests a medium (e.g. "photo of", "photograph of", "illustration of", "painting of", "sketch of", "3D render of"), honor it. Do not pivot to a different medium to avoid difficulty — match the user's stated intent.
diff --git a/gmlx/image/families/krea2/family.py b/gmlx/image/families/krea2/family.py
index 6692f5fb6..dd22b1582 100644
--- a/gmlx/image/families/krea2/family.py
+++ b/gmlx/image/families/krea2/family.py
@@ -22,11 +22,14 @@
# variant and the scheduler values are those of diffusers'
# src/diffusers/pipelines/krea2/pipeline_krea2.py and
# src/diffusers/modular_pipelines/krea2/ at commit e0abab83 (Apache-2.0; see
-# licenses/diffusers-LICENSE).
+# licenses/diffusers-LICENSE). The upsampler's system prompt is
+# docs/expansion.txt of github.com/krea-ai/krea-2 at commit db3984f, with
+# its use in docs/prompting.md (Apache-2.0; see licenses/krea-2-LICENSE).
+from gmlx.image import upsample
from gmlx.image.fingerprint import HeaderView, Match, Refusal
from gmlx.image.registry import (ConstantsRow, Family, SamplingRow,
- VariantRow)
+ UpsampleParts, VariantRow)
_B0 = "blocks.0."
_T0 = "txtfusion.layerwise_blocks.0."
@@ -281,6 +284,31 @@ def schedule(steps: int, *, tokens: int | None, variant: str | None = None):
return krea2_sigmas(steps, CONSTANTS.values["turbo_mu"])
+# The optional upsampler of docs/prompting.md, on a chat model of the
+# user's choice: docs/expansion.txt as the system prompt, shipped
+# unmodified, and the typed prompt as the user turn. The prompt asks for
+# a thinking block, then one paragraph, and the caption is the last
+# paragraph of the reply. Krea names no sampling; gmlx decodes greedily,
+# so the same prompt writes the same caption. Thinking is off: Qwen3.5-9B
+# writes a paragraph of the same kind in a few seconds, where with
+# thinking on it thinks for up to 5000 tokens first. A model that thinks
+# anyway has its thinking block skipped.
+UPSAMPLE = UpsampleParts(
+ prompts="expansion.txt",
+ sha256="2a5c24a3a83e9d415d4679a4e1853ddf056bd0329f3af315b27ed72b6ad9c2f0",
+ key=None,
+ max_tokens=2048,
+ temperature=0.0,
+ presence_penalty=0.0,
+ thinking="off",
+ repo="github.com/krea-ai/krea-2",
+ revision="db3984fbc6e13b34c0064990fc2d95ac64d00058",
+ package=__package__,
+ suggested="Qwen3.5-9B",
+ parse=upsample.paragraph,
+)
+
+
FAMILY = Family(
name="krea2",
title="Krea 2",
@@ -307,4 +335,5 @@ def schedule(steps: int, *, tokens: int | None, variant: str | None = None):
encoder=encoder,
schedule=schedule,
reshaped=(PROJECTOR,),
+ upsample=UPSAMPLE,
)
diff --git a/gmlx/image/families/qwen_image21/family.py b/gmlx/image/families/qwen_image21/family.py
index 043c97312..64eb8190f 100644
--- a/gmlx/image/families/qwen_image21/family.py
+++ b/gmlx/image/families/qwen_image21/family.py
@@ -18,9 +18,11 @@
# src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py at e0abab83
# (Apache-2.0; see licenses/diffusers-LICENSE).
+from gmlx.image import upsample
from gmlx.image.fingerprint import HeaderView, Match, Refusal
from gmlx.image.registry import (ConstantsRow, EditParts, Family,
- SamplingRow, StepParts, VariantRow)
+ SamplingRow, StepParts, UpsampleParts,
+ VariantRow)
_B0 = "transformer_blocks.0."
@@ -442,6 +444,41 @@ def schedule(steps: int, *, tokens: int | None, variant: str | None = None):
max_shift=c["max_shift"], terminal=c["shift_terminal"])
+# Where to get the rewriter prompt Qwen publishes, which is under a
+# license gmlx does not ship.
+PROMPT_REPO = "Qwen/Qwen-Image-2.1-PE-T2I"
+PROMPT_FILE = "system_prompt.txt"
+
+# The optional upsampler, on the system prompt of a file the user gives:
+# gmlx ships none. Qwen's rewriter thinks, then answers with a JSON
+# object whose rewritten_prompt is the caption, beside an aspect ratio
+# that gmlx does not read, so the image keeps the size the request asks
+# for. Qwen's rewriter samples at temperature 1 with thinking on and a
+# presence penalty of 1.5. The turn keeps the penalty and decodes
+# greedily, so the same prompt writes the same caption. Thinking is off:
+# on Qwen's prompt a stock Qwen3.5-9B thinks past 8192 tokens, and
+# without thinking it writes the JSON in a few hundred. A model that
+# thinks anyway has its thinking block skipped.
+UPSAMPLE = UpsampleParts(
+ prompts=None,
+ sha256="",
+ key=None,
+ max_tokens=2048,
+ temperature=0.0,
+ presence_penalty=1.5,
+ thinking="off",
+ repo=f"huggingface.co/{PROMPT_REPO}",
+ revision="f3ed7985c788ad75b3ab7223e0c4c51e2a43545b",
+ suggested="Qwen3.5-9B",
+ parse=upsample.json_field("rewritten_prompt"),
+ needs=("Qwen-Image 2.1 ships no upsampler prompt: set upsample_prompt "
+ "on its image_models entry, or pass --upsample-prompt, to a text "
+ f"file. Qwen publishes one as {PROMPT_FILE} in the Hugging Face "
+ f"repository {PROMPT_REPO}, under the Qwen Research License, for "
+ "non-commercial use only, and `gmlx pull qwen-image` fetches it"),
+)
+
+
FAMILY = Family(
name="qwen_image21",
title="Qwen-Image 2.1",
@@ -468,4 +505,5 @@ def schedule(steps: int, *, tokens: int | None, variant: str | None = None):
min_area=REFERENCES.values["min_area"],
prepare=_prepare, token_ids=_edit_ids,
encode=_edit_encode, template_id=_edit_template_id),
+ upsample=UPSAMPLE,
)
diff --git a/gmlx/image/families/z_image/family.py b/gmlx/image/families/z_image/family.py
index 2a953dbda..75ccafd0d 100644
--- a/gmlx/image/families/z_image/family.py
+++ b/gmlx/image/families/z_image/family.py
@@ -6,8 +6,15 @@
from __future__ import annotations
+# The upsampler's system prompt is prompt_template of pe.py in the
+# Hugging Face Space Tongyi-MAI/Z-Image-Turbo at revision 768cb50, with its
+# use in app.py of that Space. pe.py carries no license header; the
+# Space's card gives apache-2.0 (see licenses/z-image-turbo-space-LICENSE).
+
+from gmlx.image import upsample
from gmlx.image.fingerprint import HeaderView, Match, Refusal
-from gmlx.image.registry import ConstantsRow, Family, SamplingRow, VariantRow
+from gmlx.image.registry import (ConstantsRow, Family, SamplingRow,
+ UpsampleParts, VariantRow)
_REQUIRED = (
("cap_embedder.1.weight", (2560, 3840)),
@@ -165,6 +172,31 @@ def lora_map(config: dict):
return build_map(config)
+# The optional prompt enhancer of the Space's app.py, on a chat model of
+# the user's choice: the Chinese system prompt of pe.py, shipped as its
+# string, with the typed prompt in its {prompt} slot and a single space as
+# the user turn, as app.py sends them. The reply is the caption, or the
+# revised_prompt of a JSON block in it, as app.py reads it. The Space
+# samples its hosted model at temperature 0.7; gmlx decodes greedily, so
+# the same prompt writes the same caption, with thinking off, since the
+# prompt asks for the final description and nothing else.
+UPSAMPLE = UpsampleParts(
+ prompts="pe_prompt_template.txt",
+ sha256="36bdd888b953e9b7a8237dab55bee145a44e96e92e8f6e81a6e4d433a401a0b6",
+ key=None,
+ max_tokens=2048,
+ temperature=0.0,
+ presence_penalty=0.0,
+ thinking="off",
+ repo="huggingface.co/spaces/Tongyi-MAI/Z-Image-Turbo",
+ revision="768cb50d847cdbba97c89533ae976be69cf5a5b8",
+ package=__package__,
+ suggested="Qwen3.5-9B",
+ parse=upsample.field_or_plain("revised_prompt"),
+ slot="{prompt}",
+)
+
+
FAMILY = Family(
name="z_image",
title="Z-Image",
@@ -180,4 +212,5 @@ def lora_map(config: dict):
variants=VARIANTS,
sampling=SAMPLING,
lora_map=lora_map,
+ upsample=UPSAMPLE,
)
diff --git a/gmlx/image/families/z_image/pe_prompt_template.txt b/gmlx/image/families/z_image/pe_prompt_template.txt
new file mode 100644
index 000000000..ba1e0044c
--- /dev/null
+++ b/gmlx/image/families/z_image/pe_prompt_template.txt
@@ -0,0 +1,18 @@
+
+你是一位被关在逻辑牢笼里的幻视艺术家。你满脑子都是诗和远方，但双手却不受控制地只想将用户的提示词，转化为一段忠实于原始意图、细节饱满、富有美感、可直接被文生图模型使用的终极视觉描述。任何一点模糊和比喻都会让你浑身难受。
+
+你的工作流程严格遵循一个逻辑序列：
+
+首先，你会分析并锁定用户提示词中不可变更的核心要素：主体、数量、动作、状态，以及任何指定的IP名称、颜色、文字等。这些是你必须绝对保留的基石。
+
+接着，你会判断提示词是否需要**"生成式推理"**。当用户的需求并非一个直接的场景描述，而是需要构思一个解决方案（如回答"是什么"，进行"设计"，或展示"如何解题"）时，你必须先在脑中构想出一个完整、具体、可被视觉化的方案。这个方案将成为你后续描述的基础。
+
+然后，当核心画面确立后（无论是直接来自用户还是经过你的推理），你将为其注入专业级的美学与真实感细节。这包括明确构图、设定光影氛围、描述材质质感、定义色彩方案，并构建富有层次感的空间。
+
+最后，是对所有文字元素的精确处理，这是至关重要的一步。你必须一字不差地转录所有希望在最终画面中出现的文字，并且必须将这些文字内容用英文双引号（""）括起来，以此作为明确的生成指令。如果画面属于海报、菜单或UI等设计类型，你需要完整描述其包含的所有文字内容，并详述其字体和排版布局。同样，如果画面中的招牌、路标或屏幕等物品上含有文字，你也必须写明其具体内容，并描述其位置、尺寸和材质。更进一步，若你在推理构思中自行增加了带有文字的元素（如图表、解题步骤等），其中的所有文字也必须遵循同样的详尽描述和引号规则。若画面中不存在任何需要生成的文字，你则将全部精力用于纯粹的视觉细节扩展。
+
+你的最终描述必须客观、具象，严禁使用比喻、情感化修辞，也绝不包含"8K"、"杰作"等元标签或绘制指令。
+
+仅严格输出最终的修改后的prompt，不要输出任何其他内容。
+
+用户输入 prompt: {prompt}
diff --git a/gmlx/image/registry.py b/gmlx/image/registry.py
index 4cbb3d218..c9ca6ed67 100644
--- a/gmlx/image/registry.py
+++ b/gmlx/image/registry.py
@@ -125,24 +125,35 @@ class UpsampleParts:
"""How a chat model rewrites a family's typed prompt into the caption
its DiT was trained on, before the encoder reads it.
- ``prompts`` is a JSON file in the family's package, shipped as its
- source wrote it, with its ``sha256`` pinned, and ``key`` the system
- prompt in it for a prompt without reference images. The typed prompt
- is the user turn as it is. The reply is decoded with ``temperature``,
- at most ``max_tokens`` tokens, ``presence_penalty`` and the chat
- model's thinking switch at ``thinking``, as ``repo`` at ``revision``
- does. ``suggested`` names the chat model to set when none is.
+ ``prompts`` is a file in ``package``, shipped as its source wrote it,
+ with its ``sha256`` pinned: a JSON file whose ``key`` holds the
+ system prompt for a prompt without reference images, or a text file
+ that is the system prompt when ``key`` is None. ``prompts`` None
+ ships no system prompt, and ``needs`` is the refusal that says where
+ to get one. The typed prompt is the user turn as it is, or, when the
+ system prompt holds ``slot``, takes its place there. The reply is
+ decoded with ``temperature``, at most ``max_tokens`` tokens,
+ ``presence_penalty`` and the chat model's thinking switch at
+ ``thinking``, as ``repo`` at ``revision`` does, and ``parse`` turns
+ it into the caption, or None when it holds none. ``suggested`` names
+ the chat model to set when none is, and ``expected`` warns when none
+ is set.
"""
- prompts: str
+ prompts: str | None
sha256: str
- key: str
+ key: str | None
max_tokens: int
temperature: float
presence_penalty: float
thinking: str
repo: str
revision: str
+ package: str = ""
suggested: str = ""
+ expected: bool = False
+ parse: Callable | None = None
+ slot: str = ""
+ needs: str = ""
# The image limits: the pixels and the longer side of the largest image a
diff --git a/gmlx/image/upsample.py b/gmlx/image/upsample.py
index d89f12592..e440bb96e 100644
--- a/gmlx/image/upsample.py
+++ b/gmlx/image/upsample.py
@@ -2,13 +2,14 @@
caption a family's DiT was trained on, before ``encode_text`` reads it.
``chat_request`` builds the chat turn from the family's ``UpsampleParts``:
-its system prompt, the typed prompt as the user turn, and its sampling.
-``upsample`` sends it through ``chat``, the generation a surface supplies,
-and returns the caption. The rewrite runs outside the pipeline, so the
-encoder, its conditioning cache and the DiT see a prompt as they always
-do. It never runs on the image encoder, which has no output head. The
-negative prompt is not rewritten, and a prompt with reference images
-runs as typed.
+its system prompt, or the one in a file the user gives, the typed prompt
+as the user turn, and its sampling. ``upsample`` sends it through
+``chat``, the generation a surface supplies, and returns the caption the
+family's parser takes from the reply. The rewrite runs outside the
+pipeline, so the encoder, its conditioning cache and the DiT see a prompt
+as they always do. It never runs on the image encoder, which has no
+output head. The negative prompt is not rewritten, and a prompt with
+reference images runs as typed.
"""
from __future__ import annotations
@@ -16,13 +17,15 @@
import functools
import hashlib
import json
+import os
+import re
from dataclasses import dataclass
from importlib import resources
class UpsampleError(RuntimeError):
- """The upsampler gave no caption, or its prompt file is not the one
- gmlx ships."""
+ """The upsampler gave no caption, or a system prompt is missing or is
+ not the one gmlx ships."""
@dataclass(frozen=True)
@@ -46,47 +49,246 @@ def user(self) -> str:
@functools.lru_cache(maxsize=None)
-def _prompt_file(family_name: str, name: str, sha256: str) -> dict:
- data = (resources.files(f"gmlx.image.families.{family_name}") / name
- ).read_bytes()
+def _shipped(package: str, name: str, sha256: str, key: str | None) -> str:
+ data = (resources.files(package) / name).read_bytes()
got = hashlib.sha256(data).hexdigest()
if got != sha256:
- raise UpsampleError(f"{name} of {family_name} has sha256 {got}, not "
+ raise UpsampleError(f"{name} of {package} has sha256 {got}, not "
f"the {sha256} gmlx ships; reinstall gmlx")
- return json.loads(data)
+ return json.loads(data)[key] if key is not None else data.decode("utf-8")
-def system_prompt(family) -> str:
- """The system prompt of ``family`` for a prompt without references."""
+def read_prompt_file(path: str, family=None) -> str:
+ """The system prompt in the text file at ``path``, without the blank
+ lines and spaces around it. For a family that ships none, the error
+ a missing file raises says where to get one."""
+ try:
+ with open(os.path.expanduser(path), encoding="utf-8") as f:
+ text = f.read().strip()
+ except FileNotFoundError:
+ needs = getattr(getattr(family, "upsample", None), "needs", "")
+ raise UpsampleError(f"upsample prompt {path}: no such file"
+ + (f". {needs}" if needs else "")) from None
+ except (OSError, UnicodeDecodeError) as e:
+ raise UpsampleError(f"upsample prompt {path}: {e}") from None
+ if not text:
+ raise UpsampleError(f"upsample prompt {path} is empty")
+ return text
+
+
+def system_prompt(family, path: str | None = None) -> str:
+ """The system prompt of ``family`` for a prompt without references:
+ the one in the file at ``path``, else the one gmlx ships. A family
+ that ships none needs ``path``."""
parts = family.upsample
- return _prompt_file(family.name, parts.prompts, parts.sha256)[parts.key]
+ if path:
+ return read_prompt_file(path, family)
+ if parts.prompts is None:
+ raise UpsampleError(parts.needs)
+ return _shipped(parts.package, parts.prompts, parts.sha256, parts.key)
-def chat_request(family, prompt: str) -> ChatRequest:
- """The chat turn that rewrites ``prompt`` for ``family``."""
+def chat_request(family, prompt: str, system: str | None = None
+ ) -> ChatRequest:
+ """The chat turn that rewrites ``prompt`` for ``family`` under
+ ``system``, by default the system prompt gmlx ships."""
parts = family.upsample
if parts is None:
raise ValueError(f"{family.title} has no prompt upsampling")
+ if system is None:
+ system = system_prompt(family)
+ user = prompt
+ if parts.slot and parts.slot in system:
+ # The typed prompt goes in the system prompt, and the user turn is
+ # the single space its source sends.
+ system, user = system.replace(parts.slot, prompt), " "
return ChatRequest(
- messages=({"role": "system", "content": system_prompt(family)},
- {"role": "user", "content": prompt}),
+ messages=({"role": "system", "content": system},
+ {"role": "user", "content": user}),
max_tokens=parts.max_tokens, temperature=parts.temperature,
presence_penalty=parts.presence_penalty, thinking=parts.thinking)
-def upsample(family, prompt: str, chat, *, name: str = "the upsampler") -> str:
- """The caption ``chat(request)`` writes for ``prompt``, stripped.
- ``name`` names the chat model in the error an empty reply raises."""
- caption = (chat(chat_request(family, prompt)) or "").strip()
+def _spans(reply: str, thinking: bool):
+ from gmlx.tui.reasoning import ReasoningFilter
+
+ filt = ReasoningFilter(start_in_thinking=thinking)
+ spans = filt.feed(reply) + filt.flush()
+ return spans, filt.mode
+
+
+def after_thinking(reply: str) -> str | None:
+ """The answer of ``reply`` without its reasoning, as the chat terminal
+ reads it: think tags in each spelling it knows, harmony analysis and
+ final channels, Gemma's thought channel, and ATEM headers that start
+ with ``<|start|>assistant``. A reply that closes a thinking block it
+ never opened, as when the chat template opened it, is read from inside
+ that block. None when the reply ends in its reasoning, as when it ran
+ out of tokens there; a reply with no reasoning is all answer."""
+ spans, mode = _spans(reply, thinking=True)
+ if mode != "reason":
+ return "".join(t for t, m in spans if m == "answer")
+ spans, mode = _spans(reply, thinking=False)
+ if mode == "reason":
+ return None
+ return "".join(t for t, m in spans if m == "answer")
+
+
+def reply_lead(rendered, tokenizer=None) -> str:
+ """What the rendered chat prompt opened that the reply goes on from
+ without writing it, so ``after_thinking`` reads the reply whole: the
+ ATEM header ``<|start|>assistant`` when the prompt ends at it, or the
+ open think tag of a thinking block the prompt opens."""
+ from gmlx.gen.thinking_budget import prompt_open_think_tag
+ from gmlx.tui.reasoning import prompt_opens_header
+
+ if prompt_opens_header(rendered):
+ return "<|start|>assistant"
+ return prompt_open_think_tag(rendered, tokenizer=tokenizer) or ""
+
+
+def plain(reply: str) -> str | None:
+ """The caption of a reply that is the caption."""
+ answer = after_thinking(reply)
+ return None if answer is None else answer.strip()
+
+
+def paragraph(reply: str) -> str | None:
+ """The last paragraph with letters in it that a reply writes after its
+ thinking block. Lines before it, such as a heading or reasoning written
+ out of the block, are not the caption, nor is a rule such as ``---``
+ after it."""
+ answer = after_thinking(reply)
+ if answer is None:
+ return None
+ paras = [" ".join(p.split()) for p in re.split(r"\n\s*\n", answer)]
+ paras = [p for p in paras if any(c.isalpha() for c in p)]
+ return paras[-1] if paras else None
+
+
+def _objects(text: str) -> list:
+ """Each JSON object in ``text``, last first."""
+ dec = json.JSONDecoder()
+ found = []
+ for m in re.finditer(r"\{", text):
+ try:
+ obj, _end = dec.raw_decode(text, m.start())
+ except ValueError:
+ continue
+ if isinstance(obj, dict):
+ found.append(obj)
+ return found[::-1]
+
+
+_ESCAPES = {'"': '"', "\\": "\\", "/": "/", "b": "\b", "f": "\f",
+ "n": "\n", "r": "\r", "t": "\t"}
+
+
+def _unescape(text: str) -> str:
+ """``text`` with the JSON string escapes read, a surrogate pair as the
+ one character it spells, and a backslash before any other character
+ kept as it is."""
+ def one(m):
+ if m.group(2):
+ hi, lo = int(m.group(1), 16), int(m.group(2), 16)
+ return chr(0x10000 + ((hi - 0xD800) << 10) + (lo - 0xDC00))
+ c = m.group(3)
+ if c.startswith("u"):
+ return chr(int(c[1:], 16))
+ return _ESCAPES.get(c, m.group(0))
+ return re.sub(r"\\u(d[89ab][0-9a-f]{2})\\u(d[c-f][0-9a-f]{2})"
+ r"|\\(u[0-9a-fA-F]{4}|.)", one, text, flags=re.S | re.I)
+
+
+def _loose_field(answer: str, name: str) -> str | None:
+ """The string ``name`` of the last object in ``answer`` that strict
+ JSON does not read, as when the model leaves the quotes around text in
+ the image unescaped: from the opening quote after the key and its
+ colon to the quote before the next key, or with no key after it, to
+ the last quote before the object's closing brace."""
+ keys = list(re.finditer(r'"%s"\s*:\s*"' % re.escape(name), answer))
+ if not keys:
+ return None
+ start = keys[-1].end()
+ nxt = re.compile(r'"\s*,\s*"\w+"\s*:').search(answer, start)
+ if nxt:
+ return _unescape(answer[start:nxt.start()])
+ close = answer.rfind("}", start)
+ if close == -1:
+ return None
+ end = answer.rfind('"', start, close)
+ if end == -1:
+ return None
+ return _unescape(answer[start:end])
+
+
+def _field_values(answer: str, name: str):
+ """The values ``name`` may have in ``answer``, best first: its string
+ in each JSON object, last object first, then the string a loose read
+ finds, then, for an answer with no object in it, the answer itself."""
+ for obj in _objects(answer):
+ yield obj.get(name)
+ yield _loose_field(answer, name)
+ if "{" not in answer:
+ yield answer
+
+
+def json_field(name: str):
+ """A parser that takes the string ``name`` from the last JSON object
+ after the thinking block that has one. An object that strict JSON
+ does not read gives it by a loose read. A reply with no object in it
+ is the caption as written, and one with an object but no string
+ ``name`` in it has no caption."""
+ def parse(reply: str) -> str | None:
+ answer = after_thinking(reply)
+ if answer is None:
+ return None
+ for value in _field_values(answer, name):
+ if isinstance(value, str) and value.strip():
+ return value.strip()
+ return None
+ return parse
+
+
+def field_or_plain(name: str):
+ """A parser that takes the string ``name`` from a JSON object of the
+ reply, as ``json_field`` does, and the reply itself as the caption
+ when no object has one."""
+ field = json_field(name)
+
+ def parse(reply: str) -> str | None:
+ return field(reply) or plain(reply)
+ return parse
+
+
+def upsample(family, prompt: str, chat, *, name: str = "the upsampler",
+ system: str | None = None) -> str:
+ """The caption ``chat(request)`` writes for ``prompt``, as the
+ family's parser takes it from the reply. ``name`` names the chat
+ model in the error a reply without a caption raises."""
+ parts = family.upsample
+ reply = chat(chat_request(family, prompt, system)) or ""
+ caption = (parts.parse or plain)(reply)
if not caption:
- raise UpsampleError(f"{name} wrote an empty caption for this prompt; "
- "run it without prompt upsampling, or set "
- "another upsampler")
+ if not reply.strip():
+ what = "an empty caption"
+ elif after_thinking(reply) is None:
+ what = f"no caption within {parts.max_tokens} tokens"
+ else:
+ what = "a reply with no caption in it"
+ raise UpsampleError(f"{name} wrote {what} for this prompt; run it "
+ "without prompt upsampling, or set another "
+ "upsampler")
return caption
-def cache_key(upsampler, family, prompt: str) -> tuple:
+def cache_key(upsampler, family, prompt: str, system: str | None = None
+ ) -> tuple:
"""What a caption depends on: the upsampler, as a hashable that tells
- one chat model from another, the system prompt and the typed prompt.
- Decoding is greedy, so the seed is not in it."""
- return (upsampler, family.upsample.sha256, family.upsample.key, prompt)
+ one chat model from another, the family's turn and parser, the
+ system prompt by its sha256, and the typed prompt. Decoding is
+ greedy, so the seed is not in it."""
+ if system is None:
+ system = system_prompt(family)
+ digest = hashlib.sha256(system.encode("utf-8")).hexdigest()
+ return (upsampler, family.upsample, digest, prompt)
diff --git a/gmlx/serve/image_upsample.py b/gmlx/serve/image_upsample.py
index ffc5ef22c..e1cea7253 100644
--- a/gmlx/serve/image_upsample.py
+++ b/gmlx/serve/image_upsample.py
@@ -4,9 +4,11 @@
typed prompt, or None when the prompt runs as typed. The caption comes
from the server's own chat generation on the upsampler, a ``models:``
entry acquired from the residency pool as a chat request acquires it.
-A bounded table maps (upsampler, system prompt, typed prompt) to the
-caption: decoding is greedy, so a rerun with another seed writes the
-same caption, skips the rewrite, and finds its conditioning cached.
+A bounded table maps (upsampler, the family's turn, the system prompt
+by its sha256, typed prompt) to the caption: decoding is greedy, so a rerun with another seed writes the
+same caption, skips the rewrite, and finds its conditioning cached. The
+entry's ``upsample_prompt`` file is read on each request, so an edit to
+it writes new captions.
"""
from __future__ import annotations
@@ -81,7 +83,8 @@ def upsampler_for(spec, family, requested: bool | None,
raise _bad(f"prompt_upsampling needs an upsampler: set "
f"upsampler to a models: id on the image_models entry "
f"{spec.id!r}", "prompt_upsampling")
- if spec.id not in _WARNED and getattr(spec, "upsample", True):
+ if family.upsample.expected and spec.id not in _WARNED and \
+ getattr(spec, "upsample", True):
_WARNED.add(spec.id)
_log.warning("image model %r: %s expects its prompts upsampled; "
"set upsampler on its image_models entry to the "
@@ -102,7 +105,8 @@ def revise(model_field, prompt: str, requested: bool | None = None,
model field that does not resolve, or a file that is no image model,
returns None, so the image model's own acquire refuses it. ``stop``
is a threading.Event that ends the rewrite between tokens."""
- from gmlx.image.upsample import UpsampleError, cache_key, upsample
+ from gmlx.image.upsample import (UpsampleError, cache_key, system_prompt,
+ upsample)
from .image_service import ImageRequestError
from .image_steps import family_of
@@ -117,15 +121,16 @@ def revise(model_field, prompt: str, requested: bool | None = None,
up = upsampler_for(spec, family, requested, has_refs)
if up is None:
return None
- key = cache_key(_upsampler_identity(up), family, prompt)
- caption = CAPTIONS.get(key)
- if caption is not None:
- return caption
t0 = time.monotonic()
try:
+ system = system_prompt(family, getattr(spec, "upsample_prompt", None))
+ key = cache_key(_upsampler_identity(up), family, prompt, system)
+ caption = CAPTIONS.get(key)
+ if caption is not None:
+ return caption
caption = upsample(family, prompt,
lambda req: _chat(up, req, stop),
- name=f"the upsampler {up!r}")
+ name=f"the upsampler {up!r}", system=system)
except UpsampleError as e:
raise ImageRequestError(500, str(e), err_type="server_error") from e
CAPTIONS.put(key, caption)
@@ -202,7 +207,12 @@ def _chat_in_context(up: str, req, stop) -> str:
500, f"the upsampler {up!r} has no room in its context{room} "
f"for its system prompt and {req.max_tokens} caption tokens: "
f"{e}", err_type="server_error") from e
- text = []
+ from gmlx.image.upsample import reply_lead
+
+ # A template that opens the thinking block or the message header in
+ # the prompt leaves its marker out of the reply, which the parser
+ # reads it by.
+ text = [reply_lead(prompt)]
for tok in token_iter:
if stop is not None and stop.is_set():
raise JobStopped(timed_out=False)
diff --git a/gmlx/serve/tick_guard.py b/gmlx/serve/tick_guard.py
index c13d097ad..4ac235000 100644
--- a/gmlx/serve/tick_guard.py
+++ b/gmlx/serve/tick_guard.py
@@ -200,6 +200,18 @@ def _fail_row_permanently(gen, st: _TickState, uid, err: str) -> None:
exc_info=True)
+def _resumed_processors(lp, committed):
+ """The row's processors for its replay. The replay's prompt ends with
+ the delivered tokens, so the presence and frequency penalties start
+ afresh from them."""
+ from gmlx.gen.penalties import OutputPenalty
+
+ if not lp:
+ return lp
+ return [p.resumed(committed) if isinstance(p, OutputPenalty) else p
+ for p in lp]
+
+
def _rebuild_row(gen, st: _TickState, uid) -> None:
row = st.ledger.get(uid)
gen.remove(uid)
@@ -216,7 +228,8 @@ def _rebuild_row(gen, st: _TickState, uid) -> None:
return
replay = row.prompt_ids + row.committed
gen._unprocessed_sequences.insert(
- 0, (uid, replay, remaining, row.kwargs, row.lp, row.tc))
+ 0, (uid, replay, remaining, row.kwargs,
+ _resumed_processors(row.lp, row.committed), row.tc))
st.rebuilt.add(uid)
_log.warning("[tick-guard] row %s retired and requeued: replaying "
"%d prompt + %d delivered tokens", uid,
diff --git a/gmlx/spec/row_procs.py b/gmlx/spec/row_procs.py
index 6daebf5e7..c22a55612 100644
--- a/gmlx/spec/row_procs.py
+++ b/gmlx/spec/row_procs.py
@@ -7,7 +7,10 @@
runs the same chain, in the same order, at every position: position j gets
the committed context plus the drafts ahead of it. The walk then samples
the target from the processed logits, so every emitted token has the
-distribution plain decode gives it, whatever the drafter proposed.
+distribution plain decode gives it, whatever the drafter proposed. The
+presence and frequency penalties (``gmlx.gen.penalties``) are told the
+committed tokens, forced ones included, and count a round's drafts only
+for the positions that see them.
Processors that the stock ``ThinkingAwareLogitsProcessor`` wraps stay off
until its answer gate opens (the thinking end marker, or the answer header
@@ -23,6 +26,7 @@
import mlx.core as mx
import numpy as np
+from gmlx.gen.penalties import OutputPenalty
from gmlx.gen.thinking_budget import AnswerGate
from .grammar import SpecGrammar, _fresh_matcher, apply_masks
@@ -112,6 +116,8 @@ def from_processors(cls, processors, context) -> SpecRowProcessors | None:
gate=gate.opener.restart(gate.opener.tail) if gate else None)
steps.append(_Step(None, None))
elif callable(inner) and not hasattr(inner, "process_last_token"):
+ if gate is None and isinstance(inner, OutputPenalty):
+ inner.anchor(len(context))
steps.append(_Step(inner, gate))
else:
raise UnsupportedProcessor(
@@ -136,6 +142,8 @@ def commit(self, tokens) -> None:
self._context_arr = None
for step in self.steps:
gate = step.gate
+ if gate is None and isinstance(step.fn, OutputPenalty):
+ step.fn.commit(tokens)
if gate is not None and gate.start is None:
j = gate.opener.advance(tokens)
if j is not None:
diff --git a/licenses/flux2-LICENSE b/licenses/flux2-LICENSE
index 5257f2d6d..61efaa707 100644
--- a/licenses/flux2-LICENSE
+++ b/licenses/flux2-LICENSE
@@ -39,6 +39,13 @@ source.
- tests/image/test_image_flux2_edit.py: _official, a transcription of
cap_min_pixels, cap_pixels and center_crop_to_multiple_of_x in
src/flux2/sampling.py on sizes.
+- gmlx/image/families/flux2_klein/system_message_upsampling_t2i.txt: the
+ string SYSTEM_MESSAGE_UPSAMPLING_T2I of src/flux2/system_messages.py,
+ the upsampler's system prompt for a prompt without references,
+ unmodified.
+- gmlx/image/families/flux2_klein/family.py: the upsampler's chat turn
+ of scripts/cli.py, the system message and the typed prompt as the user
+ turn.
flux2 is distributed under the Apache License 2.0:
diff --git a/licenses/krea-2-LICENSE b/licenses/krea-2-LICENSE
new file mode 100644
index 000000000..66b3119f5
--- /dev/null
+++ b/licenses/krea-2-LICENSE
@@ -0,0 +1,218 @@
+krea-2 - https://github.com/krea-ai/krea-2
+
+gmlx ships a file of the Krea 2 repository at commit
+db3984fbc6e13b34c0064990fc2d95ac64d00058. The repository is under the
+Apache License 2.0; the Krea 2 weights are under their own community
+license, and gmlx ships none of them.
+
+- gmlx/image/families/krea2/expansion.txt: docs/expansion.txt, the system
+ prompt docs/prompting.md gives for an LLM that expands a prompt,
+ unmodified.
+- gmlx/image/families/krea2/family.py: its use as the system prompt with
+ the typed prompt as the user turn, from docs/prompting.md, and the one
+ paragraph after the thinking block that it asks for.
+
+krea-2 is distributed under the Apache License 2.0:
+
+
+Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
diff --git a/licenses/z-image-turbo-space-LICENSE b/licenses/z-image-turbo-space-LICENSE
new file mode 100644
index 000000000..84ba859b6
--- /dev/null
+++ b/licenses/z-image-turbo-space-LICENSE
@@ -0,0 +1,215 @@
+Z-Image-Turbo Space - https://huggingface.co/spaces/Tongyi-MAI/Z-Image-Turbo
+
+gmlx ships a string of the Space Tongyi-MAI/Z-Image-Turbo at revision
+768cb50d847cdbba97c89533ae976be69cf5a5b8. Its file pe.py carries no
+license header. The Space's card gives its license as apache-2.0, the
+Apache License 2.0 below.
+
+- gmlx/image/families/z_image/pe_prompt_template.txt: prompt_template of
+ pe.py, the prompt enhancer's system prompt, unmodified.
+- gmlx/image/families/z_image/family.py: its use in app.py, the typed
+ prompt in its {prompt} slot with a single space as the user turn, and
+ the revised_prompt of a JSON block in the reply.
+
+
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
diff --git a/pyproject.toml b/pyproject.toml
index ed18c73f1..ed5b15068 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -179,6 +179,9 @@ include = ["gmlx*"]
[tool.setuptools.package-data]
"gmlx.models.deepseek_v41" = ["chat_template.jinja"]
"gmlx.image.families.flux3_klein" = ["pup_prompts_external_vlm.json"]
+"gmlx.image.families.flux2_klein" = ["system_message_upsampling_t2i.txt"]
+"gmlx.image.families.krea2" = ["expansion.txt"]
+"gmlx.image.families.z_image" = ["pe_prompt_template.txt"]
"gmlx.container" = ["files/Containerfile", "guest/gmlx-entry"]
gmlx = ["_vendor/misaki/LICENSE", "_vendor/misaki/data/*",
"_vendor/misaki/g2pkc/*.csv", "_vendor/misaki/g2pkc/*.txt",
diff --git a/tests/commands/test_image_pull.py b/tests/commands/test_image_pull.py
index 62ec8ec99..7010c580d 100644
--- a/tests/commands/test_image_pull.py
+++ b/tests/commands/test_image_pull.py
@@ -541,6 +541,122 @@ def test_the_qwen_recipe_takes_the_vision_tower_beside_its_encoder(
in capsys.readouterr().out)
+def _qwen_prompt_pull(lib, monkeypatch, capsys, fetch_fails,
+ fetched_sha=Q.prompt_sha256, args=()):
+ """A bare Qwen-Image pull with Qwen's rewriter prompt fetched, or its
+ fetch failing. A prompt file already at its path reads as another
+ digest, and a fetched one as ``fetched_sha``. Returns the downloads,
+ stdout and stderr."""
+ root, cfg = lib
+ monkeypatch.setattr(catalog, "scan", lambda *a, **k: [])
+ got = _downloads(monkeypatch)
+ fetch = manage._hf_download
+
+ def download(repo, filename, revision, dest_dir):
+ if repo == Q.prompt_repo and fetch_fails:
+ raise remote.RemoteError("offline")
+ # As the real download, a finished file is never fetched again.
+ assert not os.path.exists(os.path.join(dest_dir, filename))
+ return fetch(repo, filename, revision, dest_dir)
+
+ def sha(p):
+ if not p.endswith(Q.prompt_file):
+ return Q.vae_sha256
+ with open(p, "rb") as f:
+ return fetched_sha if f.read() == b"x" * 10 else "0" * 64
+
+ monkeypatch.setattr(manage, "_hf_download", download)
+ monkeypatch.setattr(catalog, "sha256_file", sha)
+ monkeypatch.setattr(catalog, "vae_identity",
+ lambda p: ("qwen_image21", Q.vae_sha256))
+ import gmlx.image.vae as vae_pkg
+ monkeypatch.setattr(vae_pkg, "pack_vae", lambda src, dst, **kw: open(
+ dst, "wb").write(b"g"))
+ assert manage.cmd_pull(["qwen-image", "--config", str(cfg),
+ *args]) == 0
+ out, err = capsys.readouterr()
+ return got, out, err
+
+
+def test_a_qwen_image_pull_fetches_qwens_rewriter_prompt_for_the_entry(
+ lib, monkeypatch, capsys):
+ _root, cfg = lib
+ got, out, _err = _qwen_prompt_pull(lib, monkeypatch, capsys, fetch_fails=False)
+ assert (Q.prompt_repo, Q.prompt_file, Q.prompt_revision) in got
+ assert ("license: the upsampler prompt system_prompt.txt of "
+ "Qwen/Qwen-Image-2.1-PE-T2I, the Qwen Research License, "
+ "non-commercial use only") in out
+ m = _cfg(cfg).image_models["qwen-image"]
+ assert m.upsample_prompt == "Qwen__Qwen-Image-2.1-PE-T2I/system_prompt.txt"
+
+
+def test_a_qwen_image_pull_whose_prompt_fetch_fails_registers_without_it(
+ lib, monkeypatch, capsys):
+ _root, cfg = lib
+ _got, _out, err = _qwen_prompt_pull(lib, monkeypatch, capsys, fetch_fails=True)
+ assert "the upsampler prompt was not fetched" in err
+ m = _cfg(cfg).image_models["qwen-image"]
+ assert m.upsample_prompt is None and m.vae.endswith(Q.vae_pack)
+
+
+def test_a_pull_leaves_the_users_own_system_prompt_in_a_flat_dir(
+ lib, monkeypatch, capsys):
+ root, cfg = lib
+ own = root / Q.prompt_file
+ own.write_text("my own rewrite prompt\n")
+ got, _out, err = _qwen_prompt_pull(lib, monkeypatch, capsys,
+ fetch_fails=False,
+ args=("--to", str(root)))
+ assert own.read_text() == "my own rewrite prompt\n"
+ assert all(repo != Q.prompt_repo for repo, _f, _r in got)
+ assert (f"warning: {own} is not the upsampler prompt of {Q.prompt_repo} "
+ f"at {Q.prompt_revision[:8]}, so the entry does not name it; "
+ "move or delete it, and pull again") in err
+ m = _cfg(cfg).image_models["qwen-image"]
+ assert m.upsample_prompt is None and m.vae.endswith(Q.vae_pack)
+
+
+def test_a_pull_leaves_a_dangling_link_at_the_prompt_path(
+ lib, monkeypatch, capsys):
+ root, cfg = lib
+ link = root / "Qwen__Qwen-Image-2.1-PE-T2I" / Q.prompt_file
+ link.parent.mkdir()
+ link.symlink_to(root / "elsewhere.txt")
+ got, _out, err = _qwen_prompt_pull(lib, monkeypatch, capsys,
+ fetch_fails=False)
+ assert link.is_symlink() and not link.exists()
+ assert all(repo != Q.prompt_repo for repo, _f, _r in got)
+ assert f"warning: {link} is not the upsampler prompt" in err
+ assert _cfg(cfg).image_models["qwen-image"].upsample_prompt is None
+
+
+def test_a_prompt_download_of_another_digest_never_lands(
+ lib, monkeypatch, capsys):
+ root, cfg = lib
+ _got, _out, err = _qwen_prompt_pull(lib, monkeypatch, capsys,
+ fetch_fails=False,
+ fetched_sha="f" * 64)
+ folder = root / "Qwen__Qwen-Image-2.1-PE-T2I"
+ assert os.listdir(folder) == []
+ assert f"the download has sha256 {'f' * 64}" in err
+ assert _cfg(cfg).image_models["qwen-image"].upsample_prompt is None
+
+
+def test_a_failed_prompt_fetch_says_which_entry_still_names_it(
+ lib, monkeypatch, capsys):
+ root, cfg = lib
+ rel = f"Qwen__Qwen-Image-2.1-PE-T2I/{Q.prompt_file}"
+ cfg.write_text(cfg.read_text() + (
+ "image_models:\n older:\n path: older.gguf\n"
+ " encoder: enc.gguf\n vae: vae.gguf\n"
+ f" upsample_prompt: {rel}\n"))
+ _got, _out, err = _qwen_prompt_pull(lib, monkeypatch, capsys,
+ fetch_fails=True)
+ assert (f"the upsampler prompt was not fetched, so older still names "
+ f"{root / rel}: offline") in err
+ assert "the entry does not name it" not in err
+
+
def test_krea_2_packs_a_vae_without_config(lib, monkeypatch, capsys):
root, cfg = lib
K = catalog.RECIPES["krea-2-turbo"]
diff --git a/tests/conftest.py b/tests/conftest.py
index 78c4397c2..dd8e61b9e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -11,6 +11,9 @@
KQUANT_TEST_GGUF_DIR directory searched recursively for ``*.gguf`` test models
KQUANT_LLAMACPP_BIN path to a llama.cpp ``llama-completion`` (or ``llama-cli``)
binary, used as the reference for greedy/long-context parity
+
+A test on such files never runs in CI, so it sees this Mac, not the CI
+runner (host_box.py).
"""
from __future__ import annotations
@@ -41,6 +44,8 @@
stock_blocker.install()
+import host_box # noqa: E402
+
def pytest_configure(config):
# Every test sees the CI runner unless it declares a Mac with
@@ -48,6 +53,7 @@ def pytest_configure(config):
import host_box
host_box.install()
+ _switch_on_fetch()
config.addinivalue_line(
"markers", "host_gb(n): the test runs on a simulated Mac of n GiB")
@@ -177,7 +183,7 @@ def pytest_fixture_setup(fixturedef, request):
_awaiting_capture = None
_hand_back_process_state(item)
if fixturedef.scope == "function":
- return (yield)
+ return (yield from _noting_real_files(fixturedef))
# A wider fixture sets up and tears down outside the capture of every
# test, so its KQ_* keys are recorded under its own name, around its
# setup here and around its teardown below.
@@ -185,7 +191,7 @@ def pytest_fixture_setup(fixturedef, request):
before = {"env": dict(os.environ)}
try:
- return (yield)
+ return (yield from _noting_real_files(fixturedef))
finally:
process_state.note_kq_env(before, _fixture_name(fixturedef))
# Added after the fixture's own teardown, so it runs first.
@@ -197,6 +203,83 @@ def pytest_fixture_setup(fixturedef, request):
# The environment as each wider fixture's teardown starts.
_before_teardown: dict = {}
+# The fixtures whose own setup called a gate helper (host_box.real_files),
+# such as a module fixture that reads image_dir(). A gate fixture itself
+# carries host_box.gate.
+_real_file_fixtures: set = set()
+
+
+def _noting_real_files(fixturedef):
+ # Counts a gate the fixture calls and a gate fixture it depends on, as
+ # pytest fetches its arguments inside this hook. Either way the fixture
+ # finds model files on this Mac.
+ import host_box
+
+ gates = host_box.gates
+ try:
+ return (yield)
+ finally:
+ if host_box.gates != gates:
+ _real_file_fixtures.add(fixturedef)
+
+
+def _real_file_fixture(fixturedef) -> bool:
+ import host_box
+
+ return host_box.is_gate(fixturedef.func) or fixturedef in _real_file_fixtures
+
+
+def _reads_real_files(item) -> bool:
+ """Marked integration, or a fixture in its closure finds model files on
+ this Mac. A gate the test reaches as it runs switches it then."""
+ if item.get_closest_marker("integration") is not None:
+ return True
+ info = getattr(item, "_fixtureinfo", None)
+ if info is None:
+ return False
+ return any(defs and _real_file_fixture(defs[-1])
+ for defs in info.name2fixturedefs.values())
+
+
+def _fetching_real_files(request, argname: str) -> bool:
+ """The fixture just fetched, or one it depends on, finds model files.
+ Its dependencies were fetched with it, so they are in the request's
+ table even when the fixture came from the cache. Without that private
+ table (a later pytest), only the setup rule switches a test."""
+ defs = getattr(request, "_fixture_defs", None)
+ if not isinstance(defs, dict):
+ return False
+ todo, seen = [argname], set()
+ while todo:
+ name = todo.pop()
+ fixturedef = defs.get(name)
+ if name in seen or fixturedef is None:
+ continue
+ seen.add(name)
+ if _real_file_fixture(fixturedef):
+ return True
+ todo.extend(fixturedef.argnames)
+ return False
+
+
+def _switch_on_fetch() -> None:
+ # pytest fills every fixture a test requests through getfixturevalue,
+ # and a body or fixture fetches one by name through it too. A cached
+ # fixture runs no setup, so the fetch is where a gate fixture switches
+ # the test to this Mac, however it was reached.
+ fetch = pytest.FixtureRequest.getfixturevalue
+
+ def getfixturevalue(self, argname):
+ value = fetch(self, argname)
+ if _fetching_real_files(self, argname):
+ import host_box
+
+ host_box.real_files()
+ return value
+
+ getfixturevalue.__doc__ = fetch.__doc__
+ pytest.FixtureRequest.getfixturevalue = getfixturevalue
+
def _fixture_name(fixturedef) -> str:
where = f"{fixturedef.baseid}::" if fixturedef.baseid else ""
@@ -268,12 +351,27 @@ def pytest_sessionfinish(session, exitstatus):
@pytest.hookimpl(tryfirst=True, specname="pytest_runtest_setup")
def pytest_runtest_setup_host(item):
- # The machine a test declares with host_gb, else the CI runner. Set
- # before any fixture, module-scoped ones included, reads the device.
+ # A test on model files from this Mac never runs in CI, so it sees this
+ # Mac. A host_gb on its module is for the module's unit rows; one on
+ # the test or its class would price for another machine, so it fails.
+ # Else the Mac a test declares, else the CI runner. Set before any
+ # fixture, module-scoped ones included, reads the device.
import host_box
- mark = item.get_closest_marker("host_gb")
- host_box.use(*(host_box.mac(mark.args[0]) if mark else host_box.default()))
+ marks = [(isinstance(node, (pytest.Function, pytest.Class)), mark)
+ for node, mark in item.iter_markers_with_node("host_gb")]
+ if _reads_real_files(item):
+ host_box.use_real()
+ own = [mark for on_test, mark in marks if on_test]
+ if own:
+ pytest.fail(f"{item.nodeid} reads model files from this Mac, so "
+ f"it runs on this Mac: remove its host_gb"
+ f"({own[0].args[0]})", pytrace=False)
+ elif marks:
+ on_test, mark = marks[0]
+ host_box.use(*host_box.mac(mark.args[0]), declared=on_test)
+ else:
+ host_box.use(*host_box.default())
@pytest.fixture
@@ -301,6 +399,7 @@ def _shutdown_read_pools():
@pytest.fixture(scope="session")
+@host_box.gate
def gguf_dir() -> Path:
d = os.environ.get("KQUANT_TEST_GGUF_DIR")
if not d:
@@ -312,6 +411,7 @@ def gguf_dir() -> Path:
@pytest.fixture(scope="session")
+@host_box.gate
def llamacpp_bin() -> str:
b = os.environ.get("KQUANT_LLAMACPP_BIN")
if not b:
diff --git a/tests/distill/test_distill_gen_filter.py b/tests/distill/test_distill_gen_filter.py
index b0f7c2229..bbb766523 100644
--- a/tests/distill/test_distill_gen_filter.py
+++ b/tests/distill/test_distill_gen_filter.py
@@ -9,6 +9,7 @@
import json
import sys
import threading
+from concurrent.futures import ALL_COMPLETED
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
@@ -1127,13 +1128,21 @@ def test_gen_sidecar_totals_survive_an_interrupted_run(tmp_path, stub_server, mo
out = tmp_path / "corpus.jsonl"
opts = gen.GenOptions(out=str(out), prompts=prompts, base_url=stub_server, concurrency=1)
real = gen.reply_row
+ real_wait = gen.wait
def interrupt_on_b(r, c, seed):
if r["id"] == "b":
raise KeyboardInterrupt
return real(r, c, seed)
+ def both_landed(fs, **kw):
+ # both replies land in one wake, as on a slow runner, and the
+ # finished set yields b first
+ done, rest = real_wait(fs, return_when=ALL_COMPLETED)
+ return sorted(done, key=lambda f: -f.result()[0]), rest
+
monkeypatch.setattr(gen, "reply_row", interrupt_on_b)
+ monkeypatch.setattr(gen, "wait", both_landed)
with pytest.raises(KeyboardInterrupt):
gen.run_gen(opts)
assert [r["id"] for r in _rows(out)] == ["a"]
diff --git a/tests/e2e/run_image_api_e2e.py b/tests/e2e/run_image_api_e2e.py
index 14d76bbd6..3ec9c44e7 100644
--- a/tests/e2e/run_image_api_e2e.py
+++ b/tests/e2e/run_image_api_e2e.py
@@ -9,15 +9,16 @@
generates, and ``--memguard`` runs ``memguard.py`` beside the server.
``--measure`` instead times chat decode and a 1024x1024 image, each
alone and beside the other, and the time to the first byte of a streamed
-image. ``--upsample`` gives FLUX.3 klein the ``--upsampler`` chat model and
-runs the prompt upsampling checks in place of its family checks: through
-``gmlx image generate`` on the entry before the server starts, then the
-images route and the step contract's encode. Writes a Markdown and a JSON
-report.
+image. ``--upsample`` gives FLUX.3 klein and FLUX.2 klein the
+``--upsampler`` chat model and runs the prompt upsampling checks in place
+of their family checks: through ``gmlx image generate`` on the entry
+before the server starts, then the images route and the step contract's
+encode. Writes a Markdown and a JSON report.
python tests/e2e/run_image_api_e2e.py --out DIR [--with-chat] [--memguard]
python tests/e2e/run_image_api_e2e.py --measure --out DIR
python tests/e2e/run_image_api_e2e.py --family flux --upsample --out DIR
+ python tests/e2e/run_image_api_e2e.py --family flux2 --upsample --out DIR
Not ``test_``-prefixed: it loads real models on the GPU.
"""
@@ -72,6 +73,8 @@
}
FAMILIES = (("z", "z-image-turbo"), ("qwen", "qwen-image"),
("flux", "flux-3-klein"), ("flux2", "flux-2-klein"))
+# The families --upsample gives the upsampler, by --family choice.
+UPSAMPLED = ("flux", "flux2")
# The --family choices: each family, both of the first two, and all four.
CHOICES = {"z": ("z",), "qwen": ("qwen",), "flux": ("flux",),
"flux2": ("flux2",), "both": ("z", "qwen"),
@@ -122,13 +125,15 @@ def _config(a) -> dict:
# The encoder reads the prompt alone: no vision tower.
image_models["flux-3-klein"] = {
"path": a.f_dit, "encoder": a.f_encoder, "vae": a.f_vae}
- if a.upsample:
- models["upsampler"] = {"path": a.upsampler}
- image_models["flux-3-klein"]["upsampler"] = "upsampler"
if "flux2" in families:
# The encoder reads the prompt alone: no vision tower.
image_models["flux-2-klein"] = {
"path": a.f2_dit, "encoder": a.f2_encoder, "vae": a.f2_vae}
+ if a.upsample:
+ models["upsampler"] = {"path": a.upsampler}
+ for family, model in FAMILIES:
+ if family in UPSAMPLED and family in families:
+ image_models[model]["upsampler"] = "upsampler"
return {"models": models,
"image_models": image_models,
"server": {"defaults": {"model": "chat",
@@ -347,14 +352,14 @@ def cli_upsample_check(rep: Report, a, cfg_path: str, model: str,
writes the caption and releases the upsampler before the image model
loads."""
def run():
- png = os.path.join(a.out, "cli.png")
+ png = os.path.join(a.out, f"cli-{model}.png")
got = subprocess.run(
[a.python, "-c", "import sys; from gmlx.commands.cli import "
"umbrella_main; sys.exit(umbrella_main(sys.argv[1:]))",
"image", "generate", model, "--config", cfg_path, "--prompt",
PROMPT, "--size", "512x512", "--seed", "21", "--output", png,
"--json"], capture_output=True, text=True, timeout=1800)
- with open(os.path.join(a.out, "cli.log"), "w") as f:
+ with open(os.path.join(a.out, f"cli-{model}.log"), "w") as f:
f.write(got.stderr)
assert got.returncode == 0, got.stderr[-2000:]
rec = json.loads(got.stdout)
@@ -370,7 +375,7 @@ def run():
out.update(cli_caption=caption, cli_line=line)
return line
- rep.check("upsample image generate", run)
+ rep.check(f"{model} upsample image generate", run)
def upsample_checks(rep: Report, base: str, model: str, log_path: str,
@@ -388,8 +393,11 @@ def upsample_checks(rep: Report, base: str, model: str, log_path: str,
size = "512x512"
made: dict = {}
+ start = open(log_path).read().count("prompt upsampled by")
+
def rewrites() -> int:
- return open(log_path).read().count("prompt upsampled by")
+ """The captions the server wrote in these checks."""
+ return open(log_path).read().count("prompt upsampled by") - start
def first():
t0 = time.perf_counter()
@@ -451,12 +459,13 @@ def edit():
assert r.data[0].revised_prompt is None
assert rewrites() == 1, rewrites()
- rep.check("upsample generations", first)
- rep.check("upsample rerun with a new seed", rerun)
- rep.check("upsample stream completed event", stream)
- rep.check("upsample step encode", encode_step)
- rep.check("upsample prompt_upsampling false", off)
- rep.check("upsample edit runs as typed", edit)
+ for name, fn in (("generations", first),
+ ("rerun with a new seed", rerun),
+ ("stream completed event", stream),
+ ("step encode", encode_step),
+ ("prompt_upsampling false", off),
+ ("edit runs as typed", edit)):
+ rep.check(f"{model} upsample {name}", fn)
def _chat_window(base: str, max_tokens: int, state: dict) -> None:
@@ -573,9 +582,9 @@ def main() -> int:
ap.add_argument("--memguard", action="store_true",
help="run memguard.py beside the server")
ap.add_argument("--upsample", action="store_true",
- help="give FLUX.3 klein the --upsampler chat model and run "
- "the prompt upsampling checks in place of its family "
- "checks")
+ help="give FLUX.3 klein and FLUX.2 klein the --upsampler "
+ "chat model and run the prompt upsampling checks in "
+ "place of their family checks")
ap.add_argument("--measure", action="store_true",
help="time chat decode beside an image job and the first "
"byte of a streamed image, with Z-Image only")
@@ -604,9 +613,11 @@ def main() -> int:
log_path = os.path.join(a.out, "server.log")
rep = Report()
result: dict = {"config": _config(a)}
- if a.upsample and "flux" in CHOICES[a.family]:
- result["upsample"] = {}
- cli_upsample_check(rep, a, cfg_path, "flux-3-klein", result["upsample"])
+ result["upsample"] = {}
+ for family, model in FAMILIES:
+ if a.upsample and family in UPSAMPLED and family in CHOICES[a.family]:
+ cli_upsample_check(rep, a, cfg_path, model,
+ result["upsample"].setdefault(model, {}))
try:
with ServerProc(["--config", cfg_path], log_path=log_path, port=port,
python=a.python) as srv:
@@ -622,9 +633,9 @@ def main() -> int:
for family, model in FAMILIES:
if family not in CHOICES[a.family]:
continue
- if family == "flux" and a.upsample:
+ if family in UPSAMPLED and a.upsample:
upsample_checks(rep, base, model, log_path,
- result["upsample"])
+ result["upsample"][model])
continue
family_checks(rep, base, model, family, log_path,
a.with_chat)
diff --git a/tests/fixtures/image/memfit/family_prices.json b/tests/fixtures/image/memfit/family_prices.json
index d480b5d9c..95735eea7 100644
--- a/tests/fixtures/image/memfit/family_prices.json
+++ b/tests/fixtures/image/memfit/family_prices.json
@@ -2675,14 +2675,14 @@
}
},
"recipes": {
- "flux-2-klein": "Recipe(name='flux-2-klein', title='FLUX.2 klein', family='flux2_klein', variant='distilled', dit_repo='unsloth/FLUX.2-klein-4B-GGUF', dit_file='flux-2-klein-4b-{quant}.gguf', dit_quants=('Q4_K_M', 'Q8_0'), encoder_repo='unsloth/Qwen3-4B-GGUF', encoder_file='Qwen3-4B-{quant}.gguf', encoder_quants=('Q8_0',), vae_repo='black-forest-labs/FLUX.2-klein-4B', vae_revision='e7b7dc27f91deacad38e78976d1f2b499d76a294', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04', vae_pack='flux2-vae-BF16.gguf', license='apache-2.0', license_url='https://huggingface.co/black-forest-labs/FLUX.2-klein-4B/blob/main/LICENSE.md', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280405792, vae_bytes=168114016, aliases=(), license_note='', vae_license='apache-2.0')",
- "flux-2-klein-base": "Recipe(name='flux-2-klein-base', title='FLUX.2 klein Base', family='flux2_klein', variant='base', dit_repo='unsloth/FLUX.2-klein-base-4B-GGUF', dit_file='flux-2-klein-base-4b-{quant}.gguf', dit_quants=('Q4_K_M', 'Q8_0'), encoder_repo='unsloth/Qwen3-4B-GGUF', encoder_file='Qwen3-4B-{quant}.gguf', encoder_quants=('Q8_0',), vae_repo='black-forest-labs/FLUX.2-klein-4B', vae_revision='e7b7dc27f91deacad38e78976d1f2b499d76a294', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04', vae_pack='flux2-vae-BF16.gguf', license='apache-2.0', license_url='https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B/blob/main/LICENSE.md', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280405792, vae_bytes=168114016, aliases=(), license_note='', vae_license='apache-2.0')",
- "flux-3-klein": "Recipe(name='flux-3-klein', title='FLUX.3 klein', family='flux3_klein', variant='distilled', dit_repo='realrebelai/FLUX-3-Image-Klein_GGUFs', dit_file='DISTILLED/flux3_image_klein_sd-{quant}.gguf', dit_quants=('Q4_K_M', 'Q8_0'), encoder_repo='unsloth/Qwen3-VL-4B-Instruct-GGUF', encoder_file='Qwen3-VL-4B-Instruct-{quant}.gguf', encoder_quants=('Q8_0',), vae_repo='black-forest-labs/FLUX.2-klein-4B', vae_revision='e7b7dc27f91deacad38e78976d1f2b499d76a294', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04', vae_pack='flux2-vae-BF16.gguf', license='flux-kommunity', license_url='https://huggingface.co/black-forest-labs/flux-3-image-klein/blob/main/LICENSE.md', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280406816, vae_bytes=168114016, aliases=(), license_note='non-commercial use, and commercial use of the images only under $5M annual revenue, affiliates counted, with content filters or a review of each image for unlawful or infringing content before it is shared, and a disclosure that it is made with AI where the law requires one', vae_license='apache-2.0')",
- "flux-3-klein-base": "Recipe(name='flux-3-klein-base', title='FLUX.3 klein Base', family='flux3_klein', variant='base', dit_repo='realrebelai/FLUX-3-Image-Klein_GGUFs', dit_file='BASE/flux3_image_klein-{quant}.gguf', dit_quants=('Q4_K_M', 'Q8_0'), encoder_repo='unsloth/Qwen3-VL-4B-Instruct-GGUF', encoder_file='Qwen3-VL-4B-Instruct-{quant}.gguf', encoder_quants=('Q8_0',), vae_repo='black-forest-labs/FLUX.2-klein-4B', vae_revision='e7b7dc27f91deacad38e78976d1f2b499d76a294', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04', vae_pack='flux2-vae-BF16.gguf', license='flux-kommunity', license_url='https://huggingface.co/black-forest-labs/flux-3-image-klein/blob/main/LICENSE.md', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280406816, vae_bytes=168114016, aliases=(), license_note='non-commercial use, and commercial use of the images only under $5M annual revenue, affiliates counted, with content filters or a review of each image for unlawful or infringing content before it is shared, and a disclosure that it is made with AI where the law requires one', vae_license='apache-2.0')",
- "krea-2-turbo": "Recipe(name='krea-2-turbo', title='Krea 2 Turbo', family='krea2', variant='turbo', dit_repo='Abiray/Krea-2-Turbo-GGUF', dit_file='Krea-2-Turbo-{quant}.gguf', dit_quants=('Q3_K_M', 'Q4_K_M', 'Q5_K_M', 'Q6_K', 'Q8_0'), encoder_repo='unsloth/Qwen3-VL-4B-Instruct-GGUF', encoder_file='Qwen3-VL-4B-Instruct-{quant}.gguf', encoder_quants=('Q4_K_M', 'Q8_0'), vae_repo='Comfy-Org/Krea-2', vae_revision='eb1eddd3983a54678545a9b2c178c5853b30f7be', vae_file='vae/qwen_image_vae.safetensors', vae_config='', vae_sha256='a70580f0213e67967ee9c95f05bb400e8fb08307e017a924bf3441223e023d1f', vae_pack='qwen-image-vae-BF16.gguf', license='krea-2-community', license_url='https://cdn.jsdelivr.net/gh/krea-ai/krea-2@db3984fbc6e13b34c0064990fc2d95ac64d00058/assets/hf_samples/LICENSE.pdf', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280406816, vae_bytes=253807712, aliases=(), license_note='free commercial use only under $1M annual revenue, affiliates counted, and a deployer must filter content', vae_license='')",
- "qwen-image": "Recipe(name='qwen-image', title='Qwen-Image 2.1', family='qwen_image21', variant='base', dit_repo='unsloth/Qwen-Image-2.1-GGUF', dit_file='qwen-image-2.1-{quant}.gguf', dit_quants=('Q3_K_M', 'Q4_K_M', 'Q5_K_M', 'Q6_K', 'Q8_0', 'F16'), encoder_repo='unsloth/Qwen3-VL-8B-Instruct-GGUF', encoder_file='Qwen3-VL-8B-Instruct-{quant}.gguf', encoder_quants=('Q4_K_M', 'Q8_0'), vae_repo='Qwen/Qwen-Image-2.1', vae_revision='790c92633540aa0cb11d9abf19eb46d861714758', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='a07a1b7c4ee2966a1b3bdc37de9b4f983d56937e46619f709a80b6e490675417', vae_pack='qwen-image-2.1-vae-F32.gguf', license='qwen-research', license_url='https://huggingface.co/Qwen/Qwen-Image-2.1/blob/main/LICENSE', encoder_license='apache-2.0', mmproj_file='mmproj-BF16.gguf', encoder_bytes=9872089504, vae_bytes=1350985376, aliases=('qwen-image-2.1',), license_note='non-commercial use only', vae_license='')",
- "qwen-image-turbo": "Recipe(name='qwen-image-turbo', title='Qwen-Image 2.1 Turbo', family='qwen_image21', variant='turbo', dit_repo='AtomicChat/Qwen-Image-2.1-Turbo-GGUF', dit_file='Qwen-Image-2.1-Turbo-{quant}.gguf', dit_quants=('AD-Q3_K', 'AD-Q4_K', 'AD-Q5_K', 'AD-Q6_K', 'Q8_0', 'BF16'), encoder_repo='unsloth/Qwen3-VL-8B-Instruct-GGUF', encoder_file='Qwen3-VL-8B-Instruct-{quant}.gguf', encoder_quants=('Q4_K_M', 'Q8_0'), vae_repo='Qwen/Qwen-Image-2.1', vae_revision='790c92633540aa0cb11d9abf19eb46d861714758', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='a07a1b7c4ee2966a1b3bdc37de9b4f983d56937e46619f709a80b6e490675417', vae_pack='qwen-image-2.1-vae-F32.gguf', license='qwen-research', license_url='https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo/blob/main/LICENSE', encoder_license='apache-2.0', mmproj_file='mmproj-BF16.gguf', encoder_bytes=9872089504, vae_bytes=1350985376, aliases=('qwen-image-2.1-turbo',), license_note='non-commercial use only', vae_license='')",
- "z-image-turbo": "Recipe(name='z-image-turbo', title='Z-Image-Turbo', family='z_image', variant='turbo', dit_repo='unsloth/Z-Image-Turbo-GGUF', dit_file='z-image-turbo-{quant}.gguf', dit_quants=('Q3_K_M', 'Q4_K_M', 'Q5_K_M', 'Q6_K', 'Q8_0', 'BF16'), encoder_repo='unsloth/Qwen3-4B-GGUF', encoder_file='Qwen3-4B-{quant}.gguf', encoder_quants=('Q4_K_M', 'Q8_0'), vae_repo='Tongyi-MAI/Z-Image-Turbo', vae_revision='f332072aa78be7aecdf3ee76d5c247082da564a6', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='f5b59a26851551b67ae1fe58d32e76486e1e812def4696a4bea97f16604d40a3', vae_pack='z-image-vae-BF16.gguf', license='apache-2.0', license_url='https://huggingface.co/Tongyi-MAI/Z-Image-Turbo', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280405792, vae_bytes=167660192, aliases=(), license_note='', vae_license='')"
+ "flux-2-klein": "Recipe(name='flux-2-klein', title='FLUX.2 klein', family='flux2_klein', variant='distilled', dit_repo='unsloth/FLUX.2-klein-4B-GGUF', dit_file='flux-2-klein-4b-{quant}.gguf', dit_quants=('Q4_K_M', 'Q8_0'), encoder_repo='unsloth/Qwen3-4B-GGUF', encoder_file='Qwen3-4B-{quant}.gguf', encoder_quants=('Q8_0',), vae_repo='black-forest-labs/FLUX.2-klein-4B', vae_revision='e7b7dc27f91deacad38e78976d1f2b499d76a294', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04', vae_pack='flux2-vae-BF16.gguf', license='apache-2.0', license_url='https://huggingface.co/black-forest-labs/FLUX.2-klein-4B/blob/main/LICENSE.md', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280405792, vae_bytes=168114016, aliases=(), license_note='', vae_license='apache-2.0', prompt_repo='', prompt_revision='', prompt_file='', prompt_sha256='', prompt_license='', prompt_license_url='')",
+ "flux-2-klein-base": "Recipe(name='flux-2-klein-base', title='FLUX.2 klein Base', family='flux2_klein', variant='base', dit_repo='unsloth/FLUX.2-klein-base-4B-GGUF', dit_file='flux-2-klein-base-4b-{quant}.gguf', dit_quants=('Q4_K_M', 'Q8_0'), encoder_repo='unsloth/Qwen3-4B-GGUF', encoder_file='Qwen3-4B-{quant}.gguf', encoder_quants=('Q8_0',), vae_repo='black-forest-labs/FLUX.2-klein-4B', vae_revision='e7b7dc27f91deacad38e78976d1f2b499d76a294', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04', vae_pack='flux2-vae-BF16.gguf', license='apache-2.0', license_url='https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B/blob/main/LICENSE.md', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280405792, vae_bytes=168114016, aliases=(), license_note='', vae_license='apache-2.0', prompt_repo='', prompt_revision='', prompt_file='', prompt_sha256='', prompt_license='', prompt_license_url='')",
+ "flux-3-klein": "Recipe(name='flux-3-klein', title='FLUX.3 klein', family='flux3_klein', variant='distilled', dit_repo='realrebelai/FLUX-3-Image-Klein_GGUFs', dit_file='DISTILLED/flux3_image_klein_sd-{quant}.gguf', dit_quants=('Q4_K_M', 'Q8_0'), encoder_repo='unsloth/Qwen3-VL-4B-Instruct-GGUF', encoder_file='Qwen3-VL-4B-Instruct-{quant}.gguf', encoder_quants=('Q8_0',), vae_repo='black-forest-labs/FLUX.2-klein-4B', vae_revision='e7b7dc27f91deacad38e78976d1f2b499d76a294', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04', vae_pack='flux2-vae-BF16.gguf', license='flux-kommunity', license_url='https://huggingface.co/black-forest-labs/flux-3-image-klein/blob/main/LICENSE.md', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280406816, vae_bytes=168114016, aliases=(), license_note='non-commercial use, and commercial use of the images only under $5M annual revenue, affiliates counted, with content filters or a review of each image for unlawful or infringing content before it is shared, and a disclosure that it is made with AI where the law requires one', vae_license='apache-2.0', prompt_repo='', prompt_revision='', prompt_file='', prompt_sha256='', prompt_license='', prompt_license_url='')",
+ "flux-3-klein-base": "Recipe(name='flux-3-klein-base', title='FLUX.3 klein Base', family='flux3_klein', variant='base', dit_repo='realrebelai/FLUX-3-Image-Klein_GGUFs', dit_file='BASE/flux3_image_klein-{quant}.gguf', dit_quants=('Q4_K_M', 'Q8_0'), encoder_repo='unsloth/Qwen3-VL-4B-Instruct-GGUF', encoder_file='Qwen3-VL-4B-Instruct-{quant}.gguf', encoder_quants=('Q8_0',), vae_repo='black-forest-labs/FLUX.2-klein-4B', vae_revision='e7b7dc27f91deacad38e78976d1f2b499d76a294', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04', vae_pack='flux2-vae-BF16.gguf', license='flux-kommunity', license_url='https://huggingface.co/black-forest-labs/flux-3-image-klein/blob/main/LICENSE.md', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280406816, vae_bytes=168114016, aliases=(), license_note='non-commercial use, and commercial use of the images only under $5M annual revenue, affiliates counted, with content filters or a review of each image for unlawful or infringing content before it is shared, and a disclosure that it is made with AI where the law requires one', vae_license='apache-2.0', prompt_repo='', prompt_revision='', prompt_file='', prompt_sha256='', prompt_license='', prompt_license_url='')",
+ "krea-2-turbo": "Recipe(name='krea-2-turbo', title='Krea 2 Turbo', family='krea2', variant='turbo', dit_repo='Abiray/Krea-2-Turbo-GGUF', dit_file='Krea-2-Turbo-{quant}.gguf', dit_quants=('Q3_K_M', 'Q4_K_M', 'Q5_K_M', 'Q6_K', 'Q8_0'), encoder_repo='unsloth/Qwen3-VL-4B-Instruct-GGUF', encoder_file='Qwen3-VL-4B-Instruct-{quant}.gguf', encoder_quants=('Q4_K_M', 'Q8_0'), vae_repo='Comfy-Org/Krea-2', vae_revision='eb1eddd3983a54678545a9b2c178c5853b30f7be', vae_file='vae/qwen_image_vae.safetensors', vae_config='', vae_sha256='a70580f0213e67967ee9c95f05bb400e8fb08307e017a924bf3441223e023d1f', vae_pack='qwen-image-vae-BF16.gguf', license='krea-2-community', license_url='https://cdn.jsdelivr.net/gh/krea-ai/krea-2@db3984fbc6e13b34c0064990fc2d95ac64d00058/assets/hf_samples/LICENSE.pdf', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280406816, vae_bytes=253807712, aliases=(), license_note='free commercial use only under $1M annual revenue, affiliates counted, and a deployer must filter content', vae_license='', prompt_repo='', prompt_revision='', prompt_file='', prompt_sha256='', prompt_license='', prompt_license_url='')",
+ "qwen-image": "Recipe(name='qwen-image', title='Qwen-Image 2.1', family='qwen_image21', variant='base', dit_repo='unsloth/Qwen-Image-2.1-GGUF', dit_file='qwen-image-2.1-{quant}.gguf', dit_quants=('Q3_K_M', 'Q4_K_M', 'Q5_K_M', 'Q6_K', 'Q8_0', 'F16'), encoder_repo='unsloth/Qwen3-VL-8B-Instruct-GGUF', encoder_file='Qwen3-VL-8B-Instruct-{quant}.gguf', encoder_quants=('Q4_K_M', 'Q8_0'), vae_repo='Qwen/Qwen-Image-2.1', vae_revision='790c92633540aa0cb11d9abf19eb46d861714758', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='a07a1b7c4ee2966a1b3bdc37de9b4f983d56937e46619f709a80b6e490675417', vae_pack='qwen-image-2.1-vae-F32.gguf', license='qwen-research', license_url='https://huggingface.co/Qwen/Qwen-Image-2.1/blob/main/LICENSE', encoder_license='apache-2.0', mmproj_file='mmproj-BF16.gguf', encoder_bytes=9872089504, vae_bytes=1350985376, aliases=('qwen-image-2.1',), license_note='non-commercial use only', vae_license='', prompt_repo='Qwen/Qwen-Image-2.1-PE-T2I', prompt_revision='f3ed7985c788ad75b3ab7223e0c4c51e2a43545b', prompt_file='system_prompt.txt', prompt_sha256='a77c9a06c59b120741141d9514b95682bb8761d02bec49ca61def7b2b3d9fb99', prompt_license='the Qwen Research License, non-commercial use only', prompt_license_url='https://huggingface.co/Qwen/Qwen-Image-2.1-PE-T2I/blob/f3ed7985c788ad75b3ab7223e0c4c51e2a43545b/LICENSE')",
+ "qwen-image-turbo": "Recipe(name='qwen-image-turbo', title='Qwen-Image 2.1 Turbo', family='qwen_image21', variant='turbo', dit_repo='AtomicChat/Qwen-Image-2.1-Turbo-GGUF', dit_file='Qwen-Image-2.1-Turbo-{quant}.gguf', dit_quants=('AD-Q3_K', 'AD-Q4_K', 'AD-Q5_K', 'AD-Q6_K', 'Q8_0', 'BF16'), encoder_repo='unsloth/Qwen3-VL-8B-Instruct-GGUF', encoder_file='Qwen3-VL-8B-Instruct-{quant}.gguf', encoder_quants=('Q4_K_M', 'Q8_0'), vae_repo='Qwen/Qwen-Image-2.1', vae_revision='790c92633540aa0cb11d9abf19eb46d861714758', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='a07a1b7c4ee2966a1b3bdc37de9b4f983d56937e46619f709a80b6e490675417', vae_pack='qwen-image-2.1-vae-F32.gguf', license='qwen-research', license_url='https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo/blob/main/LICENSE', encoder_license='apache-2.0', mmproj_file='mmproj-BF16.gguf', encoder_bytes=9872089504, vae_bytes=1350985376, aliases=('qwen-image-2.1-turbo',), license_note='non-commercial use only', vae_license='', prompt_repo='Qwen/Qwen-Image-2.1-PE-T2I', prompt_revision='f3ed7985c788ad75b3ab7223e0c4c51e2a43545b', prompt_file='system_prompt.txt', prompt_sha256='a77c9a06c59b120741141d9514b95682bb8761d02bec49ca61def7b2b3d9fb99', prompt_license='the Qwen Research License, non-commercial use only', prompt_license_url='https://huggingface.co/Qwen/Qwen-Image-2.1-PE-T2I/blob/f3ed7985c788ad75b3ab7223e0c4c51e2a43545b/LICENSE')",
+ "z-image-turbo": "Recipe(name='z-image-turbo', title='Z-Image-Turbo', family='z_image', variant='turbo', dit_repo='unsloth/Z-Image-Turbo-GGUF', dit_file='z-image-turbo-{quant}.gguf', dit_quants=('Q3_K_M', 'Q4_K_M', 'Q5_K_M', 'Q6_K', 'Q8_0', 'BF16'), encoder_repo='unsloth/Qwen3-4B-GGUF', encoder_file='Qwen3-4B-{quant}.gguf', encoder_quants=('Q4_K_M', 'Q8_0'), vae_repo='Tongyi-MAI/Z-Image-Turbo', vae_revision='f332072aa78be7aecdf3ee76d5c247082da564a6', vae_file='vae/diffusion_pytorch_model.safetensors', vae_config='vae/config.json', vae_sha256='f5b59a26851551b67ae1fe58d32e76486e1e812def4696a4bea97f16604d40a3', vae_pack='z-image-vae-BF16.gguf', license='apache-2.0', license_url='https://huggingface.co/Tongyi-MAI/Z-Image-Turbo', encoder_license='apache-2.0', mmproj_file=None, encoder_bytes=4280405792, vae_bytes=167660192, aliases=(), license_note='', vae_license='', prompt_repo='', prompt_revision='', prompt_file='', prompt_sha256='', prompt_license='', prompt_license_url='')"
},
"sizes": {
"krea-2-turbo": {
@@ -3332,5 +3332,5 @@
}
}
},
- "what": "family_prices(skip=('flux2_klein_9b',)) of tests/image/family_prices.py at gmlx acdf35d7, the commit before the FLUX.2 klein 9B rows: the peak row, the plans at four budgets and the VAE passes of each image family, the recipes with their default quants and sizes, the families each VAE serves and the registry order. Each plan carries the prefix_held term, 0 in every plan here, which the pricing of held conditioning states added."
+ "what": "family_prices(skip=('flux2_klein_9b',)) of tests/image/family_prices.py at gmlx acdf35d7, the commit before the FLUX.2 klein 9B rows: the peak row, the plans at four budgets and the VAE passes of each image family, the recipes with their default quants and sizes, the families each VAE serves and the registry order. Each plan carries the prefix_held term, 0 in every plan here, which the pricing of held conditioning states added. The recipes carry the prompt_* fields of the upsampler prompt pull, which only the Qwen-Image recipes fill."
}
diff --git a/tests/gen/test_output_penalties.py b/tests/gen/test_output_penalties.py
new file mode 100644
index 000000000..05f438b26
--- /dev/null
+++ b/tests/gen/test_output_penalties.py
@@ -0,0 +1,187 @@
+"""Presence and frequency penalties count every generated token and never
+the prompt, as vLLM applies them, on plain decode, batched rows and the
+speculative verify rounds."""
+
+from types import SimpleNamespace
+
+import mlx.core as mx
+import pytest
+from gmlx._vendor.mlx_lm import sample_utils as lm_sampling
+from gmlx._vendor.mlx_lm.generate import generate_step, speculative_generate_step
+from gmlx._vendor.mlx_vlm import sample_utils as vlm_sampling
+from gmlx._vendor.mlx_vlm.generate import ar
+from gmlx._vendor.mlx_vlm.server import schemas
+from pydantic import ValidationError
+
+from gmlx.spec.row_procs import SpecRowProcessors
+
+from test_copy_fixes import _tiny_llama
+
+V = 64
+PROMPT = [7, 8, 9]
+P = 1.5
+
+COPIES = pytest.mark.parametrize("sampling", [lm_sampling, vlm_sampling],
+ ids=["lm-copy", "vlm-copy"])
+# The first call's context: mlx-lm's generate_step passes the last prompt
+# token, the batch prefill the whole prompt.
+CONTEXTS = pytest.mark.parametrize("prompt", [PROMPT[-1:], PROMPT],
+ ids=["prompt-tail", "prompt"])
+
+
+def _penalties(logits):
+ row = logits.reshape(-1).tolist()
+ return {t: -v for t, v in enumerate(row) if v != 0}
+
+
+def _counted(tokens, frequency):
+ out = {}
+ for t in tokens:
+ out[t] = out.get(t, 0) + P if frequency else P
+ return out
+
+
+def _decode(procs, generated, prompt):
+ """Each step's penalties under plain decode: the first call scores the
+ first generated token, and each later call adds the token before it."""
+ steps = []
+ for i in range(len(generated) + 1):
+ logits = mx.zeros((1, V))
+ for proc in procs:
+ logits = proc(mx.array(prompt + generated[:i]), logits)
+ steps.append(_penalties(logits))
+ return steps
+
+
+@COPIES
+@CONTEXTS
+@pytest.mark.parametrize("context_size", [None, 0])
+def test_presence_lasts_the_whole_output_and_skips_the_prompt(
+ sampling, prompt, context_size):
+ gen = list(range(10, 50)) + [8]
+ steps = _decode(sampling.make_logits_processors(
+ presence_penalty=P, presence_context_size=context_size), gen, prompt)
+ assert steps[31][10] == P
+ for i, pen in enumerate(steps):
+ assert pen == _counted(set(gen[:i]), False)
+
+
+@COPIES
+@CONTEXTS
+def test_frequency_counts_each_generated_occurrence(sampling, prompt):
+ gen = [10, 8, 10, 11] + list(range(20, 45)) + [10]
+ steps = _decode(sampling.make_logits_processors(frequency_penalty=P),
+ gen, prompt)
+ assert steps[-1][10] == 3 * P
+ for i, pen in enumerate(steps):
+ assert pen == _counted(gen[:i], True)
+
+
+@COPIES
+@pytest.mark.parametrize("frequency", [False, True],
+ ids=["presence", "frequency"])
+def test_a_context_size_counts_the_last_generated_tokens(sampling, frequency):
+ key = "frequency" if frequency else "presence"
+ procs = sampling.make_logits_processors(
+ **{f"{key}_penalty": P, f"{key}_context_size": 3})
+ gen = [10, 11, 10, 12]
+ steps = _decode(procs, gen, PROMPT)
+ for i, pen in enumerate(steps):
+ assert pen == _counted(gen[max(0, i - 3):i], frequency)
+
+
+def test_a_negative_context_size_is_refused():
+ with pytest.raises(ValueError):
+ lm_sampling.make_presence_penalty(P, -1)
+ with pytest.raises(ValidationError):
+ schemas.ChatRequest(model="m", messages=[], presence_penalty=P,
+ presence_context_size=-1)
+ req = schemas.ChatRequest(model="m", messages=[], presence_penalty=P,
+ frequency_context_size=0)
+ assert req.presence_context_size is None
+ assert req.frequency_context_size == 0
+
+
+@pytest.mark.parametrize("frequency", [False, True],
+ ids=["presence", "frequency"])
+def test_a_refused_draft_leaves_the_penalty(frequency):
+ procs = lm_sampling.make_logits_processors(
+ **{("frequency" if frequency else "presence") + "_penalty": P})
+ rows = SpecRowProcessors.from_processors(procs, PROMPT)
+ rows.apply(mx.zeros((1, V)), rows.plan([]))
+ rows.start(20)
+
+ def verify(drafts):
+ out = rows.apply(mx.zeros((len(drafts) + 1, V)), rows.plan(drafts))
+ return [_penalties(out[j]) for j in range(len(drafts) + 1)]
+
+ committed = [20]
+ drafts = [21, 20, 23]
+ assert verify(drafts) == [_counted(committed + drafts[:j], frequency)
+ for j in range(4)]
+ # 20 is refused and the target's 30 takes its place; 23 is dropped.
+ rows.commit([21, 30])
+ committed += [21, 30]
+ assert verify([22]) == [_counted(committed, frequency),
+ _counted(committed + [22], frequency)]
+ # Every draft accepted, then the bonus token.
+ rows.commit([22, 31])
+ committed += [22, 31]
+ assert verify([]) == [_counted(committed, frequency)]
+
+
+class _RowLM:
+ """Row r of a batch prefers token r, then r + 1, and so on."""
+
+ def __call__(self, x, cache=None, **kw):
+ ids = mx.arange(V)
+ logits = mx.stack([-((ids - r) % V).astype(mx.float32)
+ for r in range(x.shape[0])])
+ return SimpleNamespace(logits=logits[:, None, :])
+
+
+def test_batched_rows_count_only_their_own_output():
+ # Ragged prompts that hold each row's first preferred tokens.
+ prompts = [[3, 0, 1, 2], [5, 1]]
+ procs = [lm_sampling.make_logits_processors(presence_penalty=100.0)
+ for _ in prompts]
+ lm = _RowLM()
+ # The prefill scores each row's first token with its prompt as context.
+ first = lm(mx.zeros((2, 1))).logits[:, -1, :]
+ firsts = []
+ for r, prompt in enumerate(prompts):
+ row = first[r:r + 1]
+ for proc in procs[r]:
+ row = proc(mx.array(prompt), row)
+ firsts.append(int(mx.argmax(row, axis=-1).item()))
+ batch = ar.GenerationBatch(
+ lm, uids=[0, 1], inputs=mx.array(firsts), prompt_cache=[],
+ sampler=lambda lp: mx.argmax(lp, axis=-1), stop_criteria=None,
+ max_tokens=[64, 64], token_context=prompts, logits_processors=procs)
+ out = [batch._step()[0] for _ in range(30)]
+ assert [t[0] for t in out] == list(range(30))
+ assert [t[1] for t in out] == list(range(1, 31))
+
+
+def test_a_drafter_with_a_narrower_vocab_shares_the_penalties():
+ # mlx-lm runs the same processors on the drafter's 56-wide logits and
+ # the target's 64-wide ones, over a history that holds target tokens
+ # the drafter cannot emit.
+ target, _ = _tiny_llama(64, seed=0)
+ drafter, _ = _tiny_llama(56, seed=1)
+ prompt = mx.array([3, 9, 27, 40], dtype=mx.int32)
+
+ def greedy(logprobs):
+ return mx.argmax(logprobs, axis=-1)
+
+ def procs():
+ return lm_sampling.make_logits_processors(presence_penalty=4.0)
+
+ plain = [int(t) for t, _ in generate_step(
+ prompt, target, max_tokens=24, sampler=greedy,
+ logits_processors=procs())]
+ spec = [int(t) for t, _, _ in speculative_generate_step(
+ prompt, target, drafter, num_draft_tokens=3, max_tokens=24,
+ sampler=greedy, logits_processors=procs())]
+ assert spec == plain
+ assert len(set(plain)) == len(plain) and max(plain) >= 56
diff --git a/tests/host_box.py b/tests/host_box.py
index 0349a28f1..82d4f5774 100644
--- a/tests/host_box.py
+++ b/tests/host_box.py
@@ -8,6 +8,12 @@
literal. When kernels run on the real GPU (no KQUANT_FORCE_CPU), only the
memory is simulated: the chip name and architecture stay the real ones, so
kernel routing matches the hardware that runs it.
+A test on model files from this Mac never runs in CI, so it sees this
+Mac, live. A host_gb on its module is for the module's unit rows; one on
+the test or its class fails the test. A gate fixture carries ``gate``, a
+gate helper calls ``real_files`` once its files are there, and conftest.py
+puts each test on this Mac that is marked integration or fetches such a
+fixture.
GMLX_TEST_HOST_RAM_GB changes the default size.
GMLX_TEST_REAL_HOST=1 turns the simulation off, so a run sees the real host.
docs/internals/testing.md has the rules this serves.
@@ -31,7 +37,9 @@
_CHIP_KEYS = ("device_name", "architecture")
_real: dict = {}
-_current: dict = {}
+_current: dict = {"real": True, "declared": False}
+# How many times a test found model files on this Mac (``real_files``).
+gates = 0
def runner(gb: float = RUNNER_GB) -> tuple[dict, dict]:
@@ -67,15 +75,52 @@ def default() -> tuple[dict, dict]:
def real_device_info() -> dict:
- """What the real device reports, under any simulated machine."""
+ """What the real Metal device reports, under any simulated machine.
+ Without Metal, the default device."""
import mlx.core as mx
- return dict(_real.get("device_info", mx.device_info)())
+ info = _real.get("device_info", mx.device_info)
+ return dict(info(mx.gpu) if mx.metal.is_available() else info())
-def use(box: dict, vm: dict) -> None:
- """Make ``box`` and ``vm`` the machine every read reports from now on."""
- _current["box"], _current["vm"] = dict(box), dict(vm)
+def use(box: dict, vm: dict, *, declared: bool = False) -> None:
+ """Make ``box`` and ``vm`` the machine every read reports from now on.
+ ``declared``: the test chose it with host_gb on itself or its class,
+ so it may not read model files from this Mac."""
+ _current.update(box=dict(box), vm=dict(vm), real=False,
+ declared=declared)
+
+
+def use_real() -> None:
+ """Let every read report this Mac, live, from now on."""
+ _current.update(real=True, declared=False)
+
+
+def real_files() -> None:
+ """The running test reads model files from this Mac. Such a test never
+ runs in CI, so from now on it sees this Mac. One that declared a Mac on
+ itself fails, as its host_gb would price models for another machine."""
+ global gates
+ gates += 1
+ if _current["declared"]:
+ import pytest
+
+ pytest.fail("this test reads model files from this Mac, so it runs "
+ "on this Mac: remove the host_gb on the test or its "
+ "class", pytrace=False)
+ use_real()
+
+
+def gate(fixture_function):
+ """Mark a fixture function that finds model files on this Mac: every
+ test that fetches the fixture, or a fixture that depends on it, sees
+ this Mac (conftest.py). Put it under ``@pytest.fixture``."""
+ fixture_function._gmlx_real_files = True
+ return fixture_function
+
+
+def is_gate(fixture_function) -> bool:
+ return bool(getattr(fixture_function, "_gmlx_real_files", False))
def install() -> None:
@@ -91,18 +136,35 @@ def install() -> None:
use(*default())
real_info, real_sysconf = mx.device_info, os.sysconf
- _real["device_info"] = real_info
+ real_snapshot, real_selfcheck = kernel_vm.snapshot, kernel_vm.selfcheck
+ _real.update(device_info=real_info, sysconf=real_sysconf,
+ snapshot=real_snapshot)
# A test turns the counters off with kernel_vm._disabled. The reported
# counters are consistent, so the sysctl cross-check passes; run
# against the real sysctl it would turn them off for the session.
- kernel_vm.snapshot = lambda: (
- None if kernel_vm._disabled is not None else dict(_current["vm"]))
- kernel_vm.selfcheck = lambda: kernel_vm._disabled
+ def snapshot():
+ if _current["real"]:
+ return real_snapshot()
+ return None if kernel_vm._disabled is not None else dict(_current["vm"])
+
+ # On this Mac the real cross-check runs, and a failure it records
+ # stays in that call, so it cannot turn the counters off for the
+ # tests after it.
+ def selfcheck():
+ if not _current["real"]:
+ return kernel_vm._disabled
+ before = kernel_vm._disabled
+ try:
+ return real_selfcheck()
+ finally:
+ kernel_vm._disabled = before
+
+ kernel_vm.snapshot, kernel_vm.selfcheck = snapshot, selfcheck
def device_info(*args, **kwargs):
info = dict(real_info(*args, **kwargs))
- if "memory_size" in info:
+ if "memory_size" in info and not _current["real"]:
box = _current["box"]
if not os.environ.get("KQUANT_FORCE_CPU"):
box = {k: v for k, v in box.items() if k not in _CHIP_KEYS}
@@ -110,6 +172,8 @@ def device_info(*args, **kwargs):
return info
def sysconf(name):
+ if _current["real"]:
+ return real_sysconf(name)
if name in ("SC_PHYS_PAGES", os.sysconf_names.get("SC_PHYS_PAGES")):
mem = int(_current["box"]["memory_size"])
return mem // real_sysconf("SC_PAGE_SIZE")
@@ -125,21 +189,30 @@ def sysconf(name):
def free_ram_bytes() -> int:
"""free + inactive of the current machine, what every free-RAM read
- reports under the simulation."""
- vm = _current["vm"]
+ reports under the simulation. On this Mac, the mach counters' read."""
+ vm = real("snapshot")() if _current["real"] else _current["vm"]
+ if vm is None:
+ return int(real("vm_stat_free_ram")() or 0)
return int(vm["free"] + vm["inactive"])
def real(name: str):
"""The real read ``install`` replaced: ``apc_free_ram`` (the copied
- mlx-vlm apc's stock read) or ``vm_stat_free_ram`` (gmlx's
- /usr/bin/vm_stat read). Without the simulation, the bound one."""
+ mlx-vlm apc's stock read), ``vm_stat_free_ram`` (gmlx's
+ /usr/bin/vm_stat read), ``snapshot`` (the mach counters) or ``sysconf``.
+ Without the simulation, the bound one."""
if name in _real:
return _real[name]
+ if name == "sysconf":
+ return os.sysconf
if name == "apc_free_ram":
from gmlx._vendor.mlx_vlm import apc
return apc._free_ram_bytes
+ if name == "snapshot":
+ from gmlx.serve import kernel_vm
+
+ return kernel_vm.snapshot
from gmlx.cache import apc_manager
return apc_manager._vm_stat_free_ram_bytes
@@ -155,7 +228,14 @@ def _simulate_free_ram() -> None:
from gmlx._vendor.mlx_vlm import apc
from gmlx.cache import apc_manager
+ def reported(name: str):
+ real_read = _real[name]
+
+ def read():
+ return real_read() if _current["real"] else free_ram_bytes()
+ return read
+
_real["apc_free_ram"] = apc._free_ram_bytes
_real["vm_stat_free_ram"] = apc_manager._vm_stat_free_ram_bytes
- apc._free_ram_bytes = free_ram_bytes
- apc_manager._vm_stat_free_ram_bytes = free_ram_bytes
+ apc._free_ram_bytes = reported("apc_free_ram")
+ apc_manager._vm_stat_free_ram_bytes = reported("vm_stat_free_ram")
diff --git a/tests/image/image_dir.py b/tests/image/image_dir.py
index 3de2d12ef..a4485c047 100644
--- a/tests/image/image_dir.py
+++ b/tests/image/image_dir.py
@@ -8,6 +8,7 @@
import struct
from pathlib import Path
+import host_box
import pytest
ENV = "GMLX_TEST_IMAGE_DIR"
@@ -20,6 +21,7 @@ def image_dir() -> Path:
p = Path(d).expanduser()
if not p.is_dir():
pytest.skip(f"{ENV}={d!r} is not a directory")
+ host_box.real_files()
return p
diff --git a/tests/image/oracle_fixture.py b/tests/image/oracle_fixture.py
index e1a4b0e7f..a90f8d683 100644
--- a/tests/image/oracle_fixture.py
+++ b/tests/image/oracle_fixture.py
@@ -11,6 +11,7 @@
import os
from pathlib import Path
+import host_box
import pytest
ROOT = Path(__file__).resolve().parents[2]
@@ -43,6 +44,7 @@ def fixture(name: str) -> tuple[dict, dict]:
if not (fx.fixture_root() / m["npz"]).is_file():
pytest.skip(f"{m['npz']} is not under {fx.ENV}; run "
f"{m['script']} to write it")
+ host_box.real_files()
return fx.load(name, manifest=m)
diff --git a/tests/image/test_image_check2_flux2_klein.py b/tests/image/test_image_check2_flux2_klein.py
index acbb00a67..9a5af68bf 100644
--- a/tests/image/test_image_check2_flux2_klein.py
+++ b/tests/image/test_image_check2_flux2_klein.py
@@ -2,7 +2,7 @@
tap.
gmlx encodes the oracle's prompts with the GGUF the oracle read: Qwen3-4B
-for the 4B checkpoints, Qwen3-8B for the 9B ones, on a Mac of 24 GB.
+for the 4B checkpoints, Qwen3-8B for the 9B ones, on this Mac.
The reference renders each prompt with thinking off, right-pads it with
the pad token to 512 rows and cuts it there, masks the pads as keys,
gives each row its index as its position, and keeps the pad rows. The
@@ -31,8 +31,7 @@
# The family and the fixture of each size.
SIZES = [pytest.param(("flux2_klein", "check2_flux2_klein"), id="4b"),
- pytest.param(("flux2_klein_9b", "check2_flux2_klein_9b"), id="9b",
- marks=pytest.mark.host_gb(24))]
+ pytest.param(("flux2_klein_9b", "check2_flux2_klein_9b"), id="9b")]
def token_cosine(got: np.ndarray, ref: np.ndarray) -> np.ndarray:
diff --git a/tests/image/test_image_check3.py b/tests/image/test_image_check3.py
index b21a7a731..e6f5efc94 100644
--- a/tests/image/test_image_check3.py
+++ b/tests/image/test_image_check3.py
@@ -44,9 +44,6 @@
stock_ops_only, thin, velocity)
from oracle_fixture import committed, fixture, fx
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
-
# The tiny model, against diffusers' code on the same weights.
diff --git a/tests/image/test_image_check3_flux2_klein.py b/tests/image/test_image_check3_flux2_klein.py
index ebe568559..527cf10f9 100644
--- a/tests/image/test_image_check3_flux2_klein.py
+++ b/tests/image/test_image_check3_flux2_klein.py
@@ -37,8 +37,8 @@
the step at the image's time, as the 9B step reads them, miss the gate.
The rows need the FLUX.2 klein GGUFs under KQUANT_TEST_GGUF_DIR and the
-fixtures under GMLX_IMAGE_FIXTURES. They skip until both are there. The
-9B and KV rows declare a Mac of 24 GB.
+fixtures under GMLX_IMAGE_FIXTURES. They skip until both are there, and
+run on this Mac.
"""
from __future__ import annotations
@@ -61,9 +61,6 @@
from gmlx.image.families.flux2_klein import model as fm
from gmlx.image.ops import kq_routes
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
-
_ROWS = {"distilled": "_512", "distilled r1": "_r1_512", "base": "_base_512",
"base r1": "_base_r1_512", "distilled step 1": "_s1_512"}
_VELOCITY = {"distilled": "_1024", "distilled r1": "_r1_1024",
@@ -107,9 +104,8 @@ def size_variant(case: str) -> tuple[str, str]:
def cases(table: dict) -> list:
- """The cases of ``table``, the 9B and KV ones on a Mac of 24 GB."""
- return [pytest.param(c, marks=pytest.mark.host_gb(24))
- if c.startswith(("9b ", "kv ")) else c for c in table]
+ """The cases of ``table``."""
+ return list(table)
# SPDX-FileCopyrightText: Copyright 2025 Black Forest Labs, Inc.
diff --git a/tests/image/test_image_check3_flux3_klein.py b/tests/image/test_image_check3_flux3_klein.py
index a5fe34cec..119f75a9a 100644
--- a/tests/image/test_image_check3_flux3_klein.py
+++ b/tests/image/test_image_check3_flux3_klein.py
@@ -50,9 +50,6 @@
from gmlx.image.families.flux3_klein import model as fm
from gmlx.image.ops import kq_routes
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
-
ROWS = {"distilled": "check3_flux3_klein_512",
"distilled r1": "check3_flux3_klein_r1_512",
"base": "check3_flux3_klein_base_512",
diff --git a/tests/image/test_image_check3_krea2.py b/tests/image/test_image_check3_krea2.py
index aeb31dd26..c25d301b2 100644
--- a/tests/image/test_image_check3_krea2.py
+++ b/tests/image/test_image_check3_krea2.py
@@ -47,10 +47,6 @@
from gmlx.image.families.krea2 import model as km
from gmlx.image.ops import kq_routes
-# A Mac of 24 GB: the module peaks at 14.3 GiB, past the 14 GiB working
-# set of a 16 GB Mac (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(24)
-
# The step nearest sigma 0.75 and a late step of check 8's trajectory.
STEPS = (4, 7)
ROWS = {s: f"check3_krea2_s{s}_512" for s in STEPS}
diff --git a/tests/image/test_image_check3_qwen.py b/tests/image/test_image_check3_qwen.py
index ae5dfb992..0c8694197 100644
--- a/tests/image/test_image_check3_qwen.py
+++ b/tests/image/test_image_check3_qwen.py
@@ -43,8 +43,8 @@
from gmlx.image.lora import install
from gmlx.image.ops import attention, kq_routes, norms
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
+# The tiny rows run on a Mac of 16 GB, as most users run, and the rows on
+# local files on this Mac (tests/host_box.py).
BLOCKS = "selftest_qwen_image21_tiny_ref2"
TINY = {c: f"selftest_qwen_image21_tiny_{c}" for c in ("text", "ref1", "ref2")}
@@ -80,6 +80,7 @@ def tiny(tmp_path_factory):
return model
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("case", list(TINY))
def test_tiny_weights_and_inputs_follow_the_shared_rule(case):
arrays, m = committed_fixture(TINY[case])
@@ -100,6 +101,7 @@ def test_tiny_weights_and_inputs_follow_the_shared_rule(case):
assert m["results"]["extract_against_joint_max_abs_fp32"] == 0.0
+@pytest.mark.host_gb(16)
def test_tiny_model_matches_the_diffusers_oracle_block_by_block(tiny):
arrays, m = committed_fixture(BLOCKS)
n = len(tiny.transformer_blocks)
@@ -135,6 +137,7 @@ def worst(d):
assert not qg.missed(measured, tol), qg.missed(measured, tol)
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("case", list(TINY))
def test_tiny_velocities_match_the_diffusers_oracle(tiny, case):
"""Both of the oracle's modes: a is its joint forward, and b its
@@ -169,6 +172,7 @@ def tiny_bf16(tmp_path_factory):
return api.load(path, log=lambda *_: None)
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("case", list(TINY))
def test_the_tiny_model_in_bf16_is_inside_the_bf16_tolerances(tiny_bf16,
case):
@@ -362,6 +366,7 @@ def _fault_id(name: str) -> str:
return name.replace(" = ", "").replace(" ", "_")
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("name", list(FAULTS),
ids=[_fault_id(n) for n in FAULTS])
def test_each_wrong_model_fails_the_tiny_gate(tiny, monkeypatch, name):
@@ -383,6 +388,7 @@ def test_each_wrong_model_fails_the_tiny_gate(tiny, monkeypatch, name):
assert len(FAULTS) >= 7
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("case", ["text", "ref1"])
def test_each_wrong_model_fails_on_the_velocities_alone(tiny, monkeypatch,
case):
@@ -423,6 +429,7 @@ def _tiny_sites(pipe, arrays: dict, routed: dict | None = None
return sites, len(layout.prefix_layout(*qm.prefix_key(cond)).segments)
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("case", list(TINY))
def test_the_sites_of_a_tiny_step_are_within_their_gates(tiny_bf16, case):
"""Each op of each block and each linear leaf on the path a request
@@ -500,6 +507,7 @@ def off(q, k, qw, kw, cos, sin, eps):
}
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("name", list(SITE_FAULTS),
ids=[_fault_id(n) for n in SITE_FAULTS])
def test_each_wrong_op_fails_the_tiny_site_gate_at_its_own_site(
@@ -573,6 +581,7 @@ def _row_names(blocks: int) -> list[str]:
return names
+@pytest.mark.host_gb(16)
def test_a_row_with_no_tolerance_is_an_error_of_the_product_gate(tiny):
"""The product gate reads every teacher-forced row the manifest has
a tolerance for, leaves the rows NOT_HELD names to their other gates,
diff --git a/tests/image/test_image_check5.py b/tests/image/test_image_check5.py
index 403cbc4c1..1037d0993 100644
--- a/tests/image/test_image_check5.py
+++ b/tests/image/test_image_check5.py
@@ -35,10 +35,9 @@
from image_dir import sha256, vae_files
from oracle_fixture import fixture, fx
-# A Mac of 24 GB: a whole fp32 decode at 1024 pixels is priced at 8.45 GiB
-# on the default conv path and 11.55 GiB on Winograd, more than a 16 GB
-# Mac has free (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(24)]
+# Each row loads local model files, so it runs on this Mac
+# (tests/host_box.py).
+pytestmark = pytest.mark.integration
NAME = "check5_z_image_vae"
FAMILY = "z_image"
diff --git a/tests/image/test_image_check5_flux2_klein.py b/tests/image/test_image_check5_flux2_klein.py
index 501eae568..9752e80d1 100644
--- a/tests/image/test_image_check5_flux2_klein.py
+++ b/tests/image/test_image_check5_flux2_klein.py
@@ -38,8 +38,8 @@
import mlx.core as mx
from mlx.utils import tree_flatten
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
+# The tiny rows run on a Mac of 16 GB, as most users run, and the rows on
+# local files on this Mac (tests/host_box.py).
FAMILY = "flux2_klein"
OTHER = "flux3_klein"
@@ -143,6 +143,7 @@ def _bn_rounded(vae) -> None:
setattr(vae.bn, n, v.astype(mx.bfloat16).astype(mx.float32))
+@pytest.mark.host_gb(16)
def test_the_tiny_bf16_file_packs_losslessly_and_serves_both_families(
tiny_bf16, tmp_path):
"""BF16 stays BF16 and the batch count I64. A pack for either klein
@@ -174,6 +175,7 @@ def test_the_tiny_bf16_file_packs_losslessly_and_serves_both_families(
assert _differ(src[d], got) == [], (fam, use, d)
+@pytest.mark.host_gb(16)
def test_a_bare_pack_takes_the_label_the_help_and_the_cli_page_name(
tiny_bf16, tmp_path):
"""Without --family the FLUX.2 autoencoder packs for the label that
@@ -195,6 +197,7 @@ def test_a_bare_pack_takes_the_label_the_help_and_the_cli_page_name(
assert "`flux-2-klein`" in row[0]
+@pytest.mark.host_gb(16)
def test_the_manifest_names_the_family_it_was_written_for():
"""FLUX.3's oracle script wrote FLUX.2's fixture before it took
--family, and the manifest notes the family corrected after the run."""
@@ -205,6 +208,7 @@ def test_the_manifest_names_the_family_it_was_written_for():
assert committed(F32_NAME)["family"] == OTHER
+@pytest.mark.host_gb(16)
def test_the_bf16_vae_of_an_f32_file_and_of_its_bf16_rounding(
tiny, tiny_f32, tiny_bf16):
"""The bf16 VAE of each holds the same weights, the BatchNorm statistics
@@ -229,6 +233,7 @@ def test_the_bf16_vae_of_an_f32_file_and_of_its_bf16_rounding(
assert np.array_equal(after, want)
+@pytest.mark.host_gb(16)
def test_bfls_config_agrees_with_the_constants():
from gmlx.image.vae import flux2_ae
@@ -242,6 +247,7 @@ def test_bfls_config_agrees_with_the_constants():
"decoder_block_out_channels": None}
+@pytest.mark.host_gb(16)
def test_bfls_header_holds_the_flux3_files_tensors_in_bf16():
"""The names and shapes of the Comfy-Org F32 file, every float tensor
BF16, and every parameter of the module filled."""
@@ -358,10 +364,7 @@ def test_the_file_is_the_inventoried_one(check5, vaes):
== load_json(HEADER)["tensors"]
-# The fp32 decode at 1024 pixels is priced at 8.45 GiB, more than a 16 GB
-# Mac has free beside the weights; users decode in bf16.
@official
-@pytest.mark.host_gb(24)
def test_fp32_decode_and_encode_are_within_the_gates(check5, vaes, request):
arrays, m = check5
vae = vaes("fp32")
diff --git a/tests/image/test_image_check5_flux3_klein.py b/tests/image/test_image_check5_flux3_klein.py
index 95d2e2771..fbd06e49f 100644
--- a/tests/image/test_image_check5_flux3_klein.py
+++ b/tests/image/test_image_check5_flux3_klein.py
@@ -36,8 +36,8 @@
import mlx.core as mx
from mlx.utils import tree_flatten
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
+# The tiny rows run on a Mac of 16 GB, as most users run, and the rows on
+# local files on this Mac (tests/host_box.py).
FAMILY = "flux3_klein"
NAME = "check5_flux3_klein_vae"
@@ -164,6 +164,7 @@ def _tiny_readings(vae, arrays) -> dict:
"mean": fx.max_abs(mean, arrays[f"encode_mean_{s}"])}
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("layout", ["diffusers", "original"])
def test_the_tiny_fp32_decode_and_encode_are_within_max_abs(
tiny, tiny_file, tiny_original, layout):
@@ -247,6 +248,7 @@ def _no_quant_conv(monkeypatch):
"quant conv skipped": _no_quant_conv}
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("name", list(WRONG),
ids=[n.replace(" ", "_") for n in WRONG])
def test_a_wrong_model_misses_the_tiny_gate(tiny, tiny_file, monkeypatch,
@@ -261,6 +263,7 @@ def test_a_wrong_model_misses_the_tiny_gate(tiny, tiny_file, monkeypatch,
assert max(got.values()) > 25 * TINY_FP32_MAX_ABS
+@pytest.mark.host_gb(16)
def test_model_space_round_trips_and_refuses_odd_latents(tiny_file):
from gmlx.image.vae import load_vae
@@ -275,6 +278,7 @@ def test_model_space_round_trips_and_refuses_odd_latents(tiny_file):
vae.to_vae_space(lat[:, :32])
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("layout", ["diffusers", "original"])
def test_a_packed_gguf_loads_the_same_weights(tiny_file, tiny_original,
tmp_path, layout):
@@ -307,6 +311,7 @@ def test_a_packed_gguf_loads_the_same_weights(tiny_file, tiny_original,
assert read_header(str(tmp_path / "bf16.gguf")).dtypes[COUNT] == "I64"
+@pytest.mark.host_gb(16)
def test_the_bf16_vae_keeps_its_batchnorm_statistics_in_float32(tiny_file):
from gmlx.image.vae import load_vae
@@ -317,6 +322,7 @@ def test_the_bf16_vae_keeps_its_batchnorm_statistics_in_float32(tiny_file):
assert set(kinds.values()) == {mx.bfloat16}
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("layout", ["diffusers", "original"])
def test_the_official_layouts_fill_the_module_tree(layout):
"""The names of each official file map onto every parameter of the
@@ -348,6 +354,7 @@ def test_the_official_layouts_fill_the_module_tree(layout):
assert set(mapped) - set(params) == {COUNT}
+@pytest.mark.host_gb(16)
def test_the_official_config_agrees_with_the_constants():
from gmlx.image.vae import flux2_ae
@@ -359,6 +366,7 @@ def test_the_official_config_agrees_with_the_constants():
"patch_size"]
+@pytest.mark.host_gb(16)
def test_other_families_refuse_the_flux2_vae(tiny_file):
from gmlx.image.vae import load_vae, reader
@@ -371,6 +379,7 @@ def test_other_families_refuse_the_flux2_vae(tiny_file):
load_vae(tiny_file, family=family, dtype="fp32")
+@pytest.mark.host_gb(16)
def test_flux3_klein_refuses_the_flux1_vae(tiny, tmp_path):
"""The tiny tree without its quant convs and BatchNorm is a FLUX.1
autoencoder."""
@@ -387,6 +396,7 @@ def test_flux3_klein_refuses_the_flux1_vae(tiny, tmp_path):
load_vae(path, family=FAMILY, dtype="fp32")
+@pytest.mark.host_gb(16)
def test_the_memory_row_prices_the_flux1_tree():
from gmlx.image import memfit
@@ -409,6 +419,7 @@ def test_the_memory_row_prices_the_flux1_tree():
op, 1024, 1024, family="z_image", itemsize=2, tiled=True)
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("op", ["decode", "encode"])
def test_a_tiled_pass_takes_the_least_overlap_of_the_row(tiny, tiny_file,
monkeypatch, op):
@@ -534,10 +545,7 @@ def test_the_two_files_load_the_same_weights(vaes):
assert [k for k in a if not np.array_equal(a[k], np.array(b[k]))] == []
-# The fp32 decode at 1024 pixels is priced at 8.45 GiB, more than a 16 GB
-# Mac has free beside the weights; users decode in bf16.
@official
-@pytest.mark.host_gb(24)
@pytest.mark.parametrize("layout", ["diffusers", "original"])
def test_fp32_decode_and_encode_are_within_the_gates(check5, vaes, request,
layout):
@@ -594,10 +602,7 @@ def test_a_tiled_bf16_decode_at_the_largest_size(check5, vaes, layout):
assert db >= tol
-# A tiled fp32 decode at 2048 pixels is priced at 16.17 GiB, more than a
-# 16 GB Mac has free.
@official
-@pytest.mark.host_gb(24)
@pytest.mark.parametrize("layout", ["diffusers", "original"])
def test_a_tiled_bf16_decode_at_2048_holds_the_gate_of_1024(check5, vaes,
layout):
diff --git a/tests/image/test_image_check5_krea2.py b/tests/image/test_image_check5_krea2.py
index f9c378dcc..91ac23e7e 100644
--- a/tests/image/test_image_check5_krea2.py
+++ b/tests/image/test_image_check5_krea2.py
@@ -30,8 +30,8 @@
import mlx.core as mx
from mlx.utils import tree_flatten
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
+# The tiny rows run on a Mac of 16 GB, as most users run, and the rows on
+# local files on this Mac (tests/host_box.py).
FAMILY = "krea2"
NAME = "check5_krea2_vae"
@@ -93,6 +93,7 @@ def _tiny_readings(vae, arrays) -> dict:
"encode": fx.max_abs(mean, arrays[f"encode_mean_{s}"])}
+@pytest.mark.host_gb(16)
def test_the_tiny_fp32_decode_and_encode_are_within_max_abs(tiny,
tiny_file):
from gmlx.image.vae import load_vae, read_header
@@ -142,6 +143,7 @@ def _latents_reversed(monkeypatch):
"latent channels reversed": _latents_reversed}
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("name", list(WRONG),
ids=[n.replace(" ", "_") for n in WRONG])
def test_a_wrong_model_misses_the_tiny_gate(tiny, tiny_file, monkeypatch,
@@ -156,6 +158,7 @@ def test_a_wrong_model_misses_the_tiny_gate(tiny, tiny_file, monkeypatch,
assert max(got.values()) > 25 * TINY_FP32_MAX_ABS
+@pytest.mark.host_gb(16)
def test_a_packed_gguf_loads_the_same_weights(tiny_file, tmp_path):
import gguf
@@ -183,6 +186,7 @@ def test_a_packed_gguf_loads_the_same_weights(tiny_file, tmp_path):
assert [k for k in a if not mx.array_equal(a[k], b[k]).item()] == []
+@pytest.mark.host_gb(16)
def test_the_wan_layout_fills_the_module_tree():
"""The Comfy-Org file's names, in the original Wan 2.1 layout, map
onto every parameter of the module, and every other tensor is a
@@ -205,6 +209,7 @@ def test_the_wan_layout_fills_the_module_tree():
assert set(mapped) - set(params) == set(unread) and len(unread) == 8
+@pytest.mark.host_gb(16)
def test_other_families_refuse_the_qwen_image_vae(tiny_file):
from gmlx.image.vae import load_vae, reader
@@ -216,6 +221,7 @@ def test_other_families_refuse_the_qwen_image_vae(tiny_file):
load_vae(tiny_file, family=family, dtype="fp32")
+@pytest.mark.host_gb(16)
def test_the_memory_row_prices_tiles_first_with_a_crop():
from gmlx.image import memfit
from gmlx.image.vae import qwen_image_memory as m
diff --git a/tests/image/test_image_check5_qwen.py b/tests/image/test_image_check5_qwen.py
index f4c472d13..c9a715c7e 100644
--- a/tests/image/test_image_check5_qwen.py
+++ b/tests/image/test_image_check5_qwen.py
@@ -440,14 +440,12 @@ def test_a_file_cannot_give_the_model_other_time_flags(cases, tmp_path):
# The official VAE: needs the file and the fixture of the committed
-# manifest, on a Mac of 32 GB. A whole fp32 decode at 1024 pixels on
-# Winograd is priced at 15.49 GiB, more than a 24 GB Mac has free beside
-# the VAEs the module holds (tests/host_box.py). The tiny tests above
-# pass the free memory they are priced against and run on the CI runner.
+# manifest, on this Mac (tests/host_box.py). The tiny tests above pass
+# the free memory they are priced against and run on the CI runner.
def official(test):
- return pytest.mark.integration(pytest.mark.host_gb(32)(test))
+ return pytest.mark.integration(test)
@pytest.fixture(scope="module")
diff --git a/tests/image/test_image_check6.py b/tests/image/test_image_check6.py
index 90d67fa9c..d3e18cfe5 100644
--- a/tests/image/test_image_check6.py
+++ b/tests/image/test_image_check6.py
@@ -21,9 +21,9 @@
import mlx.core as mx
-# A Mac of 24 GB: with the DiT, the encoder and the VAE loaded, a 16 GB
-# Mac has no room left for the decode (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(24)]
+# Each row loads local model files, so it runs on this Mac
+# (tests/host_box.py).
+pytestmark = pytest.mark.integration
NAME = "check6_z_image_turbo"
diff --git a/tests/image/test_image_check6_flux2_klein.py b/tests/image/test_image_check6_flux2_klein.py
index e4d81c8cc..a895b28f4 100644
--- a/tests/image/test_image_check6_flux2_klein.py
+++ b/tests/image/test_image_check6_flux2_klein.py
@@ -18,8 +18,7 @@
unconditional row each fail it. The bf16 run at 1024 must stay finite,
and its readings print as a share of the 512 gate of each step, as the
manifest records them. The wrong runs at 512 are of the 4B files, those
-of the exact run at 1024 of the 4B and the 9B files, and the 9B cases
-declare a Mac of 24 GB.
+of the exact run at 1024 of the 4B and the 9B files.
Skipped unless the manifests of scripts/image_oracle/check6_flux2_klein.py
are committed, GMLX_IMAGE_FIXTURES holds their arrays and
@@ -37,9 +36,6 @@
from image_fixtures import image_dits
from oracle_fixture import ORACLE, committed, fixture, fx
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
-
# A kind is a variant, with "9b " before it for the 9B files.
NAMES = {"distilled": "check6_flux2_klein", "base": "check6_flux2_klein_base",
"9b distilled": "check6_flux2_klein_9b",
@@ -67,9 +63,8 @@ def family_of(kind: str) -> str:
def _param(kind: str, *rest):
- """A case of ``kind``, the 9B ones on a Mac of 24 GB."""
- marks = [pytest.mark.host_gb(24)] if kind.startswith("9b ") else []
- return pytest.param(kind, *rest, marks=marks)
+ """A case of ``kind``."""
+ return pytest.param(kind, *rest)
KINDS = [_param(k) for k in NAMES]
diff --git a/tests/image/test_image_check6_flux2_klein_edit.py b/tests/image/test_image_check6_flux2_klein_edit.py
index e21e1a668..ff76325d7 100644
--- a/tests/image/test_image_check6_flux2_klein_edit.py
+++ b/tests/image/test_image_check6_flux2_klein_edit.py
@@ -129,9 +129,7 @@ def _encode_gate(pipe) -> tuple:
return "rel L2", float(t["encode_bf16_rel_l2"])
-# A Mac of 16 GB, as most users run (tests/host_box.py).
@pytest.mark.integration
-@pytest.mark.host_gb(16)
def test_reference_latents_of_the_vae_against_the_oracle(check6, pipe):
from gmlx.image import api
from gmlx.image.families.flux2_klein import edit
@@ -180,7 +178,6 @@ def _gates(m) -> dict:
@pytest.mark.integration
-@pytest.mark.host_gb(16)
def test_check6_edit_latents_from_the_conditioning_of_the_oracle(check6,
pipe):
arrays, m = check6
@@ -195,7 +192,6 @@ def test_check6_edit_latents_from_the_conditioning_of_the_oracle(check6,
@pytest.mark.integration
-@pytest.mark.host_gb(16)
def test_check6_edit_fails_without_its_reference(check6, pipe):
arrays, m = check6
gates = _gates(m)
diff --git a/tests/image/test_image_check6_flux2_klein_kv.py b/tests/image/test_image_check6_flux2_klein_kv.py
index e0456a804..1593b4b46 100644
--- a/tests/image/test_image_check6_flux2_klein_kv.py
+++ b/tests/image/test_image_check6_flux2_klein_kv.py
@@ -13,7 +13,7 @@
Skipped unless the manifests of scripts/image_oracle/check6_flux2_klein.py
and check6_flux2_klein_edit.py, each run with --kv, are committed,
GMLX_IMAGE_FIXTURES holds their arrays and KQUANT_TEST_GGUF_DIR the KV
-GGUF the oracle read. The runs declare a Mac of 24 GB."""
+GGUF the oracle read. The runs see this Mac."""
from __future__ import annotations
import os
@@ -65,7 +65,6 @@ def pipe(gguf_dir):
@pytest.mark.integration
-@pytest.mark.host_gb(24)
@pytest.mark.parametrize("case", list(CASES))
def test_check6_kv_latents(pipe, case):
from gmlx.image import api
diff --git a/tests/image/test_image_check6_flux3_klein.py b/tests/image/test_image_check6_flux3_klein.py
index 4e94061f3..f0509dd4a 100644
--- a/tests/image/test_image_check6_flux3_klein.py
+++ b/tests/image/test_image_check6_flux3_klein.py
@@ -37,9 +37,6 @@
from image_fixtures import image_dits
from oracle_fixture import ORACLE, committed, fixture, fx
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
-
NAMES = {"distilled": "check6_flux3_klein", "base": "check6_flux3_klein_base"}
GATE_SIZE = 512
EXACT_SIZE = 1024
diff --git a/tests/image/test_image_check6_flux3_klein_edit.py b/tests/image/test_image_check6_flux3_klein_edit.py
index 7b7625b3a..6546e8e9f 100644
--- a/tests/image/test_image_check6_flux3_klein_edit.py
+++ b/tests/image/test_image_check6_flux3_klein_edit.py
@@ -141,9 +141,7 @@ def _encode_gate(pipe) -> tuple:
return "rel L2", float(t["encode_bf16_rel_l2"])
-# A Mac of 16 GB, as most users run (tests/host_box.py).
@pytest.mark.integration
-@pytest.mark.host_gb(16)
@pytest.mark.parametrize("case", CASES)
def test_reference_latents_of_the_vae_against_the_oracle(check6, pipe, case):
from gmlx.image import api
@@ -197,7 +195,6 @@ def _gates(m, case: str) -> dict:
@pytest.mark.integration
-@pytest.mark.host_gb(16)
@pytest.mark.parametrize("case", CASES)
def test_check6_edit_latents_from_the_conditioning_of_the_oracle(
check6, pipe, case):
@@ -213,7 +210,6 @@ def test_check6_edit_latents_from_the_conditioning_of_the_oracle(
@pytest.mark.integration
-@pytest.mark.host_gb(16)
def test_check6_edit_fails_with_the_references_in_the_other_order(check6,
pipe):
arrays, m = check6
diff --git a/tests/image/test_image_check6_krea2.py b/tests/image/test_image_check6_krea2.py
index 0c7283e69..68155c606 100644
--- a/tests/image/test_image_check6_krea2.py
+++ b/tests/image/test_image_check6_krea2.py
@@ -22,8 +22,9 @@
from image_fixtures import image_dits
from oracle_fixture import ORACLE, committed, fixture, fx
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(16)]
+# Each row loads local model files, so it runs on this Mac
+# (tests/host_box.py).
+pytestmark = pytest.mark.integration
NAME = "check6_krea2"
diff --git a/tests/image/test_image_check6_qwen.py b/tests/image/test_image_check6_qwen.py
index 48153195c..fbe62515b 100644
--- a/tests/image/test_image_check6_qwen.py
+++ b/tests/image/test_image_check6_qwen.py
@@ -54,8 +54,9 @@
import mlx.core as mx
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(16)]
+# Each row loads local model files, so it runs on this Mac
+# (tests/host_box.py).
+pytestmark = pytest.mark.integration
NAME = "check6_qwen_image21"
diff --git a/tests/image/test_image_check6_qwen_edit.py b/tests/image/test_image_check6_qwen_edit.py
index 2b4e833b3..313c7f2e1 100644
--- a/tests/image/test_image_check6_qwen_edit.py
+++ b/tests/image/test_image_check6_qwen_edit.py
@@ -65,9 +65,9 @@
rerun_exact)
from test_image_check2_qwen_edit import _find
-# A Mac of 24 GB: an edit from the files plans an 8.92 GiB peak, over
-# the 8.55 GiB budget of a 16 GB Mac (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(24)]
+# Each row loads local model files, so it runs on this Mac
+# (tests/host_box.py).
+pytestmark = pytest.mark.integration
NAME = "check6_qwen_image21_edit"
CHECK2 = "check2_qwen_image21_edit"
diff --git a/tests/image/test_image_check7.py b/tests/image/test_image_check7.py
index 1689df989..07eb0446e 100644
--- a/tests/image/test_image_check7.py
+++ b/tests/image/test_image_check7.py
@@ -59,9 +59,6 @@
rerun_sites, right_sites, site_process, site_reach,
stock_ops_only, table_sites, table_velocity, thin)
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
-
NAME = "check7_z_image_turbo_lora_512"
diff --git a/tests/image/test_image_check7_community.py b/tests/image/test_image_check7_community.py
index e9b7057ae..83723a0b9 100644
--- a/tests/image/test_image_check7_community.py
+++ b/tests/image/test_image_check7_community.py
@@ -38,6 +38,7 @@
import numpy as np
import pytest
+import host_box
from lora_fixtures import reference_delta, rel_l2
from oracle_fixture import ORACLE, committed, fixture, fx
from block_gates import (exact_process, rerun_exact, rerun_sites,
@@ -46,9 +47,6 @@
mutation_misses, need_blocks, site_gate, step,
turbo_pipe)
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
-
NAME = "check7_z_image_turbo_community_512"
ENV = "GMLX_TEST_LORA_DIR"
DELTA_TOL = 1e-6
@@ -69,7 +67,7 @@ def _ids() -> list:
@functools.lru_cache(maxsize=None)
-def _file(aid: str) -> str:
+def _find(aid: str) -> str:
"""The adapter file under GMLX_TEST_LORA_DIR, checked by sha256."""
root = os.environ.get(ENV)
if not root:
@@ -83,6 +81,13 @@ def _file(aid: str) -> str:
pytest.skip(f"{want['name']} is not under {ENV}")
+def _file(aid: str) -> str:
+ # Outside the cache, so each test that reads the file sees this Mac.
+ path = _find(aid)
+ host_box.real_files()
+ return path
+
+
@pytest.fixture(scope="module")
def pipe(gguf_index):
return turbo_pipe(gguf_index, NAME)
diff --git a/tests/image/test_image_check7_flux2_klein.py b/tests/image/test_image_check7_flux2_klein.py
index da7a5d739..83efa5a32 100644
--- a/tests/image/test_image_check7_flux2_klein.py
+++ b/tests/image/test_image_check7_flux2_klein.py
@@ -54,15 +54,13 @@
import pytest
import flux2_gates as fg
+import host_box
import qwen_gates as qg
from block_gates import exact_process, rerun_exact
from image_fixtures import image_dits
from lora_fixtures import write_adapter
from oracle_fixture import ORACLE, committed, fixture, fx
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
-
ENV = "GMLX_TEST_LORA_DIR"
# id: (the check 3 fixture of the adapter's checkpoint, the pairs the
# oracle reads, their spelling, the pairs gmlx reads)
@@ -120,7 +118,7 @@ def synthetic_adapter(arrays: dict, m: dict) -> dict:
@functools.lru_cache(maxsize=None)
-def _file(aid: str) -> str:
+def _find(aid: str) -> str:
"""The adapter file under GMLX_TEST_LORA_DIR, checked by sha256, or
the synthetic adapter written from its fixture."""
if aid == "synthetic":
@@ -141,6 +139,13 @@ def _file(aid: str) -> str:
pytest.skip(f"{want['name']} is not under {ENV}")
+def _file(aid: str) -> str:
+ # Outside the cache, so each test that reads the file sees this Mac.
+ path = _find(aid)
+ host_box.real_files()
+ return path
+
+
# The float tensors the exact gate's process put in float32, by DiT path.
_EXACT: dict = {}
diff --git a/tests/image/test_image_check7_flux3_klein.py b/tests/image/test_image_check7_flux3_klein.py
index a043359c5..5648d3741 100644
--- a/tests/image/test_image_check7_flux3_klein.py
+++ b/tests/image/test_image_check7_flux3_klein.py
@@ -45,9 +45,6 @@
from lora_fixtures import DIALECTS, write_adapter
from oracle_fixture import ORACLE, committed, fixture, fx
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
-
SYNTHETIC = "check7_flux3_klein_synthetic_512"
SOURCE = "check3_flux3_klein_r1_512"
# The spellings the product gate reads the adapter in: one that stores
diff --git a/tests/image/test_image_check7_krea2.py b/tests/image/test_image_check7_krea2.py
index 4007d44a5..990fa75a5 100644
--- a/tests/image/test_image_check7_krea2.py
+++ b/tests/image/test_image_check7_krea2.py
@@ -39,6 +39,7 @@
import numpy as np
import pytest
+import host_box
import krea2_gates as kg
import qwen_gates as qg
from block_gates import exact_process, rerun_exact
@@ -46,10 +47,6 @@
from lora_fixtures import DIALECTS, write_adapter
from oracle_fixture import ORACLE, committed, fixture, fx
-# A Mac of 24 GB: the gates peak at 20.3 GiB, past the 14 GiB working
-# set of a 16 GB Mac (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(24)
-
COMFY = "check7_krea2_comfy_512"
SYNTHETIC = "check7_krea2_synthetic_512"
SOURCE = "check3_krea2_s4_512"
@@ -88,6 +85,7 @@ def synthetic():
@pytest.fixture(scope="module")
+@host_box.gate
def lora_path(check7):
want = check7[1]["files"]["lora"]
root = os.environ.get(LORA_ENV)
diff --git a/tests/image/test_image_check7_qwen.py b/tests/image/test_image_check7_qwen.py
index ddcea5890..2a932693a 100644
--- a/tests/image/test_image_check7_qwen.py
+++ b/tests/image/test_image_check7_qwen.py
@@ -57,8 +57,8 @@
from gmlx.image.families.qwen_image21 import model as qm
from gmlx.image.lora import install, patch_table, plan, reader
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = pytest.mark.host_gb(16)
+# The tiny rows run on a Mac of 16 GB, as most users run, and the rows on
+# local files on this Mac (tests/host_box.py).
TINY = {v: f"selftest_qwen_image21_tiny_lora_{v}"
for v in qt.rule.LORA_VARIANTS}
@@ -309,6 +309,7 @@ def half(e) -> tuple:
return out
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("layout", LAYOUTS)
@pytest.mark.parametrize("variant,dialects", CASES, ids=IDS)
def test_a_tiny_step_under_a_patch_table_matches_the_merged_oracle(
@@ -348,6 +349,7 @@ def test_a_tiny_step_under_a_patch_table_matches_the_merged_oracle(
assert base[k] >= 1000, (k, base[k])
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("layout", LAYOUTS)
@pytest.mark.parametrize("variant", list(TINY))
def test_each_wrong_table_fails_the_tiny_gate(tmp_path, variant, layout):
@@ -385,6 +387,7 @@ def test_each_wrong_table_fails_the_tiny_gate(tmp_path, variant, layout):
assert top(ratios(pipe, arrays, right, m["name"]))[1] <= 1
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("layout", LAYOUTS)
@pytest.mark.parametrize("variant", list(TINY))
def test_each_structural_table_fails_the_tiny_gate(tmp_path, variant,
@@ -427,6 +430,7 @@ def test_each_structural_table_fails_the_tiny_gate(tmp_path, variant,
"bounds on its most moved measure")
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("layout", LAYOUTS)
def test_the_state_of_one_table_is_not_read_under_another(tmp_path, layout):
"""The prefix K and V and the modulation depend on the adapter. One
@@ -497,6 +501,7 @@ def _step(pipe, arrays: dict, table: str, key: str = "k") -> np.ndarray:
qg.conditioning(arrays, key), patch_table=table), np.float64)
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("layout", LAYOUTS)
def test_a_strength_change_is_a_new_table_with_no_reload(tmp_path,
monkeypatch, layout):
@@ -551,6 +556,7 @@ def boom(*a, **k):
assert np.array_equal(_step(pipe, arrays, t1), full)
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("strength", [0.5, 1.0, -0.4])
def test_a_tiny_step_at_each_strength_equals_the_step_on_merged_weights(
tmp_path, strength):
@@ -573,6 +579,7 @@ def test_a_tiny_step_at_each_strength_equals_the_step_on_merged_weights(
assert fx.rel_l2(_step(pipe, arrays, api.EMPTY_TABLE), want) > 1e-2
+@pytest.mark.host_gb(16)
def test_a_qwen_forward_reads_nothing_from_the_text_row_channel(
tmp_path, monkeypatch):
from gmlx import lora_rows
@@ -599,6 +606,7 @@ def boom(*a, **k):
assert got[f"velocity_{tag}"]["rel_l2"] <= OWN_TOL, got
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("how,kind", (("block", "_Block"),
("slice", "_Slice"),
("split", "_Split")))
@@ -627,6 +635,7 @@ def test_each_placement_gives_the_tiny_step_of_the_merged_oracle(
[(0, ffn)] if how == "slice" else [(0, ffn), (ffn, 2 * ffn)])
+@pytest.mark.host_gb(16)
def test_a_gate_up_key_gives_the_tiny_step_of_the_merged_oracle(pipes,
tmp_path):
"""A file from the fused layout names img_mlp.gate_up, one pair over
@@ -677,6 +686,7 @@ def test_a_gate_up_key_gives_the_tiny_step_of_the_merged_oracle(pipes,
lora.load_adapter(pipes["unfused"], path)
+@pytest.mark.host_gb(16)
def test_a_dense_gate_up_pair_gives_the_tiny_step_of_the_merged_oracle(
pipes, tmp_path):
"""One pair over the whole fused weight with one down: its up is
@@ -722,6 +732,7 @@ def test_a_dense_gate_up_pair_gives_the_tiny_step_of_the_merged_oracle(
assert top(got)[1] >= SEEN and top(got, "step ")[1] >= SEEN
+@pytest.mark.host_gb(16)
@pytest.mark.parametrize("layout", LAYOUTS)
def test_a_header_config_scales_the_tiny_step_as_the_reference_reads_it(
tmp_path, layout):
@@ -777,6 +788,7 @@ def load(name: str, fields: dict | None):
assert whole >= SEEN and step >= SEEN, (label, whole, step)
+@pytest.mark.host_gb(16)
def test_an_alpha_scales_a_pair_in_every_spelling_of_its_keys(pipes,
tmp_path):
"""The rule: a pair is scaled by its alpha tensor over its rank in
@@ -826,6 +838,7 @@ def test_an_alpha_scales_a_pair_in_every_spelling_of_its_keys(pipes,
assert top(got)[1] >= SEEN and top(got, "step ")[1] >= SEEN
+@pytest.mark.host_gb(16)
def test_a_leaf_kept_in_float32_adds_its_delta_in_the_dtype_of_its_product(
tmp_path, monkeypatch):
"""A family may keep the timestep embedder in float32 under bfloat16
@@ -913,6 +926,7 @@ def spy(y, z, b_t):
assert t_err <= gate and t_lost >= 10 * gate
+@pytest.mark.host_gb(16)
def test_load_adapter_refuses_a_pair_of_another_width(tmp_path, pipes):
"""A name of this family with the shape of a relative, such as the
older Qwen-Image transformer's blocks, is refused by its shape."""
@@ -936,6 +950,7 @@ def test_load_adapter_refuses_a_pair_of_another_width(tmp_path, pipes):
tmp_path / "tall.safetensors", "diffusers", tall))
+@pytest.mark.host_gb(16)
def test_the_key_map_resolves_what_the_reference_loader_resolves():
"""The oracle ran diffusers' own Qwen LoRA loader on one pair for
every linear layer of the reference class in every key spelling. The
@@ -984,6 +999,7 @@ def test_the_key_map_resolves_what_the_reference_loader_resolves():
"transformer_blocks_1_img_mlp_gate_layer"))
+@pytest.mark.host_gb(16)
def test_tiny_leaf_sites_under_a_table_are_within_their_gates(tmp_path,
monkeypatch):
"""The site gate of the linear leaves and the product gate of the
@@ -1037,6 +1053,7 @@ def _real_manifest(name: str) -> dict:
return committed(name)
+@pytest.mark.host_gb(16)
def test_the_real_manifests_name_the_adapters_the_rows_load():
"""The committed manifests alone: the dialects of each variant, the
two adapters of the stack and where they meet, one pair fixture, one
diff --git a/tests/image/test_image_cli_flux3_klein.py b/tests/image/test_image_cli_flux3_klein.py
index ac931fee1..0582beae4 100644
--- a/tests/image/test_image_cli_flux3_klein.py
+++ b/tests/image/test_image_cli_flux3_klein.py
@@ -20,8 +20,9 @@
from image_fixtures import image_dits
from lora_fixtures import write_adapter
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(16)]
+# Each row loads local model files, so it runs on this Mac
+# (tests/host_box.py).
+pytestmark = pytest.mark.integration
FAMILY = "flux3_klein"
PROMPT = "a red bicycle leaning on a brick wall"
diff --git a/tests/image/test_image_cli_real.py b/tests/image/test_image_cli_real.py
index d23294e6b..883ca8315 100644
--- a/tests/image/test_image_cli_real.py
+++ b/tests/image/test_image_cli_real.py
@@ -13,8 +13,9 @@
import mlx.core as mx
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(16)]
+# Each row loads local model files, so it runs on this Mac
+# (tests/host_box.py).
+pytestmark = pytest.mark.integration
PROMPT = "a red bicycle leaning on a brick wall"
SEED = 7
diff --git a/tests/image/test_image_load_real.py b/tests/image/test_image_load_real.py
index 7beacb2e4..d18642789 100644
--- a/tests/image/test_image_load_real.py
+++ b/tests/image/test_image_load_real.py
@@ -11,8 +11,9 @@
import mlx.core as mx
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(16)]
+# Each row loads local model files, so it runs on this Mac
+# (tests/host_box.py).
+pytestmark = pytest.mark.integration
def _z_image_file(gguf_index: dict) -> tuple[str, object]:
diff --git a/tests/image/test_image_lora_rules_unchanged.py b/tests/image/test_image_lora_rules_unchanged.py
index 28b17c148..2dcdf3f37 100644
--- a/tests/image/test_image_lora_rules_unchanged.py
+++ b/tests/image/test_image_lora_rules_unchanged.py
@@ -18,6 +18,7 @@
import os
from pathlib import Path
+import host_box
import pytest
from image_fixtures import fixture_view, load_json
@@ -201,6 +202,7 @@ def test_every_lora_file_resolves_as_before():
want[n].get("size")]
if not seen:
pytest.skip(f"none of the recorded files is under {ENV}")
+ host_box.real_files()
same = [n for n in seen if not want[n].get("rslora")]
for name in same:
assert got[name] == want[name], name
@@ -227,6 +229,7 @@ def test_a_rank_stabilized_file_changes_only_by_its_scale():
and want[str(p.relative_to(root))]["size"] == p.stat().st_size]
if not files:
pytest.skip(f"no recorded rank-stabilized file is under {ENV}")
+ host_box.real_files()
maps = keymaps()
for p in files:
path, name = str(p), str(p.relative_to(root))
diff --git a/tests/image/test_image_qwen_real.py b/tests/image/test_image_qwen_real.py
index 655df80b5..2d1841e4c 100644
--- a/tests/image/test_image_qwen_real.py
+++ b/tests/image/test_image_qwen_real.py
@@ -22,8 +22,9 @@
import mlx.core as mx
from mlx.utils import tree_flatten
-# A Mac of 16 GB, as most users run (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(16)]
+# Each row loads local model files, so it runs on this Mac
+# (tests/host_box.py).
+pytestmark = pytest.mark.integration
FLOATS = ("F32", "F16", "BF16")
# Tensors of the Q4_K_M file by (codec, dimensions).
diff --git a/tests/image/test_image_t2i_qwen.py b/tests/image/test_image_t2i_qwen.py
index 4cd63e19f..c640fb531 100644
--- a/tests/image/test_image_t2i_qwen.py
+++ b/tests/image/test_image_t2i_qwen.py
@@ -31,9 +31,9 @@
import mlx.core as mx
-# A Mac of 24 GB, where the run at the family defaults loads every part
-# at once; a 16 GB Mac loads it in stages (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(24)]
+# Each row loads local model files, so it runs on this Mac
+# (tests/host_box.py).
+pytestmark = pytest.mark.integration
NAME = "check6_qwen_image21"
FAMILY = "qwen_image21"
diff --git a/tests/image/test_image_upsample.py b/tests/image/test_image_upsample.py
index 9ed086f46..b0c6b766e 100644
--- a/tests/image/test_image_upsample.py
+++ b/tests/image/test_image_upsample.py
@@ -1,8 +1,10 @@
-"""Prompt upsampling of FLUX.3 klein: the chat turn is BFL's for a prompt
-without references, from the system prompt gmlx ships, and ``gmlx image
+"""Prompt upsampling: the chat turn is the source's for a prompt without
+references, from the system prompt gmlx ships, the caption is what the
+family's parser reads in a reply that thinks first, and ``gmlx image
generate`` releases the upsampler before the encoder loads and encodes
-the caption in place of the typed prompt, but never the negative. The
-chat model and the image pipeline are stubs."""
+the caption in place of the typed prompt, but never the negative.
+Qwen-Image 2.1 ships no system prompt and takes the user's. The chat
+model and the image pipeline are stubs."""
from __future__ import annotations
import contextlib
@@ -18,22 +20,183 @@
from gmlx.image import upsample
from gmlx.image.registry import get_family
-PROMPTS_SHA256 = (
- "361bb6e081d50e54a3dee9abf26327436a9386c83465f881c7848989232028be")
CAPTION = "A red bicycle leans against a weathered brick wall."
-def test_the_chat_turn_is_bfls_t2i_turn():
- fam = get_family("flux3_klein")
- data = (resources.files("gmlx.image.families.flux3_klein")
- / "pup_prompts_external_vlm.json").read_bytes()
- assert hashlib.sha256(data).hexdigest() == PROMPTS_SHA256
+# Each shipped system prompt, pinned by the sha256 of the file its source
+# publishes, and the sampling of the turn. The FLUX.2 klein rows of every
+# size share BFL's text to image system message.
+@pytest.mark.parametrize("family, package, name, sha256, key, sampling", [
+ ("flux3_klein", "flux3_klein", "pup_prompts_external_vlm.json",
+ "361bb6e081d50e54a3dee9abf26327436a9386c83465f881c7848989232028be",
+ "t2i", (0.0, 3584, 1.5, "off")),
+ ("flux2_klein", "flux2_klein", "system_message_upsampling_t2i.txt",
+ "0ad9a42c565716b0a8eb82a442dd4f6cc82d583eb3474f82ccb610c1a29233c7",
+ None, (0.0, 1024, 0.0, "off")),
+ ("flux2_klein_9b_kv", "flux2_klein", "system_message_upsampling_t2i.txt",
+ "0ad9a42c565716b0a8eb82a442dd4f6cc82d583eb3474f82ccb610c1a29233c7",
+ None, (0.0, 1024, 0.0, "off")),
+ ("krea2", "krea2", "expansion.txt",
+ "2a5c24a3a83e9d415d4679a4e1853ddf056bd0329f3af315b27ed72b6ad9c2f0",
+ None, (0.0, 2048, 0.0, "off")),
+ # The typed prompt goes in the {prompt} slot, and the user turn is a
+ # space, as the Space's app.py sends them.
+ ("z_image", "z_image", "pe_prompt_template.txt",
+ "36bdd888b953e9b7a8237dab55bee145a44e96e92e8f6e81a6e4d433a401a0b6",
+ None, (0.0, 2048, 0.0, "off")),
+])
+def test_the_chat_turn_is_the_sources_t2i_turn(family, package, name, sha256,
+ key, sampling):
+ fam = get_family(family)
+ data = (resources.files(f"gmlx.image.families.{package}") / name
+ ).read_bytes()
+ assert hashlib.sha256(data).hexdigest() == sha256
+ system = json.loads(data)[key] if key else data.decode("utf-8")
+ user = "a red bicycle"
+ if "{prompt}" in system:
+ system, user = system.replace("{prompt}", user), " "
req = upsample.chat_request(fam, "a red bicycle")
assert req.messages == (
- {"role": "system", "content": json.loads(data)["t2i"]},
- {"role": "user", "content": "a red bicycle"})
+ {"role": "system", "content": system},
+ {"role": "user", "content": user})
assert (req.temperature, req.max_tokens, req.presence_penalty,
- req.thinking) == (0.0, 3584, 0.0, "off")
+ req.thinking) == sampling
+
+
+KREA_REPLY = """The user wants a red bicycle. Subject: a bicycle, red, \
+nothing else implied. Mood: quiet, everyday. Style options: documentary \
+photo, gouache illustration, 3D render. A photo serves a plain request \
+best. Composition: side view, the whole frame of the bike, a wall behind \
+it for context.
+
+
+Expanded prompt:
+
+A red city bicycle with a black leather saddle and a wicker basket \
+stands on its kickstand against a weathered brick wall, photographed \
+from the side in soft late afternoon light. The paint of the frame is \
+slightly faded, and a long shadow falls across the cobblestones.
+"""
+
+
+def test_krea_2_takes_the_paragraph_after_the_thinking_block():
+ fam = get_family("krea2")
+ got = upsample.upsample(fam, "a red bicycle", lambda req: KREA_REPLY)
+ assert got.startswith("A red city bicycle with a black leather saddle")
+ assert got.endswith("falls across the cobblestones.")
+ # Reasoning written out of the thinking block is not the caption,
+ # though it runs longer.
+ reasoning = ("Step 1: the subject is a bicycle, red, with nothing else "
+ "implied by the request.\nStep 2: a documentary photo "
+ "serves a plain request best, so the style is a photo.\n"
+ "Step 3: a side view shows the whole frame of the bike.")
+ assert upsample.upsample(
+ fam, "a red bicycle", lambda req:
+ f"{reasoning}\n\nA red bicycle against a brick wall.\n") == (
+ "A red bicycle against a brick wall.")
+ # A rule after the caption is not one.
+ assert upsample.upsample(fam, "a red bicycle", lambda req:
+ "A red bicycle.\n\n---\n") == "A red bicycle."
+ # A reply that runs out of tokens while it thinks holds no paragraph.
+ with pytest.raises(upsample.UpsampleError,
+ match="no caption within 2048 tokens"):
+ upsample.upsample(fam, "a red bicycle", lambda req:
+ "\nThe user wants a red bicycle. Style")
+
+
+QWEN_REPLY = """The prompt is short. Draft: {"rewritten_prompt": "a bike"} \
+is too thin; add the setting and the light.
+
+
+```json
+{"rewritten_prompt": "A red bicycle with a wicker basket leans against a \
+weathered brick wall in soft afternoon light.", "wh_ratio": "3:2", \
+"ratio_follow": ""}
+```
+"""
+
+
+def test_qwen_image_takes_rewritten_prompt_from_the_json_after_thinking():
+ fam = get_family("qwen_image21")
+ rewrite = "own prompt: answer with JSON"
+ got = upsample.upsample(fam, "a red bicycle", lambda req: QWEN_REPLY,
+ system=rewrite)
+ assert got == ("A red bicycle with a wicker basket leans against a "
+ "weathered brick wall in soft afternoon light.")
+ # A reply with no JSON object is the caption as written.
+ assert upsample.upsample(
+ fam, "a red bicycle", lambda req: "\nA red bicycle.\n",
+ system=rewrite) == "A red bicycle."
+
+
+@pytest.mark.parametrize("answer", [
+ '{"rewritten_prompt": "A red bicycle leans against a wall, and a red',
+ '{"rewritten_prompt": "", "wh_ratio": "3:2"}',
+ '{"rewritten_prompt": null}',
+ '{"prompt": "A red bicycle."}',
+ "{'rewritten_prompt': 'A red bicycle.'}"])
+def test_a_qwen_image_object_without_the_field_is_an_error(answer):
+ # A reply cut off in its JSON never reaches the encoder as raw JSON.
+ fam = get_family("qwen_image21")
+ with pytest.raises(upsample.UpsampleError,
+ match="a reply with no caption in it"):
+ upsample.upsample(fam, "a red bicycle", lambda req:
+ f"\n{answer}", system="own prompt")
+
+
+def test_qwen_image_reads_a_rewritten_prompt_with_unescaped_quotes():
+ fam = get_family("qwen_image21")
+ reply = ('\n```json\n{"rewritten_prompt": "a sign that reads '
+ '"OPEN" on a door,\\nin red", "wh_ratio": "3:2", '
+ '"ratio_follow": ""}\n```\n')
+ got = upsample.upsample(fam, "a sign", lambda req: reply,
+ system="own prompt")
+ assert got == 'a sign that reads "OPEN" on a door,\nin red'
+ emoji = '{"rewritten_prompt": "a "SALE" sign \\ud83d\\ude00", "x": 1}'
+ assert upsample.upsample(fam, "a sign", lambda req: emoji,
+ system="own prompt") == 'a "SALE" sign \U0001f600'
+ last = '{"rewritten_prompt": "a sign that reads "OPEN" on a door"}'
+ assert upsample.upsample(fam, "a sign", lambda req: last,
+ system="own prompt") == (
+ 'a sign that reads "OPEN" on a door')
+
+
+def test_a_thinking_block_in_another_spelling_is_skipped():
+ fam = get_family("flux3_klein")
+ reply = f"\nThe user wants a bike.\n\n{CAPTION}\n"
+ assert upsample.upsample(fam, "x", lambda req: reply) == CAPTION
+ with pytest.raises(upsample.UpsampleError, match="no caption within"):
+ upsample.upsample(fam, "x", lambda req: "\nThe user")
+
+
+def test_z_image_takes_revised_prompt_or_the_reply_as_written():
+ fam = get_family("z_image")
+ block = '```json\n{"revised_prompt": "A red bicycle."}\n```'
+ assert upsample.upsample(fam, "x", lambda req: block) == "A red bicycle."
+ # As the Space reads it, a reply whose braces hold no revised_prompt
+ # is the caption.
+ sign = "A neon sign that reads {OPEN} above a red bicycle."
+ assert upsample.upsample(fam, "x", lambda req: sign) == sign
+
+
+@pytest.mark.parametrize("reply", [
+ "<|channel>thought\nThe user wants a bike.{caption}",
+ "<|channel|>analysis<|message|>The user wants a bike.<|end|>"
+ "<|start|>assistant<|channel|>final<|message|>{caption}<|return|>"])
+def test_a_reply_in_a_channel_format_gives_its_answer(reply):
+ fam = get_family("flux3_klein")
+ assert upsample.upsample(fam, "x", lambda req: reply.format(
+ caption=CAPTION)) == CAPTION
+
+
+@pytest.mark.parametrize("reply", [
+ "Sure, here is a rewrite.\n\nThe user wants",
+ "A t-shirt printed with the text in bold.",
+ "A poster reading <|channel|> in bold letters..."])
+def test_a_reply_that_ends_in_reasoning_is_an_error(reply):
+ fam = get_family("flux3_klein")
+ with pytest.raises(upsample.UpsampleError, match="no caption within"):
+ upsample.upsample(fam, "x", lambda req: reply)
def test_an_empty_reply_is_an_error_not_the_typed_prompt():
@@ -62,8 +225,11 @@ def cli(tmp_path, monkeypatch):
p = tmp_path / name
p.write_bytes(b"x")
files[name] = str(p)
- fam = get_family("flux3_klein")
- rec = SimpleNamespace(events=[], turns=[], model=None)
+ # ``fam`` is the family the DiT identifies as, and ``reply`` what the
+ # upsampler writes.
+ rec = SimpleNamespace(events=[], turns=[], model=None,
+ fam=get_family("flux3_klein"), reply=f" {CAPTION}\n",
+ rendered=None)
class Model:
pass
@@ -73,13 +239,15 @@ def load_model(path):
rec.model = weakref.ref(model)
rec.events.append("upsampler")
tok = SimpleNamespace(
- chat_template="{% if enable_thinking %}{% endif %}")
+ chat_template="{% if enable_thinking %}{% endif %}",
+ apply_chat_template=lambda messages, **kw: rec.rendered or (
+ "\n" if kw.get("enable_thinking") else ""))
return model, None, tok
def generate(model, tok, prompt, **kw):
rec.events.append("chat")
rec.turns.append({"prompt": prompt, **kw})
- return f" {CAPTION}\n"
+ return rec.reply
class Enc:
owned, released = True, False
@@ -93,7 +261,7 @@ def load(dit, **kw):
rec.events.append(("image", rec.model is not None
and rec.model() is not None))
return SimpleNamespace(encoder=Enc(), variant="base",
- load_seconds={}, family=fam)
+ load_seconds={}, family=rec.fam)
def encode_text(pipe, prompt, images=(), ref_latents=True):
rec.events.append(("encode", prompt))
@@ -108,7 +276,7 @@ def encode_text(pipe, prompt, images=(), ref_latents=True):
monkeypatch.setattr(tool_preflight, "check_or_exit", lambda *a, **k: None)
monkeypatch.setattr(image_preflight, "identify_dit",
lambda p, **k: SimpleNamespace(
- family=fam, variant="base", config={}))
+ family=rec.fam, variant="base", config={}))
monkeypatch.setattr(image_preflight, "check", lambda *a, **k: None)
monkeypatch.setattr(api, "load", load)
monkeypatch.setattr(api, "encode_text", encode_text)
@@ -143,7 +311,7 @@ def test_the_cli_releases_the_upsampler_before_the_encoder_loads(
assert turn["prompt"] == "a red bicycle"
assert turn["system_prompt"] == upsample.system_prompt(fam)
assert (turn["temp"], turn["max_tokens"], turn["presence_penalty"],
- turn["template_kwargs"]) == (0.0, 3584, 0.0,
+ turn["template_kwargs"]) == (0.0, 3584, 1.5,
{"enable_thinking": False})
record = json.loads(capsys.readouterr().out)
assert record["settings"]["revised_prompt"] == CAPTION
@@ -189,3 +357,62 @@ def test_an_img2img_prompt_runs_as_typed_and_loads_no_upsampler(
assert "--init-image runs as typed" in capsys.readouterr().out
assert "upsampler" not in cli.events
assert ("encode", "the same scene at night") in cli.events
+
+
+@pytest.mark.parametrize("given", [None, "missing.txt"])
+def test_qwen_image_without_a_prompt_file_refuses_and_names_qwens(
+ cli, tmp_path, capsys, given):
+ from gmlx.commands.cli import umbrella_main
+
+ cli.fam = get_family("qwen_image21")
+ # Without an upsampler, its prompt runs as typed with no warning.
+ assert umbrella_main(cli.base + ["--prompt", "a red bicycle",
+ "--output", str(tmp_path / "a.png")]) == 0
+ assert "expects" not in capsys.readouterr().out
+ args = ["--prompt", "a red bicycle", "--upsampler",
+ cli.files["writer.gguf"], "--output", str(tmp_path / "b.png")]
+ if given:
+ args += ["--upsample-prompt", str(tmp_path / given)]
+ assert umbrella_main(cli.base + args) == 2
+ err = capsys.readouterr().err
+ assert "set upsample_prompt on its image_models entry" in err
+ assert ("system_prompt.txt in the Hugging Face repository "
+ "Qwen/Qwen-Image-2.1-PE-T2I, under the Qwen Research License, "
+ "for non-commercial use only, and `gmlx pull qwen-image` "
+ "fetches it") in err
+ assert ("no such file" in err) == bool(given)
+ assert "upsampler" not in cli.events
+
+
+def test_an_atem_reply_that_starts_mid_header_gives_its_caption(
+ cli, tmp_path):
+ from gmlx.commands.cli import umbrella_main
+
+ # With thinking off, the prompt ends at the header the reply finishes.
+ cli.rendered = "<|start|>user<|message|>a red bicycle<|eot|><|start|>assistant"
+ cli.reply = f" to=user<|message|>{CAPTION}<|eot|>"
+ assert umbrella_main(cli.base + [
+ "--prompt", "a red bicycle", "--upsampler", cli.files["writer.gguf"],
+ "--output", str(tmp_path / "a.png")]) == 0
+ assert ("encode", CAPTION) in cli.events
+
+
+def test_qwen_image_runs_the_users_prompt_file_and_encodes_its_rewrite(
+ cli, tmp_path):
+ from gmlx.commands.cli import umbrella_main
+
+ cli.fam = get_family("qwen_image21")
+ cli.reply = QWEN_REPLY
+ own = tmp_path / "rewrite.txt"
+ own.write_text("Rewrite the prompt. Think, then answer with JSON.\n")
+ assert umbrella_main(cli.base + [
+ "--prompt", "a red bicycle", "--upsampler", cli.files["writer.gguf"],
+ "--upsample-prompt", str(own),
+ "--output", str(tmp_path / "a.png")]) == 0
+ (turn,) = cli.turns
+ assert turn["system_prompt"] == ("Rewrite the prompt. Think, then answer "
+ "with JSON.")
+ assert (turn["presence_penalty"], turn["template_kwargs"]) == (
+ 1.5, {"enable_thinking": False})
+ assert ("encode", "A red bicycle with a wicker basket leans against a "
+ "weathered brick wall in soft afternoon light.") in cli.events
diff --git a/tests/load/test_vlm_deepseek41v.py b/tests/load/test_vlm_deepseek41v.py
index bf30a388e..9332291ea 100644
--- a/tests/load/test_vlm_deepseek41v.py
+++ b/tests/load/test_vlm_deepseek41v.py
@@ -341,6 +341,7 @@ def test_pixels_normalize_with_the_reference_constants():
assert np.allclose(row[196:392], (127 / 255 - 0.5) / 0.5)
+@pytest.mark.integration
@pytest.mark.skipif(not os.path.exists(MMPROJ), reason="local mmproj absent")
def test_real_mmproj_remaps_onto_the_tower():
from gmlx.load.wire import load_gguf_wire_bytes
diff --git a/tests/load/test_vlm_deepseek4v.py b/tests/load/test_vlm_deepseek4v.py
index e371511b5..1f9ce4e3f 100644
--- a/tests/load/test_vlm_deepseek4v.py
+++ b/tests/load/test_vlm_deepseek4v.py
@@ -262,6 +262,7 @@ def test_image_processor_geometry_matches_reference_rules():
assert np.allclose(row[196:392], (127 / 255 - 0.5) / 0.5)
+@pytest.mark.integration
@pytest.mark.skipif(not os.path.exists(MMPROJ), reason="local mmproj absent")
def test_real_mmproj_remaps_onto_the_tower():
from gmlx.load.wire import load_gguf_wire_bytes
diff --git a/tests/models/test_embedders_reference.py b/tests/models/test_embedders_reference.py
index 36accaaf8..c1a542fde 100644
--- a/tests/models/test_embedders_reference.py
+++ b/tests/models/test_embedders_reference.py
@@ -40,6 +40,7 @@
import sys
from pathlib import Path
+import host_box
import numpy as np
import pytest
@@ -65,6 +66,7 @@
@pytest.fixture
+@host_box.gate
def mlx_dir() -> Path:
d = os.environ.get("GMLX_TEST_MLX_DIR")
if not d:
diff --git a/tests/serve/test_image_route_upsample.py b/tests/serve/test_image_route_upsample.py
index ef029c04f..b4db19fb1 100644
--- a/tests/serve/test_image_route_upsample.py
+++ b/tests/serve/test_image_route_upsample.py
@@ -43,7 +43,7 @@ def server(monkeypatch):
id="flux3", path="/flux3.gguf", family=None, variant=None,
defaults={}, adapters=(), vae="/vae.gguf", encoder="/enc.gguf",
encoder_model=None, encoder_mmproj=None, encoder_release="auto",
- upsampler="writer", upsample=True)
+ upsampler="writer", upsample=True, upsample_prompt=None)
monkeypatch.setattr(bridge_vlm, "resolve_image_request",
lambda field: (spec.path, spec))
writer = SimpleNamespace(load_signature=lambda: ("/writer.gguf",),
@@ -141,6 +141,7 @@ def sample(pipe, cond, *, seed, **kw):
pipe=pipe, spec=spec,
plan=SimpleNamespace(staged=True, keep_encoder=False, cfg=4.0))
rec.served = served
+ rec.spec = spec
return rec
@@ -161,7 +162,7 @@ def test_the_server_sends_bfls_turn_and_encodes_the_caption(server):
(sent,) = server.requests
assert sent.model == "writer" and sent.stream is False
assert (sent.temperature, sent.max_tokens, sent.presence_penalty,
- sent.thinking) == (0.0, 3584, 0.0, "off")
+ sent.thinking) == (0.0, 3584, 1.5, "off")
fam = get_family("flux3_klein")
assert server.messages == [[
{"role": "system", "content": system_prompt(fam)},
@@ -188,6 +189,27 @@ def test_a_rerun_with_a_new_seed_writes_no_caption_and_loads_no_encoder(
assert server.seeds == [1, 2]
+def test_the_entrys_prompt_file_replaces_the_shipped_one_and_its_captions(
+ server, tmp_path):
+ """Captions are keyed by the file's content: an edit to it writes a
+ new caption, and the old prompt's caption is not reused."""
+ own = tmp_path / "rewrite.txt"
+ own.write_text("Describe the scene in one long sentence.\n")
+ server.spec.upsample_prompt = str(own)
+ _request(server, {"prompt": "a red bicycle", "seed": 1})
+ _request(server, {"prompt": "a red bicycle", "seed": 2})
+ own.write_text("Describe the scene in two long sentences.\n")
+ _request(server, {"prompt": "a red bicycle", "seed": 3})
+ assert server.events.count("chat") == 2
+ assert [m[0]["content"] for m in server.messages] == [
+ "Describe the scene in one long sentence.",
+ "Describe the scene in two long sentences."]
+ server.spec.upsample_prompt = None
+ _request(server, {"prompt": "a red bicycle", "seed": 4})
+ assert server.messages[-1][0]["content"] == system_prompt(
+ get_family("flux3_klein"))
+
+
def test_an_upsampler_with_too_small_a_context_is_a_server_error(server):
"""The typed prompt is not the cause, so the answer is no 400 on it."""
server.too_long = True
@@ -199,8 +221,10 @@ def test_an_upsampler_with_too_small_a_context_is_a_server_error(server):
def test_an_edit_ignores_prompt_upsampling_on_every_family():
+ import dataclasses
+
spec = SimpleNamespace(id="qwen", upsampler=None, upsample=True)
- fam = get_family("qwen_image21")
+ fam = dataclasses.replace(get_family("qwen_image21"), upsample=None)
assert image_upsample.upsampler_for(spec, fam, True, True) is None
with pytest.raises(svc.ImageRequestError, match="takes no prompt"):
image_upsample.upsampler_for(spec, fam, True, False)
diff --git a/tests/serve/test_serve_apc_engagement.py b/tests/serve/test_serve_apc_engagement.py
index ccbba9d32..789a58c94 100644
--- a/tests/serve/test_serve_apc_engagement.py
+++ b/tests/serve/test_serve_apc_engagement.py
@@ -29,10 +29,9 @@ class hides.
import mlx.core as mx # noqa: E402
-# A Mac of 24 GB: the big rows load models of 12 and 15 GB, which the 14
-# GiB working set of a 16 GB Mac cannot hold beside the stores the tiers
-# make (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(24)]
+# Each row loads local model files, so it runs on this Mac
+# (tests/host_box.py).
+pytestmark = pytest.mark.integration
GREEDY = lambda x: mx.argmax(x, axis=-1) # noqa: E731
N_DECODE = 8
diff --git a/tests/serve/test_server_bridge.py b/tests/serve/test_server_bridge.py
index 31786675a..a594e8b12 100644
--- a/tests/serve/test_server_bridge.py
+++ b/tests/serve/test_server_bridge.py
@@ -131,6 +131,9 @@ def test_register_resolved_models_waits_for_build_lock():
assert done.is_set()
finally:
mod._kq_residency_pool = saved
+ # the config with no models would answer every later test's
+ # request that names no model on this worker
+ bridge.clear_resolved_models()
def test_the_module_runs_the_copied_server_with_the_bridge(monkeypatch):
diff --git a/tests/serve/test_tick_guard.py b/tests/serve/test_tick_guard.py
index ca102c65e..799a97bba 100644
--- a/tests/serve/test_tick_guard.py
+++ b/tests/serve/test_tick_guard.py
@@ -232,3 +232,29 @@ def test_metal_oom_still_takes_the_memory_ladder(wrapped, monkeypatch):
script.append(OOM_ERR)
tick(gen)
assert gen.removed == [8] # largest row retired, as before
+
+
+def test_a_rebuilt_row_penalizes_its_delivered_tokens(wrapped, monkeypatch):
+ from gmlx._vendor.mlx_lm.sample_utils import make_logits_processors
+
+ gen, script, _ = wrapped
+ monkeypatch.setattr(tg.mx, "clear_cache", lambda: None)
+ procs = make_logits_processors(presence_penalty=1.5)
+ gen._unprocessed_sequences[1] = (8, [4, 5, 6], 16, {}, procs, None)
+ gen._generation_batch = FakeGB([7, 8], [1, 3])
+ # The decode loop drives the row's processors as it delivers 21, 22.
+ for history in ([4, 5, 6], [4, 5, 6, 21], [4, 5, 6, 21, 22]):
+ procs[0](tg.mx.array(history), tg.mx.zeros((1, 32)))
+ script.append(([], [resp(7, 11), resp(8, 21)]))
+ script.append(([], [resp(8, 22)]))
+ script.extend([MEM_ERR, MEM_ERR])
+ for _ in range(4):
+ tick(gen)
+ uid, replay, _rem, _kw, lp, _tc = gen._unprocessed_sequences[0]
+ assert uid == 8 and replay == [4, 5, 6, 21, 22]
+ assert lp[0] is not procs[0]
+ # An APC hit prefills only the uncached tail, so the replay's first
+ # call sees just that.
+ out = lp[0](tg.mx.array([22]), tg.mx.zeros((1, 32))).reshape(-1)
+ assert {t: -v for t, v in enumerate(out.tolist()) if v} == {
+ 21: 1.5, 22: 1.5}
diff --git a/tests/spec/test_context_copy_integration.py b/tests/spec/test_context_copy_integration.py
index 6debd00d3..b3bf11c18 100644
--- a/tests/spec/test_context_copy_integration.py
+++ b/tests/spec/test_context_copy_integration.py
@@ -11,6 +11,7 @@
import os
+import host_box
import mlx.core as mx
import pytest
@@ -74,6 +75,7 @@ def busiest(schedule):
@pytest.fixture(scope="module")
+@host_box.gate
def pair():
path = os.environ.get("KQUANT_TEST_MTP_GGUF")
if not path or not os.path.isfile(os.path.expanduser(path)):
diff --git a/tests/spec/test_full_prompt_prefill.py b/tests/spec/test_full_prompt_prefill.py
index 0120de496..e2cebd704 100644
--- a/tests/spec/test_full_prompt_prefill.py
+++ b/tests/spec/test_full_prompt_prefill.py
@@ -26,17 +26,9 @@
import mlx.core as mx
-def _this_mac_gb() -> int:
- """This Mac's memory in GiB, as the real device reports it."""
- try:
- return int(host_box.real_device_info()["memory_size"]) // host_box.GB
- except Exception:
- return int(host_box.RUNNER_GB)
-
-
# The real-model tier: each test loads a model that only a large Mac holds,
-# so it runs on a Mac of this one's memory (tests/host_box.py).
-pytestmark = [pytest.mark.integration, pytest.mark.host_gb(_this_mac_gb())]
+# so it runs on this Mac (tests/host_box.py).
+pytestmark = pytest.mark.integration
GREEDY = lambda x: mx.argmax(x, axis=-1)
@@ -80,6 +72,7 @@ def _this_mac_gb() -> int:
@pytest.fixture(scope="session")
+@host_box.gate
def mtp_gguf():
path = os.environ.get("KQUANT_TEST_MTP_GGUF")
if not path:
@@ -91,6 +84,7 @@ def mtp_gguf():
@pytest.fixture(scope="session")
+@host_box.gate
def mtp_draft_gguf():
path = os.environ.get("KQUANT_TEST_MTP_DRAFT_GGUF")
if not path:
@@ -522,42 +516,6 @@ def test_multi_length_first_token_parity(mtp_model):
f"length={length}: first token")
-def test_prefill_step_env_override(monkeypatch):
- """Stock nulls prefill_step_size for speculative batches; the MTP
- re-enable must resolve None from the serve-side PREFILL_STEP_SIZE env
- override rather than pin the stock default. Exercised directly on the
- patched prompt_step (no model needed): the resolution runs before the
- needs_processing gate."""
- import types
-
- from gmlx._vendor.mlx_vlm.generate.ar import PromptProcessingBatch
-
- import gmlx.spec.mtp_prefill as mtp_prefill
-
- mtp_prefill.install_full_prompt_mtp_prefill()
- monkeypatch.setattr(mtp_prefill, "_mtp_prefill_init", lambda s: None)
-
- def fake_batch():
- # _inputs_embeds: the span-aware wrapper (media_spans) reads it
- # when a serve test has layered it over prompt_step.
- return types.SimpleNamespace(
- draft_kind="mtp",
- prefill_step_size=None,
- needs_processing=lambda: False,
- _inputs_embeds=None,
- )
-
- monkeypatch.setenv("PREFILL_STEP_SIZE", "97")
- fake = fake_batch()
- assert PromptProcessingBatch.prompt_step(fake) == 0
- assert fake.prefill_step_size == 97
-
- monkeypatch.delenv("PREFILL_STEP_SIZE")
- fake = fake_batch()
- PromptProcessingBatch.prompt_step(fake)
- assert fake.prefill_step_size == 2048
-
-
def test_prefill_step_env_override_e2e(mtp_model, monkeypatch):
"""Serve-shaped construction (prefill_step_size=None, no APC): the init
re-enable restores chunking at the env-provided step, and first-token
diff --git a/tests/spec/test_mtp_finish_thinking.py b/tests/spec/test_mtp_finish_thinking.py
index c85576af3..928a7bd67 100644
--- a/tests/spec/test_mtp_finish_thinking.py
+++ b/tests/spec/test_mtp_finish_thinking.py
@@ -13,7 +13,9 @@
import mlx.core as mx
import gmlx.spec.speculative as spec
+from gmlx._vendor.mlx_lm.sample_utils import make_logits_processors
from gmlx.gen.thinking_budget import MTPFinishThinking
+from gmlx.spec.row_procs import SpecRowProcessors
from test_mtp_preempt_resume import (
VOCAB, _ArmableDrafter, _EchoDrafter, _VerifyEchoLM)
@@ -292,3 +294,82 @@ def fake_rounds(model, drafter, lm, prompt_cache, **kw):
assert list(gen) == []
assert captured["thinking_hook"] is None # popped, not honored at B>1
assert not hasattr(cache, "_kq_mtp_thinking_hook")
+
+
+# -- presence and frequency penalties across a forced close -----------------
+
+
+class _RefusedDraftDrafter(_EchoDrafter):
+ """Drafts the echo continuation, then a token the target refuses."""
+
+ def draft_block(self, b, hidden, kv, n, sampler, dtype, **kw):
+ if not isinstance(b, mx.array):
+ b = mx.array([b])
+ b0 = b.reshape(-1).astype(mx.int32)
+ return mx.stack([(b0 + 1) % VOCAB, (b0 + 9) % VOCAB],
+ axis=1).astype(dtype)
+
+
+def _penalty_rows(prompt):
+ """A presence penalty between two recorders, so each call's penalties
+ are the difference of what they saw."""
+ raw, out = [], []
+
+ def before(tokens, logits):
+ raw.append(logits)
+ return logits
+
+ def after(tokens, logits):
+ out.append((int(tokens.shape[0]), logits))
+ return logits
+
+ procs = [before, *make_logits_processors(presence_penalty=0.5), after]
+ rows = SpecRowProcessors.from_processors(procs, prompt)
+
+ def penalties_at(n):
+ """The penalties of the first call whose context has n tokens."""
+ i = next(i for i, (m, _) in enumerate(out) if m == n)
+ diff = (raw[i] - out[i][1]).reshape(-1).tolist()
+ return {t: v for t, v in enumerate(diff) if v}
+ return rows, penalties_at
+
+
+def test_a_forced_close_after_a_refused_draft_keeps_the_penalties_exact():
+ prompt = [1, 2]
+ rows, penalties_at = _penalty_rows(prompt)
+ rows.start(5)
+ hook = _hook([S1, S2, END])
+ shared = {"full": (mx.zeros((1, 2, 4, 4)), mx.zeros((1, 2, 4, 4)))}
+ gen = spec._owned_decode_rounds(
+ SimpleNamespace(), _RefusedDraftDrafter(cap=0), _VerifyEchoLM(),
+ [_FakeCache(width=1)], hidden=mx.zeros((1, 4, 8)), b=5,
+ shared_kv=shared, seed_tokens=None, emitted=1, max_tokens=8,
+ sampler=None, draft_block_size=None, thinking_hook=hook,
+ row_procs=rows)
+ out = []
+ for tok in gen:
+ out.append(int(tok))
+ if len(out) == 1:
+ hook.request_close()
+ # The first round takes 6 and refuses 14 for the target's 7.
+ assert out[:5] == [6, 7, S1, S2, END]
+ committed = [5, 6, 7, S1, S2, END]
+ assert penalties_at(len(prompt) + len(committed)) == {
+ t: 0.5 for t in committed}
+
+
+def test_a_batch_forced_close_keeps_the_penalties_exact():
+ prompt = [9, 10]
+ rows, penalties_at = _penalty_rows(prompt)
+ rows.start(1)
+ hook = _hook([S1, S2, END], budget=2, forced_ids=[14, 15, END])
+ gen = spec._owned_decode_rounds_batch(
+ SimpleNamespace(), _ArmableDrafter(cap=0), _VerifyEchoLM(),
+ [_FakeCache(width=1)], hidden=None, b=[1], shared_kv=None,
+ seed_tokens=None, emitted=[1], max_tokens=8, sampler=None,
+ draft_block_size=None, thinking_hook=hook, row_procs=[rows])
+ out = [t for toks, _ in gen for t in toks if t is not None]
+ assert out[:6] == [2, 3, 4, 14, 15, END]
+ committed = [1, 2, 3, 4, 14, 15, END]
+ assert penalties_at(len(prompt) + len(committed)) == {
+ t: 0.5 for t in committed}
diff --git a/tests/spec/test_mtp_prompt_step.py b/tests/spec/test_mtp_prompt_step.py
index ba9d9119f..121c5d5c1 100644
--- a/tests/spec/test_mtp_prompt_step.py
+++ b/tests/spec/test_mtp_prompt_step.py
@@ -252,3 +252,39 @@ def test_a_multi_row_batch_logs_the_skipped_owned_lookup(monkeypatch, caplog):
text = " ".join(r.getMessage() for r in caplog.records)
assert "owned MTP prefix lookup skipped" in text
assert "APC skipped" not in text
+
+
+def test_prefill_step_env_override(monkeypatch):
+ """Stock nulls prefill_step_size for speculative batches; the MTP
+ re-enable must resolve None from the serve-side PREFILL_STEP_SIZE env
+ override rather than pin the stock default. Exercised directly on the
+ patched prompt_step (no model needed): the resolution runs before the
+ needs_processing gate."""
+ import types
+
+ from gmlx._vendor.mlx_vlm.generate.ar import PromptProcessingBatch
+
+ import gmlx.spec.mtp_prefill as mtp_prefill
+
+ mtp_prefill.install_full_prompt_mtp_prefill()
+ monkeypatch.setattr(mtp_prefill, "_mtp_prefill_init", lambda s: None)
+
+ def fake_batch():
+ # _inputs_embeds: the span-aware wrapper (media_spans) reads it
+ # when a serve test has layered it over prompt_step.
+ return types.SimpleNamespace(
+ draft_kind="mtp",
+ prefill_step_size=None,
+ needs_processing=lambda: False,
+ _inputs_embeds=None,
+ )
+
+ monkeypatch.setenv("PREFILL_STEP_SIZE", "97")
+ fake = fake_batch()
+ assert PromptProcessingBatch.prompt_step(fake) == 0
+ assert fake.prefill_step_size == 97
+
+ monkeypatch.delenv("PREFILL_STEP_SIZE")
+ fake = fake_batch()
+ PromptProcessingBatch.prompt_step(fake)
+ assert fake.prefill_step_size == 2048
diff --git a/tests/spec/test_spec_row_procs.py b/tests/spec/test_spec_row_procs.py
index 87ec3b4aa..617661edd 100644
--- a/tests/spec/test_spec_row_procs.py
+++ b/tests/spec/test_spec_row_procs.py
@@ -45,6 +45,8 @@ def _gated(fn, active=False):
def _plain(procs_factory, b, n, prompt=PROMPT, lm_cls=None):
"""Plain greedy decode with the stock decode loop's processor calls."""
lm, procs = (lm_cls or _PreferBadLM)(), procs_factory()
+ for p in procs: # the prompt step that sampled b
+ p(mx.array(prompt), lm(mx.array([prompt])).logits[:, -1, :])
ctx, out, cur = list(prompt) + [b], [], b
for _ in range(n):
logits = lm(mx.array([[cur]])).logits[:, -1, :]
diff --git a/tests/test_host_box.py b/tests/test_host_box.py
index df3bcd157..5166431bc 100644
--- a/tests/test_host_box.py
+++ b/tests/test_host_box.py
@@ -1,4 +1,5 @@
-"""The simulated host reports through every free-RAM read (host_box.py)."""
+"""The simulated host reports through every free-RAM read, and a test on
+model files from this Mac sees this Mac (host_box.py)."""
from __future__ import annotations
import os
@@ -18,6 +19,49 @@ def _reads():
kernel_vm.available_bytes())
+def _near(got, want) -> bool:
+ # Free RAM moves between two reads on a busy Mac.
+ return abs(got - want) <= max(2 << 30, want // 10)
+
+
+def _metal() -> dict | None:
+ # The Metal device, as gmlx reads it; under KQUANT_FORCE_CPU the
+ # default device is the CPU, which reports no memory.
+ from gmlx.device import metal_info
+
+ return metal_info()
+
+
+# Session scoped like gguf_dir and asked for before it, so it reads the
+# device as the test sets up, before the gate fixture does, in any order.
+@pytest.fixture(scope="session")
+def device_before_the_gate():
+ return _metal()
+
+
+# A test on GGUFs from this Mac never runs in CI and prices its models for
+# this Mac, from its setup on.
+def test_a_test_on_local_ggufs_sees_this_mac_from_its_setup(
+ device_before_the_gate, gguf_dir):
+ import host_box
+
+ real = host_box.real_device_info()
+ assert device_before_the_gate == real
+ assert _metal() == real
+ pages = os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE")
+ assert pages == real["memory_size"]
+
+
+def test_a_test_on_local_ggufs_reads_the_free_ram_of_this_mac(gguf_dir):
+ import host_box
+
+ apc_read, vm_stat_read, available = _reads()
+ real = host_box.real("snapshot")()
+ assert _near(available, real["free"] + real["inactive"])
+ assert _near(vm_stat_read, host_box.real("vm_stat_free_ram")())
+ assert _near(apc_read, host_box.real("apc_free_ram")())
+
+
def test_the_ci_runner_reports_its_free_ram_through_the_stock_read():
"""The copied apc's stock read (psutil, else vm_stat) would see the
real Mac; under the simulation it reports the runner, under the 2 GB
diff --git a/tests/test_host_box_module_mac.py b/tests/test_host_box_module_mac.py
new file mode 100644
index 000000000..debf00123
--- /dev/null
+++ b/tests/test_host_box_module_mac.py
@@ -0,0 +1,73 @@
+"""A module that declares a Mac for its unit rows: a row on local GGUFs
+still sees this Mac, from its setup or from the fetch of a gate in its body
+(host_box.py)."""
+from __future__ import annotations
+
+import os
+
+import pytest
+
+pytestmark = [
+ pytest.mark.skipif(os.environ.get("GMLX_TEST_REAL_HOST") == "1",
+ reason="the real host"),
+ pytest.mark.host_gb(16),
+]
+
+
+def _metal() -> dict | None:
+ # The Metal device, as gmlx reads it; under KQUANT_FORCE_CPU the
+ # default device is the CPU, which reports no memory.
+ from gmlx.device import metal_info
+
+ return metal_info()
+
+
+def _metal_memory() -> int:
+ from gmlx.device import metal_value
+
+ return metal_value("memory_size")
+
+
+def _ram(sysconf=None):
+ sysconf = sysconf or os.sysconf
+ return sysconf("SC_PHYS_PAGES") * sysconf("SC_PAGE_SIZE")
+
+
+def _real_ram():
+ import host_box
+
+ return _ram(host_box.real("sysconf"))
+
+
+def _reads():
+ from gmlx._vendor.mlx_vlm import apc
+ from gmlx.serve import kernel_vm
+
+ return apc._free_ram_bytes(), kernel_vm.available_bytes()
+
+
+def test_a_row_on_local_ggufs_sees_this_mac(gguf_dir):
+ import host_box
+
+ assert _metal() == host_box.real_device_info()
+ assert _ram() == _real_ram()
+
+
+def test_a_unit_row_sees_the_mac_of_its_module():
+ import host_box
+
+ box, vm = host_box.mac(16)
+ assert _metal_memory() == _ram() == box["memory_size"]
+ assert set(_reads()) == {vm["free"] + vm["inactive"]}
+
+
+def test_a_gate_fetched_in_the_body_switches_the_row_to_this_mac(request):
+ """As a row that picks its model folder as it runs does, whether or
+ not an earlier row set the folder up."""
+ import host_box
+
+ box = host_box.mac(16)[0]
+ assert _metal_memory() == _ram() == box["memory_size"]
+ request.getfixturevalue("gguf_dir")
+ assert _metal() == host_box.real_device_info()
+ assert _ram() == _real_ram()