feat(image): flux.3 klein prompt upsampling

CHANGELOG.md

@@ -8,6 +8,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ### Added
 
+- 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`.
 - `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

THIRD_PARTY_NOTICES.md

@@ -21,7 +21,7 @@ in [`licenses/`](licenses/).
 | [system-one-adapter-python](https://github.com/typesafe-ai/system-one-adapter-python) | MIT | The official Jev confidence formulas in `gmlx/systemone/confidence.py`, from `src/system_one_adapter/_utils/confidence_metrics.py` at commit `e1d4cc9`. |
 | [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`. |
+| [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`. |
 | [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. |

docs/cli.md

@@ -1068,6 +1068,8 @@ 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. |
+| `--no-upsample` | Off | Run the prompt as typed, without the upsampler. |
 | `--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. |

docs/config.md

@@ -271,6 +271,8 @@ entries.
 | <a id="image_modelsfamily"></a>`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. |
 | <a id="image_modelsvariant"></a>`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. |
 | <a id="image_modelsadapters"></a>`adapters` | none | LoRA files applied to every request, each a path or `{path, strength}` with strength 1 by default |
+| <a id="image_modelsupsampler"></a>`upsampler` | none | The `models` id of the chat model that rewrites each FLUX.3 klein prompt, as [Prompt upsampling](images.md#prompt-upsampling) describes |
+| <a id="image_modelsupsample"></a>`upsample` | `true` | Rewrite prompts with the `upsampler`. `false` runs them as typed, unless a request sends `prompt_upsampling: true`. |
 | <a id="image_modelsdefaults"></a>`defaults` | the family's | Request defaults: `size` as `WxH`, `steps`, `guidance`, `sampler` (`euler` or `uni_pc_bh2`) and `negative_prompt` |
 | <a id="image_modelsencoder_release"></a>`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 |
 | <a id="image_modelspin"></a>`pin` | `false` | Load at start and stay loaded, as a model's [`pin`](#modelspin) |

docs/image-api.md

@@ -47,7 +47,8 @@ model, and no context figures.
 | `steps` | [Steps](glossary.md#steps), 1 to 200, in place of `quality`. |
 | `seed` | The seed of the first image, random without it. |
 | `guidance` | [Guidance](glossary.md#guidance), 1 to 30. 1 turns it off. |
-| `negative_prompt` | What to steer away from, read when `guidance` is above 1. |
+| `negative_prompt` | What to steer away from, read when `guidance` is above 1. It is never rewritten. |
+| `prompt_upsampling` | `false` runs the prompt as typed on an entry with an [upsampler](images.md#prompt-upsampling), and `true` rewrites it on one with `upsample: false`. |
 | `sampler` | [Sampler](glossary.md#sampler), `euler` or `uni_pc_bh2`. |
 | `scheduler` | `default` only, the model's own schedule. |
 | `strength` | How far an edit or a variation departs from its image, from 0.01 to 1. |
@@ -73,6 +74,11 @@ ignored, and the log names them in a warning. The server writes no file
 that a request names, so `output_path`, `output_dir` and `response_format:
 "path"` get 400.
 
+On an entry with an [upsampler](images.md#prompt-upsampling), each
+image of the answer carries the caption the prompt was rewritten to as
+`revised_prompt`, as does a stream's `completed` event. An edit or a
+variation runs its prompt as typed and ignores `prompt_upsampling`.
+
 ## Edits and variations
 
 An edit takes its images as files of a multipart form, `image` for one and
@@ -152,6 +158,9 @@ for event in stream:
 | More reference images than fit in memory | 400 naming the count and the reference size that fit |
 | 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'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` |
 | The job waits for its turn, or for one step, longer than [`token_queue_timeout_s`](config.md#servertoken_queue_timeout_s) | 504, and the log marks the request `timeout` |

docs/images.md

@@ -76,6 +76,7 @@ has the fix.
 
 - [Call it from an app](#call-it-from-an-app)
 - [Use a LoRA](#use-a-lora)
+- [Prompt upsampling](#prompt-upsampling)
 - [ComfyUI](#comfyui)
 - [Models and quants](#models-and-quants)
 - [Licenses](#licenses)
@@ -163,6 +164,61 @@ LoHa, LoKr, DoRA and the other kinds that are not plain LoRA pairs are
 refused by name, as is a LoRA of another model. ComfyUI takes LoRAs as
 [LoRAs](comfyui.md#loras) describes.
 
+## 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:
+
+```yaml
+models:
+  qwen3.5-9b:
+    path: Qwen3.5-9B-Q6_K.gguf
+image_models:
+  flux-3-klein:
+    path: flux3_image_klein_sd-Q4_K_M.gguf
+    encoder: Qwen3-VL-4B-Instruct-Q8_0.gguf
+    vae: flux2-vae-BF16.gguf
+    upsampler: qwen3.5-9b
+```
+
+With an upsampler set, each prompt is rewritten:
+
+- `gmlx image generate flux-3-klein` loads the chat model, writes the
+  caption, and unloads the chat model before the image model loads.
+  `--upsampler` names a chat model, a GGUF or a `models` id, for one run.
+- `gmlx serve` writes the caption with the chat model, which loads and
+  unloads as it does for a chat request, and returns the caption as
+  `revised_prompt` with the image. A request with the same prompt and
+  another seed reuses the caption. With the
+  [prompt cache](config.md#prompt-cache) on, a new prompt reuses the
+  system prompt's prefill. The first caption also loads the chat model,
+  which `pin: true` on its `models` entry keeps loaded.
+- The ComfyUI pack rewrites the prompt of its text encode node, and
+  ComfyUI's log shows the caption.
+
+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.
+
+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`.
+
 ## ComfyUI
 
 [ComfyUI](comfyui.md) builds image workflows from nodes in a browser. The

docs/internals/image-steps.md

@@ -93,6 +93,11 @@ with sides a multiple of 32, and an image outside that area gets 400. An
 encode with `refs: false` leaves out the VAE latents of its reference
 images, for a client that encodes them with a VAE of its own.
 
+An encode that upsamples the prompt returns the caption as
+`revised_prompt` and its conditioning as `cond`, and the prompt as typed
+as `cond_typed`. A client that adds reference images after the encode,
+as ComfyUI does, samples an edit and a negative on `cond_typed`.
+
 ## Steps and runs
 
 A step names its conditionings and its patch table by hash. When the

gmlx/comfyui_pack/client.py

@@ -546,8 +546,9 @@ def encode(self, model: str, text: str, images: list | None = None) -> tuple:
             if s != 200:
                 self._fail(s, d, "the text encode")
             out, fields = unpack(d)
-        if fields.get("cond"):
-            self.known_tensors.add(fields["cond"])
+        for k in ("cond", "cond_typed"):
+            if fields.get(k):
+                self.known_tensors.add(fields[k])
         for h in fields.get("ref_latents") or []:
             self.known_tensors.add(h)
         return out, fields

gmlx/comfyui_pack/clip.py

@@ -33,6 +33,34 @@
 _SPACE_FOR_EMPTY = frozenset({"qwen_image", "qwen_image21", "krea2", "flux3_klein"})
 
 
+# The conditioning key of the prompt as typed beside an upsampled one.
+TYPED = "gmlx_typed"
+
+
+class TypedPrompt:
+    """The conditioning of a caption and that of the prompt as typed. The
+    text encode node sees no reference images, so the model samples an
+    edit and a negative on the typed one."""
+
+    __slots__ = ("caption", "typed")
+
+    def __init__(self, caption: torch.Tensor, typed: torch.Tensor):
+        self.caption = caption
+        self.typed = typed
+
+    def holds(self, cond) -> bool:
+        """Whether ``cond`` is still the caption, which a node such as
+        ConditioningZeroOut or ConditioningConcat replaces."""
+        c = self.caption
+        return cond is c or (isinstance(cond, torch.Tensor) and cond.shape == c.shape
+                             and torch.equal(cond.to(c.device, c.dtype), c))
+
+
+def _batched(t: torch.Tensor) -> torch.Tensor:
+    t = t.float()
+    return t[None] if t.ndim == 2 else t
+
+
 class _StubEncoder(torch.nn.Module):
     """A text encoder with no weights; ComfyUI reads its options and state dict."""
 
@@ -141,13 +169,16 @@ def encode_from_tokens(self, tokens, return_pooled=False, return_dict=False):
         if p.get("prevent_empty_text") and text == "" and self.gmlx_family in _SPACE_FOR_EMPTY:
             text = " "
         out, fields = self.gmlx_client.encode(self.gmlx_model, text, images)
-        cond = out["cond"].float()
-        if cond.ndim == 2:
-            cond = cond[None]
+        if fields.get("revised_prompt"):
+            log.info("gmlx: the server encoded the upsampled prompt: %s",
+                     fields["revised_prompt"])
+        cond = _batched(out["cond"])
         if return_dict:
             d = {"cond": cond, "pooled_output": None}
             if images and fields.get("image_slots"):
                 d["image_slots"] = [int(s) for s in fields["image_slots"]]
+            if "cond_typed" in out:
+                d[TYPED] = TypedPrompt(cond, _batched(out["cond_typed"]))
             self.add_hooks_to_dict(d)
             return d
         if return_pooled: