From 139be2f2b13f6a7b350cbcde5337ed66854ef49c Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 19:45:35 -0500 Subject: [PATCH 01/20] docs(flux3): FLUX 3 Image Klein audit and catalog design Launch-day audit of the vendor and diffusers sources, and the design for the distilled core: probes P1/P2 on lem (bf16 fits 24 GB at 1024^2 with every component placed; the pup runs through text_generation but needs enable_thinking off), and the from_hub system-prompt reference. Co-Authored-By: Claude Opus 5.5 --- .../audits/2026-10-15-flux-3-klein-audit.md | 271 ++++++++++++++++++ .../2026-10-15-flux-3-klein-catalog-design.md | 128 +++++++++ 2 files changed, 399 insertions(+) create mode 100644 docs/proposals/audits/2026-10-15-flux-3-klein-audit.md create mode 100644 docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md diff --git a/docs/proposals/audits/2026-10-15-flux-3-klein-audit.md b/docs/proposals/audits/2026-10-15-flux-3-klein-audit.md new file mode 100644 index 00000000..3982a690 --- /dev/null +++ b/docs/proposals/audits/2026-10-15-flux-3-klein-audit.md @@ -0,0 +1,271 @@ +# FLUX 3 Image Klein: vendor knowledge for onboarding + +Researched 2026-10-15, the vendor's launch day. Every source below is less than 72 hours old, and the vendor repos were still taking README commits at 22:58Z. Re-check the volatile items (marked **[volatile]**) before you pin anything. + +Naming: the vendor calls the model "FLUX 3 Image Klein", not "FLUX.3 [klein]". Diffusers calls the family "FLUX.3", with classes named `Flux3*`. + +## Sources + +| URL | Type | Date | Contributes | +|---|---|---|---| +| https://huggingface.co/black-forest-labs/flux-3-image-klein (+ README via `/resolve/main/README.md`, readable without auth) | primary | repo created 2026-10-13, last modified 2026-10-15T22:57Z | Variant table (16 steps, no guidance), "do not use without upsampler", Diffusers `ModularPipeline` snippet, VRAM table, box format, license | +| https://huggingface.co/api/models/black-forest-labs/flux-3-image-klein | primary | 2026-10-15 | `gated: auto`, `library_name: diffusers`, file list (no `model_index.json`; has `modular_model_index.json`), 5,548,681,984 BF16 params | +| https://huggingface.co/black-forest-labs/flux-3-image-klein-base | primary | created 2026-10-13 | Undistilled model: 50 steps, CFG 3.5, for LoRA and fine-tuning; FP8 VRAM 29 GB, NVFP4 26 GB | +| https://huggingface.co/black-forest-labs/flux-3-image-klein-pup | primary | created 2026-10-14, sha 3f3edb52 | Prompt upsampler: Qwen3.5-9B fine-tune, Apache-2.0, not gated, 18.8 GB; `system_prompts.json` (`t2i`, `i2i`, `t2i_bbox`, `i2i_bbox`); decoding settings | +| https://huggingface.co/black-forest-labs/flux-3-image-klein-comfy | primary | created 2026-10-14 | int8/NVFP4/BF16 single files; int8 runs on 8 GB GPUs; VAE file is `flux2-vae.safetensors`; sampler defaults | +| https://github.com/black-forest-labs/flux-3-image-klein (commit bbfe317d) | primary | created 2026-10-15 | Reference code: `src/flux3/pipeline.py` constants (MAX_REFERENCES=10, sides 64-4096, alpha 6.93), CLI defaults, example JSONs holding the upsampled captions, base-vs-distilled timing | +| https://bfl.ai/blog/flux-3-image-klein | primary | 2026-10-15 | Announcement: 5.6B, 10 refs, 4K, Elo against FLUX.2 [klein] 9B, int8 on 8 GB RTX, LoRA via Diffusers / AI Toolkit / NeMo | +| https://docs.bfl.ml/flux_3/flux3_image_klein.md | primary | undated (launch) | Spec table (lists distilled "Guidance (CFG) 1"), LoRA training config, LoRA-on-distilled scale 0.75 | +| https://docs.bfl.ml/flux_3/flux3_image_bounding_boxes.md | primary | undated | Box format: `[y0,x0,y1,x1]` on a 0-1000 grid; edit rows `from`/`src_bbox`/`tgt_bbox`; kinds new/move/anchor | +| https://docs.bfl.ml/guides/prompting_summary.md and `prompting_unified_basics`, `_building`, `_style`, `_technical`, `prompting_layout`, `prompting_editing_*` | primary | undated | Prompting guide. **Written for the FLUX 3 Image API model**, which upsamples on the server side; covers quoting text, refs by position, no negative prompts, aspect ratios | +| https://docs.bfl.ml/api-reference/utility/generate-an-image-with-flux-3-klein.md | primary | undated | API: `steps` 1-16 (default 16), `seed`, `aspect_ratio` (default `auto`), up to 10 `images`; "pixel restoration" is API-only | +| https://github.com/huggingface/diffusers/pull/15046 (merge commit 7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941) | primary (library) | merged 2026-10-15T16:30Z | Adds `src/diffusers/modular_pipelines/flux3/`, `Flux3Transformer2DModel`, `BFLQuantizationConfig`, `Flux3LoraLoaderMixin`, DreamBooth LoRA scripts | +| https://github.com/huggingface/diffusers/blob/7ecb1e1b4f/docs/source/en/api/pipelines/flux3.md (rendered at https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux3) | primary (library) | 2026-10-15 | Diffusers docs: defaults, guider, refs ``, memory, single-file and quantized loading | +| https://github.com/huggingface/diffusers/blob/7ecb1e1b4f/docs/source/en/quantization/bfl.md | primary (library) | 2026-10-15 | fp8r needs CC >= 8.9 to run fused; NVFP4 needs Blackwell + flashinfer; on other devices weights are dequantized per forward pass | +| https://github.com/huggingface/diffusers/pull/15058 | primary (library), **open** | opened 2026-10-15T22:32Z | `Flux3PromptUpsampleStep`, plus the constants `SYSTEM_PROMPTS` and `PROMPT_UPSAMPLER_REPO` in `diffusers/modular_pipelines/flux3/prompt_upsampler.py` (not merged) | +| https://huggingface.co/api/models?filter=base_model:adapter:black-forest-labs/flux-3-image-klein-base | secondary | 2026-10-15 | First community LoRAs, e.g. `Muinez/flux3-klein-distil-4step` and `diffusers-modular/flux3-tarot-lora` | +| Web search "FLUX 3 Image Klein" | secondary | 2026-10-15 | **No independent coverage yet.** Every result is about FLUX.2 [klein] (Jan 2026). | + +Fetch failures: +- `raw/main/model_index.json`: 401. The file does not exist; the repo ships `modular_model_index.json` instead. +- `modular_model_index.json`, `transformer/config.json` and the other config files: 401. They are gated, and there is no local HF token. +- `github.com/black-forest-labs/flux-3`, linked from the diffusers docs (flux3.md line 175): does not resolve. The real repo is `flux-3-image-klein`. + +## Model facts + +- **Parameters.** The transformer is 5.6B by the vendor's count, 5,548,681,984 by the Hub safetensors count, and 5.5B in the diffusers docs (flux3.md:203). The Qwen3-VL-4B-Instruct text encoder (~4B) is not counted in that figure. +- **Variants** (HF card table): + + | Variant | Repo | Weights | Steps | Guidance | Built for | + |---|---|---|---|---|---| + | Klein (distilled) | `flux-3-image-klein` | step-distilled flow map, bf16 | 16 | none | generation and editing; the model behind the BFL API | + | Klein Base | `flux-3-image-klein-base` | undistilled | 50 | CFG 3.5 | fine-tuning, LoRA, research | + + Both repos carry `variants/fp8r` and `variants/nvfp4`. The vendor names no size other than 5.6B. +- **Siblings on the Hub (org listing):** + - `flux-3-image-klein-pup`: the prompt upsampler. + - `flux-3-image-klein-comfy`: ComfyUI repack. + - `flux-3-action-{base,droid,so101}`: a 7B robot-action model released 2026-09-22. It is irrelevant here. + - "FLUX 3 Image", the full model, is API-only. No weights. +- **License.** Weights are under the FLUX Kommunity License v1.0: non-commercial use, plus commercial *output* use for Qualifying Users under US$5M gross annual revenue. The text encoder and the upsampler are Apache-2.0. The reference code is Apache-2.0. +- **Gating.** Both weight repos are gated with `gated: "auto"`: you click-accept and access is granted automatically. The pup repo is not gated. `hf auth login` plus accepting the terms on both the Klein and Klein Base repos is required. +- **Text encoder.** Qwen3-VL-4B-Instruct, class `transformers.Qwen3VLForConditionalGeneration` with an `AutoProcessor` tokenizer. Embeddings stack hidden layers (4, 8, …, 32). +- **VAE.** `AutoencoderKLFlux2`, the same FLUX.2 autoencoder: the comfy file is `vae/flux2-vae.safetensors`, and diffusers flux3.md:17 agrees. 8x downsample with 2x2 packing gives a 16 px token. Latent channels 32 x 4 = 128. +- **Tasks.** One model covers text-to-image, single-reference edit, multi-reference composition (up to 10 refs), bounding-box layout and box edits, and text rendering. +- **Resolution.** "Up to 4K" (GitHub README, BFL docs). Reference code: `MIN_SIDE, MAX_SIDE = 64, 4096` (`src/flux3/pipeline.py:26`). Default output is 1024x1024. +- **VRAM, as the vendor states it:** + - B200, 1024², upsampler off: distilled 31 GB bf16 / 26 GB FP8 / 25 GB NVFP4; base 31 / 29 / 26 GB. + - The upsampler adds ~28 GB. + - BFL docs say it "fits on consumer GPUs, down to an RTX 3090". + - Blog and comfy card: int8 runs on 8 GB RTX (3070) in ComfyUI. + - Diffusers: ~20 GB of bf16 weights in total, with the VAE in fp32 (flux3.md:203). +- **Speed.** H200 at 1 MP: about 3 s distilled vs 12 s base (GitHub README). RTX 4070 int8 distilled, 1024², 16 steps: 8.9 s (blog chart). +- **Quantized checkpoints from the vendor:** + - fp8r (rowwise FP8) and NVFP4 (ModelOpt layout), under `variants/` in both repos. + - Comfy int8 "convrot" files for the DiT, the text encoder and the upsampler. + - NVFP4 needs Blackwell. fp8r runs fused only on CC >= 8.9 (Ada and newer); the RTX 3090 (8.6) falls back to per-forward dequantization (bfl.md). + +## Pipeline facts (diffusers @ 7ecb1e1b4f, PR #15046) + +**Local install.** `venv/.../diffusers` is 0.41.0.dev0 at commit 1e004eb6. It has `pipelines/flux`, `pipelines/flux2` (`Flux2Pipeline`, `Flux2KleinPipeline`, `…KleinInpaint`, `…KleinKV`) and `modular_pipelines/flux2`. It has **no flux3**. Reinstall diffusers from main at or after 7ecb1e1b. Local transformers 5.19.0 already has `Qwen3VLForConditionalGeneration` and `Qwen3_5ForConditionalGeneration`. + +**Klein does NOT load through `Flux2Pipeline` or `Flux2KleinPipeline`.** It needs a new transformer (`Flux3Transformer2DModel`) and a new text encoder. It ships **only as a modular pipeline**: there is no `pipelines/flux3` and no `DiffusionPipeline` class. (The HF sidebar snippet's `DiffusionPipeline.from_pretrained` is the Hub's auto-generated snippet; the card itself uses `ModularPipeline`.) + +**Classes and modules:** +- `src/diffusers/modular_pipelines/flux3/modular_pipeline.py:24`: `Flux3ModularPipeline(ModularPipeline, Flux3LoraLoaderMixin)`. Its default blocks are `Flux3AutoBlocks`. +- `modular_pipeline.py:52`: `Flux3DistilledModularPipeline`. Its default blocks are `Flux3DistilledAutoBlocks` (`modular_blocks_flux3_distilled.py:256`). +- `modular_pipelines/modular_pipeline.py` maps the `"flux3"` and `"flux3-distilled"` model names. Per diffusers docs flux3.md:161, the repo's `modular_model_index.json` resolves to the distilled classes. +- Load with `ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein")`, then `pipe.load_components(dtype=torch.bfloat16)`. + +**trust_remote_code.** Inferred: not needed. Every class named is in the diffusers library, and the file list has no `.py` files. The `modular_model_index.json` itself could not be read (gated). Verify after you accept the gate. + +**Components** (`modular_blocks_flux3_distilled.py:265-268`): + +| Component | Class | +|---|---| +| `text_encoder` | `Qwen3VLForConditionalGeneration` | +| `tokenizer` | `AutoProcessor` | +| `image_processor` | `Flux3ImageProcessor` (from_config) | +| `vae` | `AutoencoderKLFlux2` | +| `transformer` | `Flux3Transformer2DModel` (`models/transformers/transformer_flux3.py:341`, `flow_map` config at :406) | +| `scheduler` | `FlowMatchEulerDiscreteScheduler` (from_config, shift 6.93) | +| `guider` | `ClassifierFreeGuidance`, base only | + +**Call inputs and defaults** (`Flux3DistilledAutoBlocks` docstring, `modular_blocks_flux3_distilled.py:270-304`): +- `prompt`: a **str only**. Batched prompts raise (`encoders.py:229`). Use `num_images_per_prompt` (default 1) for several images. +- `max_sequence_length`: 8192 (`encoders.py:196`). The text window is right-padded to a multiple of 80 tokens (`encoders.py:166`). The padding "is part of the model recipe". +- `text_encoder_out_layers`: `(4, 8, 12, 16, 20, 24, 28, 32)`. +- `image`: a PIL image or a list of them. Non-PIL input raises (`encoders.py:84`). +- `height` / `width`: + - Text-to-image default is 1024 x 1024 (`before_denoise.py:186`). + - With `image`, the default is 1 MP at the **first reference's** aspect ratio, rounded to 16 (`inputs.py:246, 302-309`). + - Both must be divisible by 16 or the call raises (`before_denoise.py:232`). +- `num_inference_steps`: 16 distilled (`modular_blocks_flux3_distilled.py:48,116`), 50 base (`before_denoise.py:95`). The default sigmas are `np.linspace(1.0, 1/n, n)` with scheduler `shift=6.93` (`before_denoise.py:95,151`). +- Also accepted: `timesteps` and `sigmas` (overrides), `latents`, `generator`, `output_type` (`pil`). +- `output="images"` selects the return value. + +**Guidance:** +- Distilled: no guider, no `negative_prompt`. The transformer is a flow map that gets the current and the target timestep at each step (flux3.md:161; `transformer_flux3.py:546`). +- Base: real CFG through the `guider` component, `ClassifierFreeGuidance(guidance_scale=3.5)` (`denoise.py:48`). **There is no `guidance_scale` call argument.** You change it with `pipe.guider = pipe.guider.new(guidance_scale=…)` (flux3.md:151). The base model takes `negative_prompt`. +- Diffusers says the distilled model "is trained for 4 to 16 steps" (flux3.md:171). The API caps `steps` at 1-16. + +**References:** +- Each reference keeps its own resolution. It is downscaled only above 4096x4096 and center-cropped to the 16 px grid (`encoders.py:88`). +- In prompts, refer to references as ``, ``, … (flux3.md:108). +- **Diffusers enforces no maximum count.** The reference code enforces `MAX_REFERENCES = 10` (`src/flux3/pipeline.py:27`), and the vendor states 10 everywhere. + +**Reference handling differs between diffusers and the vendor code:** +- **Vendor, single reference:** pads the image to a canvas of at most 4 MP (`MATCH_MAX_PIXELS`), then crops the output back. The output size defaults to the reference's own size. +- **Vendor, multi-reference:** each reference is resized to at most 1 MP. +- **Diffusers:** keeps references at up to 16 MP and defaults the output to 1 MP. + +**Prompt upsampler / enhancer helper:** +- **Not in merged diffusers.** Open PR #15058 adds `Flux3PromptUpsampleStep` and the module constants `SYSTEM_PROMPTS` (a dict of `t2i`/`i2i`/`t2i_bbox`/`i2i_bbox`) and `PROMPT_UPSAMPLER_REPO = "black-forest-labs/flux-3-image-klein-pup"`, in `diffusers/modular_pipelines/flux3/prompt_upsampler.py`. The PR copies `system_prompts.json` at pup commit 8b069751. The pup head is now 3f3edb52 **[volatile]**. +- Until that merges, the vendor text lives only in the pup repo's `system_prompts.json` and in the reference repo's `src/flux3/pup_prompts.json`. Pin to the diffusers constant once it lands, rather than copying the text. + +**Quantized loading:** +- `BFLQuantizationConfig(compute_dtype=torch.bfloat16)` works **only through `Flux3Transformer2DModel.from_single_file(…/variants/{fp8r,nvfp4}/model.safetensors, config=, subfolder="transformer")`**, followed by `pipe.update_components(transformer=…)`. Pipeline-level loading of these checkpoints is unsupported (bfl.md). +- Load the other components by name first. A bare `load_components()` would also load the bf16 transformer. + +**Other library facts:** +- **LoRA:** `Flux3LoraLoaderMixin` is on the pipeline. DreamBooth LoRA scripts are at `examples/dreambooth/train_dreambooth_lora_flux3_klein{,_img2img}.py`. +- **`_repeated_blocks`:** `["Flux3ModeTransformerBlock", "Flux3SingleTransformerBlock"]`, which is usable for regional compile. +- **Memory:** flux3.md:208 suggests `pipe.enable_auto_cpu_offload(device="cuda")`. The base-model docs use VAE fp32 (`dtype={"vae": torch.float32, "default": torch.bfloat16}`); the distilled card uses bf16 throughout. + +## Vendor recommended settings + +| Variant | Steps | Guidance | Resolution | Source | +|---|---|---|---|---| +| Klein (distilled) | 16 (trained for 4-16; API 1-16) | none. BFL docs list "CFG 1" (equivalent to none), and the reference code raises on anything but None/1.0 | 1024² default; up to 4K; sides 64-4096, multiple of 16 | HF card table; `src/flux3/generate.py`; docs.bfl.ml klein page; flux3.md:171 | +| Klein Base | 50 | CFG 3.5 | same | HF card; CLI `--model` help; `pipeline.py:20` | +| LoRA on distilled | 16 | CFG 1 | — | LoRA scale 0.75 (0.5-1.0), docs.bfl.ml klein page | +| LoRA training | rank/alpha 32/32, AdamW (0.9, 0.95), lr 1e-4 with 200 warmup steps then linear decay, 4000 steps, ~1 MP native aspect ratio, train on **Base** | — | — | docs.bfl.ml klein page | +| Upsampler decoding | greedy, `max_new_tokens=3584`, `enable_thinking=False` | — | — | pup card (note below) | + +The upsampler decoding note: the pup card also says "top_p 0.95, top_k 20, presence penalty 1.5 for t2i / 0 for edits". That reads as contradictory alongside greedy decoding; the card's own snippet uses `do_sample=False`. + +Vendor example sizes are 1024x1024, 832x1248 (portrait) and 1248x832 (landscape) (`examples/README.md`). The vendor does not state a scheduler shift; it is alpha 6.93 in `pipeline.py:63` and in diffusers. + +## Prompt format + +- **Mandatory upsampling.** "Do not use the model without a prompt upsampler" (HF card; GitHub README `[!IMPORTANT]`). The model is trained on the captions the upsampler writes, and short prompts give generic output. + - **Text-to-image:** the upsampler (`t2i` system prompt) expands the request into one English scene caption, about 200 words in the vendor examples (e.g. `examples/distilled/man_dog.json`). The caption covers medium/style, framing, subject with colors, materials and positions, background and lighting. + - **Edits** (`i2i`): the user message is the reference images followed by a three-line text: "Reference images follow in order (the first image, the second image, ...).\nUser request: …\nOutput the caption." The output is a complete caption of the reference with the edit folded in. Pass the same images to both the upsampler and the pipeline. + - **Any other LLM** may stand in, given these system prompts. +- **Natural language, not tags.** Describe attributes (palette, stroke, texture) rather than naming an artwork (HF card *Prompting*). +- **Text rendering.** Quote the exact words, then say where they go and how they look (docs.bfl.ml `prompting_unified_basics`; the upsampler keeps quoted text verbatim, per the pup card). +- **No negative prompts.** Rephrase each negation as a positive description (docs.bfl.ml `prompting_unified_technical`). The distilled pipeline also rejects `negative_prompt`. +- **Layout / JSON.** + - Write the caption with `` tags, then append a JSON array of `{"id","bbox":[top,left,bottom,right] on 0-1000,"desc"}`. + - For box edits, add `from: "ref_image_N"`, `src_bbox` and `tgt_bbox`. A `null` target removes the element; equal boxes anchor it. + - Boxes are a hint, not a hard mask. + - Sources: HF card; https://docs.bfl.ml/flux_3/flux3_image_bounding_boxes; flux3.md:122-147. + - The reference CLI's `boxes.json` uses 0-1 floats and converts them (`bbox.to_canvas`). +- **Naming references.** Diffusers and the vendor box format use `` (zero-based). The BFL API guide uses "image 1", "image 2" (one-based). The upsampler input uses "the first image, the second image". Use the convention of whichever layer the prompt is going to. +- **Scope of the API guide.** Its prompting guide (https://docs.bfl.ml/guides/prompting_summary.md) is written for the FLUX 3 Image *API*, where a short prompt is expanded on the server side ("A short prompt works"). For local Klein, that advice holds only once an upsampler is in front of the model. + +## Missing knowledge / caveats + +- **Diffusers support is hours old.** It merged 2026-10-15 16:30Z and is not in any release. The upsampler block (PR #15058) is unmerged. The diffusers docs link a nonexistent `github.com/black-forest-labs/flux-3`. +- **Gated configs unread.** I could not read `modular_model_index.json`, `transformer/config.json` or `variants/*/hf_quant_config.json` (gated, no token). The trust_remote_code status and the exact block mapping are inferred, not observed. +- **No guider support in the repo.** A grep of `docs/WORKFLOW_GUIDE.md` and `dw/pipeline_processors/` found no handling for modular `guider` components, so a Base CFG other than 3.5 has no knob today. The distilled template is unaffected. +- **The vendor's quantized variants need repo support** for `from_single_file` + `update_components` on a modular pipeline. That is unverified for this engine. The fallbacks are SDNQ/TorchAO through `load_components.quantization_config`, which is how the H3 templates quantize. +- **24 GB fit.** ~11.1 GB transformer + ~8.9 GB text encoder + VAE ≈ 20 GB bf16 is tight on 24 GB at 1 MP. The upsampler (18.8 GB bf16; 9.8 GB int8 in the comfy repo) cannot be resident alongside. It must run as a separate prior step, or as a remote LLM. `enable_auto_cpu_offload` (`components_manager`) is CUDA-only per the repo docs. +- **MPS** is completely unstated by the vendor. The fp8r and NVFP4 kernels are CUDA-only; elsewhere they dequantize per forward pass. bf16 on MPS needs ~20 GB of unified memory plus activations. +- **Known failure modes** (HF card *Limitations*): + - Rendered text can be inaccurate. + - Boxes are not masks; small boxed elements may not appear, and elements can extend past their box. + - Edits can change content outside the requested region. API "pixel restoration" fixes this, but it is API-only. + - A one-line edit instruction sent raw "reframes the shot". + - The upsampler resizes references to ≤1 MP, so it can misdescribe small details. It never refuses, and its output is English only. +- **Seeds.** Seeds reproduce only when noise is drawn as the vendor draws it (flux3.md:175-185). +- **Text encoder version quirk.** Qwen3-VL padding position ids differ between transformers 4.57 and 5.x. Diffusers passes them explicitly (`encoders.py:269-273`), so do not bypass the pipeline's text encoder step. +- **LoRA.** Train on Base; the LoRA runs on distilled without conversion at scale ~0.75. Trainers: Diffusers DreamBooth scripts, Ostris AI Toolkit (`arch: flux3`), NVIDIA NeMo AutoModel. A LoRA adds ~5% to generation time (blog). Community LoRAs exist from day one, including a "4-step distil" LoRA by Muinez, not vendor-verified. +- **FLUX.2 vs FLUX.3.** FLUX.2 [klein] (4B Apache / 9B non-commercial, released 2026-01-15) is the previous release. + - It has its own diffusers classes (`Flux2KleinPipeline`, already local). The blog claims Klein beats FLUX.2 [klein] 9B by +91 Elo t2i and +70 edit with 38% fewer parameters (internal evaluation). + - Shared with FLUX.2: the VAE, 10 references, and "up to 4MP"-style editing. + - Different from FLUX.2: the transformer, the text encoder (Qwen3-VL-4B), mandatory upsampling, bbox prompts, and the license (Kommunity rather than Apache/non-commercial). + - Do not carry FLUX.2 [klein] settings or guides over; the BFL docs' FLUX.2 pages (training, HEX/JSON prompting use cases) are about the earlier release. +- **No independent secondary coverage** exists yet. Every claim above rests on vendor or diffusers sources alone. + +## What carries over from `workflows/models/flux2-dev.json` + +- **Carries over:** the result block (`image/jpeg`, `embed_metadata`), `seed: variable:seed`, `cost_drivers`, `configures: templates/text-to-image`, and the `hf auth login` gating pattern. +- **Differs:** + - `component_type` becomes `ModularPipeline` (as in the H3 templates), not `Flux2Pipeline`. + - There is no `offload`; use `components_manager` / on-demand residency instead. + - Per-component dtype and quantization go through `load_components`. + - There is no `guidance_scale` argument (distilled has none; base uses the guider). + - `output: "images"`. + - A prompt-upsampler step must precede the generation step. + - Steps are 16, not 40. + - The text encoder is a transformers Qwen3-VL model and is not pre-quantized by the vendor except in comfy int8. +- **Repo conventions** (RECIPES_24GB, ACCELERATION, QUANTIZATION): + - Unquantized FLUX uses `offload: "model"`, which modular pipelines cannot use. + - `mag_ratios: "flux"` is a FLUX.1 preset and is not valid for Flux3; calibrate instead. + - `is_guidance_distilled` (FasterCache) is FLUX.1 guidance embedding, which does not apply here. The distilled Klein is a flow map with no CFG batch. + +## What the templates should teach (reading order) + +1. **text-to-image (distilled)** + - Setup: `black-forest-labs/flux-3-image-klein`, 16 steps, 1024², bf16, no guidance. + - The prompt is already a full caption, or comes from an upsampler step. + - Vendor example: "A man walking with his dog, afternoon sun", seed 7 (`examples/README.md`); the expanded caption is in `examples/distilled/man_dog.json`. +2. **upsample-then-generate** + - Two steps: a `text_generation` step with `flux-3-image-klein-pup` (`Qwen3_5ForConditionalGeneration`, system prompt `t2i`, greedy, `max_new_tokens` 3584, `enable_thinking=False`), whose output is passed by `previous_result:` as `prompt`. + - Unload between the steps on 24 GB. + - Pin the system prompt to diffusers `SYSTEM_PROMPTS` once PR #15058 merges. + - Source: HF card *Prompt upsampling*. +3. **image edit (single reference)** + - `image` = one asset; the `i2i` upsampler caption gets the same image; 16 steps. + - Output defaults to 1 MP at the reference's aspect ratio. + - Vendor examples: "put a red bandana around the dog's neck" (seed 22), "fill the mug with coffee and put a croissant next to it" (seed 21). +4. **multi-reference** + - 2-10 images, referred to as ``; the first image sets the aspect ratio. + - Vendor example: "print the dog from the second image onto the mug" (seed 23). Diffusers example: fox + bench at 1360x768. +5. **layout with bounding boxes** + - Caption with `` tags plus a JSON box table on 0-1000. + - Vendor examples: the silhouette example (HF card) and the balloons boxes (`examples/inputs/balloons_boxes.json`). Docs: docs.bfl.ml bounding-box page. +6. **typography** + - Quoted text, 832x1248 portrait. + - Vendor example: the gig poster, seed 1 (`examples/typography/`). +7. **base + CFG** (optional, for LoRA users) + - `flux-3-image-klein-base`, 50 steps, guider 3.5, `negative_prompt` allowed. + - The repo needs guider support to vary the CFG. +8. **LoRA on distilled** + - A Base-trained LoRA at scale 0.75, 16 steps (docs.bfl.ml klein page). +9. **low-VRAM 24 GB CUDA** + - Option A, the vendor fp8r `from_single_file` + `BFLQuantizationConfig`: fused on Ada+ only, dequant fallback on a 3090. Needs engine support. + - Option B, SDNQ/TorchAO int8 through `load_components.quantization_config` on the transformer and text encoder (H3 pattern). + - With either option, `components_manager` auto-offload and the upsampler as a separate step. Measure; the vendor gives no 24 GB recipe. +10. **Apple MPS** + - bf16 with no vendor quantization (the kernels are CUDA-only), 1024² or lower. + - Upsampler run separately or replaced by another LLM. + - The vendor says nothing about MPS; it must be measured. + +## Assessment + +**Trust as-is.** These facts are consistent across the HF card, the reference code, the BFL docs and the diffusers source: +- Distilled: 16 steps, no CFG. Base: 50 steps, CFG 3.5. Shift 6.93. +- 1024² default; sides a multiple of 16, at most 4096. +- Up to 10 references (vendor-enforced; diffusers does not enforce it). +- Qwen3-VL-4B text encoder and the FLUX.2 VAE. +- `max_sequence_length` 8192; prompt must be a single string. +- No negative prompt on distilled. +- The bbox format. +- The mandatory upsampler and its decoding settings. +- The LoRA recipe (train on Base, scale 0.75 on distilled). +- The class names: `Flux3DistilledModularPipeline` / `Flux3ModularPipeline`, loaded via `ModularPipeline.from_pretrained`. + +**Uncertain:** +- The upsampler system prompts are still moving (pup repo updated hours ago; the diffusers PR pins an older commit). Do not pin text yet. +- Diffusers and the vendor code prepare references differently (1 MP vs native-size output default; 16 MP vs 4 MP/1 MP reference cap). Expect diffusers output to differ from the vendor's for edits. +- The distilled guidance label differs: "none" (card) vs "CFG 1" (BFL docs). They are equivalent. +- 24 GB and MPS fit are unmeasured. The vendor's VRAM figures are B200 peaks. +- trust_remote_code: inferred not needed; the gated index is unread. + +**Blockers:** +1. **The local diffusers has no flux3.** Diffusers must be updated to main ≥ 7ecb1e1b (merged today; unreleased). +2. **Gated weights.** Accept the Kommunity license on the Klein and Base repos, and set an HF token on the server. +3. **The upsampler.** A faithful template needs either a 9B upsampler step (18.8 GB bf16) or a different LLM given BFL's system prompts. The library constant for those prompts is not merged yet. +4. **Engine gaps.** The vendor fp8r/NVFP4 variants and a non-default Base CFG require engine support for modular `from_single_file`/`update_components` and the `guider`. Neither is present today. diff --git a/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md b/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md new file mode 100644 index 00000000..a48c44fc --- /dev/null +++ b/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md @@ -0,0 +1,128 @@ +# FLUX 3 Image Klein: catalog onboarding design + +Date: 2026-10-15. Audit: [2026-10-15-flux-3-klein-audit.md](../../proposals/audits/2026-10-15-flux-3-klein-audit.md). +Status: draft, probes on lem in progress. + +## Decisions taken (Don, 2026-10-15) + +- **Scope: the distilled core.** Text-to-image, upsample-then-generate, single edit, + multi-reference, bounding-box layout, and a 24 GB recipe. Base + CFG (needs modular + `guider` support), LoRA and typography are dated follow-ups. +- **The upsampler is a prior step** on `black-forest-labs/flux-3-image-klein-pup` + (Qwen3.5-9B, Apache-2.0, ungated), released before the generation step loads. BFL: + "do not use the model without a prompt upsampler". +- **lem moved after the loop finished:** diffusers 7ecb1e1b (PR #15046) installed via + `update_diffusers` and `dw-serve` restarted, 2026-10-15 18:58 CDT. + +## Verified on lem + +- `modular_model_index.json` (read with lem's token): `Qwen3VLForConditionalGeneration`, + `Qwen3VLProcessor`, `Flux3Transformer2DModel`, `AutoencoderKLFlux2`. No `.py` in the + repo, so no `trust_remote_code`. The catalog structure test's refusal does not bite. +- A `ModularPipeline` step on `black-forest-labs/flux-3-image-klein` validates on lem. +- `update_diffusers` alone leaves a running `dw.serve` with a half-old diffusers: the lazy + `modular_pipeline` import raises `RuntimeError` in `component_types`/`component_names` + validation until the server is restarted (the known admission gap from the + Kandinsky 6 onboarding). +- The first in-worker download looked hung (3.0 GB, no traffic for a 5 s sample) and P1 + was cancelled, but that blob had completed a minute earlier, so it was probably + between files: not a confirmed hang. The weights were then fetched with `hf download` + on the box (~34 MB/s). + +## Probes (results to fill in) + +| Probe | Job | Result | +|---|---|---| +| P1a as below, text encoder + VAE `residency: on_demand` | 8c88debc90ab | Failed: `index is on cpu` in the text encoder | +| P1b transformer and text encoder unplaced, VAE on demand | 19ece0b9522a | Failed: `Cannot generate a cpu tensor from a generator of type cuda` | +| P1d distilled t2i bf16, 1024², 16 steps, vendor `man_dog` caption, seed 7, transformer + text encoder + VAE `device: cuda` | 9cc1ce441555 | **Passed.** 24 s total: load 10.1 s, encode 0.4 s, denoise 12.1 s (0.8 s/step), decode 0.4 s. 19.2 GB allocated, 22.1 GB reserved peak at decode (1.7 GB free). The image follows the caption element by element | +| P2 pup via `text_generation` as-is, `t2i` system prompt, `man_dog` request | a5471561de18 | Ran (22 s, cuda:1); vendor-style caption, but it opens mid-phrase in lowercase ("medium-wide outdoor photograph ...") | +| P2b chat template check | (script) | The template's default opens `\n` (thinking on); `enable_thinking=False` gives `\n\n\n\n`. P2 ran in thinking mode | +| P3 single edit / multi-reference | | | + +What P1 settled: + +- `ModularPipeline._execution_device` is the first model component's device unless a + group-offload or accelerate hook says otherwise. An unplaced component stays on the + CPU, so the pipeline encodes and draws noise on the CPU. **Every Klein component needs + an explicit `device`.** `residency: on_demand` on the text encoder breaks for the same + reason: the encoder step moves `input_ids` to the execution device before the wrapper + moves the weights. On the VAE it is fine once the transformer is placed. +- bf16 with everything resident fits a 24 GB card at 1024², just. Larger outputs and + edits (references add tokens) want the VAE on demand or the components manager. + +## Templates (`workflows/templates/flux3/`) + +One idea each, in reading order; every setting from the audit's *Vendor recommended +settings*: + +1. `text-to-image.json` - distilled, 16 steps, no guidance, 1024², `width`/`height` + modulus-16 constraints (the pipeline raises after load), prompt is a full caption + (stored prompt = the vendor's `man_dog` upsampled caption). +2. `upsample-then-generate.json` - pup `t2i` step -> `previous_result:` prompt; + `release_models` between them. +3. `image-edit.json` - one reference; the `i2i` upsampler sees the same image; output + defaults to 1 MP at the reference's aspect ratio (diffusers' rule, not the vendor + code's - stated in the description). +4. `multi-reference.json` - 2-10 references named `` in the caption; the + 10 cap is a `variable_constraints` rule because diffusers does not enforce it. +5. `layout-boxes.json` - `` caption + 0-1000 box table (`t2i_bbox`). +6. A 24 GB recipe, from the P1 measurement: bf16 with on-demand residency if it fits, + else SDNQ int8 on the transformer and text encoder (the H3 pattern). + +Plus `workflows/models/flux3-klein.json` and the family README. + +## Engine work (generic, no per-model Python) + +`text_generation` (`dw/tasks/text_generation.py`) loads and runs the pup as it is (P2). It +needs two things to run it as BFL specifies: + +1. **`chat_template_kwargs`** passed to the chat template (`enable_thinking: false`). + Without it the pup runs in thinking mode (P2b). +2. **A list of images** - an edit sends the upsampler every reference, in order. + +`repetition_penalty` is already overridable (1.15 is only the vision default), and +`generate_kwargs` already passes through, so the vendor's greedy decoding needs no +change: the edit templates set `repetition_penalty: 1.0`. + +## Decision: the upsampler's system prompts are read from the Hub (Don, 2026-10-15) + +BFL's four system prompts (`t2i`, `i2i`, `t2i_bbox`, `i2i_bbox`) are in the pup repo's +`system_prompts.json` and still changing (pup head 3f3edb52; diffusers PR #15058, open, +copies an older 8b069751). The skill rule is to point at vendor text, not copy it, so +nothing is transcribed. Considered and not taken: a pinned copy as stored prompts (drifts), +and waiting for #15058 (no upsampler templates until it merges). + +`text_generation` takes `system_prompt` as today, or a generic reference to a JSON file in +a Hub repo: + +```json +"system_prompt": {"from_hub": {"repo": "black-forest-labs/flux-3-image-klein-pup", + "file": "system_prompts.json", + "revision": "3f3edb52b2c55de0d49feab28598b68f7f8e4830", + "key": "t2i"}} +``` + +- Fetched with `hf_hub_download` into the normal HF cache, so it's one small download, + then offline. +- `revision` is required and must be a full 40-hex commit, so the text a template runs + can't move under it. +- `repo` matches the Hub id pattern, and `file` is a relative path with no `..` or + absolute part. `key` is optional; the value it selects (or the file, without one) must + be a string. +- Model-agnostic: any vendor that ships prompts as a repo file can use it. When #15058 + merges, the skill can cite the diffusers constant beside it. A test compares the + constant with the file only if it pins the same revision. + +## Follow-ups (dated 2026-10-15) + +- Base + CFG: modular `guider` support (`pipe.guider.new(guidance_scale=...)`). +- Vendor fp8r/NVFP4 via `Flux3Transformer2DModel.from_single_file` + `BFLQuantizationConfig`; + fused only on Ada+/Blackwell, so of little use on the 3090s. +- LoRA on distilled (scale 0.75), typography template, MPS measurement on mini-ai. +- Re-pin the system prompts when #15058 merges. +- dw gap: `validate_workflow`'s `plan.downloads_required` was empty for both uncached + repos (the modular Klein repo, ~20 GB, and the pup as a `text_generation` + `model_name`, 18.8 GB), so the cost quote misses the download. Once a partial + cache existed it swung the other way: 40.3 GB, the whole repo including the + `variants/` the pipeline never loads. From 67ab2a02bf0fa41c3f1260d06b4ee0dd029afd76 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 19:50:41 -0500 Subject: [PATCH 02/20] docs(flux3): Klein catalog implementation plan Co-Authored-By: Claude Opus 5.5 --- .../plans/2026-10-15-flux3-klein-catalog.md | 1158 +++++++++++++++++ .../2026-10-15-flux-3-klein-catalog-design.md | 16 +- 2 files changed, 1171 insertions(+), 3 deletions(-) create mode 100644 docs/superpowers/plans/2026-10-15-flux3-klein-catalog.md diff --git a/docs/superpowers/plans/2026-10-15-flux3-klein-catalog.md b/docs/superpowers/plans/2026-10-15-flux3-klein-catalog.md new file mode 100644 index 00000000..f244d2f1 --- /dev/null +++ b/docs/superpowers/plans/2026-10-15-flux3-klein-catalog.md @@ -0,0 +1,1158 @@ +# FLUX 3 Image Klein Catalog Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Onboard FLUX 3 Image Klein (distilled) into the dw catalog: upsampler support in `text_generation`, two upsampler builtins, five templates, a model entry, a family README, GPU-measured costs, and the `flux3-klein` composition skill. + +**Architecture:** Model knowledge stays in JSON, prompts and the skill. The engine gains three model-agnostic `text_generation` options (`chat_template_kwargs`, `revision`, a list for `image`) and a `system_prompt` that can name a JSON file in a Hub repo at a pinned commit. Two builtins (`flux3_klein_upsample`, `flux3_klein_upsample_edit`) carry BFL's upsampler settings once, and every template that upsamples composes one of them. + +**Tech Stack:** Python 3.12+, transformers 5.19 pipelines, huggingface_hub `hf_hub_download`, diffusers `ModularPipeline` at or after 7ecb1e1b (PR #15046), pytest. + +**Spec:** [docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md](../specs/2026-10-15-flux-3-klein-catalog-design.md), audit [docs/proposals/audits/2026-10-15-flux-3-klein-audit.md](../../proposals/audits/2026-10-15-flux-3-klein-audit.md). + +## Global Constraints + +- Branch `feat/flux3-klein`. Commits end with `Co-Authored-By: Claude Opus 5.5 `. +- No per-model Python. Nothing in `dw/` names Flux3, Klein or BFL. +- diffusers must be at or after `7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941` (the Flux3 merge). Install it locally in Task 0; lem already has it. +- Klein repo `black-forest-labs/flux-3-image-klein`; upsampler repo `black-forest-labs/flux-3-image-klein-pup` pinned at `3f3edb52b2c55de0d49feab28598b68f7f8e4830` for both its weights and its `system_prompts.json`. +- Distilled settings: `num_inference_steps` 16, no guidance argument, no `negative_prompt`, `width`/`height` multiples of 16, default 1024x1024, at most 10 references. +- Every Klein pipeline step places every component explicitly: `"components": {"transformer": {"device": "cuda"}, "text_encoder": {"device": "cuda"}, "vae": {"device": "cuda"}}`. An unplaced component leaves `_execution_device` on the CPU (spec, P1a/P1b). Never put `residency: on_demand` on the text encoder. +- Upsampler decoding: greedy (the task's default), `max_new_tokens` 3584, `chat_template_kwargs: {"enable_thinking": false}`, `repetition_penalty: 1.0` on the image path. +- No vendor prompt text copied into the repo. The system prompts are read with `from_hub`. The edit scaffold (three short lines from BFL's `src/flux3/pup.py` at `bbfe317dd81c825c30d76090a67d3d0db51da3cc`) is the one quoted vendor string, and the builtin's description cites that file and commit. +- Descriptions name rules without single quotes (the drift check reads a single-quoted name as a variable or step). +- Security: validate the Hub repo id with `dw.locations.validate_hub_repo_id`; the `file` is a relative POSIX path with no `..`, no absolute part and no backslash; the `revision` is exactly 40 lowercase hex characters. No `eval`, `exec` or `shell=True`. +- Unit tests: `venv/bin/python -m pytest tests/ -v` locally (the venv has torch). The full suite runs with `-n auto` from `pytest.ini`. +- `tests/test_catalog_structure.py::test_the_compact_listing_fits_the_budget` may fail once the templates land. Raising that budget needs Don's OK: stop and ask, don't edit the number. + +## Review Focus + +1. **`from_hub` naming a key that isn't in the file, or a value that isn't text.** Expect a `ValueError` that lists the keys the file has, raised before any model loads, not a `KeyError` traceback. Test in Task 2. +2. **`image` given as `[]` or `[null]`.** That's an optional reference list left empty by a template. Expect the text-only path, exactly as with no image. Test in Task 1. +3. **A short `revision` (`3f3edb52`) or a branch name (`main`) in `from_hub`.** Expect a refusal naming the 40-hex rule. Without one, the prompt text could move under a pinned template. Test in Task 2. +4. **The Klein step with a component left unplaced**, as a cold agent copying the template might write it. Expect a catalog test to refuse it: `test_every_flux3_pipeline_step_places_every_component` in Task 4. +5. **More than 10 references.** Diffusers doesn't refuse it, and the model was trained on 10. Expect the multi-reference description and the skill to state the cap; an enforcing check is a dated follow-up. Test that the text says it in Task 6. + +--- + +### Task 0: Local diffusers with Flux3 + +**Files:** none (environment) + +- [ ] **Step 1: Install diffusers at the Flux3 merge or later** + +Run: `venv/bin/pip install -q --force-reinstall --no-deps git+https://github.com/huggingface/diffusers@7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941` + +- [ ] **Step 2: Confirm the modular package imports** + +Run: `venv/bin/python -c "from diffusers.modular_pipelines.flux3 import modular_blocks_flux3_distilled as m; print(m.__file__)"` +Expected: a path ending in `modular_pipelines/flux3/modular_blocks_flux3_distilled.py`. + +--- + +### Task 1: `text_generation` takes chat-template settings, a pinned revision and a list of images + +**Files:** +- Modify: `dw/tasks/text_generation.py` (`_image_part`, `_build_messages`, `generate_text`, `_generate`) +- Test: `tests/test_text_generation.py` + +**Interfaces:** +- Produces: `generate_text(prompt, device="cpu", **kwargs)` now also reads `chat_template_kwargs: dict | None`, `revision: str | None`, and `image: PIL | str | list[PIL | str] | None`. The return value doesn't change. + +- [ ] **Step 1: Write the failing tests** (append to `tests/test_text_generation.py`, inside `TestTextGeneration`) + +```python + @patch("dw.tasks.text_generation.hf_pipeline") + def test_chat_template_kwargs_reach_the_text_template(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text( + "a cat", device="cpu", model_name="m/text-chat-kwargs", + chat_template_kwargs={"enable_thinking": False}, + ) + + self.assertEqual( + pipe.call_args[1]["tokenizer_encode_kwargs"], {"enable_thinking": False} + ) + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_chat_template_kwargs_reach_the_vision_template(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text( + "a cat", device="cpu", model_name="m/vision-chat-kwargs", + image="https://example.com/a.png", + chat_template_kwargs={"enable_thinking": False}, + ) + + self.assertIs(pipe.call_args[1]["enable_thinking"], False) + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_no_chat_template_kwargs_passes_nothing_extra(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text("a cat", device="cpu", model_name="m/no-chat-kwargs") + + self.assertNotIn("tokenizer_encode_kwargs", pipe.call_args[1]) + + def test_chat_template_kwargs_must_be_a_mapping(self): + with self.assertRaises(ValueError): + generate_text("a cat", device="cpu", chat_template_kwargs=["enable_thinking"]) + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_revision_reaches_the_loader(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text("a cat", device="cpu", model_name="m/revision", revision="a" * 40) + + self.assertEqual(mock_pipeline.call_args[1]["revision"], "a" * 40) + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_a_list_of_images_becomes_one_part_each_in_order(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text( + "edit", device="cpu", model_name="m/image-list", + image=["https://example.com/1.png", "https://example.com/2.png"], + ) + + user = pipe.call_args[1]["text"][-1]["content"] + self.assertEqual( + [part.get("url") for part in user[:2]], + ["https://example.com/1.png", "https://example.com/2.png"], + ) + self.assertEqual(user[2], {"type": "text", "text": "edit"}) + self.assertEqual(mock_pipeline.call_args[0][0], "image-text-to-text") + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_an_empty_or_all_null_image_list_is_the_text_path(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + for empty in ([], [None]): + generate_text("a cat", device="cpu", model_name=f"m/empty-{len(empty)}", image=empty) + self.assertEqual(mock_pipeline.call_args[0][0], "text-generation") +``` + +Each test uses its own `model_name` because `cached_model` keys on it and a shared name would return an earlier test's mock. + +- [ ] **Step 2: Run them and watch them fail** + +Run: `venv/bin/python -m pytest tests/test_text_generation.py -v -n0` +Expected: the seven new tests FAIL (`KeyError: 'tokenizer_encode_kwargs'`, `'enable_thinking'`, `'revision'`, no `ValueError` raised, and the list tests on the message shape or pipeline task). + +- [ ] **Step 3: Implement** + +In `dw/tasks/text_generation.py`, replace `_build_messages` with the version below (keep `_image_part` as it is): + +```python +def _images(image): + """The images a call carries, as a list in order, or None for none. + + A template's optional reference list arrives as [] or [None] when the + caller supplied nothing, and that is the same statement as no image. + """ + if image is None: + return None + images = list(image) if isinstance(image, (list, tuple)) else [image] + images = [i for i in images if i is not None and i != ""] + return images or None + + +def _build_messages(prompt, system_prompt, images): + """Chat messages in the shape the chosen pipeline expects. + + Text-generation models take plain string content. Vision models take a + list of typed parts, because the images are part of the message rather + than something alongside it - every image first, in order, then the text. + """ + messages = [] + if images is None: + if system_prompt is not None: + messages.append({"role": "system", "content": system_prompt}) + messages.append({"role": "user", "content": prompt}) + else: + if system_prompt is not None: + messages.append( + {"role": "system", "content": [{"type": "text", "text": system_prompt}]} + ) + messages.append( + { + "role": "user", + "content": [_image_part(i) for i in images] + + [{"type": "text", "text": prompt}], + } + ) + return messages +``` + +In `generate_text`, document the three new kwargs in the docstring: + +``` + image: Optional PIL image, URL or path, or a list of them - every + one goes to the model, in order. Its presence is what selects + the vision pipeline; an empty list is no image. + chat_template_kwargs: Settings for the model's chat template, + such as {"enable_thinking": false} for a Qwen3 model that + must not reason before answering. + revision: The model repo's commit, branch or tag to load. +``` + +Then change the body: + +```python + images = _images(kwargs.get("image", None) or None) + system_prompt = kwargs.get("system_prompt", None) + max_new_tokens = int(kwargs.get("max_new_tokens", 500)) + revision = kwargs.get("revision", None) + chat_template_kwargs = kwargs.get("chat_template_kwargs") or {} + if not isinstance(chat_template_kwargs, dict): + raise ValueError( + "text_generation: chat_template_kwargs must be an object of settings " + f"for the chat template, got {type(chat_template_kwargs).__name__}" + ) + + if images is None: + pipeline_task = "text-generation" + model_name = kwargs.get("model_name", _DEFAULT_MODEL) + else: + pipeline_task = "image-text-to-text" + model_name = kwargs.get("model_name", DEFAULT_VISION_MODEL) +``` + +`load_pipe` passes the revision only when one is set: + +```python + def load_pipe(): + logger.info(f"Generating text with {model_name} on {device}") + placement = hf_pipeline_placement(device) + if revision is not None: + placement = {**placement, "revision": revision} + return hf_pipeline( + pipeline_task, + model=model_name, + torch_dtype=dtype, + **placement, + ) +``` + +The cache key gains the revision, so two pins of one repo don't share a model: + +```python + pipe = cached_model( + ("text_generation", pipeline_task, model_name, str(revision), str(device), str(dtype)), + load_pipe, +``` + +Pass `images` where the body used `image`: `_build_messages(prompt, system_prompt, images)`, `if images is not None:` around the repetition penalty, and `_generate(pipe, messages, images, max_new_tokens, generation, chat_template_kwargs)`. + +`_generate` hands the template settings to each path the way that pipeline accepts them: + +```python +def _generate(pipe, messages, images, max_new_tokens, generation, chat_template_kwargs): + """Run the pipeline, which each path calls differently.""" + if images is None: + # This pipeline collects anything it does not name into the arguments it + # forwards to generate(), so settings go in as plain keywords. Chat + # template settings have their own argument, which it hands to + # apply_chat_template + extra = {"tokenizer_encode_kwargs": chat_template_kwargs} if chat_template_kwargs else {} + results = pipe( + messages, + max_new_tokens=max_new_tokens, + return_full_text=False, + **generation, + **extra, + ) + else: + # (keep the existing comment) ... Keywords this pipeline does not name + # go to the processor's apply_chat_template, which is where the chat + # template settings belong + results = pipe( + text=messages, + max_new_tokens=max_new_tokens, + return_full_text=False, + generate_kwargs=generation, + **chat_template_kwargs, + ) +``` + +- [ ] **Step 4: Run the file and the callers' tests** + +Run: `venv/bin/python -m pytest tests/test_text_generation.py tests/test_image_to_text.py tests/test_judge.py -v -n0` (skip any of those names that don't exist; `ls tests | grep -i "image_to_text\|judge"`) +Expected: all PASS, including the existing tests, which call `_build_messages`'s callers with a single image. + +- [ ] **Step 5: Commit** + +```bash +git add dw/tasks/text_generation.py tests/test_text_generation.py +git commit -m "feat(text_generation): chat_template_kwargs, revision and a list of images + +A Qwen3.5 upsampler must run with enable_thinking off, an edit sends every +reference, and a pinned prompt wants pinned weights. + +Co-Authored-By: Claude Opus 5.5 " +``` + +--- + +### Task 2: `system_prompt` can name a JSON file in a Hub repo at a pinned commit + +**Files:** +- Modify: `dw/tasks/text_generation.py` (new `resolve_system_prompt`, called at the top of `generate_text`) +- Modify: `docs/TASKS.md` (the `text_generation` section) +- Test: `tests/test_text_generation.py` (new class `TestSystemPromptFromHub`) + +**Interfaces:** +- Consumes: `dw.locations.validate_hub_repo_id(name, what=...)`. +- Produces: `resolve_system_prompt(value) -> str | None`. `value` is `None`, a `str`, or `{"from_hub": {"repo": str, "file": str, "revision": str, "key": str (optional)}}`. Tasks 3-4 write the dict form in JSON. + +- [ ] **Step 1: Write the failing tests** (append to `tests/test_text_generation.py`) + +```python +import json as _json + +from dw.tasks.text_generation import resolve_system_prompt + +PIN = "3f3edb52b2c55de0d49feab28598b68f7f8e4830" + + +def _spec(**over): + spec = {"repo": "black-forest-labs/flux-3-image-klein-pup", + "file": "system_prompts.json", "revision": PIN, "key": "t2i"} + spec.update(over) + return {"from_hub": spec} + + +class TestSystemPromptFromHub(unittest.TestCase): + def setUp(self): + import tempfile, os + self.dir = tempfile.mkdtemp() + self.json_path = os.path.join(self.dir, "system_prompts.json") + with open(self.json_path, "w") as f: + _json.dump({"t2i": "expand the request", "i2i": "edit it", "n": 3}, f) + + def test_a_string_or_none_passes_through(self): + self.assertEqual(resolve_system_prompt("plain"), "plain") + self.assertIsNone(resolve_system_prompt(None)) + + @patch("dw.tasks.text_generation.hf_hub_download") + def test_the_key_s_text_is_returned_from_the_pinned_file(self, download): + download.return_value = self.json_path + self.assertEqual(resolve_system_prompt(_spec()), "expand the request") + download.assert_called_once_with( + "black-forest-labs/flux-3-image-klein-pup", "system_prompts.json", revision=PIN + ) + + @patch("dw.tasks.text_generation.hf_hub_download") + def test_a_missing_key_names_the_keys_the_file_has(self, download): + download.return_value = self.json_path + with self.assertRaises(ValueError) as caught: + resolve_system_prompt(_spec(key="t2v")) + self.assertIn("i2i", str(caught.exception)) + self.assertIn("t2i", str(caught.exception)) + + @patch("dw.tasks.text_generation.hf_hub_download") + def test_a_value_that_is_not_text_is_refused(self, download): + download.return_value = self.json_path + with self.assertRaises(ValueError): + resolve_system_prompt(_spec(key="n")) + + def test_a_short_hash_or_branch_is_refused(self): + for revision in ("3f3edb52", "main", PIN.upper()): + with self.assertRaises(ValueError) as caught: + resolve_system_prompt(_spec(revision=revision)) + self.assertIn("40", str(caught.exception)) + + def test_a_path_that_escapes_the_repo_is_refused(self): + for file in ("../secrets.json", "/etc/passwd", "a\\b.json", "", "x/../../y.json"): + with self.assertRaises(ValueError): + resolve_system_prompt(_spec(file=file)) + + def test_a_bad_repo_id_is_refused(self): + from dw.security import InvalidInputError + with self.assertRaises((ValueError, InvalidInputError)): + resolve_system_prompt(_spec(repo="https://evil.example/x")) + + def test_an_unknown_field_or_shape_is_refused(self): + for bad in ({"from_hub": {**_spec()["from_hub"], "branch": "main"}}, + {"from_url": "x"}, ["t2i"]): + with self.assertRaises(ValueError): + resolve_system_prompt(bad) + + @patch("dw.tasks.text_generation.hf_pipeline") + @patch("dw.tasks.text_generation.hf_hub_download") + def test_generate_text_sends_the_resolved_text(self, download, mock_pipeline): + download.return_value = self.json_path + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text("a dog", device="cpu", model_name="m/from-hub", system_prompt=_spec()) + + messages = pipe.call_args[0][0] + self.assertEqual(messages[0], {"role": "system", "content": "expand the request"}) +``` + +- [ ] **Step 2: Run them and watch them fail** + +Run: `venv/bin/python -m pytest tests/test_text_generation.py -v -n0 -k FromHub` +Expected: ImportError on `resolve_system_prompt`. + +- [ ] **Step 3: Implement** + +At the top of `dw/tasks/text_generation.py`, add module-level imports (a module-level `hf_hub_download` is what the tests patch): + +```python +import json +import re +from pathlib import PurePosixPath + +from huggingface_hub import hf_hub_download + +from ..locations import validate_hub_repo_id +``` + +If importing `..locations` here creates an import cycle (run the tests to see), import `validate_hub_repo_id` inside `resolve_system_prompt` instead. + +Add, above `generate_text`: + +```python +# A pinned prompt has to stay pinned: a short hash can become ambiguous and a +# branch moves, so only a full commit is accepted +_FULL_COMMIT = re.compile(r"^[0-9a-f]{40}$") +_FROM_HUB_FIELDS = {"repo", "file", "revision", "key"} + + +def _hub_file(file): + """A file path inside a repo, refused if it could name anything outside it.""" + if not isinstance(file, str) or not file or "\\" in file: + raise ValueError(f"system_prompt.from_hub.file must be a path inside the repo, got {file!r}") + path = PurePosixPath(file) + if path.is_absolute() or ".." in path.parts: + raise ValueError(f"system_prompt.from_hub.file must be a path inside the repo, got {file!r}") + return file + + +def resolve_system_prompt(value): + """The system prompt text, reading it from a Hub repo when one is named. + + Some vendors ship the instructions their model was trained against as a + file in its repo, and change them. Naming the file and a full commit runs + exactly that text without copying it into a workflow, where it would + drift from the source. + + Args: + value: None, the text itself, or {"from_hub": {"repo", "file", + "revision", "key"}}. `revision` is a full 40-character commit. + A .json file takes `key` to pick one string out of a top-level + object; any other file is read whole as text + + Raises: + ValueError: On any other shape, a revision that is not a full commit, + a file path that leaves the repo, a missing key, or a value that + is not text + """ + if value is None or isinstance(value, str): + return value + if not isinstance(value, dict) or set(value) != {"from_hub"} or not isinstance(value["from_hub"], dict): + raise ValueError( + 'system_prompt must be text or {"from_hub": {"repo", "file", "revision", "key"}}' + ) + spec = value["from_hub"] + unknown = set(spec) - _FROM_HUB_FIELDS + if unknown: + raise ValueError(f"system_prompt.from_hub does not take {sorted(unknown)}") + repo = validate_hub_repo_id(spec.get("repo"), what="system_prompt.from_hub.repo") + file = _hub_file(spec.get("file")) + revision = spec.get("revision") + if not isinstance(revision, str) or not _FULL_COMMIT.match(revision): + raise ValueError( + "system_prompt.from_hub.revision must be a full 40-character lowercase " + f"commit hash, so the text cannot change under the workflow; got {revision!r}" + ) + key = spec.get("key") + + path = hf_hub_download(repo, file, revision=revision) + with open(path, encoding="utf-8") as f: + if file.endswith(".json"): + data = json.load(f) + else: + data = f.read() + if key is not None: + if not isinstance(data, dict) or key not in data: + keys = sorted(data) if isinstance(data, dict) else [] + raise ValueError( + f"system_prompt.from_hub: {file} at {revision[:8]} has no key {key!r}; " + f"it has {keys}" + ) + data = data[key] + if not isinstance(data, str) or not data.strip(): + raise ValueError( + f"system_prompt.from_hub: {file}{'[' + repr(key) + ']' if key else ''} " + "is not text" + ) + return data +``` + +In `generate_text`, replace `system_prompt = kwargs.get("system_prompt", None)` with: + +```python + # Resolved before anything loads, so a bad reference costs nothing + system_prompt = resolve_system_prompt(kwargs.get("system_prompt", None)) +``` + +In `docs/TASKS.md`'s `text_generation` section, add the three Task 1 arguments and this paragraph with an example: + +```markdown +`system_prompt` may name a file in a Hugging Face repo instead of carrying the text, +for a vendor that ships its prompt instructions there: +`{"from_hub": {"repo": "owner/name", "file": "system_prompts.json", "revision": "<40-hex commit>", "key": "t2i"}}`. +The revision must be a full commit, so the text can't move under the workflow. A `.json` +file takes `key` to pick one string; any other file is read whole. +``` + +- [ ] **Step 4: Run the tests** + +Run: `venv/bin/python -m pytest tests/test_text_generation.py -v -n0` +Expected: all PASS. + +- [ ] **Step 5: Commit** + +```bash +git add dw/tasks/text_generation.py tests/test_text_generation.py docs/TASKS.md +git commit -m "feat(text_generation): system_prompt from a Hub repo file at a pinned commit + +Co-Authored-By: Claude Opus 5.5 " +``` + +--- + +### Task 3: The upsampler builtins + +**Files:** +- Create: `dw/workflows/flux3_klein_upsample.json` +- Create: `dw/workflows/flux3_klein_upsample_edit.json` +- Test: `tests/test_flux3_catalog.py` (new) + +**Interfaces:** +- Consumes: `text_generation` arguments from Tasks 1-2; the `compose_text` task (`parts`, `separator`). +- Produces: `builtin:flux3_klein_upsample.json` (variables `prompt`, `device`) and `builtin:flux3_klein_upsample_edit.json` (variables `instruction`, `images`, `device`). Each returns the caption text from its last step. + +- [ ] **Step 1: Write the failing test** (`tests/test_flux3_catalog.py`) + +```python +"""The FLUX 3 Image Klein catalog: the upsampler builtins, the templates' settings +and their placement. The vendor values here are the ones the design states +(docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md).""" + +import json +import os + +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +BUILTINS = os.path.join(ROOT, "dw", "workflows") +PUP = "black-forest-labs/flux-3-image-klein-pup" +PUP_PIN = "3f3edb52b2c55de0d49feab28598b68f7f8e4830" + + +def load(path): + with open(path) as f: + return json.load(f) + + +def generation_steps(definition): + return [s for s in definition["steps"] if s.get("task", {}).get("command") == "text_generation"] + + +@pytest.mark.parametrize("name,key", [ + ("flux3_klein_upsample.json", "t2i"), + ("flux3_klein_upsample_edit.json", "i2i"), +]) +def test_each_upsampler_runs_the_pinned_pup_as_bfl_does(name, key): + definition = load(os.path.join(BUILTINS, name)) + (step,) = generation_steps(definition) + args = step["task"]["arguments"] + assert args["model_name"] == PUP + assert args["revision"] == PUP_PIN + assert args["system_prompt"] == {"from_hub": { + "repo": PUP, "file": "system_prompts.json", "revision": PUP_PIN, "key": key}} + assert args["chat_template_kwargs"] == {"enable_thinking": False} + assert args["max_new_tokens"] == 3584 + + +def test_the_edit_upsampler_wraps_the_instruction_in_bfl_s_scaffold(): + definition = load(os.path.join(BUILTINS, "flux3_klein_upsample_edit.json")) + compose = next(s for s in definition["steps"] if s["task"]["command"] == "compose_text") + parts = compose["task"]["arguments"]["parts"] + assert parts == [ + "Reference images follow in order (the first image, the second image, ...).\nUser request: ", + "variable:instruction", + "\nOutput the caption.", + ] + assert compose["task"]["arguments"]["separator"] == "" + (step,) = generation_steps(definition) + assert step["task"]["arguments"]["image"] == "variable:images" + assert step["task"]["arguments"]["repetition_penalty"] == 1.0 + assert "bbfe317d" in definition["description"] +``` + +- [ ] **Step 2: Run it and watch it fail** + +Run: `venv/bin/python -m pytest tests/test_flux3_catalog.py -v -n0` +Expected: FAIL, `FileNotFoundError`. + +- [ ] **Step 3: Write the two builtins** + +`dw/workflows/flux3_klein_upsample.json`: + +```json +{ + "id": "flux3_klein_upsample", + "description": "Expands a short request into the caption FLUX 3 Image Klein was trained on, with Black Forest Labs' own prompt upsampler, flux-3-image-klein-pup, a Qwen3.5-9B fine-tune. BFL says not to use Klein without it. The weights and the text-to-image system prompt are both pinned to pup commit 3f3edb52, and the prompt is read from the repo's system_prompts.json rather than copied. Thinking is switched off and decoding is greedy, as BFL runs it. BFL also applies a presence penalty of 1.5 for text-to-image, which transformers' generate() has no equivalent for, so it is left out. About 18 GB in bf16.", + "variables": { + "prompt": "A man walking with his dog, afternoon sun", + "device": "cuda" + }, + "steps": [ + { + "name": "upsample", + "task": { + "command": "text_generation", + "arguments": { + "prompt": "variable:prompt", + "model_name": "black-forest-labs/flux-3-image-klein-pup", + "revision": "3f3edb52b2c55de0d49feab28598b68f7f8e4830", + "system_prompt": {"from_hub": { + "repo": "black-forest-labs/flux-3-image-klein-pup", + "file": "system_prompts.json", + "revision": "3f3edb52b2c55de0d49feab28598b68f7f8e4830", + "key": "t2i" + }}, + "chat_template_kwargs": {"enable_thinking": false}, + "max_new_tokens": 3584, + "device": "variable:device" + } + }, + "result": {"content_type": "text/plain", "save": false} + } + ] +} +``` + +`dw/workflows/flux3_klein_upsample_edit.json` has the same shape, with these differences: +- `id` `flux3_klein_upsample_edit`. +- `variables`: `{"instruction": "put a red bandana around the dog's neck", "images": [], "device": "cuda"}`. +- A first step `scaffold` with `"command": "compose_text"`, `"arguments": {"parts": [], "separator": ""}`, and `"result": {"content_type": "text/plain", "save": false}`. +- The `upsample` step uses `"prompt": "previous_result:scaffold"`, `"image": "variable:images"`, `"repetition_penalty": 1.0` and `"key": "i2i"`. +- The description says the upsampler sees every reference in order, as the pipeline will, and that the scaffold's three lines are BFL's `I2I_SCAFFOLD` from `src/flux3/pup.py` at commit bbfe317d in github.com/black-forest-labs/flux-3-image-klein. It also says BFL resizes each reference to a pixel budget (1 MP for one reference, 724x724 for two, 512x512 for three or more) before the upsampler sees it, while transformers' processor applies its own sizing, so the two can describe small details differently. + +- [ ] **Step 4: Run the test and the library tests** + +Run: `venv/bin/python -m pytest tests/test_flux3_catalog.py tests/test_library_path.py tests/test_library_sources.py -v -n0` +Expected: PASS. + +Then validate each builtin against the schema: `venv/bin/python -m dw.validate dw/workflows/flux3_klein_upsample.json` and the same for `flux3_klein_upsample_edit.json`. Expected: valid. + +- [ ] **Step 5: Commit** + +```bash +git add dw/workflows/flux3_klein_upsample.json dw/workflows/flux3_klein_upsample_edit.json tests/test_flux3_catalog.py +git commit -m "feat(flux3): Klein prompt-upsampler builtins on the pinned pup + +Co-Authored-By: Claude Opus 5.5 " +``` + +--- + +### Task 4: Text-to-image template, model entry and stored prompt + +**Files:** +- Create: `prompts/flux3/man_dog.json` +- Create: `workflows/templates/flux3/text-to-image.json` +- Create: `workflows/models/flux3-klein.json` +- Test: `tests/test_flux3_catalog.py` + +**Interfaces:** +- Produces: the Klein pipeline step shape every later template copies: + +```json +{ + "name": "text_to_image", + "pipeline": { + "configuration": { + "component_type": "ModularPipeline", + "load_components": {"dtype": "torch.bfloat16"}, + "components": { + "transformer": {"device": "cuda"}, + "text_encoder": {"device": "cuda"}, + "vae": {"device": "cuda"} + } + }, + "from_pretrained_arguments": {"model_name": "black-forest-labs/flux-3-image-klein"}, + "arguments": { + "prompt": "variable:prompt", + "width": "variable:width", + "height": "variable:height", + "num_inference_steps": 16, + "output": "images" + } + }, + "result": {"content_type": "image/jpeg", "embed_metadata": true, "subfolder": "final"} +} +``` + +- [ ] **Step 1: Write the failing tests** (append to `tests/test_flux3_catalog.py`) + +```python +TEMPLATES = os.path.join(ROOT, "workflows", "templates", "flux3") +KLEIN = "black-forest-labs/flux-3-image-klein" + + +def klein_steps(definition): + return [s for s in definition["steps"] + if s.get("pipeline", {}).get("from_pretrained_arguments", {}).get("model_name") == KLEIN] + + +def flux3_definitions(): + paths = [os.path.join(TEMPLATES, n) for n in sorted(os.listdir(TEMPLATES)) if n.endswith(".json")] + paths.append(os.path.join(ROOT, "workflows", "models", "flux3-klein.json")) + return paths + + +@pytest.mark.parametrize("path", flux3_definitions()) +def test_every_flux3_pipeline_step_places_every_component(path): + steps = klein_steps(load(path)) + assert steps, f"{path} runs no Klein step" + for step in steps: + components = step["pipeline"]["configuration"]["components"] + for name in ("transformer", "text_encoder", "vae"): + assert components.get(name, {}).get("device") == "cuda", (path, name) + assert components["text_encoder"].get("residency") is None + + +@pytest.mark.parametrize("path", flux3_definitions()) +def test_every_klein_step_runs_the_distilled_schedule_without_guidance(path): + for step in klein_steps(load(path)): + args = step["pipeline"]["arguments"] + assert args["num_inference_steps"] == 16 + assert "guidance_scale" not in args and "negative_prompt" not in args + assert args["output"] == "images" + + +@pytest.mark.parametrize("path", flux3_definitions()) +def test_size_variables_are_held_to_the_16_pixel_grid(path): + definition = load(path) + constraints = definition.get("variable_constraints", {}) + for name in ("width", "height"): + if name in definition.get("variables", {}): + assert constraints[name]["modulus"] == 16 and constraints[name].get("remainder", 0) == 0 + + +def test_the_16_pixel_grid_and_16_step_default_are_the_pipeline_s(): + import inspect + from diffusers.modular_pipelines.flux3 import before_denoise, modular_blocks_flux3_distilled + assert "num_inference_steps" in inspect.getsource(modular_blocks_flux3_distilled) + assert "16" in inspect.getsource(modular_blocks_flux3_distilled) + assert "vae_scale_factor * 2" in inspect.getsource(before_denoise) +``` + +- [ ] **Step 2: Run them and watch them fail** + +Run: `venv/bin/python -m pytest tests/test_flux3_catalog.py -v -n0` +Expected: collection error (the `flux3` templates folder doesn't exist). + +- [ ] **Step 3: Write the files** + +`prompts/flux3/man_dog.json`: `{"text": ""}`. Copy the `prompt` field of `examples/distilled/man_dog.json` in github.com/black-forest-labs/flux-3-image-klein at bbfe317d byte for byte (`gh api "repos/black-forest-labs/flux-3-image-klein/contents/examples/distilled/man_dog.json?ref=bbfe317dd81c825c30d76090a67d3d0db51da3cc" --jq .content | base64 -d`). It's the vendor's example caption, already the length and form the model was trained on. Match the other files' shape in `prompts/` (`cat prompts/kandinsky6/mountain_steam_train.json`). + +`workflows/templates/flux3/text-to-image.json`: + +```json +{ + "id": "Flux3KleinT2I", + "description": "FLUX 3 Image Klein (distilled, 5.6B) text to image: the family's baseline. The prompt here is already a full caption - BFL's own upsampled caption for their man walking a dog example - because Klein was trained on long upsampler captions and turns short prompts into generic pictures; upsample-then-generate shows how to get one from a short request. Sixteen steps with no guidance: the distilled model is a flow map, so it takes no guidance scale and no negative prompt. Width and height must be multiples of 16, up to 4096. Every component is placed on the GPU explicitly, about 19 GB of bf16 weights, which fits a 24 GB card at 1024x1024 with under 2 GB to spare. The weights are gated under the FLUX Kommunity License: accept it on the model page and log in with hf auth login first.", + "summary": "", + "cost_drivers": ["width", "height"], + "variables": { + "prompt": "prompt:flux3/man_dog", + "width": 1024, + "height": 1024, + "seed": 7 + }, + "variable_constraints": { + "width": {"modulus": 16, "remainder": 0, "reason": "FLUX 3 Image Klein's pipeline refuses a width not divisible by 16 (the VAE's 8x compression times the 2x2 latent packing), and only after the weights have loaded, so it is checked here first"}, + "height": {"modulus": 16, "remainder": 0, "reason": "FLUX 3 Image Klein's pipeline refuses a height not divisible by 16 (the VAE's 8x compression times the 2x2 latent packing), and only after the weights have loaded, so it is checked here first"} + }, + "seed": "variable:seed", + "steps": [""] +} +``` + +`workflows/models/flux3-klein.json`: the same Klein step with `"configures": "templates/text-to-image"`. Copy `variables` and `cost_drivers` from `workflows/templates/text-to-image.json` (`prompt` there is a short prompt, and the description says so: generic output without the upsampler). Add a `width`/`height` pair and their constraints only if that template declares them. Its description names the gated license and points at `templates/flux3/upsample-then-generate` for full captions. + +Fill `summary` to pass `test_every_template_has_a_summary_that_fits`. + +- [ ] **Step 4: Run the catalog tests** + +Run: `venv/bin/python -m pytest tests/test_flux3_catalog.py tests/test_catalog_structure.py tests/test_catalog_shape.py tests/test_template_defaults.py tests/test_template_subfolders.py tests/test_prompt_library.py tests/test_prompt_references.py -v -n0` +Expected: all PASS. If only `test_the_compact_listing_fits_the_budget` fails, stop and report its numbers (Global Constraints). + +Then: `venv/bin/python -m dw.validate workflows/templates/flux3/text-to-image.json` and the same for `workflows/models/flux3-klein.json`. Expected: valid. + +- [ ] **Step 5: Commit** + +```bash +git add prompts/flux3/man_dog.json workflows/templates/flux3/text-to-image.json workflows/models/flux3-klein.json tests/test_flux3_catalog.py +git commit -m "feat(flux3): Klein text-to-image template and model entry + +Co-Authored-By: Claude Opus 5.5 " +``` + +--- + +### Task 5: Upsample-then-generate template + +**Files:** +- Create: `workflows/templates/flux3/upsample-then-generate.json` +- Test: `tests/test_flux3_catalog.py` + +**Interfaces:** +- Consumes: `builtin:flux3_klein_upsample.json` (Task 3) and the Klein step (Task 4). + +- [ ] **Step 1: Write the failing test** + +```python +def test_upsample_then_generate_feeds_the_builtin_caption_to_klein(): + definition = load(os.path.join(TEMPLATES, "upsample-then-generate.json")) + first, second = definition["steps"] + assert first["workflow"]["path"] == "builtin:flux3_klein_upsample.json" + assert first["release_models"] is True + assert second["pipeline"]["arguments"]["prompt"] == f"previous_result:{first['name']}" +``` + +- [ ] **Step 2: Run it and watch it fail** + +Run: `venv/bin/python -m pytest tests/test_flux3_catalog.py -v -n0 -k upsample_then` +Expected: FAIL, file not found. + +- [ ] **Step 3: Write the template** + +Steps: `upsample` (`"workflow": {"path": "builtin:flux3_klein_upsample.json", "arguments": {"prompt": "variable:request"}}`, `"release_models": true`, result `text/plain`, `subfolder` `intermediate`), then the Klein step (Task 4 Interfaces) with `"prompt": "previous_result:upsample"`. Variables: `request` ("A man walking with his dog, afternoon sun"), `width`, `height`, `seed` 7, with the Task 4 constraints. The description says why `release_models` is there: the 18 GB upsampler and the 19 GB generator don't fit a 24 GB card together, so the upsampler is released before Klein loads. It also says the caption is saved as an intermediate so it can be reused. + +- [ ] **Step 4: Run the catalog tests and validate** + +Run: `venv/bin/python -m pytest tests/test_flux3_catalog.py tests/test_catalog_structure.py -v -n0 && venv/bin/python -m dw.validate workflows/templates/flux3/upsample-then-generate.json` +Expected: PASS and valid. + +- [ ] **Step 5: Commit** + +```bash +git add workflows/templates/flux3/upsample-then-generate.json tests/test_flux3_catalog.py +git commit -m "feat(flux3): upsample-then-generate template + +Co-Authored-By: Claude Opus 5.5 " +``` + +--- + +### Task 6: Probe P3, then the edit and multi-reference templates + +**Files:** +- Create: `workflows/templates/flux3/image-edit.json` +- Create: `workflows/templates/flux3/multi-reference.json` +- Modify: `docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md` (P3 row) +- Test: `tests/test_flux3_catalog.py` + +**Interfaces:** +- Consumes: `builtin:flux3_klein_upsample_edit.json` (Task 3) and the Klein step (Task 4). + +- [ ] **Step 1: Probe how a list of images reaches the pipeline (lem, through the dw MCP)** + +The Klein pipeline's `image` takes a PIL image or a list of them (`encoders.py:84` refuses anything else). Before writing the templates, find out whether a variable holding `[{"location": ""}, {"location": ""}]` arrives as a list of PIL images. Validate and run, on lem, the Task 4 step with `"image": "variable:images"` added, `images` set to two image URLs (any two public JPEGs), and a prompt naming `` and ``. Quote the plan's estimate and get Don's go-ahead first. Record the job id and the result in the spec's P3 row, along with peak memory from `get_job_events(kinds=["memory"])`. If the list doesn't load, stop and report: the fix is engine work outside this plan. + +- [ ] **Step 2: Write the failing tests** + +```python +@pytest.mark.parametrize("name", ["image-edit.json", "multi-reference.json"]) +def test_the_edit_templates_show_the_upsampler_the_pipeline_s_references(name): + definition = load(os.path.join(TEMPLATES, name)) + upsample = next(s for s in definition["steps"] if "workflow" in s) + assert upsample["workflow"]["path"] == "builtin:flux3_klein_upsample_edit.json" + (klein,) = klein_steps(definition) + assert upsample["workflow"]["arguments"]["images"] == klein["pipeline"]["arguments"]["image"] + assert klein["pipeline"]["arguments"]["prompt"] == f"previous_result:{upsample['name']}" + + +def test_multi_reference_states_the_ten_reference_cap_and_the_naming(): + description = load(os.path.join(TEMPLATES, "multi-reference.json"))["description"] + assert "10 references" in description + assert "" in description +``` + +- [ ] **Step 3: Run them and watch them fail** + +Run: `venv/bin/python -m pytest tests/test_flux3_catalog.py -v -n0 -k "edit_templates or ten_reference"` +Expected: FAIL, files not found. + +- [ ] **Step 4: Write the templates** + +`image-edit.json`: variables `instruction` ("put a red bandana around the dog's neck", BFL's example), `image` (a `{"location": ...}` reference; use the same public dog photo P3 used), `seed` 22, no width/height (diffusers defaults the output to 1 MP at the reference's aspect ratio, rounded to 16, and the description says so). Steps: `upsample` → `builtin:flux3_klein_upsample_edit.json` with `"images": ["variable:image"]` and `"instruction": "variable:instruction"`, `release_models: true`; then the Klein step with `"image": ["variable:image"]` and `"prompt": "previous_result:upsample"`. + +`multi-reference.json`: the same, with `images` a list variable of two references, `instruction` "print the dog from the second image onto the mug" (BFL's example), seed 23. Both the builtin and the Klein step read `"variable:images"`. The description states: references are named ``, ``, … (zero-based) in a caption written by hand, while the upsampler's own wording is "the first image, the second image"; the first reference sets the output's aspect ratio; and BFL trained Klein on up to 10 references, a cap the pipeline doesn't enforce, so more than 10 runs but is outside what the model learned. + +Add `"cost_drivers"` naming the image variable(s). + +- [ ] **Step 5: Run the catalog tests and validate both** + +Run: `venv/bin/python -m pytest tests/test_flux3_catalog.py tests/test_catalog_structure.py -v -n0 && venv/bin/python -m dw.validate workflows/templates/flux3/image-edit.json && venv/bin/python -m dw.validate workflows/templates/flux3/multi-reference.json` +Expected: PASS and valid. + +- [ ] **Step 6: Commit** + +```bash +git add workflows/templates/flux3/image-edit.json workflows/templates/flux3/multi-reference.json tests/test_flux3_catalog.py docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md +git commit -m "feat(flux3): image-edit and multi-reference templates + +Co-Authored-By: Claude Opus 5.5 " +``` + +--- + +### Task 7: Layout-with-boxes template + +**Files:** +- Create: `workflows/templates/flux3/layout-boxes.json` +- Create: `prompts/flux3/balloons_layout.json` +- Test: `tests/test_flux3_catalog.py` + +Scope note: BFL's box upsampling (`t2i_bbox`) goes through extra scaffolding in their `bbox.plan`/`bbox.reconcile` code. This template takes a caption written in the box format directly, as the Hugging Face card shows. Upsampled boxes are a dated follow-up. + +- [ ] **Step 1: Write the failing test** + +```python +def test_the_layout_caption_carries_tags_and_a_0_to_1000_box_table(): + import re + with open(os.path.join(ROOT, "prompts", "flux3", "balloons_layout.json")) as f: + text = json.load(f)["text"] + table = json.loads(text[text.index("["):]) + assert table and all({"id", "bbox", "desc"} <= set(row) for row in table) + for row in table: + assert f"<{row['id']}>" in text + top, left, bottom, right = row["bbox"] + assert 0 <= top < bottom <= 1000 and 0 <= left < right <= 1000 + assert klein_steps(load(os.path.join(TEMPLATES, "layout-boxes.json"))) +``` + +- [ ] **Step 2: Run it and watch it fail** + +Run: `venv/bin/python -m pytest tests/test_flux3_catalog.py -v -n0 -k layout` +Expected: FAIL, file not found. + +- [ ] **Step 3: Write the prompt and template** + +`prompts/flux3/balloons_layout.json`: a caption in BFL's box format (see the audit's *Prompt format → Layout / JSON* and https://docs.bfl.ml/flux_3/flux3_image_bounding_boxes). That's prose naming each boxed element as ``, then a JSON array of `{"id", "bbox": [top, left, bottom, right] on 0-1000, "desc"}`. Base it on BFL's `examples/inputs/balloons_boxes.json` at bbfe317d, converting its 0-1 floats to 0-1000 integers. Write the prose yourself: the example file holds boxes, not the caption. + +`layout-boxes.json`: the Task 4 template shape with `"prompt": "prompt:flux3/balloons_layout"` and seed 7. The description says boxes are a hint, not a mask: small boxed elements may not appear, and elements can run past their box. + +- [ ] **Step 4: Run the catalog tests and validate** + +Run: `venv/bin/python -m pytest tests/test_flux3_catalog.py tests/test_catalog_structure.py tests/test_prompt_library.py -v -n0 && venv/bin/python -m dw.validate workflows/templates/flux3/layout-boxes.json` +Expected: PASS and valid. + +- [ ] **Step 5: Commit** + +```bash +git add workflows/templates/flux3/layout-boxes.json prompts/flux3/balloons_layout.json tests/test_flux3_catalog.py +git commit -m "feat(flux3): layout-with-boxes template + +Co-Authored-By: Claude Opus 5.5 " +``` + +--- + +### Task 8: Family README + +**Files:** +- Create: `workflows/templates/flux3/README.md` +- Test: the existing README link checks (`grep -rln "README" tests/ | head` to find them; `tests/test_plugin_skills.py` and `tests/test_catalog_structure.py` are likely) + +- [ ] **Step 1: Write the README** + +Follow `workflows/templates/kandinsky6/README.md`: +- **First paragraph:** the model (5.6B distilled flow map, 16 steps, no guidance; the FLUX.2 VAE; the Qwen3-VL-4B text encoder), the sources (the Hugging Face card, github.com/black-forest-labs/flux-3-image-klein, docs.bfl.ml, diffusers `modular_pipelines/flux3`, merged 2026-10-15 and not yet released, so diffusers must be installed from git), the gated FLUX Kommunity License, and a link to `../../../docs/proposals/audits/2026-10-15-flux-3-klein-audit.md`. +- **A prompt paragraph:** BFL says not to use Klein without its upsampler, and the upsampler builtins read BFL's system prompts from the pup repo at a pinned commit. For writing captions by hand, link BFL's prompting guide (https://docs.bfl.ml/guides/prompting_summary) and note it's written for the API model, which upsamples on the server side. +- **A reading-order table:** text-to-image, upsample-then-generate, image-edit, multi-reference, layout-boxes, one idea each. +- **A placement note** marked as dw craft: every component placed explicitly, and why. +- **The agent pointer** to the `flux3-klein` skill of the dw plugin. + +- [ ] **Step 2: Run the link and catalog tests** + +Run: `venv/bin/python -m pytest tests/test_catalog_structure.py tests/test_plugin_skills.py -v -n0` +Expected: PASS. + +- [ ] **Step 3: Commit** + +```bash +git add workflows/templates/flux3/README.md +git commit -m "docs(flux3): Klein templates README + +Co-Authored-By: Claude Opus 5.5 " +``` + +--- + +### Task 9: GPU runs and costs (lem) + +**Files:** +- Modify: every `workflows/templates/flux3/*.json` and `workflows/models/flux3-klein.json` (`cost`) +- Modify: the spec's probe table + +- [ ] **Step 1: Deploy the branch to lem** + +Ask Don to deploy `feat/flux3-klein` to lem (the auto-mode classifier refuses `deploy.sh` over ssh). Then confirm with `get_diffusers_state` that the commit is at or after 7ecb1e1b. If `update_diffusers` was used, `dw-serve` must be restarted before validation sees the new classes. + +- [ ] **Step 2: Run each template once warm** + +For each template: `validate_workflow` with `workflow_path` set to the catalog name, quote the estimate, get Don's go-ahead (one batch approval for all five is fine), then `run_workflow`, `wait_for_job`, `get_job`, and look at each output (`get_output_image`, `get_output_text` for captions). Run each one a second time to get a warm timing, since the first run loads the models. + +- [ ] **Step 3: Write `cost`** + +Add `"cost": [{"device": "cuda", "name": "RTX 3090", "vram_gb": 24, "minutes": }]` to each template and the model entry. Record each job id, warm minutes and peak reserved memory in the spec's probe table. + +- [ ] **Step 4: Run the catalog tests, then commit** + +Run: `venv/bin/python -m pytest tests/test_catalog_structure.py tests/test_flux3_catalog.py -v -n0` + +```bash +git add workflows/templates/flux3 workflows/models/flux3-klein.json docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md +git commit -m "feat(flux3): measured RTX 3090 costs + +Co-Authored-By: Claude Opus 5.5 " +``` + +--- + +### Task 10: The `flux3-klein` composition skill + +**Files:** +- Create: `plugins/dw/skills/flux3-klein/SKILL.md` +- Modify: `plugins/dw/README.md` (name the skill) +- Modify: `CLAUDE.md` (the plugin paragraph lists every skill) +- Modify: `tests/test_plugin_skills.py` (new `TestFlux3KleinSkill`) + +- [ ] **Step 1: Write the failing tests** (append to `tests/test_plugin_skills.py`) + +```python +FLUX3_SKILL = os.path.join(PLUGIN_DIR, "skills", "flux3-klein", "SKILL.md") + + +class TestFlux3KleinSkill: + """The Klein skill's numbers come from the diffusers modular blocks that + enforce them; the 10-reference cap is the vendor's, which the pipeline + does not enforce, and the skill says so.""" + + def test_the_step_count_and_grid_are_the_pipeline_s(self): + import inspect + from diffusers.modular_pipelines.flux3 import before_denoise, modular_blocks_flux3_distilled + text = skill_text(FLUX3_SKILL) + assert "16" in inspect.getsource(modular_blocks_flux3_distilled) + assert "vae_scale_factor * 2" in inspect.getsource(before_denoise) + assert "16 steps" in text + assert "multiples of 16" in text + + def test_the_distilled_model_takes_no_guidance(self): + assert "no guidance" in skill_text(FLUX3_SKILL) + + def test_the_reference_cap_is_stated_as_the_vendor_s(self): + text = skill_text(FLUX3_SKILL) + assert "10 references" in text and "" in text + + def test_the_skill_defers_prompt_format_to_bfl(self): + text = skill_text(FLUX3_SKILL) + assert "docs.bfl.ml" in text + assert "flux3_klein_upsample" in text + + def test_the_skill_starts_with_the_server(self): + text = skill_text(FLUX3_SKILL) + assert text.index("get_server_info") < text.index("list_workflows") +``` + +Read the existing `TestKandinsky6Skill` and `skill_text` helper first and match their style. If the parametrized skill tests (size cap, links, catalog names) already pick up every `SKILL.md` under `plugins/dw/skills/`, they cover the new file with no change. + +- [ ] **Step 2: Run them and watch them fail** + +Run: `venv/bin/python -m pytest tests/test_plugin_skills.py -v -n0 -k Flux3` +Expected: FAIL, file not found. + +- [ ] **Step 3: Write the skill** + +Start from a copy of `plugins/dw/skills/kandinsky-6/SKILL.md` and replace every section: +- **Frontmatter `description`:** written for triggering ("Use when a dw MCP server is connected and the user wants an image from FLUX 3 Image Klein - a picture from a request, an edit of their photo, several references combined, or a layout with boxes…"). +- **Start with the server:** `get_server_info`, then `list_workflows(shape="image")` / `shape="image-edit"`. +- **The shape decision:** short request → upsample-then-generate; full caption → text-to-image; one photo to change → image-edit; several → multi-reference; positions → layout-boxes. +- **Hard rules:** 16 steps, no guidance, multiples of 16, up to 4096, 10 references named ``, gated license and `hf auth login`, explicit placement. +- **Prompts:** use the upsampler builtins (`flux3_klein_upsample`, `flux3_klein_upsample_edit`), and for hand-written captions BFL's guide at docs.bfl.ml. Nothing transcribed. +- **Validate, quote cost, run, look.** +- **Failure modes:** generic output from short prompts without the upsampler; boxes are hints; edits can change content outside the requested region; small reference details can be misdescribed. +- **Sources with dates.** + +Keep it under the 12 KB cap. + +Add the skill to `plugins/dw/README.md`'s list and to `CLAUDE.md`'s plugin paragraph (one composition skill per model family: add `flux3-klein`). + +- [ ] **Step 4: Run the plugin tests** + +Run: `venv/bin/python -m pytest tests/test_plugin_skills.py -v -n0` +Expected: all PASS. + +- [ ] **Step 5: Commit** + +```bash +git add plugins/dw/skills/flux3-klein plugins/dw/README.md CLAUDE.md tests/test_plugin_skills.py +git commit -m "feat(plugin): flux3-klein composition skill + +Co-Authored-By: Claude Opus 5.5 " +``` + +--- + +### Task 11: Full suite, cold drill, close + +- [ ] **Step 1: Full suite and lint** + +Run: `venv/bin/python -m pytest tests/ -q` and `venv/bin/ruff check dw tests` +Expected: all pass, no lint errors. + +- [ ] **Step 2: Cold drill** + +Follow `.claude/skills/model-family-onboarding/references/cold-drill-example.md`. Use lem on this branch with `dw-serve` restarted, a fresh Claude Code session in an empty directory with the plugin installed, and the request "make me a picture of a lighthouse keeper's cat on a stormy night". Run a control without the plugin. Record both transcripts' paths and the verdict (the skill fired, upsample-then-generate was chosen, cost was quoted before the run, the image was looked at) in the spec. + +- [ ] **Step 3: Follow-ups in the spec, dated** + +The spec's *Follow-ups* section holds these; add any new ones: +- Base + CFG (`guider`) +- fp8r/NVFP4 +- LoRA at 0.75 +- typography +- MPS on mini-ai +- upsampled boxes (`t2i_bbox`) +- a list-length constraint for the 10-reference cap +- the `downloads_required` gaps +- re-pinning the system prompts and the scaffold when diffusers PR #15058 merges + +- [ ] **Step 4: Hand off** + +Use superpowers:finishing-a-development-branch. diff --git a/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md b/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md index a48c44fc..e882ade5 100644 --- a/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md +++ b/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md @@ -66,9 +66,19 @@ settings*: code's - stated in the description). 4. `multi-reference.json` - 2-10 references named `` in the caption; the 10 cap is a `variable_constraints` rule because diffusers does not enforce it. -5. `layout-boxes.json` - `` caption + 0-1000 box table (`t2i_bbox`). -6. A 24 GB recipe, from the P1 measurement: bf16 with on-demand residency if it fits, - else SDNQ int8 on the transformer and text encoder (the H3 pattern). +5. `layout-boxes.json` - `` caption + 0-1000 box table, written in the box format + directly as the HF card shows. Upsampled boxes (`t2i_bbox`) need BFL's + `bbox.plan`/`bbox.reconcile` scaffolding, so they're a follow-up. + +No separate 24 GB recipe: P1 showed bf16 with every component placed fits the card at +1024², so the templates above are the 24 GB recipe. SDNQ int8 (the H3 pattern) waits +until a larger size or an edit is measured not to fit. + +The upsampler builtins (`builtin:flux3_klein_upsample.json`, +`builtin:flux3_klein_upsample_edit.json`) carry the pinned pup settings once. The edit +one wraps the instruction in BFL's three-line `I2I_SCAFFOLD` (`src/flux3/pup.py` at +bbfe317d) with `compose_text`. BFL's presence penalty (1.5 for text-to-image) has no +transformers `generate()` equivalent and is left out. Plus `workflows/models/flux3-klein.json` and the family README. From 86459df83c2871ca532534bc34d3778b75dc8836 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 19:56:28 -0500 Subject: [PATCH 03/20] feat(text_generation): chat_template_kwargs, revision and a list of images A Qwen3.5 upsampler must run with enable_thinking off, an edit sends every reference, and a pinned prompt wants pinned weights. Co-Authored-By: Claude Opus 5.5 --- dw/tasks/text_generation.py | 82 +++++++++++++++++++++++------- tests/test_text_generation.py | 95 +++++++++++++++++++++++++++++++++++ 2 files changed, 160 insertions(+), 17 deletions(-) diff --git a/dw/tasks/text_generation.py b/dw/tasks/text_generation.py index def832bd..5b7ec41f 100644 --- a/dw/tasks/text_generation.py +++ b/dw/tasks/text_generation.py @@ -42,15 +42,28 @@ def _image_part(image): return {"type": "image", "url" if isinstance(image, str) else "image": image} -def _build_messages(prompt, system_prompt, image): +def _images(image): + """The images a call carries, as a list in order, or None for none. + + A template's optional reference list arrives as [] or [None] when the + caller supplied nothing, and that is the same statement as no image. + """ + if image is None: + return None + images = list(image) if isinstance(image, (list, tuple)) else [image] + images = [i for i in images if i is not None and i != ""] + return images or None + + +def _build_messages(prompt, system_prompt, images): """Chat messages in the shape the chosen pipeline expects. Text-generation models take plain string content. Vision models take a - list of typed parts, because the image is a part of the message rather - than something alongside it. + list of typed parts, because the images are part of the message rather + than something alongside it - every image first, in order, then the text. """ messages = [] - if image is None: + if images is None: if system_prompt is not None: messages.append({"role": "system", "content": system_prompt}) messages.append({"role": "user", "content": prompt}) @@ -62,7 +75,8 @@ def _build_messages(prompt, system_prompt, image): messages.append( { "role": "user", - "content": [_image_part(image), {"type": "text", "text": prompt}], + "content": [_image_part(i) for i in images] + + [{"type": "text", "text": prompt}], } ) return messages @@ -81,8 +95,13 @@ def generate_text(prompt, device="cpu", **kwargs): accept one - a text-only model will fail to load as one. system_prompt: Optional system instruction for the model. max_new_tokens: Max tokens to generate (default: 500). - image: Optional PIL image, URL or path. Its presence is what - selects the vision pipeline. + image: Optional PIL image, URL or path, or a list of them - every + one goes to the model, in order. Its presence is what selects + the vision pipeline; an empty list is no image. + chat_template_kwargs: Settings for the model's chat template, + such as {"enable_thinking": false} for a Qwen3 model that + must not reason before answering. + revision: The model repo's commit, branch or tag to load. repetition_penalty: Vision pipeline only (default: 1.15). Raise it if a model still repeats itself, or set 1.0 to disable. generate_kwargs: Anything else to hand the model's generate() - @@ -94,11 +113,18 @@ def generate_text(prompt, device="cpu", **kwargs): """ # A workflow declaring an optional image passes it through as null when the # caller supplies none, and an empty string is the same statement - image = kwargs.get("image", None) or None + images = _images(kwargs.get("image", None) or None) system_prompt = kwargs.get("system_prompt", None) max_new_tokens = int(kwargs.get("max_new_tokens", 500)) + revision = kwargs.get("revision", None) + chat_template_kwargs = kwargs.get("chat_template_kwargs") or {} + if not isinstance(chat_template_kwargs, dict): + raise ValueError( + "text_generation: chat_template_kwargs must be an object of settings " + f"for the chat template, got {type(chat_template_kwargs).__name__}" + ) - if image is None: + if images is None: pipeline_task = "text-generation" model_name = kwargs.get("model_name", _DEFAULT_MODEL) else: @@ -110,6 +136,8 @@ def generate_text(prompt, device="cpu", **kwargs): def load_pipe(): logger.info(f"Generating text with {model_name} on {device}") placement = hf_pipeline_placement(device) + if revision is not None: + placement = {**placement, "revision": revision} return hf_pipeline( pipeline_task, model=model_name, @@ -120,35 +148,52 @@ def load_pipe(): # The task is part of the identity - the same model name can be loaded # under either pipeline, and they are not interchangeable pipe = cached_model( - ("text_generation", pipeline_task, model_name, str(device), str(dtype)), + ( + "text_generation", + pipeline_task, + model_name, + str(revision), + str(device), + str(dtype), + ), load_pipe, ) - messages = _build_messages(prompt, system_prompt, image) + messages = _build_messages(prompt, system_prompt, images) # Decoding is greedy, so the same input returns the same text every run - # which is what a workflow wants, and why there is nothing here to seed generation = {"do_sample": False} - if image is not None: + if images is not None: generation["repetition_penalty"] = float( kwargs.get("repetition_penalty", _VISION_REPETITION_PENALTY) ) # Last, so a workflow can override anything decided above generation.update(kwargs.get("generate_kwargs") or {}) - return _generate(pipe, messages, image, max_new_tokens, generation) + return _generate( + pipe, messages, images, max_new_tokens, generation, chat_template_kwargs + ) -def _generate(pipe, messages, image, max_new_tokens, generation): +def _generate(pipe, messages, images, max_new_tokens, generation, chat_template_kwargs): """Run the pipeline, which each path calls differently.""" - if image is None: + if images is None: # This pipeline collects anything it does not name into the arguments it - # forwards to generate(), so settings go in as plain keywords + # forwards to generate(), so settings go in as plain keywords. Chat + # template settings have their own argument, which it hands to + # apply_chat_template + extra = ( + {"tokenizer_encode_kwargs": chat_template_kwargs} + if chat_template_kwargs + else {} + ) results = pipe( messages, max_new_tokens=max_new_tokens, return_full_text=False, **generation, + **extra, ) else: # Images live inside the messages, so the chat goes in as `text` - the @@ -156,12 +201,15 @@ def _generate(pipe, messages, image, max_new_tokens, generation): # settings have to go through generate_kwargs here: anything else this # pipeline does not name explicitly is forwarded to the processor and # dropped, so a bare do_sample=False would leave sampling on. Passing - # max_new_tokens both ways is an error, so it stays a direct argument + # max_new_tokens both ways is an error, so it stays a direct argument. + # Keywords this pipeline does not name go to the processor's + # apply_chat_template, which is where the chat template settings belong results = pipe( text=messages, max_new_tokens=max_new_tokens, return_full_text=False, generate_kwargs=generation, + **chat_template_kwargs, ) text = results[0]["generated_text"].strip() diff --git a/tests/test_text_generation.py b/tests/test_text_generation.py index ca1f778a..1970c462 100644 --- a/tests/test_text_generation.py +++ b/tests/test_text_generation.py @@ -104,6 +104,101 @@ def test_strips_whitespace(self, mock_pipeline): result = generate_text("test", device="cpu") self.assertEqual(result, "some text with spaces") + @patch("dw.tasks.text_generation.hf_pipeline") + def test_chat_template_kwargs_reach_the_text_template(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text( + "a cat", + device="cpu", + model_name="m/text-chat-kwargs", + chat_template_kwargs={"enable_thinking": False}, + ) + + self.assertEqual( + pipe.call_args[1]["tokenizer_encode_kwargs"], {"enable_thinking": False} + ) + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_chat_template_kwargs_reach_the_vision_template(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text( + "a cat", + device="cpu", + model_name="m/vision-chat-kwargs", + image="https://example.com/a.png", + chat_template_kwargs={"enable_thinking": False}, + ) + + self.assertIs(pipe.call_args[1]["enable_thinking"], False) + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_no_chat_template_kwargs_passes_nothing_extra(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text("a cat", device="cpu", model_name="m/no-chat-kwargs") + + self.assertNotIn("tokenizer_encode_kwargs", pipe.call_args[1]) + + def test_chat_template_kwargs_must_be_a_mapping(self): + with self.assertRaises(ValueError): + generate_text( + "a cat", device="cpu", chat_template_kwargs=["enable_thinking"] + ) + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_revision_reaches_the_loader(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text("a cat", device="cpu", model_name="m/revision", revision="a" * 40) + + self.assertEqual(mock_pipeline.call_args[1]["revision"], "a" * 40) + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_a_list_of_images_becomes_one_part_each_in_order(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text( + "edit", + device="cpu", + model_name="m/image-list", + image=["https://example.com/1.png", "https://example.com/2.png"], + ) + + user = pipe.call_args[1]["text"][-1]["content"] + self.assertEqual( + [part.get("url") for part in user[:2]], + ["https://example.com/1.png", "https://example.com/2.png"], + ) + self.assertEqual(user[2], {"type": "text", "text": "edit"}) + self.assertEqual(mock_pipeline.call_args[0][0], "image-text-to-text") + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_an_empty_or_all_null_image_list_is_the_text_path(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + for empty in ([], [None]): + generate_text( + "a cat", + device="cpu", + model_name=f"m/empty-{len(empty)}", + image=empty, + ) + self.assertEqual(mock_pipeline.call_args[0][0], "text-generation") + class TestTextGenerationWithImage(unittest.TestCase): """An image switches generation onto a vision-language pipeline.""" From a213781df726be243671f6cac7f84d56e3ac87f7 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 19:59:29 -0500 Subject: [PATCH 04/20] feat(text_generation): system_prompt from a Hub repo file at a pinned commit Co-Authored-By: Claude Opus 5.5 --- docs/TASKS.md | 10 ++++ dw/tasks/text_generation.py | 94 +++++++++++++++++++++++++++++- tests/test_text_generation.py | 104 ++++++++++++++++++++++++++++++++++ 3 files changed, 207 insertions(+), 1 deletion(-) diff --git a/docs/TASKS.md b/docs/TASKS.md index 974c80d3..e0c0fc06 100644 --- a/docs/TASKS.md +++ b/docs/TASKS.md @@ -2850,6 +2850,16 @@ Generate or expand text using a local language model. Useful for expanding short | `repetition_penalty` | No | Vision path only (default: 1.15) — see below | | `generate_kwargs` | No | Anything else to pass to the model's `generate()` — `no_repeat_ngram_size`, `top_p`, `min_new_tokens`. Merged last, so it overrides the settings above | | `max_new_tokens` | No | Maximum tokens to generate (default: 500) | +| `chat_template_kwargs` | No | Extra keyword arguments for the model's chat template (for a template that takes switches such as a thinking toggle) | +| `revision` | No | Model revision (branch, tag or commit) to load `model_name` at | + +`image` may also be a list of images, each sent as its own picture in the one user message. + +`system_prompt` may name a file in a Hugging Face repo instead of carrying the text, +for a vendor that ships its prompt instructions there: +`{"from_hub": {"repo": "owner/name", "file": "system_prompts.json", "revision": "<40-hex commit>", "key": "t2i"}}`. +The revision must be a full commit, so the text can't move under the workflow. A `.json` +file takes `key` to pick one string; any other file is read whole. ### Writing a prompt from a picture diff --git a/dw/tasks/text_generation.py b/dw/tasks/text_generation.py index 5b7ec41f..e830310a 100644 --- a/dw/tasks/text_generation.py +++ b/dw/tasks/text_generation.py @@ -11,9 +11,15 @@ workflow whose prompt agrees with its keyframe and one whose prompt fights it. """ +import json import logging +import re +from pathlib import PurePosixPath + +from huggingface_hub import hf_hub_download from transformers import pipeline as hf_pipeline from .. import preferred_task_dtype +from ..locations import validate_hub_repo_id from .model_cache import cached_model, hf_pipeline_placement logger = logging.getLogger("dw") @@ -82,6 +88,91 @@ def _build_messages(prompt, system_prompt, images): return messages +# A pinned prompt has to stay pinned: a short hash can become ambiguous and a +# branch moves, so only a full commit is accepted +_FULL_COMMIT = re.compile(r"^[0-9a-f]{40}$") +_FROM_HUB_FIELDS = {"repo", "file", "revision", "key"} + + +def _hub_file(file): + """A file path inside a repo, refused if it could name anything outside it.""" + if not isinstance(file, str) or not file or "\\" in file: + raise ValueError( + f"system_prompt.from_hub.file must be a path inside the repo, got {file!r}" + ) + path = PurePosixPath(file) + if path.is_absolute() or ".." in path.parts: + raise ValueError( + f"system_prompt.from_hub.file must be a path inside the repo, got {file!r}" + ) + return file + + +def resolve_system_prompt(value): + """The system prompt text, reading it from a Hub repo when one is named. + + Some vendors ship the instructions their model was trained against as a + file in its repo, and change them. Naming the file and a full commit runs + exactly that text without copying it into a workflow, where it would + drift from the source. + + Args: + value: None, the text itself, or {"from_hub": {"repo", "file", + "revision", "key"}}. `revision` is a full 40-character commit. + A .json file takes `key` to pick one string out of a top-level + object; any other file is read whole as text + + Raises: + ValueError: On any other shape, a revision that is not a full commit, + a file path that leaves the repo, a missing key, or a value that + is not text + """ + if value is None or isinstance(value, str): + return value + if ( + not isinstance(value, dict) + or set(value) != {"from_hub"} + or not isinstance(value["from_hub"], dict) + ): + raise ValueError( + 'system_prompt must be text or {"from_hub": {"repo", "file", "revision", "key"}}' + ) + spec = value["from_hub"] + unknown = set(spec) - _FROM_HUB_FIELDS + if unknown: + raise ValueError(f"system_prompt.from_hub does not take {sorted(unknown)}") + repo = validate_hub_repo_id(spec.get("repo"), what="system_prompt.from_hub.repo") + file = _hub_file(spec.get("file")) + revision = spec.get("revision") + if not isinstance(revision, str) or not _FULL_COMMIT.match(revision): + raise ValueError( + "system_prompt.from_hub.revision must be a full 40-character lowercase " + f"commit hash, so the text cannot change under the workflow; got {revision!r}" + ) + key = spec.get("key") + + path = hf_hub_download(repo, file, revision=revision) + with open(path, encoding="utf-8") as f: + if file.endswith(".json"): + data = json.load(f) + else: + data = f.read() + if key is not None: + if not isinstance(data, dict) or key not in data: + keys = sorted(data) if isinstance(data, dict) else [] + raise ValueError( + f"system_prompt.from_hub: {file} at {revision[:8]} has no key {key!r}; " + f"it has {keys}" + ) + data = data[key] + if not isinstance(data, str) or not data.strip(): + raise ValueError( + f"system_prompt.from_hub: {file}{'[' + repr(key) + ']' if key else ''} " + "is not text" + ) + return data + + def generate_text(prompt, device="cpu", **kwargs): """Generate text from a prompt using a local language model. @@ -114,7 +205,8 @@ def generate_text(prompt, device="cpu", **kwargs): # A workflow declaring an optional image passes it through as null when the # caller supplies none, and an empty string is the same statement images = _images(kwargs.get("image", None) or None) - system_prompt = kwargs.get("system_prompt", None) + # Resolved before anything loads, so a bad reference costs nothing + system_prompt = resolve_system_prompt(kwargs.get("system_prompt", None)) max_new_tokens = int(kwargs.get("max_new_tokens", 500)) revision = kwargs.get("revision", None) chat_template_kwargs = kwargs.get("chat_template_kwargs") or {} diff --git a/tests/test_text_generation.py b/tests/test_text_generation.py index 1970c462..8530f23f 100644 --- a/tests/test_text_generation.py +++ b/tests/test_text_generation.py @@ -1,5 +1,8 @@ """Tests for text_generation task.""" +import json as _json +import os +import tempfile import unittest from unittest.mock import patch, MagicMock @@ -8,6 +11,7 @@ _DEFAULT_MODEL, DEFAULT_VISION_MODEL, _VISION_REPETITION_PENALTY, + resolve_system_prompt, ) @@ -431,5 +435,105 @@ def test_other_devices_still_use_device_map(self, mock_pipeline): self.assertNotIn("device", kwargs) +PIN = "3f3edb52b2c55de0d49feab28598b68f7f8e4830" + + +def _spec(**over): + spec = { + "repo": "black-forest-labs/flux-3-image-klein-pup", + "file": "system_prompts.json", + "revision": PIN, + "key": "t2i", + } + spec.update(over) + return {"from_hub": spec} + + +class TestSystemPromptFromHub(unittest.TestCase): + def setUp(self): + + self.dir = tempfile.mkdtemp() + self.json_path = os.path.join(self.dir, "system_prompts.json") + with open(self.json_path, "w") as f: + _json.dump({"t2i": "expand the request", "i2i": "edit it", "n": 3}, f) + + def test_a_string_or_none_passes_through(self): + self.assertEqual(resolve_system_prompt("plain"), "plain") + self.assertIsNone(resolve_system_prompt(None)) + + @patch("dw.tasks.text_generation.hf_hub_download") + def test_the_key_s_text_is_returned_from_the_pinned_file(self, download): + download.return_value = self.json_path + self.assertEqual(resolve_system_prompt(_spec()), "expand the request") + download.assert_called_once_with( + "black-forest-labs/flux-3-image-klein-pup", + "system_prompts.json", + revision=PIN, + ) + + @patch("dw.tasks.text_generation.hf_hub_download") + def test_a_missing_key_names_the_keys_the_file_has(self, download): + download.return_value = self.json_path + with self.assertRaises(ValueError) as caught: + resolve_system_prompt(_spec(key="t2v")) + self.assertIn("i2i", str(caught.exception)) + self.assertIn("t2i", str(caught.exception)) + + @patch("dw.tasks.text_generation.hf_hub_download") + def test_a_value_that_is_not_text_is_refused(self, download): + download.return_value = self.json_path + with self.assertRaises(ValueError): + resolve_system_prompt(_spec(key="n")) + + def test_a_short_hash_or_branch_is_refused(self): + for revision in ("3f3edb52", "main", PIN.upper()): + with self.assertRaises(ValueError) as caught: + resolve_system_prompt(_spec(revision=revision)) + self.assertIn("40", str(caught.exception)) + + def test_a_path_that_escapes_the_repo_is_refused(self): + for file in ( + "../secrets.json", + "/etc/passwd", + "a\\b.json", + "", + "x/../../y.json", + ): + with self.assertRaises(ValueError): + resolve_system_prompt(_spec(file=file)) + + def test_a_bad_repo_id_is_refused(self): + from dw.security import InvalidInputError + + with self.assertRaises((ValueError, InvalidInputError)): + resolve_system_prompt(_spec(repo="https://evil.example/x")) + + def test_an_unknown_field_or_shape_is_refused(self): + for bad in ( + {"from_hub": {**_spec()["from_hub"], "branch": "main"}}, + {"from_url": "x"}, + ["t2i"], + ): + with self.assertRaises(ValueError): + resolve_system_prompt(bad) + + @patch("dw.tasks.text_generation.hf_pipeline") + @patch("dw.tasks.text_generation.hf_hub_download") + def test_generate_text_sends_the_resolved_text(self, download, mock_pipeline): + download.return_value = self.json_path + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + + generate_text( + "a dog", device="cpu", model_name="m/from-hub", system_prompt=_spec() + ) + + messages = pipe.call_args[0][0] + self.assertEqual( + messages[0], {"role": "system", "content": "expand the request"} + ) + + if __name__ == "__main__": unittest.main() From 4c6adb39fc4a595a44c37c6df4360a90f7bcd950 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:02:49 -0500 Subject: [PATCH 05/20] fix(text_generation): refuse structural keys and non-scalars in chat_template_kwargs Workflows are untrusted and the settings are splatted into the pipeline and apply_chat_template calls, so images/tools/etc. would reach transformers unvalidated. Also pins that image list elements get the URL check. Co-Authored-By: Claude Opus 5.5 --- dw/tasks/text_generation.py | 64 +++++++++++++++++++++++++++++++---- tests/test_security_ssrf.py | 14 ++++++++ tests/test_text_generation.py | 55 ++++++++++++++++++++++++++++++ 3 files changed, 127 insertions(+), 6 deletions(-) diff --git a/dw/tasks/text_generation.py b/dw/tasks/text_generation.py index e830310a..1773861a 100644 --- a/dw/tasks/text_generation.py +++ b/dw/tasks/text_generation.py @@ -173,6 +173,61 @@ def resolve_system_prompt(value): return data +# The pipelines' and apply_chat_template's own structural arguments. Chat +# template settings are splatted into those calls, so one of these in a +# workflow (which is untrusted) would not set a template variable - it would +# hand transformers an image or video to fetch that dw never validated, or +# change what the template call does (tokenize, tools, the template itself) +_RESERVED_CHAT_TEMPLATE_KEYS = frozenset( + { + "images", + "videos", + "audio", + "audios", + "text", + "timeout", + "processor_kwargs", + "generate_kwargs", + "max_new_tokens", + "return_full_text", + "return_tensors", + "tokenize", + "return_dict", + "chat_template", + "add_generation_prompt", + "continue_final_message", + "tools", + "documents", + "tokenizer_encode_kwargs", + } +) + + +def _validated_chat_template_kwargs(value): + """Chat template settings: an object of JSON scalars, none of them one of + the call's own structural arguments.""" + if not value: + return {} + if not isinstance(value, dict): + raise ValueError( + "text_generation: chat_template_kwargs must be an object of settings " + f"for the chat template, got {type(value).__name__}" + ) + for key, setting in value.items(): + if key in _RESERVED_CHAT_TEMPLATE_KEYS: + raise ValueError( + f"text_generation: chat_template_kwargs key '{key}' is not a " + "chat-template setting - it is an argument of the generation " + "call itself" + ) + if setting is not None and not isinstance(setting, (bool, int, float, str)): + raise ValueError( + f"text_generation: chat_template_kwargs value for '{key}' must be " + f"a boolean, number, string or null, got {type(setting).__name__}" + ) + return value + + def generate_text(prompt, device="cpu", **kwargs): """Generate text from a prompt using a local language model. @@ -209,12 +264,9 @@ def generate_text(prompt, device="cpu", **kwargs): system_prompt = resolve_system_prompt(kwargs.get("system_prompt", None)) max_new_tokens = int(kwargs.get("max_new_tokens", 500)) revision = kwargs.get("revision", None) - chat_template_kwargs = kwargs.get("chat_template_kwargs") or {} - if not isinstance(chat_template_kwargs, dict): - raise ValueError( - "text_generation: chat_template_kwargs must be an object of settings " - f"for the chat template, got {type(chat_template_kwargs).__name__}" - ) + chat_template_kwargs = _validated_chat_template_kwargs( + kwargs.get("chat_template_kwargs") + ) if images is None: pipeline_task = "text-generation" diff --git a/tests/test_security_ssrf.py b/tests/test_security_ssrf.py index 557a3033..28a39bac 100644 --- a/tests/test_security_ssrf.py +++ b/tests/test_security_ssrf.py @@ -786,3 +786,17 @@ def test_trust_lifts_the_scope(self, no_real_sockets, monkeypatch): monkeypatch.setenv(TRUST_WORKFLOWS_ENV_VAR, "0") transport = _encode("https://evil.example/encode", monkeypatch) assert not _carried_token(transport.sent[0]) + + +class TestImageListElementsAreValidated: + """An `image` argument may be a list (text_generation sends every one to a + model); each element goes through the same location check as a lone string.""" + + def test_a_blocked_url_in_a_list_is_refused(self, untrusted, no_real_sockets): + import pytest + + from dw.arguments import realize_args + + arguments = {"image": ["http://169.254.169.254/x.png"]} + with pytest.raises(Exception, match="Refusing to fetch an image"): + realize_args(arguments) diff --git a/tests/test_text_generation.py b/tests/test_text_generation.py index 8530f23f..01369a4b 100644 --- a/tests/test_text_generation.py +++ b/tests/test_text_generation.py @@ -203,6 +203,61 @@ def test_an_empty_or_all_null_image_list_is_the_text_path(self, mock_pipeline): ) self.assertEqual(mock_pipeline.call_args[0][0], "text-generation") + @patch("dw.tasks.text_generation.hf_pipeline") + def test_a_non_scalar_chat_template_value_is_refused(self, mock_pipeline): + with self.assertRaisesRegex(ValueError, "enable_thinking"): + generate_text( + "a cat", + device="cpu", + model_name="m/nonscalar", + chat_template_kwargs={"enable_thinking": {"a": 1}}, + ) + mock_pipeline.assert_not_called() + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_a_reserved_chat_template_key_is_refused_before_any_call( + self, mock_pipeline + ): + pipe = MagicMock() + mock_pipeline.return_value = pipe + for key in ("images", "tools", "tokenizer_encode_kwargs"): + with self.assertRaisesRegex(ValueError, f"'{key}'.*not a chat-template"): + generate_text( + "a cat", + device="cpu", + model_name="m/reserved-vision", + image="https://example.com/a.png", + chat_template_kwargs={key: "http://169.254.169.254/x"}, + ) + with self.assertRaisesRegex(ValueError, f"'{key}'.*not a chat-template"): + generate_text( + "a cat", + device="cpu", + model_name="m/reserved-text", + chat_template_kwargs={key: "x"}, + ) + pipe.assert_not_called() + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_a_legitimate_setting_passes_on_both_paths(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "out"}] + mock_pipeline.return_value = pipe + settings = {"enable_thinking": False, "effort": "low", "n": 2, "t": 0.5} + + generate_text( + "a cat", device="cpu", model_name="m/ok-text", chat_template_kwargs=settings + ) + self.assertEqual(pipe.call_args[1]["tokenizer_encode_kwargs"], settings) + generate_text( + "a cat", + device="cpu", + model_name="m/ok-vision", + image="https://example.com/a.png", + chat_template_kwargs=settings, + ) + self.assertIs(pipe.call_args[1]["enable_thinking"], False) + class TestTextGenerationWithImage(unittest.TestCase): """An image switches generation onto a vision-language pipeline.""" From 643227f99bc4a22fa210f2e8a726da7b7d210434 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:04:55 -0500 Subject: [PATCH 06/20] fix(text_generation): from_hub validates repo/file/key, public files only, size cap Co-Authored-By: Claude Opus 5.5 --- docs/TASKS.md | 4 +++- dw/tasks/text_generation.py | 38 +++++++++++++++++++--------------- tests/test_text_generation.py | 39 ++++++++++++++++++++++++++++++++++- 3 files changed, 62 insertions(+), 19 deletions(-) diff --git a/docs/TASKS.md b/docs/TASKS.md index e0c0fc06..980a0e60 100644 --- a/docs/TASKS.md +++ b/docs/TASKS.md @@ -2859,7 +2859,9 @@ Generate or expand text using a local language model. Useful for expanding short for a vendor that ships its prompt instructions there: `{"from_hub": {"repo": "owner/name", "file": "system_prompts.json", "revision": "<40-hex commit>", "key": "t2i"}}`. The revision must be a full commit, so the text can't move under the workflow. A `.json` -file takes `key` to pick one string; any other file is read whole. +file takes `key` to pick one string; any other file is read whole. Only public repos are +reachable (the download never uses the server's Hugging Face token), and a file over 1 MiB +is refused. ### Writing a prompt from a picture diff --git a/dw/tasks/text_generation.py b/dw/tasks/text_generation.py index 1773861a..6ea9418c 100644 --- a/dw/tasks/text_generation.py +++ b/dw/tasks/text_generation.py @@ -13,13 +13,13 @@ import json import logging +import os import re -from pathlib import PurePosixPath from huggingface_hub import hf_hub_download from transformers import pipeline as hf_pipeline from .. import preferred_task_dtype -from ..locations import validate_hub_repo_id +from ..locations import validate_hub_repo_id, validate_weight_name from .model_cache import cached_model, hf_pipeline_placement logger = logging.getLogger("dw") @@ -94,18 +94,8 @@ def _build_messages(prompt, system_prompt, images): _FROM_HUB_FIELDS = {"repo", "file", "revision", "key"} -def _hub_file(file): - """A file path inside a repo, refused if it could name anything outside it.""" - if not isinstance(file, str) or not file or "\\" in file: - raise ValueError( - f"system_prompt.from_hub.file must be a path inside the repo, got {file!r}" - ) - path = PurePosixPath(file) - if path.is_absolute() or ".." in path.parts: - raise ValueError( - f"system_prompt.from_hub.file must be a path inside the repo, got {file!r}" - ) - return file +# A system prompt is a page of instructions; anything bigger is not one +_MAX_PROMPT_FILE_BYTES = 1024 * 1024 def resolve_system_prompt(value): @@ -141,8 +131,13 @@ def resolve_system_prompt(value): unknown = set(spec) - _FROM_HUB_FIELDS if unknown: raise ValueError(f"system_prompt.from_hub does not take {sorted(unknown)}") - repo = validate_hub_repo_id(spec.get("repo"), what="system_prompt.from_hub.repo") - file = _hub_file(spec.get("file")) + repo = spec.get("repo") + if not isinstance(repo, str) or not repo: + raise ValueError( + f"system_prompt.from_hub.repo must be a Hub repo id string, got {repo!r}" + ) + repo = validate_hub_repo_id(repo, what="system_prompt.from_hub.repo") + file = validate_weight_name(spec.get("file"), what="system_prompt.from_hub.file") revision = spec.get("revision") if not isinstance(revision, str) or not _FULL_COMMIT.match(revision): raise ValueError( @@ -150,8 +145,17 @@ def resolve_system_prompt(value): f"commit hash, so the text cannot change under the workflow; got {revision!r}" ) key = spec.get("key") + if key is not None and not isinstance(key, str): + raise ValueError(f"system_prompt.from_hub.key must be a string, got {key!r}") - path = hf_hub_download(repo, file, revision=revision) + # token=False: the server's own Hub login must not reach a private or gated + # file on a workflow's say-so, since the text is echoed into outputs + path = hf_hub_download(repo, file, revision=revision, token=False) + if os.path.getsize(path) > _MAX_PROMPT_FILE_BYTES: + raise ValueError( + f"system_prompt.from_hub: {file} is larger than " + f"{_MAX_PROMPT_FILE_BYTES} bytes, too big for a system prompt" + ) with open(path, encoding="utf-8") as f: if file.endswith(".json"): data = json.load(f) diff --git a/tests/test_text_generation.py b/tests/test_text_generation.py index 01369a4b..b5617d3a 100644 --- a/tests/test_text_generation.py +++ b/tests/test_text_generation.py @@ -524,6 +524,7 @@ def test_the_key_s_text_is_returned_from_the_pinned_file(self, download): "black-forest-labs/flux-3-image-klein-pup", "system_prompts.json", revision=PIN, + token=False, ) @patch("dw.tasks.text_generation.hf_hub_download") @@ -547,16 +548,52 @@ def test_a_short_hash_or_branch_is_refused(self): self.assertIn("40", str(caught.exception)) def test_a_path_that_escapes_the_repo_is_refused(self): + from dw.security import InvalidInputError, PathTraversalError + for file in ( "../secrets.json", "/etc/passwd", "a\\b.json", "", "x/../../y.json", + "a\x00b.json", + "~/x", + "C:x", + "./x", + "a//b", ): - with self.assertRaises(ValueError): + with self.assertRaises((ValueError, InvalidInputError, PathTraversalError)): resolve_system_prompt(_spec(file=file)) + def test_a_missing_or_non_string_repo_is_refused(self): + omitted = _spec() + del omitted["from_hub"]["repo"] + for bad in (omitted, _spec(repo=None), _spec(repo=3)): + with self.assertRaises(ValueError) as caught: + resolve_system_prompt(bad) + self.assertIn("system_prompt.from_hub.repo", str(caught.exception)) + + def test_a_non_string_key_is_refused(self): + for key in (["t2i"], {"a": 1}, 3): + with self.assertRaises(ValueError): + resolve_system_prompt(_spec(key=key)) + + @patch("dw.tasks.text_generation.hf_hub_download") + def test_the_download_never_carries_the_server_token(self, download): + download.return_value = self.json_path + resolve_system_prompt(_spec()) + self.assertIs(download.call_args.kwargs["token"], False) + + @patch("dw.tasks.text_generation.hf_hub_download") + def test_an_oversize_file_is_refused_unread(self, download): + big = os.path.join(self.dir, "big.txt") + with open(big, "wb") as f: + f.truncate(1024 * 1024 + 1) + download.return_value = big + with patch("builtins.open", side_effect=AssertionError("read")): + with self.assertRaises(ValueError): + resolve_system_prompt(_spec(file="big.txt", key=None)) + def test_a_bad_repo_id_is_refused(self): from dw.security import InvalidInputError From 66314b63f36e59e8e835b9e28b064786bbb79d10 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:07:39 -0500 Subject: [PATCH 07/20] feat(flux3): Klein prompt-upsampler builtins on the pinned pup Co-Authored-By: Claude Opus 5.5 --- dw/workflows/flux3_klein_upsample.json | 33 ++++++++++ dw/workflows/flux3_klein_upsample_edit.json | 51 +++++++++++++++ tests/test_flux3_catalog.py | 69 +++++++++++++++++++++ 3 files changed, 153 insertions(+) create mode 100644 dw/workflows/flux3_klein_upsample.json create mode 100644 dw/workflows/flux3_klein_upsample_edit.json create mode 100644 tests/test_flux3_catalog.py diff --git a/dw/workflows/flux3_klein_upsample.json b/dw/workflows/flux3_klein_upsample.json new file mode 100644 index 00000000..9a2ef5d4 --- /dev/null +++ b/dw/workflows/flux3_klein_upsample.json @@ -0,0 +1,33 @@ +{ + "id": "flux3_klein_upsample", + "description": "Expands a short request into the caption FLUX 3 Image Klein was trained on, with Black Forest Labs' own prompt upsampler, flux-3-image-klein-pup, a Qwen3.5-9B fine-tune. BFL says not to use Klein without it. The weights and the text-to-image system prompt are both pinned to pup commit 3f3edb52, and the prompt is read from the repo's system_prompts.json rather than copied. Thinking is switched off and decoding is greedy, as BFL runs it. BFL also applies a presence penalty of 1.5 for text-to-image, which transformers' generate() has no equivalent for, so it is left out. About 18 GB in bf16.", + "variables": { + "prompt": "A man walking with his dog, afternoon sun", + "device": "cuda" + }, + "steps": [ + { + "name": "upsample", + "task": { + "command": "text_generation", + "arguments": { + "prompt": "variable:prompt", + "model_name": "black-forest-labs/flux-3-image-klein-pup", + "revision": "3f3edb52b2c55de0d49feab28598b68f7f8e4830", + "system_prompt": { + "from_hub": { + "repo": "black-forest-labs/flux-3-image-klein-pup", + "file": "system_prompts.json", + "revision": "3f3edb52b2c55de0d49feab28598b68f7f8e4830", + "key": "t2i" + } + }, + "chat_template_kwargs": {"enable_thinking": false}, + "max_new_tokens": 3584, + "device": "variable:device" + } + }, + "result": {"content_type": "text/plain", "save": false} + } + ] +} diff --git a/dw/workflows/flux3_klein_upsample_edit.json b/dw/workflows/flux3_klein_upsample_edit.json new file mode 100644 index 00000000..e6505f24 --- /dev/null +++ b/dw/workflows/flux3_klein_upsample_edit.json @@ -0,0 +1,51 @@ +{ + "id": "flux3_klein_upsample_edit", + "description": "Expands an edit instruction into the caption FLUX 3 Image Klein was trained on, with Black Forest Labs' own prompt upsampler, flux-3-image-klein-pup, a Qwen3.5-9B fine-tune. The upsampler sees every reference in order, as the pipeline will. The instruction is wrapped in the three-line scaffold that is BFL's I2I_SCAFFOLD from src/flux3/pup.py at commit bbfe317d in github.com/black-forest-labs/flux-3-image-klein. The weights and the image-edit system prompt are both pinned to pup commit 3f3edb52, and the prompt is read from the repo's system_prompts.json rather than copied. Thinking is switched off and decoding is greedy with no repetition penalty, as BFL runs it. BFL resizes each reference to a pixel budget (1 MP for one reference, 724x724 for two, 512x512 for three or more) before the upsampler sees it, while transformers' processor applies its own sizing, so the two can describe small details differently. About 18 GB in bf16.", + "variables": { + "instruction": "put a red bandana around the dog's neck", + "images": [], + "device": "cuda" + }, + "steps": [ + { + "name": "scaffold", + "task": { + "command": "compose_text", + "arguments": { + "parts": [ + "Reference images follow in order (the first image, the second image, ...).\nUser request: ", + "variable:instruction", + "\nOutput the caption." + ], + "separator": "" + } + }, + "result": {"content_type": "text/plain", "save": false} + }, + { + "name": "upsample", + "task": { + "command": "text_generation", + "arguments": { + "prompt": "previous_result:scaffold", + "image": "variable:images", + "model_name": "black-forest-labs/flux-3-image-klein-pup", + "revision": "3f3edb52b2c55de0d49feab28598b68f7f8e4830", + "system_prompt": { + "from_hub": { + "repo": "black-forest-labs/flux-3-image-klein-pup", + "file": "system_prompts.json", + "revision": "3f3edb52b2c55de0d49feab28598b68f7f8e4830", + "key": "i2i" + } + }, + "chat_template_kwargs": {"enable_thinking": false}, + "max_new_tokens": 3584, + "repetition_penalty": 1.0, + "device": "variable:device" + } + }, + "result": {"content_type": "text/plain", "save": false} + } + ] +} diff --git a/tests/test_flux3_catalog.py b/tests/test_flux3_catalog.py new file mode 100644 index 00000000..1d0bc983 --- /dev/null +++ b/tests/test_flux3_catalog.py @@ -0,0 +1,69 @@ +"""The FLUX 3 Image Klein catalog: the upsampler builtins, the templates' settings +and their placement. The vendor values here are the ones the design states +(docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md).""" + +import json +import os + +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +BUILTINS = os.path.join(ROOT, "dw", "workflows") +PUP = "black-forest-labs/flux-3-image-klein-pup" +PUP_PIN = "3f3edb52b2c55de0d49feab28598b68f7f8e4830" + + +def load(path): + with open(path) as f: + return json.load(f) + + +def generation_steps(definition): + return [ + s + for s in definition["steps"] + if s.get("task", {}).get("command") == "text_generation" + ] + + +@pytest.mark.parametrize( + "name,key", + [ + ("flux3_klein_upsample.json", "t2i"), + ("flux3_klein_upsample_edit.json", "i2i"), + ], +) +def test_each_upsampler_runs_the_pinned_pup_as_bfl_does(name, key): + definition = load(os.path.join(BUILTINS, name)) + (step,) = generation_steps(definition) + args = step["task"]["arguments"] + assert args["model_name"] == PUP + assert args["revision"] == PUP_PIN + assert args["system_prompt"] == { + "from_hub": { + "repo": PUP, + "file": "system_prompts.json", + "revision": PUP_PIN, + "key": key, + } + } + assert args["chat_template_kwargs"] == {"enable_thinking": False} + assert args["max_new_tokens"] == 3584 + + +def test_the_edit_upsampler_wraps_the_instruction_in_bfl_s_scaffold(): + definition = load(os.path.join(BUILTINS, "flux3_klein_upsample_edit.json")) + compose = next( + s for s in definition["steps"] if s["task"]["command"] == "compose_text" + ) + parts = compose["task"]["arguments"]["parts"] + assert parts == [ + "Reference images follow in order (the first image, the second image, ...).\nUser request: ", + "variable:instruction", + "\nOutput the caption.", + ] + assert compose["task"]["arguments"]["separator"] == "" + (step,) = generation_steps(definition) + assert step["task"]["arguments"]["image"] == "variable:images" + assert step["task"]["arguments"]["repetition_penalty"] == 1.0 + assert "bbfe317d" in definition["description"] From c7ca24b430bf4cbc4061ec389f8b8b8315b88c1c Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:10:40 -0500 Subject: [PATCH 08/20] fix(text_generation): from_hub fails closed unless the Hub serves the file anonymously A full commit is answered from the local cache without an auth check, so token=False alone left cached gated files readable. Ask the Hub first. Co-Authored-By: Claude Opus 5.5 --- docs/TASKS.md | 3 ++- dw/tasks/text_generation.py | 21 ++++++++++++++++++- tests/test_text_generation.py | 38 ++++++++++++++++++++++++++++++++++- 3 files changed, 59 insertions(+), 3 deletions(-) diff --git a/docs/TASKS.md b/docs/TASKS.md index 980a0e60..1d738578 100644 --- a/docs/TASKS.md +++ b/docs/TASKS.md @@ -2861,7 +2861,8 @@ for a vendor that ships its prompt instructions there: The revision must be a full commit, so the text can't move under the workflow. A `.json` file takes `key` to pick one string; any other file is read whole. Only public repos are reachable (the download never uses the server's Hugging Face token), and a file over 1 MiB -is refused. +is refused. The Hub is asked, without credentials, whether the file is public before every +use, so `from_hub` needs network access even when the file is already cached. ### Writing a prompt from a picture diff --git a/dw/tasks/text_generation.py b/dw/tasks/text_generation.py index 6ea9418c..6cc1ee5c 100644 --- a/dw/tasks/text_generation.py +++ b/dw/tasks/text_generation.py @@ -16,7 +16,7 @@ import os import re -from huggingface_hub import hf_hub_download +from huggingface_hub import get_hf_file_metadata, hf_hub_download, hf_hub_url from transformers import pipeline as hf_pipeline from .. import preferred_task_dtype from ..locations import validate_hub_repo_id, validate_weight_name @@ -148,6 +148,25 @@ def resolve_system_prompt(value): if key is not None and not isinstance(key, str): raise ValueError(f"system_prompt.from_hub.key must be a string, got {key!r}") + # hf_hub_download answers a full commit from the local cache without asking + # the Hub, so files the server fetched with its own login (gated or private + # repos) would be readable by any workflow. Ask the Hub, anonymously, first, + # and fail closed on any error rather than fall back to the cache + try: + metadata = get_hf_file_metadata( + hf_hub_url(repo, file, revision=revision), token=False + ) + except Exception as e: + raise ValueError( + f"system_prompt.from_hub: {repo}/{file} at {revision[:8]} is not a " + f"public file the Hub will serve without credentials ({type(e).__name__})" + ) from e + if metadata.size is not None and metadata.size > _MAX_PROMPT_FILE_BYTES: + raise ValueError( + f"system_prompt.from_hub: {file} is larger than " + f"{_MAX_PROMPT_FILE_BYTES} bytes, too big for a system prompt" + ) + # token=False: the server's own Hub login must not reach a private or gated # file on a workflow's say-so, since the text is echoed into outputs path = hf_hub_download(repo, file, revision=revision, token=False) diff --git a/tests/test_text_generation.py b/tests/test_text_generation.py index b5617d3a..0d873c6a 100644 --- a/tests/test_text_generation.py +++ b/tests/test_text_generation.py @@ -506,7 +506,10 @@ def _spec(**over): class TestSystemPromptFromHub(unittest.TestCase): def setUp(self): - + meta = patch("dw.tasks.text_generation.get_hf_file_metadata") + self.metadata = meta.start() + self.addCleanup(meta.stop) + self.metadata.return_value = MagicMock(size=100) self.dir = tempfile.mkdtemp() self.json_path = os.path.join(self.dir, "system_prompts.json") with open(self.json_path, "w") as f: @@ -573,6 +576,39 @@ def test_a_missing_or_non_string_repo_is_refused(self): resolve_system_prompt(bad) self.assertIn("system_prompt.from_hub.repo", str(caught.exception)) + @patch("dw.tasks.text_generation.hf_hub_download") + def test_the_hub_is_asked_anonymously_before_the_download(self, download): + download.return_value = self.json_path + order = MagicMock() + order.attach_mock(self.metadata, "metadata") + order.attach_mock(download, "download") + resolve_system_prompt(_spec()) + self.assertEqual([c[0] for c in order.mock_calls], ["metadata", "download"]) + self.assertIs(self.metadata.call_args.kwargs["token"], False) + + @patch("dw.tasks.text_generation.hf_hub_download") + def test_a_file_the_hub_will_not_serve_publicly_is_refused(self, download): + from huggingface_hub.errors import HfHubHTTPError + + class _Gated(HfHubHTTPError): + def __init__(self, message): + Exception.__init__(self, message) + + for err in (Exception("boom"), _Gated("403 gated"), OSError("offline")): + self.metadata.side_effect = err + with self.assertRaises(ValueError) as caught: + resolve_system_prompt(_spec()) + self.assertIn("public file", str(caught.exception)) + self.assertIs(caught.exception.__cause__, err) + download.assert_not_called() + + @patch("dw.tasks.text_generation.hf_hub_download") + def test_an_oversize_file_is_refused_before_downloading(self, download): + self.metadata.return_value = MagicMock(size=1024 * 1024 + 1) + with self.assertRaises(ValueError): + resolve_system_prompt(_spec()) + download.assert_not_called() + def test_a_non_string_key_is_refused(self): for key in (["t2i"], {"a": 1}, 3): with self.assertRaises(ValueError): From 68adce1650e46d6102983ecae592340ac0a9fe25 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:15:06 -0500 Subject: [PATCH 09/20] feat(flux3): Klein text-to-image template and model entry Co-Authored-By: Claude Opus 5.5 --- prompts/flux3/man_dog.json | 3 + tests/test_flux3_catalog.py | 74 ++++++++++++++++++++ tests/test_step_cache.py | 3 + workflows/models/flux3-klein.json | 70 ++++++++++++++++++ workflows/templates/flux3/text-to-image.json | 66 +++++++++++++++++ 5 files changed, 216 insertions(+) create mode 100644 prompts/flux3/man_dog.json create mode 100644 workflows/models/flux3-klein.json create mode 100644 workflows/templates/flux3/text-to-image.json diff --git a/prompts/flux3/man_dog.json b/prompts/flux3/man_dog.json new file mode 100644 index 00000000..0583568c --- /dev/null +++ b/prompts/flux3/man_dog.json @@ -0,0 +1,3 @@ +{ + "text": "Full-color outdoor photograph taken at eye level in the afternoon, with low sunlight coming from frame-right and casting long shadows across a neighborhood park. In the center, a white man in his thirties with fair skin, short brown hair, and light stubble walks along a paved path while holding a medium-sized golden retriever on a dark green leash. He wears a faded blue overshirt over a grey T-shirt, tan trousers, and worn brown walking shoes. His body faces left as he takes a step forward, looking down toward the dog with a neutral expression. The leash hangs in a loose curve between his right hand and the dog’s red collar. The golden retriever walks beside him on the left, its cream-gold coat lit along the back and ears, head lowered toward the grass and tail raised behind it. Dry leaves and small twigs lie near the edge of the path in the foreground. On the right, the sun sits just above the line of trees, producing a bright patch of glare and warm highlights on the man’s sleeve and the dog’s fur. Behind them, an open lawn extends toward scattered shrubs and mature tree trunks, with a wooden bench resting near the path. In the background, muted brick houses and parked cars form simple shapes beyond the foliage, while the pale blue sky occupies a narrow band above the treetops." +} diff --git a/tests/test_flux3_catalog.py b/tests/test_flux3_catalog.py index 1d0bc983..2a888b82 100644 --- a/tests/test_flux3_catalog.py +++ b/tests/test_flux3_catalog.py @@ -67,3 +67,77 @@ def test_the_edit_upsampler_wraps_the_instruction_in_bfl_s_scaffold(): assert step["task"]["arguments"]["image"] == "variable:images" assert step["task"]["arguments"]["repetition_penalty"] == 1.0 assert "bbfe317d" in definition["description"] + + +TEMPLATES = os.path.join(ROOT, "workflows", "templates", "flux3") +KLEIN = "black-forest-labs/flux-3-image-klein" + + +def klein_steps(definition): + return [ + s + for s in definition["steps"] + if s.get("pipeline", {}).get("from_pretrained_arguments", {}).get("model_name") + == KLEIN + ] + + +def flux3_definitions(): + paths = [ + os.path.join(TEMPLATES, n) + for n in sorted(os.listdir(TEMPLATES)) + if n.endswith(".json") + ] + paths.append(os.path.join(ROOT, "workflows", "models", "flux3-klein.json")) + return paths + + +@pytest.mark.parametrize("path", flux3_definitions()) +def test_every_flux3_pipeline_step_places_every_component(path): + steps = klein_steps(load(path)) + assert steps, f"{path} runs no Klein step" + for step in steps: + components = step["pipeline"]["configuration"]["components"] + for name in ("transformer", "text_encoder", "vae"): + assert components.get(name, {}).get("device") == "cuda", (path, name) + assert components["text_encoder"].get("residency") is None + + +@pytest.mark.parametrize("path", flux3_definitions()) +def test_every_klein_step_runs_the_distilled_schedule_without_guidance(path): + for step in klein_steps(load(path)): + args = step["pipeline"]["arguments"] + assert args["num_inference_steps"] == 16 + assert "guidance_scale" not in args and "negative_prompt" not in args + assert args["output"] == "images" + + +@pytest.mark.parametrize("path", flux3_definitions()) +def test_size_variables_are_held_to_the_16_pixel_grid(path): + definition = load(path) + constraints = definition.get("variable_constraints", {}) + for name in ("width", "height"): + if name in definition.get("variables", {}): + assert ( + constraints[name]["modulus"] == 16 + and constraints[name].get("remainder", 0) == 0 + ) + + +def test_the_16_step_default_is_the_distilled_pipeline_s(): + from diffusers.modular_pipelines.flux3.modular_blocks_flux3_distilled import ( + Flux3DistilledCoreDenoiseStep, + ) + + defaults = {p.name: p.default for p in Flux3DistilledCoreDenoiseStep().inputs} + assert defaults["num_inference_steps"] == 16 + + +def test_the_16_pixel_grid_is_the_vae_scale_times_the_2x2_packing(): + import inspect + + from diffusers.modular_pipelines.flux3 import before_denoise + + # Source check: the multiple is derived from the loaded VAE, so it cannot + # be read without weights; 8 (VAE) * 2 (packing) = 16 is the vendor's. + assert "vae_scale_factor * 2" in inspect.getsource(before_denoise) diff --git a/tests/test_step_cache.py b/tests/test_step_cache.py index e0ef2220..4823fa87 100644 --- a/tests/test_step_cache.py +++ b/tests/test_step_cache.py @@ -683,6 +683,9 @@ def test_normalized_downstream_false_when_nothing_references_it(): "ip-adapter.json": { "main": "a67c769c0ee881e36aa15b62f8226a233c1e4b93542a8caaf5eab5715cb9f556" }, + "flux3/text-to-image.json": { + "text_to_image": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" + }, "kandinsky6/generate-and-upscale.json": { "text_to_video": "3095873d2e63833c2e8c8048379035840694ca108309e544ff17a0722b4ee51a", "upscale": "838beb7eefde695da365c37927645fc388ab8feeaf8e1bc1695440cbd2b54fdd", diff --git a/workflows/models/flux3-klein.json b/workflows/models/flux3-klein.json new file mode 100644 index 00000000..1bf5457c --- /dev/null +++ b/workflows/models/flux3-klein.json @@ -0,0 +1,70 @@ +{ + "id": "flux3-klein", + "description": "FLUX 3 Image Klein (distilled, 5.6B) as a model entry, configured by the text-to-image template. The prompt here is a short one, and Klein turns a short prompt into a generic picture: for full captions run templates/flux3/upsample-then-generate. Sixteen steps with no guidance or negative prompt. The weights are gated under the FLUX Kommunity License: accept it on the model page and log in with hf auth login first.", + "configures": "templates/text-to-image", + "cost_drivers": [ + "num_images_per_prompt", + "width", + "height" + ], + "variables": { + "prompt": "an apple", + "num_images_per_prompt": 1, + "width": 1024, + "height": 1024, + "seed": null + }, + "seed": "variable:seed", + "steps": [ + { + "name": "text_to_image", + "pipeline": { + "configuration": { + "component_type": "ModularPipeline", + "load_components": { + "dtype": "torch.bfloat16" + }, + "components": { + "transformer": { + "device": "cuda" + }, + "text_encoder": { + "device": "cuda" + }, + "vae": { + "device": "cuda" + } + } + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/flux-3-image-klein" + }, + "arguments": { + "prompt": "variable:prompt", + "width": "variable:width", + "height": "variable:height", + "num_inference_steps": 16, + "output": "images", + "num_images_per_prompt": "variable:num_images_per_prompt" + } + }, + "result": { + "content_type": "image/jpeg", + "embed_metadata": true, + "subfolder": "final" + } + } + ], + "variable_constraints": { + "width": { + "modulus": 16, + "remainder": 0, + "reason": "FLUX 3 Image Klein's pipeline refuses a width not divisible by 16 (the VAE's 8x compression times the 2x2 latent packing), and only after the weights have loaded, so it is checked here first" + }, + "height": { + "modulus": 16, + "remainder": 0, + "reason": "FLUX 3 Image Klein's pipeline refuses a height not divisible by 16 (the VAE's 8x compression times the 2x2 latent packing), and only after the weights have loaded, so it is checked here first" + } + } +} diff --git a/workflows/templates/flux3/text-to-image.json b/workflows/templates/flux3/text-to-image.json new file mode 100644 index 00000000..fb8fc6fd --- /dev/null +++ b/workflows/templates/flux3/text-to-image.json @@ -0,0 +1,66 @@ +{ + "id": "Flux3KleinT2I", + "description": "FLUX 3 Image Klein (distilled, 5.6B) text to image, the baseline of the FLUX 3 templates. The prompt here is already a full caption - BFL's own upsampled caption for their man walking a dog example - because Klein was trained on long upsampler captions and turns short prompts into generic pictures; upsample-then-generate shows how to get one from a short request. Sixteen steps with no guidance: the distilled model is a flow map, so it takes no guidance scale and no negative prompt. Width and height must be multiples of 16, up to 4096. Every component is placed on the GPU explicitly, about 19 GB of bf16 weights, which fits a 24 GB card at 1024x1024 with under 2 GB to spare. The weights are gated under the FLUX Kommunity License: accept it on the model page and log in with hf auth login first.", + "cost_drivers": [ + "width", + "height" + ], + "variables": { + "prompt": "prompt:flux3/man_dog", + "width": 1024, + "height": 1024, + "seed": 7 + }, + "variable_constraints": { + "width": { + "modulus": 16, + "remainder": 0, + "reason": "FLUX 3 Image Klein's pipeline refuses a width not divisible by 16 (the VAE's 8x compression times the 2x2 latent packing), and only after the weights have loaded, so it is checked here first" + }, + "height": { + "modulus": 16, + "remainder": 0, + "reason": "FLUX 3 Image Klein's pipeline refuses a height not divisible by 16 (the VAE's 8x compression times the 2x2 latent packing), and only after the weights have loaded, so it is checked here first" + } + }, + "seed": "variable:seed", + "steps": [ + { + "name": "text_to_image", + "pipeline": { + "configuration": { + "component_type": "ModularPipeline", + "load_components": { + "dtype": "torch.bfloat16" + }, + "components": { + "transformer": { + "device": "cuda" + }, + "text_encoder": { + "device": "cuda" + }, + "vae": { + "device": "cuda" + } + } + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/flux-3-image-klein" + }, + "arguments": { + "prompt": "variable:prompt", + "width": "variable:width", + "height": "variable:height", + "num_inference_steps": 16, + "output": "images" + } + }, + "result": { + "content_type": "image/jpeg", + "embed_metadata": true, + "subfolder": "final" + } + } + ] +} From c4d0fadc116bbbba1ac262094c4e0789904dd04a Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:18:25 -0500 Subject: [PATCH 10/20] feat(flux3): upsample-then-generate template Co-Authored-By: Claude Opus 5.5 --- tests/test_flux3_catalog.py | 8 ++ tests/test_step_cache.py | 3 + .../flux3/upsample-then-generate.json | 81 +++++++++++++++++++ 3 files changed, 92 insertions(+) create mode 100644 workflows/templates/flux3/upsample-then-generate.json diff --git a/tests/test_flux3_catalog.py b/tests/test_flux3_catalog.py index 2a888b82..934e3f6e 100644 --- a/tests/test_flux3_catalog.py +++ b/tests/test_flux3_catalog.py @@ -141,3 +141,11 @@ def test_the_16_pixel_grid_is_the_vae_scale_times_the_2x2_packing(): # Source check: the multiple is derived from the loaded VAE, so it cannot # be read without weights; 8 (VAE) * 2 (packing) = 16 is the vendor's. assert "vae_scale_factor * 2" in inspect.getsource(before_denoise) + + +def test_upsample_then_generate_feeds_the_builtin_caption_to_klein(): + definition = load(os.path.join(TEMPLATES, "upsample-then-generate.json")) + first, second = definition["steps"] + assert first["workflow"]["path"] == "builtin:flux3_klein_upsample.json" + assert first["release_models"] is True + assert second["pipeline"]["arguments"]["prompt"] == f"previous_result:{first['name']}" diff --git a/tests/test_step_cache.py b/tests/test_step_cache.py index 4823fa87..776b2e4b 100644 --- a/tests/test_step_cache.py +++ b/tests/test_step_cache.py @@ -686,6 +686,9 @@ def test_normalized_downstream_false_when_nothing_references_it(): "flux3/text-to-image.json": { "text_to_image": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" }, + "flux3/upsample-then-generate.json": { + "text_to_image": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" + }, "kandinsky6/generate-and-upscale.json": { "text_to_video": "3095873d2e63833c2e8c8048379035840694ca108309e544ff17a0722b4ee51a", "upscale": "838beb7eefde695da365c37927645fc388ab8feeaf8e1bc1695440cbd2b54fdd", diff --git a/workflows/templates/flux3/upsample-then-generate.json b/workflows/templates/flux3/upsample-then-generate.json new file mode 100644 index 00000000..9bebf350 --- /dev/null +++ b/workflows/templates/flux3/upsample-then-generate.json @@ -0,0 +1,81 @@ +{ + "id": "Flux3KleinUpsampleThenGenerate", + "description": "FLUX 3 Image Klein from a short request, the way Black Forest Labs say to run it: their prompt upsampler first expands the request into the long caption Klein was trained on, then Klein draws from that caption. The upsampler is about 18 GB and Klein about 19 GB in bf16, which do not fit a 24 GB card together, so the upsample step sets release_models and the upsampler is freed before Klein loads. The caption is saved as an intermediate so it can be read, edited or reused as the prompt of text-to-image without paying for the upsampler again. Sixteen steps with no guidance or negative prompt, as the distilled model takes neither. Width and height must be multiples of 16, up to 4096. The weights are gated under the FLUX Kommunity License: accept it on the model page and log in with hf auth login first.", + "summary": "FLUX 3 Klein image from a short request, expanded by BFL's upsampler first.", + "cost_drivers": [ + "width", + "height" + ], + "variables": { + "request": "A man walking with his dog, afternoon sun", + "width": 1024, + "height": 1024, + "seed": 7 + }, + "variable_constraints": { + "width": { + "modulus": 16, + "remainder": 0, + "reason": "FLUX 3 Image Klein's pipeline refuses a width not divisible by 16 (the VAE's 8x compression times the 2x2 latent packing), and only after the weights have loaded, so it is checked here first" + }, + "height": { + "modulus": 16, + "remainder": 0, + "reason": "FLUX 3 Image Klein's pipeline refuses a height not divisible by 16 (the VAE's 8x compression times the 2x2 latent packing), and only after the weights have loaded, so it is checked here first" + } + }, + "seed": "variable:seed", + "steps": [ + { + "name": "upsample", + "release_models": true, + "workflow": { + "path": "builtin:flux3_klein_upsample.json", + "arguments": { + "prompt": "variable:request" + } + }, + "result": { + "content_type": "text/plain", + "subfolder": "intermediate" + } + }, + { + "name": "text_to_image", + "pipeline": { + "configuration": { + "component_type": "ModularPipeline", + "load_components": { + "dtype": "torch.bfloat16" + }, + "components": { + "transformer": { + "device": "cuda" + }, + "text_encoder": { + "device": "cuda" + }, + "vae": { + "device": "cuda" + } + } + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/flux-3-image-klein" + }, + "arguments": { + "prompt": "previous_result:upsample", + "width": "variable:width", + "height": "variable:height", + "num_inference_steps": 16, + "output": "images" + } + }, + "result": { + "content_type": "image/jpeg", + "embed_metadata": true, + "subfolder": "final" + } + } + ] +} From d377e721a7c9719e35e2614435e0c527ddea65f9 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:24:52 -0500 Subject: [PATCH 11/20] feat(flux3): image-edit and multi-reference templates Co-Authored-By: Claude Opus 5.5 --- .../2026-10-15-flux-3-klein-catalog-design.md | 2 +- tests/test_flux3_catalog.py | 16 ++++ tests/test_step_cache.py | 6 ++ workflows/templates/flux3/image-edit.json | 73 ++++++++++++++++++ .../templates/flux3/multi-reference.json | 74 +++++++++++++++++++ 5 files changed, 170 insertions(+), 1 deletion(-) create mode 100644 workflows/templates/flux3/image-edit.json create mode 100644 workflows/templates/flux3/multi-reference.json diff --git a/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md b/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md index e882ade5..a598e06a 100644 --- a/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md +++ b/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md @@ -38,7 +38,7 @@ Status: draft, probes on lem in progress. | P1d distilled t2i bf16, 1024², 16 steps, vendor `man_dog` caption, seed 7, transformer + text encoder + VAE `device: cuda` | 9cc1ce441555 | **Passed.** 24 s total: load 10.1 s, encode 0.4 s, denoise 12.1 s (0.8 s/step), decode 0.4 s. 19.2 GB allocated, 22.1 GB reserved peak at decode (1.7 GB free). The image follows the caption element by element | | P2 pup via `text_generation` as-is, `t2i` system prompt, `man_dog` request | a5471561de18 | Ran (22 s, cuda:1); vendor-style caption, but it opens mid-phrase in lowercase ("medium-wide outdoor photograph ...") | | P2b chat template check | (script) | The template's default opens `\n` (thinking on); `enable_thinking=False` gives `\n\n\n\n`. P2 ran in thinking mode | -| P3 single edit / multi-reference | | | +| P3 multi-reference edit: `image` a list of two `{"location": url}` refs (BFL examples/inputs 2.png mug, 3.png dog, commit bbfe317d), BFL upsampled mugprint caption, seed 23, 1024², 16 steps | cc7fb1b29091 | **Passed.** The list arrives as PIL images; the output printed the dog on the mug and kept the scene. 45 s total: encode 0.5 s, denoise 41.0 s (2.6 s/step, references add tokens), decode 0.3 s. 19.2 GB allocated, 22.1 GB reserved peak | What P1 settled: diff --git a/tests/test_flux3_catalog.py b/tests/test_flux3_catalog.py index 934e3f6e..41655b8c 100644 --- a/tests/test_flux3_catalog.py +++ b/tests/test_flux3_catalog.py @@ -149,3 +149,19 @@ def test_upsample_then_generate_feeds_the_builtin_caption_to_klein(): assert first["workflow"]["path"] == "builtin:flux3_klein_upsample.json" assert first["release_models"] is True assert second["pipeline"]["arguments"]["prompt"] == f"previous_result:{first['name']}" + + +@pytest.mark.parametrize("name", ["image-edit.json", "multi-reference.json"]) +def test_the_edit_templates_show_the_upsampler_the_pipeline_s_references(name): + definition = load(os.path.join(TEMPLATES, name)) + upsample = next(s for s in definition["steps"] if "workflow" in s) + assert upsample["workflow"]["path"] == "builtin:flux3_klein_upsample_edit.json" + (klein,) = klein_steps(definition) + assert upsample["workflow"]["arguments"]["images"] == klein["pipeline"]["arguments"]["image"] + assert klein["pipeline"]["arguments"]["prompt"] == f"previous_result:{upsample['name']}" + + +def test_multi_reference_states_the_ten_reference_cap_and_the_naming(): + description = load(os.path.join(TEMPLATES, "multi-reference.json"))["description"] + assert "10 references" in description + assert "" in description diff --git a/tests/test_step_cache.py b/tests/test_step_cache.py index 776b2e4b..b2241e39 100644 --- a/tests/test_step_cache.py +++ b/tests/test_step_cache.py @@ -683,6 +683,12 @@ def test_normalized_downstream_false_when_nothing_references_it(): "ip-adapter.json": { "main": "a67c769c0ee881e36aa15b62f8226a233c1e4b93542a8caaf5eab5715cb9f556" }, + "flux3/image-edit.json": { + "edit": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" + }, + "flux3/multi-reference.json": { + "edit": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" + }, "flux3/text-to-image.json": { "text_to_image": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" }, diff --git a/workflows/templates/flux3/image-edit.json b/workflows/templates/flux3/image-edit.json new file mode 100644 index 00000000..c94480dd --- /dev/null +++ b/workflows/templates/flux3/image-edit.json @@ -0,0 +1,73 @@ +{ + "id": "Flux3KleinImageEdit", + "description": "FLUX 3 Image Klein image edit: BFL's prompt upsampler reads the reference and the instruction and writes the caption Klein was trained on, then Klein edits the reference by that caption. The upsampler sees the same reference as the pipeline, and is freed (release_models) before Klein loads. The caption is saved as an intermediate. Width and height are left unset, so the output defaults to 1 MP at the reference's aspect ratio, rounded to the 16-pixel grid. Family rules (16 steps, no guidance, component placement, the gated license) are as in text-to-image.json.", + "summary": "FLUX 3 Klein edit of one picture, the instruction expanded by BFL's upsampler first.", + "cost_drivers": [ + "image" + ], + "variables": { + "instruction": "put a red bandana around the dog's neck", + "image": { + "location": "https://raw.githubusercontent.com/black-forest-labs/flux-3-image-klein/bbfe317dd81c825c30d76090a67d3d0db51da3cc/examples/inputs/3.png" + }, + "seed": 22 + }, + "seed": "variable:seed", + "steps": [ + { + "name": "upsample", + "release_models": true, + "workflow": { + "path": "builtin:flux3_klein_upsample_edit.json", + "arguments": { + "instruction": "variable:instruction", + "images": [ + "variable:image" + ] + } + }, + "result": { + "content_type": "text/plain", + "subfolder": "intermediate" + } + }, + { + "name": "edit", + "pipeline": { + "configuration": { + "component_type": "ModularPipeline", + "load_components": { + "dtype": "torch.bfloat16" + }, + "components": { + "transformer": { + "device": "cuda" + }, + "text_encoder": { + "device": "cuda" + }, + "vae": { + "device": "cuda" + } + } + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/flux-3-image-klein" + }, + "arguments": { + "prompt": "previous_result:upsample", + "image": [ + "variable:image" + ], + "num_inference_steps": 16, + "output": "images" + } + }, + "result": { + "content_type": "image/jpeg", + "embed_metadata": true, + "subfolder": "final" + } + } + ] +} diff --git a/workflows/templates/flux3/multi-reference.json b/workflows/templates/flux3/multi-reference.json new file mode 100644 index 00000000..76e86780 --- /dev/null +++ b/workflows/templates/flux3/multi-reference.json @@ -0,0 +1,74 @@ +{ + "id": "Flux3KleinMultiReference", + "description": "FLUX 3 Image Klein with several references: BFL's prompt upsampler reads every reference in order and the instruction, writes the caption, then Klein draws from the caption and the references. The upsampler sees the same references as the pipeline. In a caption written by hand, name them , and so on (zero-based); the upsampler's own wording is the first image, the second image. The first reference sets the output's aspect ratio, at 1 MP by default. BFL trained Klein on up to 10 references; the pipeline does not enforce that cap, so more than 10 references run but are outside what the model learned. Family rules (16 steps, no guidance, component placement, the gated license) are as in text-to-image.json.", + "summary": "FLUX 3 Klein edit from several reference pictures, expanded by BFL's upsampler first.", + "cost_drivers": [ + "images" + ], + "variables": { + "instruction": "print the dog from the second image onto the mug", + "images": [ + { + "location": "https://raw.githubusercontent.com/black-forest-labs/flux-3-image-klein/bbfe317dd81c825c30d76090a67d3d0db51da3cc/examples/inputs/2.png" + }, + { + "location": "https://raw.githubusercontent.com/black-forest-labs/flux-3-image-klein/bbfe317dd81c825c30d76090a67d3d0db51da3cc/examples/inputs/3.png" + } + ], + "seed": 23 + }, + "seed": "variable:seed", + "steps": [ + { + "name": "upsample", + "release_models": true, + "workflow": { + "path": "builtin:flux3_klein_upsample_edit.json", + "arguments": { + "instruction": "variable:instruction", + "images": "variable:images" + } + }, + "result": { + "content_type": "text/plain", + "subfolder": "intermediate" + } + }, + { + "name": "edit", + "pipeline": { + "configuration": { + "component_type": "ModularPipeline", + "load_components": { + "dtype": "torch.bfloat16" + }, + "components": { + "transformer": { + "device": "cuda" + }, + "text_encoder": { + "device": "cuda" + }, + "vae": { + "device": "cuda" + } + } + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/flux-3-image-klein" + }, + "arguments": { + "prompt": "previous_result:upsample", + "image": "variable:images", + "num_inference_steps": 16, + "output": "images" + } + }, + "result": { + "content_type": "image/jpeg", + "embed_metadata": true, + "subfolder": "final" + } + } + ] +} From 8a62a7e3f295ace911a3a63ec7af574718e47387 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:29:34 -0500 Subject: [PATCH 12/20] feat(flux3): layout-with-boxes template Co-Authored-By: Claude Opus 5.5 --- prompts/flux3/balloons_layout.json | 3 + tests/test_flux3_catalog.py | 12 ++++ tests/test_step_cache.py | 3 + workflows/templates/flux3/layout-boxes.json | 67 +++++++++++++++++++++ 4 files changed, 85 insertions(+) create mode 100644 prompts/flux3/balloons_layout.json create mode 100644 workflows/templates/flux3/layout-boxes.json diff --git a/prompts/flux3/balloons_layout.json b/prompts/flux3/balloons_layout.json new file mode 100644 index 00000000..637aef44 --- /dev/null +++ b/prompts/flux3/balloons_layout.json @@ -0,0 +1,3 @@ +{ + "text": "Full-color outdoor photograph at sunrise, with a broad dawn sky above a green valley . On the left, a large hot-air balloon floats over the valley, while a smaller blue hot-air balloon hangs farther back on the right. Along the horizon, the rising sun appears between distant mountain ridges , casting low orange light across the slopes and leaving the foreground meadow in cooler shadow. [{\"id\": \"obj_1\", \"bbox\": [80, 60, 480, 380], \"desc\": \"A large red and yellow striped hot-air balloon with alternating vertical fabric panels, a rounded envelope, dark suspension cables, and a small rectangular wicker basket hanging below; the lower edge of the envelope is partly shaded by the sunrise angle.\"}, {\"id\": \"obj_2\", \"bbox\": [200, 620, 420, 800], \"desc\": \"A small blue hot-air balloon with a rounded cobalt-blue envelope, narrow darker seams, fine suspension lines, and a tiny brown basket.\"}, {\"id\": \"sky_1\", \"bbox\": [0, 0, 650, 1000], \"desc\": \"A clear sunrise sky grading from pale orange near the horizon through peach to muted blue overhead, with several thin horizontal clouds catching warm light.\"}, {\"id\": \"sun_1\", \"bbox\": [470, 455, 535, 525], \"desc\": \"A low circular sunrise disk in bright yellow-orange, partially veiled by the distant ridge line.\"}, {\"id\": \"mountain_1\", \"bbox\": [430, 0, 690, 1000], \"desc\": \"Layered distant mountain ridges forming an uneven horizon, with blue-gray silhouettes in front and warmer brown-green slopes along their upper edges.\"}, {\"id\": \"valley_1\", \"bbox\": [560, 0, 1000, 1000], \"desc\": \"A broad green valley with rolling grassy slopes, scattered darker tree masses, and a winding river reflecting a narrow strip of orange morning light; the nearest meadow contains patches of deep green shadow and uneven vegetation.\"}]" +} diff --git a/tests/test_flux3_catalog.py b/tests/test_flux3_catalog.py index 41655b8c..a7b6b3b4 100644 --- a/tests/test_flux3_catalog.py +++ b/tests/test_flux3_catalog.py @@ -165,3 +165,15 @@ def test_multi_reference_states_the_ten_reference_cap_and_the_naming(): description = load(os.path.join(TEMPLATES, "multi-reference.json"))["description"] assert "10 references" in description assert "" in description + + +def test_the_layout_caption_carries_tags_and_a_0_to_1000_box_table(): + with open(os.path.join(ROOT, "prompts", "flux3", "balloons_layout.json")) as f: + text = json.load(f)["text"] + table = json.loads(text[text.index("["):]) + assert table and all({"id", "bbox", "desc"} <= set(row) for row in table) + for row in table: + assert f"<{row['id']}>" in text + top, left, bottom, right = row["bbox"] + assert 0 <= top < bottom <= 1000 and 0 <= left < right <= 1000 + assert klein_steps(load(os.path.join(TEMPLATES, "layout-boxes.json"))) diff --git a/tests/test_step_cache.py b/tests/test_step_cache.py index b2241e39..d01b1053 100644 --- a/tests/test_step_cache.py +++ b/tests/test_step_cache.py @@ -686,6 +686,9 @@ def test_normalized_downstream_false_when_nothing_references_it(): "flux3/image-edit.json": { "edit": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" }, + "flux3/layout-boxes.json": { + "layout_to_image": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" + }, "flux3/multi-reference.json": { "edit": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" }, diff --git a/workflows/templates/flux3/layout-boxes.json b/workflows/templates/flux3/layout-boxes.json new file mode 100644 index 00000000..3a0ec424 --- /dev/null +++ b/workflows/templates/flux3/layout-boxes.json @@ -0,0 +1,67 @@ +{ + "id": "Flux3KleinLayoutBoxes", + "summary": "FLUX 3 Klein image placed by a caption with layout boxes (BFL's balloons example).", + "description": "FLUX 3 Image Klein text to image from a caption with layout boxes: the prose tags each element as , and a JSON table gives every id a bbox as [top, left, bottom, right] on a 0 to 1000 grid plus a description. The caption here is BFL's own upsampled balloons example, used as written; upsampling boxes from a short request is not covered yet. Boxes are hints, not masks: a small boxed element may not appear, and elements can run past their box. The family rules for steps, guidance, size, placement and the license are in text-to-image.", + "cost_drivers": [ + "width", + "height" + ], + "variables": { + "prompt": "prompt:flux3/balloons_layout", + "width": 1024, + "height": 1024, + "seed": 1 + }, + "variable_constraints": { + "width": { + "modulus": 16, + "remainder": 0, + "reason": "FLUX 3 Image Klein's pipeline refuses a width not divisible by 16 (the VAE's 8x compression times the 2x2 latent packing), and only after the weights have loaded, so it is checked here first" + }, + "height": { + "modulus": 16, + "remainder": 0, + "reason": "FLUX 3 Image Klein's pipeline refuses a height not divisible by 16 (the VAE's 8x compression times the 2x2 latent packing), and only after the weights have loaded, so it is checked here first" + } + }, + "seed": "variable:seed", + "steps": [ + { + "name": "layout_to_image", + "pipeline": { + "configuration": { + "component_type": "ModularPipeline", + "load_components": { + "dtype": "torch.bfloat16" + }, + "components": { + "transformer": { + "device": "cuda" + }, + "text_encoder": { + "device": "cuda" + }, + "vae": { + "device": "cuda" + } + } + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/flux-3-image-klein" + }, + "arguments": { + "prompt": "variable:prompt", + "width": "variable:width", + "height": "variable:height", + "num_inference_steps": 16, + "output": "images" + } + }, + "result": { + "content_type": "image/jpeg", + "embed_metadata": true, + "subfolder": "final" + } + } + ] +} From 803203cc5234f6085ea0a4800df5779dde6032a9 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:31:26 -0500 Subject: [PATCH 13/20] docs(flux3): Klein templates README Co-Authored-By: Claude Opus 5.5 --- workflows/templates/flux3/README.md | 74 +++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 workflows/templates/flux3/README.md diff --git a/workflows/templates/flux3/README.md b/workflows/templates/flux3/README.md new file mode 100644 index 00000000..770df26d --- /dev/null +++ b/workflows/templates/flux3/README.md @@ -0,0 +1,74 @@ +# FLUX 3 Image Klein workflows + +Text-to-image, image editing and multi-reference generation with Black Forest Labs' +[FLUX 3 Image Klein](https://huggingface.co/black-forest-labs/flux-3-image-klein): a +5.6B distilled flow map that draws in 16 steps with no guidance scale and no negative +prompt. It pairs the FLUX.2 VAE with a Qwen3-VL-4B text encoder. The weights are gated +under the FLUX Kommunity License: accept it on the model page and run `hf auth login` +before the first download. + +Settings come from the Hugging Face card, BFL's +[flux-3-image-klein](https://github.com/black-forest-labs/flux-3-image-klein) +repository, [docs.bfl.ml](https://docs.bfl.ml) and the diffusers pipeline source +(`diffusers.modular_pipelines.flux3`, merged on 2026-10-15 and not yet in a release). +These templates need diffusers installed from git. Audited against those sources on +2026-10-15 ([the audit](../../../docs/proposals/audits/2026-10-15-flux-3-klein-audit.md)). + +Klein was trained on long captions written by BFL's prompt upsampler, and BFL says not +to use it without one: a short prompt turns into a generic picture. The two upsampler +builtins (`flux3_klein_upsample` and `flux3_klein_upsample_edit`) run BFL's +`flux-3-image-klein-pup` for you, and the system prompts they give it are BFL's own, +read from the pup repository at a pinned commit (see below). To write a caption by hand, +read BFL's [prompting guide](https://docs.bfl.ml/guides/prompting_summary). It is +written for the API model, which upsamples on the server side, so treat it as a +description of the target caption and not as advice for a short prompt. + +An agent driving this family from Claude Code has the `flux3-klein` skill of the +[dw plugin](../../../plugins/dw/README.md), which chooses between these templates and +states the family's rules. + +Read them in this order and each introduces one new idea on top of the last. + +## The templates + +| Example | What it introduces | +| ------- | ------------------ | +| [text-to-image.json](text-to-image.json) | The baseline. The prompt is already a full upsampled caption (BFL's man-walking-a-dog example), 16 steps, no guidance. Width and height are multiples of 16, up to 4096, and the pipeline only complains after the load, so the workflow declares the rule up front. 24 s on an RTX 3090 at 1024x1024 (0.8 s/step), peaking at 22.1GB reserved | +| [upsample-then-generate.json](upsample-then-generate.json) | A short request becomes a caption first: the pup upsampler writes it, `release_models` frees the upsampler (about 18GB), and Klein draws from the result. The caption is saved as an intermediate, so it can be read, edited or reused as the prompt of text-to-image | +| [image-edit.json](image-edit.json) | One reference and an instruction. The upsampler sees the same image as the pipeline. With no width or height the output is 1 MP at the reference's aspect ratio, on the 16-pixel grid | +| [multi-reference.json](multi-reference.json) | Several references, in order. A caption written by hand names them ``, `` and so on (zero-based). The first reference sets the aspect ratio. BFL trained on up to 10; the pipeline does not enforce that, so more than 10 run but are outside what the model learned. Two references ran in 45 s on an RTX 3090 (2.6 s/step, references add tokens), also peaking at 22.1GB reserved | +| [layout-boxes.json](layout-boxes.json) | Placement by box. Elements are tagged `` in the prose, and a JSON table gives each id a `[top, left, bottom, right]` box on a 0 to 1000 grid. Boxes are hints, not masks: a small boxed element may not appear, and elements can run past their box | + +## dw craft, not vendor guidance + +Three things in these templates are this repository's choices. BFL does not say them. + +**Every component is placed explicitly.** Each of the transformer, text encoder and VAE +carries a `device`. The modular pipeline's execution device is its first model +component's device, so a component left unplaced stays on the CPU, and the pipeline then +encodes the prompt and draws its noise there. The failures looked like `index is on cpu` +in the text encoder, and `Cannot generate a cpu tensor from a generator of type cuda`. +`residency: on_demand` on the text encoder breaks for the same reason, because the +encoder step moves `input_ids` to the execution device before the wrapper moves the +weights. The weights are about 19GB in bf16, so with everything resident a 24GB card +fits 1024x1024 with under 2GB to spare. A larger output, or several references, wants +the VAE on demand or a quantized transformer, and neither has been measured yet. + +**The upsampler is released before Klein loads.** The pup (Qwen3.5-9B, Apache-2.0, +ungated) is about 18GB and Klein about 19GB, which do not fit a 24GB card together. The +upsample step sets `release_models`, and the generation step loads after it. + +**The upsampler's system prompts are read from the Hub at a pinned commit.** BFL's +prompts are in the pup repository's `system_prompts.json` and are still changing, so the +builtins name the file, the key and a full 40-character revision through the +`system_prompt.from_hub` form of `text_generation`, and never carry a copy. A template +runs the same text until someone re-pins it. The builtins also turn the chat template's +thinking mode off, which BFL's usage assumes. + +## Not here yet + +These follow-ups are dated 2026-10-15 and are in the audit's reading order: +- The base (non-distilled) checkpoint with classifier-free guidance, which needs the modular `guider` to be reachable from a workflow. +- Upsampled layout boxes from a short request (BFL's `t2i_bbox`), which need BFL's box planning and reconciling code. +- LoRA, a typography template, the fp8 and NVFP4 checkpoints (fused only on Ada and Blackwell cards), and a measurement on Apple Silicon. +- Re-pinning the system prompts when diffusers PR #15058, which carries them, merges. From c3690b437ae92535d559767c4daa61b040a0d9ee Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:36:53 -0500 Subject: [PATCH 14/20] feat(plugin): flux3-klein composition skill The skill picks among the five FLUX 3 Klein templates by request shape, states the 16-step no-guidance schedule, the 16-pixel grid, the 4096 ceiling as stated rather than checked, and the 10-reference training limit the pipeline does not enforce. Prompts go through the upsampler builtins, or BFL's prompting guide for hand-written captions; the box format is in references/layout-boxes.md. Named in the plugin README and CLAUDE.md; the tests read the step default off the distilled denoise block. Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 2 +- plugins/dw/README.md | 1 + plugins/dw/skills/flux3-klein/SKILL.md | 143 ++++++++++++++++++ .../flux3-klein/references/layout-boxes.md | 31 ++++ tests/test_plugin_skills.py | 52 +++++++ 5 files changed, 228 insertions(+), 1 deletion(-) create mode 100644 plugins/dw/skills/flux3-klein/SKILL.md create mode 100644 plugins/dw/skills/flux3-klein/references/layout-boxes.md diff --git a/CLAUDE.md b/CLAUDE.md index d56341cd..b25fa046 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -72,7 +72,7 @@ The guides in `docs/`, by topic: `.claude-plugin/marketplace.json` publishes the `dw` plugin in `plugins/dw/`: one composition skill per model family (`minimax-h3`, `minimax-music3`, `ltx-2.5`, -`kandinsky-6`) that chooses a template for a request's shape and states the family's +`kandinsky-6`, `flux3-klein`) that chooses a template for a request's shape and states the family's hard rules, plus cross-cutting composition skills (`script-to-video`, `series-episodes`) - shapes above the families that orchestrate the decision trees and cast consistency across multiple diff --git a/plugins/dw/README.md b/plugins/dw/README.md index a204d6bb..8b0b4685 100644 --- a/plugins/dw/README.md +++ b/plugins/dw/README.md @@ -19,6 +19,7 @@ Install once, from Claude Code: | `minimax-music3` | MiniMax Music 3: a song, an instrumental, a score under a film, the soundtrack a music video is cut to; the ceiling semantics, the tag vocabulary, the caps and the 44.1 kHz output. Captions come from MiniMax's own `music-caption-rewriter` skill. | | `ltx-2.5` | LTX-2.5 video with a soundtrack: a single clip, first-frame and keyframe conditioning, the three-move two-stage flow, the IC-LoRA upscale, extend and chain; the distilled schedule and the frame and size rules. Prompts follow the trained-caption spec that ships inside diffusers. | | `kandinsky-6` | Kandinsky 6.0 video with a soundtrack: a five-second clip from text or from a picture on the 3B Lite-distill checkpoint; the distilled schedule, the size and frame rules the pipeline's own defaults get wrong, and why Apple Silicon is parked. Prompts follow Kandinsky Lab's own beautifier system prompts. | +| `flux3-klein` | FLUX 3 Image Klein images: a picture from a short request through BFL's prompt upsampler, from a full caption, an edit of one photo, several references combined, or a layout placed by boxes; the 16-step distilled schedule with no guidance, the 16-pixel grid, the 10-reference training limit the pipeline does not enforce, and the gated license. Prompts come from BFL's upsampler builtins or their prompting guide. | | `script-to-video` | From a prose script to finished video: decompose into shots, cast recurring characters once, pick each shot's template and family, validate cost before spending, run and diagnose, hand off to cut-and-score. The decision tree that turns a script into the shots the other skills know what to do with. | | `series-episodes` | A series rather than one generation: several episodes over one cast, each cut and scored from its own template runs. The five-beat recut-bed-match_levels-normalize-pair procedure, an optional per-shot look (one shared palette through `apply_lut`, then `film_grain`) before the recut, and the one step that must not be skipped or the cast drifts between episodes. | diff --git a/plugins/dw/skills/flux3-klein/SKILL.md b/plugins/dw/skills/flux3-klein/SKILL.md new file mode 100644 index 00000000..c95af9a4 --- /dev/null +++ b/plugins/dw/skills/flux3-klein/SKILL.md @@ -0,0 +1,143 @@ +--- +name: flux3-klein +description: Use when a dw MCP server is connected and the user wants an image from FLUX 3 Image Klein - a picture from a short request or a full caption, an edit of their photo, several reference pictures combined into one, or a layout with elements placed by box. Picks the template, states the distilled schedule, the size grid and the reference rules the pipeline does not check for you, quotes cost, and points at Black Forest Labs' upsampler and prompting guide. +--- + +# FLUX 3 Image Klein on a dw server + +FLUX 3 Image Klein (Black Forest Labs) is a 5.6B distilled flow map that +draws text-to-image, edits and multi-reference images in 16 steps. It pairs +the FLUX.2 VAE with a Qwen3-VL-4B text encoder, about 19 GB of bf16 weights. +Klein was trained on long captions written by BFL's prompt upsampler, and BFL +says not to run it without one: a short prompt gives a generic picture. + +## Before anything + +1. `get_server_info`: the device. The templates place every component on + `cuda`. Apple Silicon is unmeasured (about 20 GB of unified memory plus + activations, by the vendor's sizes); on `mps` say so before spending, and + on `cpu` stop. +2. `get_diffusers_state`, then `list_classes` for + `Flux3DistilledModularPipeline`. The modular pipeline is in diffusers + `main` only (merged 2026-10-15). If it is missing, the server needs + `update_diffusers` and a restart. That is the user's call; ask first. +3. The weights are gated under the FLUX Kommunity License (non-commercial, + plus commercial use of outputs below US$5M annual revenue). The user + accepts it on the model page and the server runs `hf auth login`; a 401 or + `GatedRepoError` in `get_job` means one of those is missing. The + upsampler is Apache-2.0 and ungated. +4. `list_workflows(shape="image")` for the text-to-image templates and + `list_workflows(shape="image-edit")` for the edit ones. Take names, + `summary` and `cost` from the listing over the names quoted below, then + `get_workflow` on the one chosen, for its variables. + +## Which shape is the request + +- **A short request** ("a man walking his dog"): + `templates/flux3/upsample-then-generate`. The upsampler writes the caption, + is released, and Klein draws from it. The caption is saved as an + `intermediate` text file: read it, and reuse or edit it as the prompt of + text-to-image without paying for the upsampler again. +- **A full caption already written** (or one the upsampler wrote earlier): + `templates/flux3/text-to-image`. +- **One photo to change**: `templates/flux3/image-edit`, an `instruction` + and an `image`. The upsampler sees the same picture as the pipeline. With + no width or height the output is 1 MP at the reference's aspect ratio. +- **Several pictures combined** ("the dog from this one on that mug"): + `templates/flux3/multi-reference`, an `instruction` and `images` in order. + The first reference sets the aspect ratio. +- **Elements at given positions**: `templates/flux3/layout-boxes`. The + caption carries the boxes; the format is in `references/layout-boxes.md`. + Boxes from a short request (BFL's box upsampler) are not in the catalog + yet, so the caption is written by hand. +- **The base (non-distilled) checkpoint, guidance, LoRA, fp8 or NVFP4**: not + in the catalog. Say so rather than improvising. + +Input media must be on the server: `upload_asset`, then `asset:` in the +argument, or `output:` for an earlier run's file. + +## Hard rules + +- 16 steps, no guidance and no negative prompt. The distilled model is a + flow map; BFL's code refuses any guidance but none. Leave + `num_inference_steps` at 16 (trained for 4-16) and add no + `guidance_scale` or `negative_prompt` to a workflow you write. +- `width` and `height` are multiples of 16 (the VAE's 8x times the 2x2 + latent packing). The pipeline only refuses an off-grid size after the + weights load, so the templates declare the rule and `validate_workflow` + catches it first. +- Sides run from 64 to 4096, as BFL states it. Nothing in the templates or + the pipeline checks that ceiling, and above 1024x1024 a 24 GB card has not + been measured: everything resident leaves under 2 GB spare at 1024x1024. +- Up to 10 references. That is BFL's training limit, and the pipeline does + not enforce it: an 11th reference runs, outside what the model learned. +- In a caption written by hand, name references ``, + `` and so on, zero-based in the order given. The upsampler's + own captions say "the first image", "the second image". +- Every component carries an explicit `device`. The modular pipeline runs + on its first component's device, so one left unplaced keeps the prompt + and the noise on the CPU (`index is on cpu`, `Cannot generate a cpu + tensor from a generator of type cuda`). For the same reason the text + encoder cannot be `residency: on_demand`. +- The upsampler (about 18 GB) and Klein (about 19 GB) do not fit a 24 GB + card together. Keep `release_models` on the upsample step and load Klein + after it. + +## Prompts + +The prompt path is the upsampler. The builtins `flux3_klein_upsample` (a +request to a caption) and `flux3_klein_upsample_edit` (an instruction plus +the references to a caption) run BFL's `flux-3-image-klein-pup` with BFL's +own system prompts, read from the pup repository at a pinned commit. The +templates call them; a workflow you compose names them as a `builtin:` step +before the Klein step. + +To write a caption by hand, read BFL's prompting guide: +https://docs.bfl.ml/guides/prompting_summary. It is written for the API +model, which upsamples on the server side, so read it as a description of +the caption Klein expects, not as advice for a short prompt. The stored +`flux3/` prompts (`list_prompts`, `get_prompt`) are BFL's own upsampled +captions, `man_dog` and `balloons_layout`: read one before writing your own. +The upsampler writes English only. + +## Run and judge + +1. `validate_workflow` first - free, and catches an off-grid size. +2. Quote `plan.estimate` and name any `downloads_required` (Klein, plus the + upsampler for every template but text-to-image and layout-boxes). Do not + quote minutes from memory: the estimate is the number. More references + cost more per step, because each one adds tokens. Get the go-ahead, then + `run_workflow` with `acknowledged_cost` set to the plan's + `{fingerprint, minutes, downloads}`. +3. `wait_for_job`, `timeout_seconds` = estimate plus margin. Call again + while `still_running`. If the call dies client-side first, call again + with `timeout_seconds` 60 and loop. +4. A saving step carries a `subfolder`: the image shown to the user is + `final`, the upsampler's caption `intermediate`. + `list_gallery(subfolder="final")` lists only deliverables, so keep both + in anything you compose. +5. Judge it yourself: `get_output_image` on the `final` image against the + ask, and `get_output_text` on the caption when the picture misses - the + caption is what Klein drew. Then `get_job` for the manifest and its + warnings, and give the user the gallery `url`. +6. Save a keeper (`get_job_workflow`, `save_workflow`) to rerun it by name. + +What goes wrong, and where to look: +- A generic, flat picture: a short prompt went to text-to-image without the + upsampler. Rerun through upsample-then-generate. +- An edit that changed more than was asked: Klein can alter content outside + the requested region (BFL's fix for that is API-only). Narrow the + instruction, or try another seed; compare with `compare_outputs`. +- A small detail of a reference rendered wrong: the upsampler sees each + reference reduced (about 1 MP or less) and can misdescribe fine detail. Read + the caption with `get_output_text`, correct it, and run text-to-image or + the edit with the corrected caption. +- A boxed element missing or spilling over: boxes are hints, not masks. + Enlarge the box or describe the element more, in the prose and the table. + +## Sources + +The flux-3-image-klein and flux-3-image-klein-pup model cards (Hugging Face), +BFL's flux-3-image-klein GitHub repository (commit bbfe317d), docs.bfl.ml, +and `diffusers.modular_pipelines.flux3`. Read 2026-10-15; audit +`docs/proposals/audits/2026-10-15-flux-3-klein-audit.md`. diff --git a/plugins/dw/skills/flux3-klein/references/layout-boxes.md b/plugins/dw/skills/flux3-klein/references/layout-boxes.md new file mode 100644 index 00000000..fa2bbe67 --- /dev/null +++ b/plugins/dw/skills/flux3-klein/references/layout-boxes.md @@ -0,0 +1,31 @@ +# Layout boxes for `templates/flux3/layout-boxes` + +Read this when a request places elements at given positions. The format is +BFL's, from the Klein model card and the box captions in their reference +repository (read 2026-10-15); the stored caption `flux3/balloons_layout` +(`get_prompt`) is a complete worked example, and the clearest one to copy. + +A layout caption is one prompt with two parts: + +1. **Prose with tags.** The usual long caption, with each placed element + followed by an id in angle brackets where it is described, such as + `` or ``. Ids are short and unique; BFL's examples use a + kind and a number. +2. **A JSON table after the prose.** A list with one object per id: + `{"id": "obj_1", "bbox": [top, left, bottom, right], "desc": "..."}`. + The box is on a 0 to 1000 grid over the image, whatever its pixel size, + with `top < bottom` and `left < right`. `desc` describes that element + alone, in the same detail as the prose. + +Every tag in the prose has a row, and every row's id is tagged in the prose. +Background regions (sky, ground) can be boxed too; boxes may overlap. + +Boxes are hints, not masks. A small boxed element may not appear at all, and +an element can run past its box. When one is missing, enlarge its box or give +it more words in both the prose and `desc`, and try another seed before +rewriting the layout. + +Writing the boxes from a short request (BFL's `t2i_bbox` upsampler mode) +needs BFL's box planning code and is not in the catalog yet, so a layout +caption is written by hand or adapted from the stored example. Keep the +width and height on the 16-pixel grid as for any Klein run. diff --git a/tests/test_plugin_skills.py b/tests/test_plugin_skills.py index 42fcb78e..be3c3cad 100644 --- a/tests/test_plugin_skills.py +++ b/tests/test_plugin_skills.py @@ -719,6 +719,58 @@ def test_the_skill_starts_with_the_server(self): assert text.index("get_server_info") < text.index("templates/kandinsky6/") +FLUX3_SKILL = os.path.join(PLUGIN_DIR, "skills", "flux3-klein", "SKILL.md") + + +def flux3_text(): + """The Klein skill with its line wraps collapsed, so a phrase the tests + pin does not break when a paragraph is rewrapped.""" + return " ".join(skill_text(FLUX3_SKILL).split()) + + +class TestFlux3KleinSkill: + """The Klein skill's numbers come from the diffusers modular blocks that + enforce them; the 10-reference cap is the vendor's training limit, which + the pipeline does not enforce, and the skill says so.""" + + def test_the_step_count_is_the_distilled_pipeline_s(self): + # The distilled denoise step, not Flux3SetTimestepsStep (which + # defaults to 50), holds the 16 the templates run. + from diffusers.modular_pipelines.flux3.modular_blocks_flux3_distilled import ( + Flux3DistilledCoreDenoiseStep, + ) + + defaults = {p.name: p.default for p in Flux3DistilledCoreDenoiseStep().inputs} + assert defaults["num_inference_steps"] == 16 + assert "16 steps" in flux3_text() + + def test_the_grid_is_the_vae_scale_times_the_2x2_packing(self): + from diffusers.modular_pipelines.flux3 import before_denoise + + # The multiple is read off the loaded VAE, so only the source shows it. + assert "vae_scale_factor * 2" in inspect.getsource(before_denoise) + assert "multiples of 16" in flux3_text() + + def test_the_distilled_model_takes_no_guidance(self): + assert "no guidance" in flux3_text() + + def test_the_reference_cap_is_stated_as_the_vendor_s(self): + text = flux3_text() + assert "10 references" in text and "" in text + assert "does not enforce" in text + + def test_the_skill_defers_prompt_format_to_bfl(self): + text = flux3_text() + assert "https://docs.bfl.ml/guides/prompting_summary" in text + assert "`flux3_klein_upsample`" in text + assert "`flux3_klein_upsample_edit`" in text + + def test_the_skill_starts_with_the_server(self): + text = flux3_text() + assert text.index("get_server_info") < text.index("list_workflows") + assert text.index("get_server_info") < text.index("templates/flux3/") + + SERIES_SKILL = os.path.join(PLUGIN_DIR, "skills", "series-episodes", "SKILL.md") SERIES_LOOK = os.path.join(os.path.dirname(SERIES_SKILL), "references", "look.md") From 81069fc58fee42a5cf69f76307545289977a5497 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:41:25 -0500 Subject: [PATCH 15/20] fix(plugin): flux3-klein skill's reference-detail failure mode The edit templates always re-upsample and take no caption, so a corrected caption is reused through a one-step workflow built from the edit template's Klein step, not by passing prompt= or falling back to text-to-image. dw's reference sizing for the upsampler is not BFL's, so BFL's budget is no longer quoted. One memory figure for Apple Silicon, the guidance rule as BFL's reference code states it, and CLAUDE.md rewrapped. Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 4 ++-- plugins/dw/skills/flux3-klein/SKILL.md | 20 +++++++++++++------- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index b25fa046..e3ad6fd7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -72,8 +72,8 @@ The guides in `docs/`, by topic: `.claude-plugin/marketplace.json` publishes the `dw` plugin in `plugins/dw/`: one composition skill per model family (`minimax-h3`, `minimax-music3`, `ltx-2.5`, -`kandinsky-6`, `flux3-klein`) that chooses a template for a request's shape and states the family's -hard rules, plus +`kandinsky-6`, `flux3-klein`) that chooses a template for a request's shape and +states the family's hard rules, plus cross-cutting composition skills (`script-to-video`, `series-episodes`) - shapes above the families that orchestrate the decision trees and cast consistency across multiple generations. Every skill the directory holds is named in `plugins/dw/README.md` and here, diff --git a/plugins/dw/skills/flux3-klein/SKILL.md b/plugins/dw/skills/flux3-klein/SKILL.md index c95af9a4..73bd71de 100644 --- a/plugins/dw/skills/flux3-klein/SKILL.md +++ b/plugins/dw/skills/flux3-klein/SKILL.md @@ -14,8 +14,8 @@ says not to run it without one: a short prompt gives a generic picture. ## Before anything 1. `get_server_info`: the device. The templates place every component on - `cuda`. Apple Silicon is unmeasured (about 20 GB of unified memory plus - activations, by the vendor's sizes); on `mps` say so before spending, and + `cuda`. Apple Silicon is unmeasured (the same 19 GB of weights in unified + memory, plus activations); on `mps` say so before spending, and on `cpu` stop. 2. `get_diffusers_state`, then `list_classes` for `Flux3DistilledModularPipeline`. The modular pipeline is in diffusers @@ -59,7 +59,8 @@ argument, or `output:` for an earlier run's file. ## Hard rules - 16 steps, no guidance and no negative prompt. The distilled model is a - flow map; BFL's code refuses any guidance but none. Leave + flow map; BFL's reference code raises on any guidance other than None + or 1.0. Leave `num_inference_steps` at 16 (trained for 4-16) and add no `guidance_scale` or `negative_prompt` to a workflow you write. - `width` and `height` are multiples of 16 (the VAE's 8x times the 2x2 @@ -128,10 +129,15 @@ What goes wrong, and where to look: - An edit that changed more than was asked: Klein can alter content outside the requested region (BFL's fix for that is API-only). Narrow the instruction, or try another seed; compare with `compare_outputs`. -- A small detail of a reference rendered wrong: the upsampler sees each - reference reduced (about 1 MP or less) and can misdescribe fine detail. Read - the caption with `get_output_text`, correct it, and run text-to-image or - the edit with the corrected caption. +- A small detail of a reference rendered wrong: the upsampler can + misdescribe fine detail in a reference, and dw sizes references for it + differently from BFL's own pipeline. Read the caption with + `get_output_text`. The edit templates always re-upsample: their + prompt is the upsampler's output and they have no caption variable. Either narrow the + instruction and rerun, or reuse the corrected caption: `get_workflow` on + the edit template, and compose a one-step workflow from its Klein step + with the caption as `prompt` and the same image list. Text-to-image is no + answer here, since it drops the references. - A boxed element missing or spilling over: boxes are hints, not masks. Enlarge the box or describe the element more, in the prose and the table. From f77e09c79eeffba2fa0cc971e9d32d4901ef47dd Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:57:30 -0500 Subject: [PATCH 16/20] fix(text_generation): gate and filter generate_kwargs; fail closed on unsized from_hub generate_kwargs reached model.generate() unfiltered, so an untrusted workflow could set custom_generate (transformers downloads and runs the repo's custom_generate/generate.py) with trust_remote_code, or pass tokenizer_encode_kwargs around the chat_template_kwargs checks. - dw/trust.py: require_trusted_generate_arguments refuses custom_generate and trust_remote_code unless workflows are trusted - generate_kwargs: the call's structural keys are refused on both paths, and values must be scalars or lists of scalars; all checked before any load - _RESERVED_CHAT_TEMPLATE_KEYS gains the calls' remaining named parameters - system_prompt.from_hub refuses a file the Hub reports no size for Co-Authored-By: Claude Opus 5.5 --- dw/tasks/text_generation.py | 93 +++++++++++++++++++++++++++++-- dw/trust.py | 33 +++++++++++ tests/test_security_trust_gate.py | 45 +++++++++++++++ tests/test_text_generation.py | 89 +++++++++++++++++++++++++++++ 4 files changed, 254 insertions(+), 6 deletions(-) diff --git a/dw/tasks/text_generation.py b/dw/tasks/text_generation.py index 6cc1ee5c..03a81063 100644 --- a/dw/tasks/text_generation.py +++ b/dw/tasks/text_generation.py @@ -20,6 +20,7 @@ from transformers import pipeline as hf_pipeline from .. import preferred_task_dtype from ..locations import validate_hub_repo_id, validate_weight_name +from ..trust import require_trusted_generate_arguments from .model_cache import cached_model, hf_pipeline_placement logger = logging.getLogger("dw") @@ -161,7 +162,15 @@ def resolve_system_prompt(value): f"system_prompt.from_hub: {repo}/{file} at {revision[:8]} is not a " f"public file the Hub will serve without credentials ({type(e).__name__})" ) from e - if metadata.size is not None and metadata.size > _MAX_PROMPT_FILE_BYTES: + # No size is no promise the file is small, so it is refused rather than + # downloaded on trust + if metadata.size is None: + raise ValueError( + f"system_prompt.from_hub: the Hub reported no size for {file}, so it " + f"cannot be checked against the {_MAX_PROMPT_FILE_BYTES}-byte limit " + "before downloading" + ) + if metadata.size > _MAX_PROMPT_FILE_BYTES: raise ValueError( f"system_prompt.from_hub: {file} is larger than " f"{_MAX_PROMPT_FILE_BYTES} bytes, too big for a system prompt" @@ -200,7 +209,9 @@ def resolve_system_prompt(value): # template settings are splatted into those calls, so one of these in a # workflow (which is untrusted) would not set a template variable - it would # hand transformers an image or video to fetch that dw never validated, or -# change what the template call does (tokenize, tools, the template itself) +# change what the template call does (tokenize, tools, the template itself). +# The rest are the calls' other named parameters: one of those is a refusal +# here, before the model loads, rather than a TypeError after it has _RESERVED_CHAT_TEMPLATE_KEYS = frozenset( { "images", @@ -222,9 +233,42 @@ def resolve_system_prompt(value): "tools", "documents", "tokenizer_encode_kwargs", + "conversation", + "padding", + "truncation", + "max_length", + "return_assistant_tokens_mask", + "sanitize_control_tokens", + "tokenizer_kwargs", + "load_audio_from_video", + "return_type", + "return_text", + "stop_sequence", + "clean_up_tokenization_spaces", + "skip_special_tokens", + "inputs", + "text_inputs", + "batch_size", + "num_workers", + "prefix", + "handle_long_generation", } ) +# generate_kwargs go into the same calls - splatted into the text pipeline's, +# handed whole to the vision pipeline's generate() - so the call's own +# arguments are refused there too, along with the text pipeline's +# preprocessing switches +_RESERVED_GENERATE_KEYS = _RESERVED_CHAT_TEMPLATE_KEYS | { + "tokenizer_encode_kwargs", + "prefix", + "handle_long_generation", +} + + +def _is_scalar(value): + return value is None or isinstance(value, (bool, int, float, str)) + def _validated_chat_template_kwargs(value): """Chat template settings: an object of JSON scalars, none of them one of @@ -243,7 +287,7 @@ def _validated_chat_template_kwargs(value): "chat-template setting - it is an argument of the generation " "call itself" ) - if setting is not None and not isinstance(setting, (bool, int, float, str)): + if not _is_scalar(setting): raise ValueError( f"text_generation: chat_template_kwargs value for '{key}' must be " f"a boolean, number, string or null, got {type(setting).__name__}" @@ -251,6 +295,39 @@ def _validated_chat_template_kwargs(value): return value +def _validated_generate_kwargs(value): + """Generation settings: an object of JSON scalars or lists of them + (stop_strings is a list), none of them one of the call's own arguments, + and none that runs code from the Hub unless workflows are trusted.""" + if not value: + return {} + if not isinstance(value, dict): + raise ValueError( + "text_generation: generate_kwargs must be an object of settings for " + f"the model's generate(), got {type(value).__name__}" + ) + # First: custom_generate's value is a plain string, so nothing below would + # stop it + require_trusted_generate_arguments(value, "text_generation") + for key, setting in value.items(): + if key in _RESERVED_GENERATE_KEYS: + raise ValueError( + f"text_generation: generate_kwargs key '{key}' is not a " + "generation setting - it is an argument of the pipeline call " + "itself" + ) + if not ( + _is_scalar(setting) + or (isinstance(setting, list) and all(_is_scalar(v) for v in setting)) + ): + raise ValueError( + f"text_generation: generate_kwargs value for '{key}' must be a " + "boolean, number, string, null or a list of those, got " + f"{type(setting).__name__}" + ) + return value + + def generate_text(prompt, device="cpu", **kwargs): """Generate text from a prompt using a local language model. @@ -274,8 +351,11 @@ def generate_text(prompt, device="cpu", **kwargs): repetition_penalty: Vision pipeline only (default: 1.15). Raise it if a model still repeats itself, or set 1.0 to disable. generate_kwargs: Anything else to hand the model's generate() - - no_repeat_ngram_size, top_p, min_new_tokens and so on. Merged - over what this function sets, so it can override those too. + no_repeat_ngram_size, top_p, min_new_tokens, stop_strings and + so on. Merged over what this function sets, so it can override + those too. Values are scalars or lists of scalars; the call's + own arguments are refused, and custom_generate and + trust_remote_code need trusted workflows. Returns: Generated text string. @@ -290,6 +370,7 @@ def generate_text(prompt, device="cpu", **kwargs): chat_template_kwargs = _validated_chat_template_kwargs( kwargs.get("chat_template_kwargs") ) + generate_kwargs = _validated_generate_kwargs(kwargs.get("generate_kwargs")) if images is None: pipeline_task = "text-generation" @@ -336,7 +417,7 @@ def load_pipe(): kwargs.get("repetition_penalty", _VISION_REPETITION_PENALTY) ) # Last, so a workflow can override anything decided above - generation.update(kwargs.get("generate_kwargs") or {}) + generation.update(generate_kwargs) return _generate( pipe, messages, images, max_new_tokens, generation, chat_template_kwargs diff --git a/dw/trust.py b/dw/trust.py index 4a922b77..9e9dcd63 100644 --- a/dw/trust.py +++ b/dw/trust.py @@ -275,3 +275,36 @@ def require_trusted_from_pretrained_arguments(arguments, what: str) -> None: f"Python - see docs/SECURITY.md. Pass --trust-workflows if you " f"trust this workflow's source." ) + + +# generate() arguments that load and run code from the Hub. transformers' +# generate() hands `custom_generate` (a repo id) to load_custom_generate, which +# downloads that repo's custom_generate/generate.py and executes it in place of +# the built-in decoding loop; `trust_remote_code` is what lets it. A task that +# passes a workflow's generate_kwargs through reaches that with no from_pretrained +# call for require_trusted_from_pretrained_arguments to see +GENERATE_REMOTE_CODE_ARGUMENTS = ("custom_generate", "trust_remote_code") + + +def require_trusted_generate_arguments(arguments, what: str) -> None: + """Refuse generate() arguments that download and execute Python from the + Hub unless the workflow is trusted. + + Args: + arguments: The generate_kwargs block a workflow supplied + what: The task or call being configured, for the error message + + Raises: + UntrustedWorkflowError: If untrusted and any of + GENERATE_REMOTE_CODE_ARGUMENTS is set + """ + if workflows_are_trusted() or not isinstance(arguments, dict): + return + for key in GENERATE_REMOTE_CODE_ARGUMENTS: + if arguments.get(key): + raise UntrustedWorkflowError( + f"Refusing to run {what}: its generate_kwargs set '{key}', which " + f"downloads and executes Python from a Hub repository, so it " + f"needs trusted workflows - see docs/SECURITY.md. Pass " + f"--trust-workflows if you trust this workflow's source." + ) diff --git a/tests/test_security_trust_gate.py b/tests/test_security_trust_gate.py index 2b0a1eac..d63ce602 100644 --- a/tests/test_security_trust_gate.py +++ b/tests/test_security_trust_gate.py @@ -779,6 +779,51 @@ def test_an_in_package_class_of_no_accepted_kind_is_refused(self, untrusted, nam # --------------------------------------------------- validate-time probes +class TestRemoteGenerateCodeUntrusted: + """text_generation's generate_kwargs reach model.generate(), and + `custom_generate` there makes transformers download a repo's + custom_generate/generate.py and run it. Untrusted, both it and + trust_remote_code are refused before the model loads or the Hub is asked.""" + + @pytest.mark.parametrize( + "settings", + [ + pytest.param( + {"custom_generate": "attacker/repo", "trust_remote_code": True}, + id="both", + ), + pytest.param({"custom_generate": "attacker/repo"}, id="custom_generate"), + pytest.param({"trust_remote_code": True}, id="trust_remote_code"), + ], + ) + @pytest.mark.parametrize("image", [None, "https://example.com/a.png"]) + def test_refused_before_the_model_loads( + self, untrusted, no_network, monkeypatch, settings, image + ): + from dw.tasks import text_generation + + loader = MagicMock() + monkeypatch.setattr(text_generation, "hf_pipeline", loader) + key = next(iter(settings)) + with pytest.raises(UntrustedWorkflowError, match=f"'{key}'.*trusted"): + text_generation.generate_text( + "a cat", + device="cpu", + model_name="m/untrusted-generate", + image=image, + generate_kwargs=settings, + ) + loader.assert_not_called() + assert list(no_network) == [] + + def test_trusted_lets_them_through(self, trusted): + from dw.trust import require_trusted_generate_arguments + + require_trusted_generate_arguments( + {"custom_generate": "a/b", "trust_remote_code": True}, "text_generation" + ) + + def _probe_recorder(info): """A `probe=` callable that records every path it is asked about. diff --git a/tests/test_text_generation.py b/tests/test_text_generation.py index 0d873c6a..9be473f0 100644 --- a/tests/test_text_generation.py +++ b/tests/test_text_generation.py @@ -452,6 +452,88 @@ def test_overrides_what_the_task_decided(self, mock_pipeline): ) +class TestGenerateKwargsAreFiltered(unittest.TestCase): + """generate_kwargs reach the pipeline call, so they are checked first. + + The untrusted refusal of custom_generate and trust_remote_code lives with + the other trust gate tests, in tests/test_security_trust_gate.py; this + suite runs trusted (tests/conftest.py). + """ + + def _mock_pipe(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "output"}] + mock_pipeline.return_value = pipe + return pipe + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_trusted_workflows_may_name_remote_generate_code(self, mock_pipeline): + pipe = self._mock_pipe(mock_pipeline) + settings = {"custom_generate": "someone/decoder", "trust_remote_code": True} + with patch("dw.trust.workflows_are_trusted", return_value=True): + generate_text( + "a cat", device="cpu", model_name="m/trusted", generate_kwargs=settings + ) + self.assertEqual(pipe.call_args[1]["custom_generate"], "someone/decoder") + self.assertIs(pipe.call_args[1]["trust_remote_code"], True) + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_a_structural_key_is_refused_before_any_load(self, mock_pipeline): + for key in ( + "tokenizer_encode_kwargs", + "prefix", + "handle_long_generation", + "max_new_tokens", + "images", + ): + for image in (None, "https://example.com/a.png"): + with self.assertRaisesRegex(ValueError, f"'{key}'.*not a generation"): + generate_text( + "a cat", + device="cpu", + model_name="m/structural", + image=image, + generate_kwargs={key: "x"}, + ) + mock_pipeline.assert_not_called() + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_a_nested_value_is_refused(self, mock_pipeline): + for bad in ({"a": 1}, [{"a": 1}], [[1]]): + with self.assertRaisesRegex(ValueError, "sequence_bias"): + generate_text( + "a cat", + device="cpu", + model_name="m/nested", + generate_kwargs={"sequence_bias": bad}, + ) + with self.assertRaisesRegex(ValueError, "must be an object"): + generate_text("a cat", device="cpu", generate_kwargs=["top_p"]) + mock_pipeline.assert_not_called() + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_legitimate_settings_pass_on_both_paths(self, mock_pipeline): + pipe = self._mock_pipe(mock_pipeline) + settings = {"no_repeat_ngram_size": 3, "stop_strings": ["\n\n"]} + + generate_text( + "a cat", device="cpu", model_name="m/legit", generate_kwargs=settings + ) + self.assertEqual(pipe.call_args[1]["no_repeat_ngram_size"], 3) + self.assertEqual(pipe.call_args[1]["stop_strings"], ["\n\n"]) + + generate_text( + "a cat", + device="cpu", + model_name="m/legit-vision", + image="https://example.com/a.png", + generate_kwargs=settings, + ) + forwarded = pipe.call_args[1]["generate_kwargs"] + self.assertEqual(forwarded["no_repeat_ngram_size"], 3) + self.assertEqual(forwarded["stop_strings"], ["\n\n"]) + + class TestTextGenerationLoadDevice(unittest.TestCase): """How the model gets onto the device it runs on. @@ -609,6 +691,13 @@ def test_an_oversize_file_is_refused_before_downloading(self, download): resolve_system_prompt(_spec()) download.assert_not_called() + @patch("dw.tasks.text_generation.hf_hub_download") + def test_a_file_of_unknown_size_is_refused_before_downloading(self, download): + self.metadata.return_value = MagicMock(size=None) + with self.assertRaisesRegex(ValueError, "no size"): + resolve_system_prompt(_spec()) + download.assert_not_called() + def test_a_non_string_key_is_refused(self): for key in (["t2i"], {"a": 1}, 3): with self.assertRaises(ValueError): From e6baee556384716fd0fd58f881246556a2b7cbe2 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 20:58:41 -0500 Subject: [PATCH 17/20] fix(flux3): state the 4096 side as unchecked, the edit upsampler's image need - text-to-image, upsample-then-generate and the README declare only the multiple-of-16 rule; BFL's 64 to 4096 range is stated as not checked - flux3_klein_upsample_edit says at least one image is required, and what an empty list produces - tests: flux3 step-cache rows in alphabetical order; the upsample step's result is asserted to be an intermediate Co-Authored-By: Claude Opus 5.5 --- dw/workflows/flux3_klein_upsample_edit.json | 2 +- tests/test_flux3_catalog.py | 17 ++++++++--- tests/test_step_cache.py | 30 +++++++++---------- workflows/templates/flux3/README.md | 2 +- workflows/templates/flux3/text-to-image.json | 2 +- .../flux3/upsample-then-generate.json | 2 +- 6 files changed, 32 insertions(+), 23 deletions(-) diff --git a/dw/workflows/flux3_klein_upsample_edit.json b/dw/workflows/flux3_klein_upsample_edit.json index e6505f24..d4b9baf6 100644 --- a/dw/workflows/flux3_klein_upsample_edit.json +++ b/dw/workflows/flux3_klein_upsample_edit.json @@ -1,6 +1,6 @@ { "id": "flux3_klein_upsample_edit", - "description": "Expands an edit instruction into the caption FLUX 3 Image Klein was trained on, with Black Forest Labs' own prompt upsampler, flux-3-image-klein-pup, a Qwen3.5-9B fine-tune. The upsampler sees every reference in order, as the pipeline will. The instruction is wrapped in the three-line scaffold that is BFL's I2I_SCAFFOLD from src/flux3/pup.py at commit bbfe317d in github.com/black-forest-labs/flux-3-image-klein. The weights and the image-edit system prompt are both pinned to pup commit 3f3edb52, and the prompt is read from the repo's system_prompts.json rather than copied. Thinking is switched off and decoding is greedy with no repetition penalty, as BFL runs it. BFL resizes each reference to a pixel budget (1 MP for one reference, 724x724 for two, 512x512 for three or more) before the upsampler sees it, while transformers' processor applies its own sizing, so the two can describe small details differently. About 18 GB in bf16.", + "description": "Expands an edit instruction into the caption FLUX 3 Image Klein was trained on, with Black Forest Labs' own prompt upsampler, flux-3-image-klein-pup, a Qwen3.5-9B fine-tune. The upsampler sees every reference in order, as the pipeline will. At least one image is required: an empty list produces a caption written without seeing the image. The instruction is wrapped in the three-line scaffold that is BFL's I2I_SCAFFOLD from src/flux3/pup.py at commit bbfe317d in github.com/black-forest-labs/flux-3-image-klein. The weights and the image-edit system prompt are both pinned to pup commit 3f3edb52, and the prompt is read from the repo's system_prompts.json rather than copied. Thinking is switched off and decoding is greedy with no repetition penalty, as BFL runs it. BFL resizes each reference to a pixel budget (1 MP for one reference, 724x724 for two, 512x512 for three or more) before the upsampler sees it, while transformers' processor applies its own sizing, so the two can describe small details differently. About 18 GB in bf16.", "variables": { "instruction": "put a red bandana around the dog's neck", "images": [], diff --git a/tests/test_flux3_catalog.py b/tests/test_flux3_catalog.py index a7b6b3b4..53898c90 100644 --- a/tests/test_flux3_catalog.py +++ b/tests/test_flux3_catalog.py @@ -148,7 +148,10 @@ def test_upsample_then_generate_feeds_the_builtin_caption_to_klein(): first, second = definition["steps"] assert first["workflow"]["path"] == "builtin:flux3_klein_upsample.json" assert first["release_models"] is True - assert second["pipeline"]["arguments"]["prompt"] == f"previous_result:{first['name']}" + assert first["result"]["subfolder"] == "intermediate" + assert ( + second["pipeline"]["arguments"]["prompt"] == f"previous_result:{first['name']}" + ) @pytest.mark.parametrize("name", ["image-edit.json", "multi-reference.json"]) @@ -157,8 +160,14 @@ def test_the_edit_templates_show_the_upsampler_the_pipeline_s_references(name): upsample = next(s for s in definition["steps"] if "workflow" in s) assert upsample["workflow"]["path"] == "builtin:flux3_klein_upsample_edit.json" (klein,) = klein_steps(definition) - assert upsample["workflow"]["arguments"]["images"] == klein["pipeline"]["arguments"]["image"] - assert klein["pipeline"]["arguments"]["prompt"] == f"previous_result:{upsample['name']}" + assert ( + upsample["workflow"]["arguments"]["images"] + == klein["pipeline"]["arguments"]["image"] + ) + assert ( + klein["pipeline"]["arguments"]["prompt"] + == f"previous_result:{upsample['name']}" + ) def test_multi_reference_states_the_ten_reference_cap_and_the_naming(): @@ -170,7 +179,7 @@ def test_multi_reference_states_the_ten_reference_cap_and_the_naming(): def test_the_layout_caption_carries_tags_and_a_0_to_1000_box_table(): with open(os.path.join(ROOT, "prompts", "flux3", "balloons_layout.json")) as f: text = json.load(f)["text"] - table = json.loads(text[text.index("["):]) + table = json.loads(text[text.index("[") :]) assert table and all({"id", "bbox", "desc"} <= set(row) for row in table) for row in table: assert f"<{row['id']}>" in text diff --git a/tests/test_step_cache.py b/tests/test_step_cache.py index d01b1053..98aa1321 100644 --- a/tests/test_step_cache.py +++ b/tests/test_step_cache.py @@ -661,6 +661,21 @@ def test_normalized_downstream_false_when_nothing_references_it(): "generate": "bfe2e09d70b6bf6e7799266cf735abd779e5a9e3c283b8287d962eab17d95425" }, "expand-prompt.json": {}, + "flux3/image-edit.json": { + "edit": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" + }, + "flux3/layout-boxes.json": { + "layout_to_image": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" + }, + "flux3/multi-reference.json": { + "edit": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" + }, + "flux3/text-to-image.json": { + "text_to_image": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" + }, + "flux3/upsample-then-generate.json": { + "text_to_image": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" + }, "generate-speech.json": {}, "image-edit.json": { "edit": "241e4cc77f00d837873d01e3ea9ec7105513af8362908d8229c740cbca997737" @@ -683,21 +698,6 @@ def test_normalized_downstream_false_when_nothing_references_it(): "ip-adapter.json": { "main": "a67c769c0ee881e36aa15b62f8226a233c1e4b93542a8caaf5eab5715cb9f556" }, - "flux3/image-edit.json": { - "edit": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" - }, - "flux3/layout-boxes.json": { - "layout_to_image": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" - }, - "flux3/multi-reference.json": { - "edit": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" - }, - "flux3/text-to-image.json": { - "text_to_image": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" - }, - "flux3/upsample-then-generate.json": { - "text_to_image": "58be5755e37a15b2780cc4fa3eead66e91e2eeec21e15f15624f60a9a072f235" - }, "kandinsky6/generate-and-upscale.json": { "text_to_video": "3095873d2e63833c2e8c8048379035840694ca108309e544ff17a0722b4ee51a", "upscale": "838beb7eefde695da365c37927645fc388ab8feeaf8e1bc1695440cbd2b54fdd", diff --git a/workflows/templates/flux3/README.md b/workflows/templates/flux3/README.md index 770df26d..63326842 100644 --- a/workflows/templates/flux3/README.md +++ b/workflows/templates/flux3/README.md @@ -33,7 +33,7 @@ Read them in this order and each introduces one new idea on top of the last. | Example | What it introduces | | ------- | ------------------ | -| [text-to-image.json](text-to-image.json) | The baseline. The prompt is already a full upsampled caption (BFL's man-walking-a-dog example), 16 steps, no guidance. Width and height are multiples of 16, up to 4096, and the pipeline only complains after the load, so the workflow declares the rule up front. 24 s on an RTX 3090 at 1024x1024 (0.8 s/step), peaking at 22.1GB reserved | +| [text-to-image.json](text-to-image.json) | The baseline. The prompt is already a full upsampled caption (BFL's man-walking-a-dog example), 16 steps, no guidance. Width and height must be multiples of 16, and the pipeline only complains after the load, so the workflow declares that rule up front. BFL states sides of 64 to 4096; that range is not checked. 24 s on an RTX 3090 at 1024x1024 (0.8 s/step), peaking at 22.1GB reserved | | [upsample-then-generate.json](upsample-then-generate.json) | A short request becomes a caption first: the pup upsampler writes it, `release_models` frees the upsampler (about 18GB), and Klein draws from the result. The caption is saved as an intermediate, so it can be read, edited or reused as the prompt of text-to-image | | [image-edit.json](image-edit.json) | One reference and an instruction. The upsampler sees the same image as the pipeline. With no width or height the output is 1 MP at the reference's aspect ratio, on the 16-pixel grid | | [multi-reference.json](multi-reference.json) | Several references, in order. A caption written by hand names them ``, `` and so on (zero-based). The first reference sets the aspect ratio. BFL trained on up to 10; the pipeline does not enforce that, so more than 10 run but are outside what the model learned. Two references ran in 45 s on an RTX 3090 (2.6 s/step, references add tokens), also peaking at 22.1GB reserved | diff --git a/workflows/templates/flux3/text-to-image.json b/workflows/templates/flux3/text-to-image.json index fb8fc6fd..2498f8c4 100644 --- a/workflows/templates/flux3/text-to-image.json +++ b/workflows/templates/flux3/text-to-image.json @@ -1,6 +1,6 @@ { "id": "Flux3KleinT2I", - "description": "FLUX 3 Image Klein (distilled, 5.6B) text to image, the baseline of the FLUX 3 templates. The prompt here is already a full caption - BFL's own upsampled caption for their man walking a dog example - because Klein was trained on long upsampler captions and turns short prompts into generic pictures; upsample-then-generate shows how to get one from a short request. Sixteen steps with no guidance: the distilled model is a flow map, so it takes no guidance scale and no negative prompt. Width and height must be multiples of 16, up to 4096. Every component is placed on the GPU explicitly, about 19 GB of bf16 weights, which fits a 24 GB card at 1024x1024 with under 2 GB to spare. The weights are gated under the FLUX Kommunity License: accept it on the model page and log in with hf auth login first.", + "description": "FLUX 3 Image Klein (distilled, 5.6B) text to image, the baseline of the FLUX 3 templates. The prompt here is already a full caption - BFL's own upsampled caption for their man walking a dog example - because Klein was trained on long upsampler captions and turns short prompts into generic pictures; upsample-then-generate shows how to get one from a short request. Sixteen steps with no guidance: the distilled model is a flow map, so it takes no guidance scale and no negative prompt. Width and height must be multiples of 16, which the workflow declares; BFL states sides of 64 to 4096, which is not checked. Every component is placed on the GPU explicitly, about 19 GB of bf16 weights, which fits a 24 GB card at 1024x1024 with under 2 GB to spare. The weights are gated under the FLUX Kommunity License: accept it on the model page and log in with hf auth login first.", "cost_drivers": [ "width", "height" diff --git a/workflows/templates/flux3/upsample-then-generate.json b/workflows/templates/flux3/upsample-then-generate.json index 9bebf350..f9d8e5e8 100644 --- a/workflows/templates/flux3/upsample-then-generate.json +++ b/workflows/templates/flux3/upsample-then-generate.json @@ -1,6 +1,6 @@ { "id": "Flux3KleinUpsampleThenGenerate", - "description": "FLUX 3 Image Klein from a short request, the way Black Forest Labs say to run it: their prompt upsampler first expands the request into the long caption Klein was trained on, then Klein draws from that caption. The upsampler is about 18 GB and Klein about 19 GB in bf16, which do not fit a 24 GB card together, so the upsample step sets release_models and the upsampler is freed before Klein loads. The caption is saved as an intermediate so it can be read, edited or reused as the prompt of text-to-image without paying for the upsampler again. Sixteen steps with no guidance or negative prompt, as the distilled model takes neither. Width and height must be multiples of 16, up to 4096. The weights are gated under the FLUX Kommunity License: accept it on the model page and log in with hf auth login first.", + "description": "FLUX 3 Image Klein from a short request, the way Black Forest Labs say to run it: their prompt upsampler first expands the request into the long caption Klein was trained on, then Klein draws from that caption. The upsampler is about 18 GB and Klein about 19 GB in bf16, which do not fit a 24 GB card together, so the upsample step sets release_models and the upsampler is freed before Klein loads. The caption is saved as an intermediate so it can be read, edited or reused as the prompt of text-to-image without paying for the upsampler again. Sixteen steps with no guidance or negative prompt, as the distilled model takes neither. Width and height must be multiples of 16, which the workflow declares; BFL states sides of 64 to 4096, which is not checked. The weights are gated under the FLUX Kommunity License: accept it on the model page and log in with hf auth login first.", "summary": "FLUX 3 Klein image from a short request, expanded by BFL's upsampler first.", "cost_drivers": [ "width", From 181554d010c209f255568cc0cfc783cf2ae683e5 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 21:01:06 -0500 Subject: [PATCH 18/20] docs: the generate_kwargs gate, the from_hub fetch, and what the spec built - TASKS.md: chat_template_kwargs and generate_kwargs filtering, the custom_generate/trust_remote_code trust rule, and that an image list is one prompt with one answer (also for judge and image_to_text) - SECURITY.md and ARCHITECTURE.md: the generate_kwargs gate and the anonymous, pinned, fail-closed from_hub fetch - Klein design spec: from_hub asks the Hub before every use and fails closed offline; the 10-reference cap is stated, not enforced; follow-ups for the other tasks' generate_kwargs and revision validation parity Co-Authored-By: Claude Opus 5.5 --- docs/ARCHITECTURE.md | 2 ++ docs/SECURITY.md | 20 +++++++++++ docs/TASKS.md | 8 +++-- .../2026-10-15-flux-3-klein-catalog-design.md | 34 ++++++++++++++++--- 4 files changed, 56 insertions(+), 8 deletions(-) diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 4649da13..d7006929 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -147,6 +147,8 @@ one to open. | Locations from a workflow | `dw/locations.py` | A local path in a workflow's media arguments (named like their media, or declared by the task's `@register_command(media_arguments=...)`, read through `TASK_MEDIA_ARGUMENTS`) is confined to a known root, a location written as a URL is one `security.validate_url` accepts, and a URL must not name a host inside the deployment. A host that does not resolve is refused. The requests themselves are `dw/outbound.py`'s. | `tests/test_locations.py::TestMediaPathContainment::test_absolute_path_outside_every_root_is_refused`, `tests/test_locations.py::TestMediaPathContainment::test_a_url_that_is_not_http_is_refused`, `tests/test_locations.py::TestMediaPathContainment::test_an_upper_case_https_url_is_not_refused_as_another_scheme`, `tests/test_security_ssrf.py::TestInternalAddressSpellings::test_a_name_that_resolves_inside`, `tests/test_locations.py::TestTaskMediaArguments::test_an_unreadable_source_is_an_error_at_its_path` | | Outbound requests | `dw/outbound.py`: `safe_get`, `safe_post` | Every outbound request is validated per hop against `dw/locations.py`'s policy, resolved once and dialed at the address the policy checked, capped and timed. No other module under `dw/` makes an HTTP request or builds a session or adapter. | `tests/test_outbound.py::TestSafeRequest::test_one_lookup_per_hop_and_the_dial_is_the_checked_address` | | Trust gate | `dw/trust.py`: `require_trusted_dotted_name` | In an untrusted workflow (the default), a dotted type outside the allowlist is a validation error before anything is imported. | `tests/test_security_trust_gate.py::TestValidationRefusesBeforeImport::test_a_dotted_type_is_a_validation_error` | +| Remote generate code | `dw/trust.py`: `require_trusted_generate_arguments`; `dw/tasks/text_generation.py`: `_validated_generate_kwargs`, `_validated_chat_template_kwargs` | Untrusted, `generate_kwargs` may not set `custom_generate` or `trust_remote_code` (transformers would download and run a repo's `custom_generate/generate.py`); both it and `chat_template_kwargs` hold only scalars (or, for `generate_kwargs`, lists of them) and none of the pipeline call's own keys. All checked before the model loads. | `tests/test_security_trust_gate.py::TestRemoteGenerateCodeUntrusted::test_refused_before_the_model_loads`, `tests/test_text_generation.py::TestGenerateKwargsAreFiltered::test_a_structural_key_is_refused_before_any_load`, `tests/test_text_generation.py::TestGenerateKwargsAreFiltered::test_a_nested_value_is_refused` | +| Hub-file system prompts | `dw/tasks/text_generation.py`: `resolve_system_prompt` | A `system_prompt.from_hub` file is pinned to a full commit, asked of the Hub anonymously before every use (failing closed on any error, offline or under `HF_HUB_OFFLINE`), downloaded with `token=False`, and refused over 1 MiB or when the Hub reports no size. | `tests/test_text_generation.py::TestSystemPromptFromHub::test_a_file_the_hub_will_not_serve_publicly_is_refused`, `tests/test_text_generation.py::TestSystemPromptFromHub::test_the_download_never_carries_the_server_token`, `tests/test_text_generation.py::TestSystemPromptFromHub::test_a_file_of_unknown_size_is_refused_before_downloading` | | CodeQL path-injection model | `.github/codeql/dw-security/DwPathSanitizers.qll`, `.github/workflows/codeql.yml` | The local pack models the `dw/security.py` validators (and `dw/locations.py`'s `validate_media_path`) as sanitizers, so a validator that moves is re-modelled in the same commit. | `.github/codeql/dw-security/DwPathInjection.ql` | | Archives never follow links | `dw/server/outputs.py`: `zip_download` | A gallery listing drops a symlink that leaves its root, and an archive skips one. | `tests/test_security_symlinks.py::TestOutputs::test_the_archive_route_does_not_follow_the_link`, `tests/test_security_symlinks.py::TestOutputs::test_the_gallery_listing_does_not_enumerate_the_link` | | Engine/server import direction | `dw/workspace.py`: `EXPORTS_SUBDIR` | No engine module imports `dw.server` except the entry point `dw/serve.py`, and a constant both sides need lives engine-side. | `docs/stabilization/baseline.json`: `import_cycles` ratchet (indirect: it fails only on an import that closes a cycle) | diff --git a/docs/SECURITY.md b/docs/SECURITY.md index fd63c0fe..382142db 100644 --- a/docs/SECURITY.md +++ b/docs/SECURITY.md @@ -120,6 +120,26 @@ diffusers' remote-code paths open. No bundled catalog entry sets either (`tests/test_catalog_structure.py` refuses one that does); a workflow that needs them needs `--trust-workflows`. +`text_generation`'s `generate_kwargs` reach the model's `generate()`, which +has a remote-code path of its own: `custom_generate` names a Hub repo whose +`custom_generate/generate.py` transformers downloads and runs, and +`trust_remote_code` lets it. Both are refused untrusted +(`require_trusted_generate_arguments` in `dw/trust.py`), before the model +loads. Its other keys must be scalars or lists of them, and none may be one of +the pipeline call's own arguments - the same filter `chat_template_kwargs` +gets, so neither can pass `tokenizer_encode_kwargs`, `images` or a template +of its own around the task's checks. + +`text_generation`'s `system_prompt` may be read from a Hub repo file +(`{"from_hub": {...}}`). That fetch is anonymous only: the Hub is asked +without credentials (`token=False`) whether it serves the file publicly at the +pinned 40-hex commit before every use, and the download carries no token +either, so the server's own login never reaches a gated or private file on a +workflow's say-so. It fails closed - any error from that ask, being offline, +or `HF_HUB_OFFLINE` refuses the run rather than fall back to the cache - and +it is capped at 1 MiB, checked against the Hub's reported size before the +download (no reported size is a refusal) and against the file after it. + ### Where a workflow may read and reach (`dw/locations.py`) Code execution is not the only thing a workflow file chooses. Its arguments diff --git a/docs/TASKS.md b/docs/TASKS.md index 1d738578..4127bb18 100644 --- a/docs/TASKS.md +++ b/docs/TASKS.md @@ -2848,12 +2848,12 @@ Generate or expand text using a local language model. Useful for expanding short | `model_name` | No | HuggingFace model ID (default: `Qwen/Qwen2.5-1.5B-Instruct`, or `HuggingFaceTB/SmolVLM-256M-Instruct` when an image is supplied) | | `image` | No | PIL Image, URL/path, or `previous_result:` reference — see below | | `repetition_penalty` | No | Vision path only (default: 1.15) — see below | -| `generate_kwargs` | No | Anything else to pass to the model's `generate()` — `no_repeat_ngram_size`, `top_p`, `min_new_tokens`. Merged last, so it overrides the settings above | +| `generate_kwargs` | No | Anything else to pass to the model's `generate()` — `no_repeat_ngram_size`, `top_p`, `min_new_tokens`, `stop_strings`. Merged last, so it overrides the settings above. Filtered like `chat_template_kwargs` — see below | | `max_new_tokens` | No | Maximum tokens to generate (default: 500) | -| `chat_template_kwargs` | No | Extra keyword arguments for the model's chat template (for a template that takes switches such as a thinking toggle) | +| `chat_template_kwargs` | No | Extra keyword arguments for the model's chat template (for a template that takes switches such as a thinking toggle). Values must be scalars (boolean, number, string or null), and the generation call's own keys (`images`, `tools`, `max_new_tokens`, `tokenizer_encode_kwargs`, ...) are refused | | `revision` | No | Model revision (branch, tag or commit) to load `model_name` at | -`image` may also be a list of images, each sent as its own picture in the one user message. +`image` may also be a list of images, each sent as its own picture in the one user message. A list is one prompt over all the images, in order, and returns one answer — not one answer per image. The same holds for `judge` and `image_to_text`, which hand their `image` to this task. An empty list, or one holding only nulls, is no image. `system_prompt` may name a file in a Hugging Face repo instead of carrying the text, for a vendor that ships its prompt instructions there: @@ -2903,6 +2903,8 @@ For anything the arguments above do not cover, `generate_kwargs` goes straight t It is merged after everything else, so it can override `repetition_penalty` and the sampling settings as well as add to them. +`generate_kwargs` is checked before the model loads, the same way `chat_template_kwargs` is. Values must be scalars or lists of scalars (`stop_strings` is a list), and the call's own keys are refused (`max_new_tokens`, `tokenizer_encode_kwargs`, `prefix`, `handle_long_generation` and the rest of the pipeline's named arguments). `custom_generate` and `trust_remote_code` need trusted workflows (`--trust-workflows`): with them, transformers downloads a repo's `custom_generate/generate.py` and runs it in place of its own decoding loop. + **Examples:** - [expand-prompt.json](../workflows/templates/expand-prompt.json) — Expand a short prompt and save as `.txt` diff --git a/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md b/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md index a598e06a..dc2707f4 100644 --- a/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md +++ b/docs/superpowers/specs/2026-10-15-flux-3-klein-catalog-design.md @@ -64,8 +64,10 @@ settings*: 3. `image-edit.json` - one reference; the `i2i` upsampler sees the same image; output defaults to 1 MP at the reference's aspect ratio (diffusers' rule, not the vendor code's - stated in the description). -4. `multi-reference.json` - 2-10 references named `` in the caption; the - 10 cap is a `variable_constraints` rule because diffusers does not enforce it. +4. `multi-reference.json` - 2-10 references named `` in the caption. The + 10 cap is stated in the description and the skill, not enforced: diffusers does not + enforce it, and `variable_constraints` are numeric only, so they cannot check a + list's length. An enforcing list-length check is a follow-up. 5. `layout-boxes.json` - `` caption + 0-1000 box table, written in the box format directly as the HF card shows. Upsampled boxes (`t2i_bbox`) need BFL's `bbox.plan`/`bbox.reconcile` scaffolding, so they're a follow-up. @@ -93,7 +95,10 @@ needs two things to run it as BFL specifies: `repetition_penalty` is already overridable (1.15 is only the vision default), and `generate_kwargs` already passes through, so the vendor's greedy decoding needs no -change: the edit templates set `repetition_penalty: 1.0`. +change: the edit templates set `repetition_penalty: 1.0`. Built later: `generate_kwargs` +is filtered like `chat_template_kwargs` (scalars or lists of them, none of the call's own +keys), and `custom_generate`/`trust_remote_code` need trusted workflows, since +transformers runs a `custom_generate` repo's code. ## Decision: the upsampler's system prompts are read from the Hub (Don, 2026-10-15) @@ -113,8 +118,15 @@ a Hub repo: "key": "t2i"}} ``` -- Fetched with `hf_hub_download` into the normal HF cache, so it's one small download, - then offline. +- Not "one download, then offline": before every use the Hub is asked, anonymously + (`get_hf_file_metadata`, `token=False`), whether it serves the file publicly at that + commit, and any error refuses the run rather than falling back to the cache. So it + fails closed offline or under `HF_HUB_OFFLINE`. Without that ask, a full commit would + be answered from the local cache, where files the server fetched with its own login + (gated or private repos) would be readable by any workflow. +- The download is public-only (`hf_hub_download(..., token=False)`) and capped at + 1 MiB: the Hub's reported size is checked first (no reported size is a refusal), + and the downloaded file again before it is read. - `revision` is required and must be a full 40-hex commit, so the text a template runs can't move under it. - `repo` matches the Hub id pattern, and `file` is a relative path with no `..` or @@ -136,3 +148,15 @@ a Hub repo: `model_name`, 18.8 GB), so the cost quote misses the download. Once a partial cache existed it swung the other way: 40.3 GB, the whole repo including the `variants/` the pipeline never loads. +- Enforce the multi-reference 10 cap with a list-length check (`variable_constraints` + are numeric only). +- `generate_kwargs` sweep of the other tasks: `text_generation` now refuses + `custom_generate`/`trust_remote_code` untrusted and filters the call's own keys and + non-scalar values. `speech_generation` still passes a workflow's `generate_kwargs` + (`dw/tasks/speech_generation.py:219`) and `forward_params` (`:190`, `:218`) to the + transformers pipeline unfiltered, and the controlnet_aux detector handlers splat + workflow kwargs into the detector call (`dw/tasks/image_utils.py:103`, `:114`, `:124`). + The repo owner decides that fix. +- `revision` validation parity (M7, deferred): `text_generation`'s model `revision` is + passed to the loader as given, while `system_prompt.from_hub.revision` must be a full + commit; whether the model revision gets the same pin rule is open. From b9b997f846a5224e7afce829add6b2b08e0734c9 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Thu, 15 Oct 2026 22:40:49 -0500 Subject: [PATCH 19/20] test(catalog): raise compact listing budget 10_850 -> 11_400 for FLUX 3 Klein (#930) Approved by Don 2026-10-15: measured 11_256 with the five Klein templates, about 85 tokens held for the Task 9 cost blocks. Co-Authored-By: Claude Opus 5.5 --- tests/test_catalog_structure.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_catalog_structure.py b/tests/test_catalog_structure.py index d040038a..8d9f157d 100644 --- a/tests/test_catalog_structure.py +++ b/tests/test_catalog_structure.py @@ -484,7 +484,10 @@ def test_no_stale_entry_in_the_allowlist(): # Then to 10_850 for the whole-object `cache` variable on the same nineteen H3 # templates (#865, stage #884, 2026-10-13), measured at 10_800: about 42 # tokens, so first_block on a long base-model schedule is an argument too. -COMPACT_BUDGET = 10_850 +# Then to 11_400 for the five FLUX 3 Klein templates (#930, 2026-10-15), +# measured at 11_256: a new family, its summaries already trimmed to the idea +# each template adds, with about 85 tokens held for its Task 9 cost blocks. +COMPACT_BUDGET = 11_400 FILTERED_BUDGET = 1_500 From ad23a852f12d6888501477bf1dfc5684802632a8 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Fri, 16 Oct 2026 06:00:55 -0500 Subject: [PATCH 20/20] test(text_generation): patch hf_pipeline and hub download once via fixtures Brings test_dw_patch_targets back under the ratchet for #930. Co-Authored-By: Claude Opus 5.5 --- tests/test_text_generation.py | 288 ++++++++++++++++------------------ 1 file changed, 133 insertions(+), 155 deletions(-) diff --git a/tests/test_text_generation.py b/tests/test_text_generation.py index 9be473f0..b477c6e4 100644 --- a/tests/test_text_generation.py +++ b/tests/test_text_generation.py @@ -6,6 +6,8 @@ import unittest from unittest.mock import patch, MagicMock +import pytest + from dw.tasks.text_generation import ( generate_text, _DEFAULT_MODEL, @@ -15,49 +17,64 @@ ) +# Each dw target is patched once here; a class opts in with usefixtures and +# reads the mock off self, so no test carries a patch of its own +@pytest.fixture +def fake_pipeline(request): + """hf_pipeline replaced for one test, as self.mock_pipeline.""" + with patch("dw.tasks.text_generation.hf_pipeline") as mock_pipeline: + request.instance.mock_pipeline = mock_pipeline + yield mock_pipeline + + +@pytest.fixture +def fake_hub_download(request): + """hf_hub_download replaced for one test, as self.download.""" + with patch("dw.tasks.text_generation.hf_hub_download") as download: + request.instance.download = download + yield download + + +@pytest.mark.usefixtures("fake_pipeline") class TestTextGeneration(unittest.TestCase): """Tests for the generate_text function.""" - @patch("dw.tasks.text_generation.hf_pipeline") - def test_returns_generated_string(self, mock_pipeline): + def test_returns_generated_string(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "an expanded detailed prompt"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe result = generate_text("a cat", device="cpu") self.assertEqual(result, "an expanded detailed prompt") - mock_pipeline.assert_called_once() + self.mock_pipeline.assert_called_once() - @patch("dw.tasks.text_generation.hf_pipeline") - def test_uses_default_model(self, mock_pipeline): + def test_uses_default_model(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "output"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe generate_text("test", device="cpu") - call_kwargs = mock_pipeline.call_args + call_kwargs = self.mock_pipeline.call_args self.assertEqual(call_kwargs[1]["model"], _DEFAULT_MODEL) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_custom_model_name(self, mock_pipeline): + def test_custom_model_name(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "output"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe generate_text( "test", device="cpu", model_name="meta-llama/Llama-3.2-1B-Instruct" ) - call_kwargs = mock_pipeline.call_args + call_kwargs = self.mock_pipeline.call_args self.assertEqual(call_kwargs[1]["model"], "meta-llama/Llama-3.2-1B-Instruct") - @patch("dw.tasks.text_generation.hf_pipeline") - def test_system_prompt_included(self, mock_pipeline): + def test_system_prompt_included(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "output"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe generate_text( "a cat", @@ -75,11 +92,10 @@ def test_system_prompt_included(self, mock_pipeline): self.assertEqual(messages[1]["role"], "user") self.assertEqual(messages[1]["content"], "a cat") - @patch("dw.tasks.text_generation.hf_pipeline") - def test_no_system_prompt(self, mock_pipeline): + def test_no_system_prompt(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "output"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe generate_text("a cat", device="cpu") @@ -88,31 +104,28 @@ def test_no_system_prompt(self, mock_pipeline): self.assertEqual(len(messages), 1) self.assertEqual(messages[0]["role"], "user") - @patch("dw.tasks.text_generation.hf_pipeline") - def test_max_new_tokens(self, mock_pipeline): + def test_max_new_tokens(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "output"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe generate_text("test", device="cpu", max_new_tokens=200) call_args = pipe.call_args self.assertEqual(call_args[1]["max_new_tokens"], 200) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_strips_whitespace(self, mock_pipeline): + def test_strips_whitespace(self): pipe = MagicMock() pipe.return_value = [{"generated_text": " some text with spaces \n"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe result = generate_text("test", device="cpu") self.assertEqual(result, "some text with spaces") - @patch("dw.tasks.text_generation.hf_pipeline") - def test_chat_template_kwargs_reach_the_text_template(self, mock_pipeline): + def test_chat_template_kwargs_reach_the_text_template(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "out"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe generate_text( "a cat", @@ -125,11 +138,10 @@ def test_chat_template_kwargs_reach_the_text_template(self, mock_pipeline): pipe.call_args[1]["tokenizer_encode_kwargs"], {"enable_thinking": False} ) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_chat_template_kwargs_reach_the_vision_template(self, mock_pipeline): + def test_chat_template_kwargs_reach_the_vision_template(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "out"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe generate_text( "a cat", @@ -141,11 +153,10 @@ def test_chat_template_kwargs_reach_the_vision_template(self, mock_pipeline): self.assertIs(pipe.call_args[1]["enable_thinking"], False) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_no_chat_template_kwargs_passes_nothing_extra(self, mock_pipeline): + def test_no_chat_template_kwargs_passes_nothing_extra(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "out"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe generate_text("a cat", device="cpu", model_name="m/no-chat-kwargs") @@ -157,21 +168,19 @@ def test_chat_template_kwargs_must_be_a_mapping(self): "a cat", device="cpu", chat_template_kwargs=["enable_thinking"] ) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_revision_reaches_the_loader(self, mock_pipeline): + def test_revision_reaches_the_loader(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "out"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe generate_text("a cat", device="cpu", model_name="m/revision", revision="a" * 40) - self.assertEqual(mock_pipeline.call_args[1]["revision"], "a" * 40) + self.assertEqual(self.mock_pipeline.call_args[1]["revision"], "a" * 40) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_a_list_of_images_becomes_one_part_each_in_order(self, mock_pipeline): + def test_a_list_of_images_becomes_one_part_each_in_order(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "out"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe generate_text( "edit", @@ -186,13 +195,12 @@ def test_a_list_of_images_becomes_one_part_each_in_order(self, mock_pipeline): ["https://example.com/1.png", "https://example.com/2.png"], ) self.assertEqual(user[2], {"type": "text", "text": "edit"}) - self.assertEqual(mock_pipeline.call_args[0][0], "image-text-to-text") + self.assertEqual(self.mock_pipeline.call_args[0][0], "image-text-to-text") - @patch("dw.tasks.text_generation.hf_pipeline") - def test_an_empty_or_all_null_image_list_is_the_text_path(self, mock_pipeline): + def test_an_empty_or_all_null_image_list_is_the_text_path(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "out"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe for empty in ([], [None]): generate_text( @@ -201,10 +209,9 @@ def test_an_empty_or_all_null_image_list_is_the_text_path(self, mock_pipeline): model_name=f"m/empty-{len(empty)}", image=empty, ) - self.assertEqual(mock_pipeline.call_args[0][0], "text-generation") + self.assertEqual(self.mock_pipeline.call_args[0][0], "text-generation") - @patch("dw.tasks.text_generation.hf_pipeline") - def test_a_non_scalar_chat_template_value_is_refused(self, mock_pipeline): + def test_a_non_scalar_chat_template_value_is_refused(self): with self.assertRaisesRegex(ValueError, "enable_thinking"): generate_text( "a cat", @@ -212,14 +219,11 @@ def test_a_non_scalar_chat_template_value_is_refused(self, mock_pipeline): model_name="m/nonscalar", chat_template_kwargs={"enable_thinking": {"a": 1}}, ) - mock_pipeline.assert_not_called() + self.mock_pipeline.assert_not_called() - @patch("dw.tasks.text_generation.hf_pipeline") - def test_a_reserved_chat_template_key_is_refused_before_any_call( - self, mock_pipeline - ): + def test_a_reserved_chat_template_key_is_refused_before_any_call(self): pipe = MagicMock() - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe for key in ("images", "tools", "tokenizer_encode_kwargs"): with self.assertRaisesRegex(ValueError, f"'{key}'.*not a chat-template"): generate_text( @@ -238,11 +242,10 @@ def test_a_reserved_chat_template_key_is_refused_before_any_call( ) pipe.assert_not_called() - @patch("dw.tasks.text_generation.hf_pipeline") - def test_a_legitimate_setting_passes_on_both_paths(self, mock_pipeline): + def test_a_legitimate_setting_passes_on_both_paths(self): pipe = MagicMock() pipe.return_value = [{"generated_text": "out"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe settings = {"enable_thinking": False, "effort": "low", "n": 2, "t": 0.5} generate_text( @@ -259,6 +262,7 @@ def test_a_legitimate_setting_passes_on_both_paths(self, mock_pipeline): self.assertIs(pipe.call_args[1]["enable_thinking"], False) +@pytest.mark.usefixtures("fake_pipeline") class TestTextGenerationWithImage(unittest.TestCase): """An image switches generation onto a vision-language pipeline.""" @@ -273,33 +277,29 @@ def _mock_pipe(self, mock_pipeline): mock_pipeline.return_value = pipe return pipe - @patch("dw.tasks.text_generation.hf_pipeline") - def test_image_selects_the_vision_pipeline(self, mock_pipeline): - self._mock_pipe(mock_pipeline) + def test_image_selects_the_vision_pipeline(self): + self._mock_pipe(self.mock_pipeline) generate_text("describe it", device="cpu", image=self._make_image()) - self.assertEqual(mock_pipeline.call_args[0][0], "image-text-to-text") + self.assertEqual(self.mock_pipeline.call_args[0][0], "image-text-to-text") - @patch("dw.tasks.text_generation.hf_pipeline") - def test_no_image_stays_on_text_generation(self, mock_pipeline): - self._mock_pipe(mock_pipeline) + def test_no_image_stays_on_text_generation(self): + self._mock_pipe(self.mock_pipeline) generate_text("a cat", device="cpu") - self.assertEqual(mock_pipeline.call_args[0][0], "text-generation") + self.assertEqual(self.mock_pipeline.call_args[0][0], "text-generation") - @patch("dw.tasks.text_generation.hf_pipeline") - def test_image_uses_the_vision_default_model(self, mock_pipeline): - self._mock_pipe(mock_pipeline) + def test_image_uses_the_vision_default_model(self): + self._mock_pipe(self.mock_pipeline) generate_text("describe it", device="cpu", image=self._make_image()) - self.assertEqual(mock_pipeline.call_args[1]["model"], DEFAULT_VISION_MODEL) + self.assertEqual(self.mock_pipeline.call_args[1]["model"], DEFAULT_VISION_MODEL) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_vision_content_is_typed_parts(self, mock_pipeline): - pipe = self._mock_pipe(mock_pipeline) + def test_vision_content_is_typed_parts(self): + pipe = self._mock_pipe(self.mock_pipeline) image = self._make_image() generate_text( @@ -321,9 +321,8 @@ def test_vision_content_is_typed_parts(self, mock_pipeline): ], ) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_string_image_is_passed_as_a_url(self, mock_pipeline): - pipe = self._mock_pipe(mock_pipeline) + def test_string_image_is_passed_as_a_url(self): + pipe = self._mock_pipe(self.mock_pipeline) generate_text("describe it", device="cpu", image="https://example.com/cat.jpg") @@ -332,9 +331,8 @@ def test_string_image_is_passed_as_a_url(self, mock_pipeline): content[0], {"type": "image", "url": "https://example.com/cat.jpg"} ) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_sampling_is_disabled_through_generate_kwargs(self, mock_pipeline): - pipe = self._mock_pipe(mock_pipeline) + def test_sampling_is_disabled_through_generate_kwargs(self): + pipe = self._mock_pipe(self.mock_pipeline) generate_text("describe it", device="cpu", image=self._make_image()) @@ -343,9 +341,8 @@ def test_sampling_is_disabled_through_generate_kwargs(self, mock_pipeline): self.assertIs(pipe.call_args[1]["generate_kwargs"]["do_sample"], False) self.assertNotIn("do_sample", pipe.call_args[1]) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_vision_defaults_to_a_repetition_penalty(self, mock_pipeline): - pipe = self._mock_pipe(mock_pipeline) + def test_vision_defaults_to_a_repetition_penalty(self): + pipe = self._mock_pipe(self.mock_pipeline) generate_text("describe it", device="cpu", image=self._make_image()) @@ -356,9 +353,8 @@ def test_vision_defaults_to_a_repetition_penalty(self, mock_pipeline): _VISION_REPETITION_PENALTY, ) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_repetition_penalty_is_overridable(self, mock_pipeline): - pipe = self._mock_pipe(mock_pipeline) + def test_repetition_penalty_is_overridable(self): + pipe = self._mock_pipe(self.mock_pipeline) generate_text( "describe it", @@ -371,9 +367,8 @@ def test_repetition_penalty_is_overridable(self, mock_pipeline): pipe.call_args[1]["generate_kwargs"]["repetition_penalty"], 1.2 ) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_text_path_takes_no_repetition_penalty(self, mock_pipeline): - pipe = self._mock_pipe(mock_pipeline) + def test_text_path_takes_no_repetition_penalty(self): + pipe = self._mock_pipe(self.mock_pipeline) generate_text("a cat", device="cpu") @@ -381,9 +376,8 @@ def test_text_path_takes_no_repetition_penalty(self, mock_pipeline): # the output of every workflow already using them self.assertNotIn("generate_kwargs", pipe.call_args[1]) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_text_and_vision_variants_cache_separately(self, mock_pipeline): - self._mock_pipe(mock_pipeline) + def test_text_and_vision_variants_cache_separately(self): + self._mock_pipe(self.mock_pipeline) # Same model name under both pipelines - they are not interchangeable, # so the cache must not hand one back for the other @@ -392,11 +386,12 @@ def test_text_and_vision_variants_cache_separately(self, mock_pipeline): "a cat", device="cpu", model_name="some/model", image=self._make_image() ) - self.assertEqual(mock_pipeline.call_count, 2) - tasks = [call[0][0] for call in mock_pipeline.call_args_list] + self.assertEqual(self.mock_pipeline.call_count, 2) + tasks = [call[0][0] for call in self.mock_pipeline.call_args_list] self.assertEqual(tasks, ["text-generation", "image-text-to-text"]) +@pytest.mark.usefixtures("fake_pipeline") class TestGenerateKwargsPassThrough(unittest.TestCase): """Anything generate() understands can be set from a workflow.""" @@ -411,9 +406,8 @@ def _mock_pipe(self, mock_pipeline): mock_pipeline.return_value = pipe return pipe - @patch("dw.tasks.text_generation.hf_pipeline") - def test_reaches_the_text_pipeline(self, mock_pipeline): - pipe = self._mock_pipe(mock_pipeline) + def test_reaches_the_text_pipeline(self): + pipe = self._mock_pipe(self.mock_pipeline) generate_text( "a cat", device="cpu", generate_kwargs={"no_repeat_ngram_size": 25} @@ -421,9 +415,8 @@ def test_reaches_the_text_pipeline(self, mock_pipeline): self.assertEqual(pipe.call_args[1]["no_repeat_ngram_size"], 25) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_reaches_the_vision_pipeline(self, mock_pipeline): - pipe = self._mock_pipe(mock_pipeline) + def test_reaches_the_vision_pipeline(self): + pipe = self._mock_pipe(self.mock_pipeline) generate_text( "describe it", @@ -436,9 +429,8 @@ def test_reaches_the_vision_pipeline(self, mock_pipeline): pipe.call_args[1]["generate_kwargs"]["no_repeat_ngram_size"], 25 ) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_overrides_what_the_task_decided(self, mock_pipeline): - pipe = self._mock_pipe(mock_pipeline) + def test_overrides_what_the_task_decided(self): + pipe = self._mock_pipe(self.mock_pipeline) generate_text( "describe it", @@ -452,6 +444,7 @@ def test_overrides_what_the_task_decided(self, mock_pipeline): ) +@pytest.mark.usefixtures("fake_pipeline") class TestGenerateKwargsAreFiltered(unittest.TestCase): """generate_kwargs reach the pipeline call, so they are checked first. @@ -466,9 +459,8 @@ def _mock_pipe(self, mock_pipeline): mock_pipeline.return_value = pipe return pipe - @patch("dw.tasks.text_generation.hf_pipeline") - def test_trusted_workflows_may_name_remote_generate_code(self, mock_pipeline): - pipe = self._mock_pipe(mock_pipeline) + def test_trusted_workflows_may_name_remote_generate_code(self): + pipe = self._mock_pipe(self.mock_pipeline) settings = {"custom_generate": "someone/decoder", "trust_remote_code": True} with patch("dw.trust.workflows_are_trusted", return_value=True): generate_text( @@ -477,8 +469,7 @@ def test_trusted_workflows_may_name_remote_generate_code(self, mock_pipeline): self.assertEqual(pipe.call_args[1]["custom_generate"], "someone/decoder") self.assertIs(pipe.call_args[1]["trust_remote_code"], True) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_a_structural_key_is_refused_before_any_load(self, mock_pipeline): + def test_a_structural_key_is_refused_before_any_load(self): for key in ( "tokenizer_encode_kwargs", "prefix", @@ -495,10 +486,9 @@ def test_a_structural_key_is_refused_before_any_load(self, mock_pipeline): image=image, generate_kwargs={key: "x"}, ) - mock_pipeline.assert_not_called() + self.mock_pipeline.assert_not_called() - @patch("dw.tasks.text_generation.hf_pipeline") - def test_a_nested_value_is_refused(self, mock_pipeline): + def test_a_nested_value_is_refused(self): for bad in ({"a": 1}, [{"a": 1}], [[1]]): with self.assertRaisesRegex(ValueError, "sequence_bias"): generate_text( @@ -509,11 +499,10 @@ def test_a_nested_value_is_refused(self, mock_pipeline): ) with self.assertRaisesRegex(ValueError, "must be an object"): generate_text("a cat", device="cpu", generate_kwargs=["top_p"]) - mock_pipeline.assert_not_called() + self.mock_pipeline.assert_not_called() - @patch("dw.tasks.text_generation.hf_pipeline") - def test_legitimate_settings_pass_on_both_paths(self, mock_pipeline): - pipe = self._mock_pipe(mock_pipeline) + def test_legitimate_settings_pass_on_both_paths(self): + pipe = self._mock_pipe(self.mock_pipeline) settings = {"no_repeat_ngram_size": 3, "stop_strings": ["\n\n"]} generate_text( @@ -534,6 +523,7 @@ def test_legitimate_settings_pass_on_both_paths(self, mock_pipeline): self.assertEqual(forwarded["stop_strings"], ["\n\n"]) +@pytest.mark.usefixtures("fake_pipeline") class TestTextGenerationLoadDevice(unittest.TestCase): """How the model gets onto the device it runs on. @@ -551,23 +541,21 @@ def _mock_pipe(mock_pipeline): mock_pipeline.return_value = pipe return pipe - @patch("dw.tasks.text_generation.hf_pipeline") - def test_mps_loads_on_cpu_then_moves(self, mock_pipeline): - self._mock_pipe(mock_pipeline) + def test_mps_loads_on_cpu_then_moves(self): + self._mock_pipe(self.mock_pipeline) generate_text("test", device="mps") - kwargs = mock_pipeline.call_args[1] + kwargs = self.mock_pipeline.call_args[1] self.assertEqual(kwargs["device"], "mps") self.assertNotIn("device_map", kwargs) - @patch("dw.tasks.text_generation.hf_pipeline") - def test_other_devices_still_use_device_map(self, mock_pipeline): - self._mock_pipe(mock_pipeline) + def test_other_devices_still_use_device_map(self): + self._mock_pipe(self.mock_pipeline) generate_text("test", device="cuda") - kwargs = mock_pipeline.call_args[1] + kwargs = self.mock_pipeline.call_args[1] self.assertEqual(kwargs["device_map"], "cuda") self.assertNotIn("device", kwargs) @@ -586,6 +574,7 @@ def _spec(**over): return {"from_hub": spec} +@pytest.mark.usefixtures("fake_pipeline", "fake_hub_download") class TestSystemPromptFromHub(unittest.TestCase): def setUp(self): meta = patch("dw.tasks.text_generation.get_hf_file_metadata") @@ -601,28 +590,25 @@ def test_a_string_or_none_passes_through(self): self.assertEqual(resolve_system_prompt("plain"), "plain") self.assertIsNone(resolve_system_prompt(None)) - @patch("dw.tasks.text_generation.hf_hub_download") - def test_the_key_s_text_is_returned_from_the_pinned_file(self, download): - download.return_value = self.json_path + def test_the_key_s_text_is_returned_from_the_pinned_file(self): + self.download.return_value = self.json_path self.assertEqual(resolve_system_prompt(_spec()), "expand the request") - download.assert_called_once_with( + self.download.assert_called_once_with( "black-forest-labs/flux-3-image-klein-pup", "system_prompts.json", revision=PIN, token=False, ) - @patch("dw.tasks.text_generation.hf_hub_download") - def test_a_missing_key_names_the_keys_the_file_has(self, download): - download.return_value = self.json_path + def test_a_missing_key_names_the_keys_the_file_has(self): + self.download.return_value = self.json_path with self.assertRaises(ValueError) as caught: resolve_system_prompt(_spec(key="t2v")) self.assertIn("i2i", str(caught.exception)) self.assertIn("t2i", str(caught.exception)) - @patch("dw.tasks.text_generation.hf_hub_download") - def test_a_value_that_is_not_text_is_refused(self, download): - download.return_value = self.json_path + def test_a_value_that_is_not_text_is_refused(self): + self.download.return_value = self.json_path with self.assertRaises(ValueError): resolve_system_prompt(_spec(key="n")) @@ -658,18 +644,16 @@ def test_a_missing_or_non_string_repo_is_refused(self): resolve_system_prompt(bad) self.assertIn("system_prompt.from_hub.repo", str(caught.exception)) - @patch("dw.tasks.text_generation.hf_hub_download") - def test_the_hub_is_asked_anonymously_before_the_download(self, download): - download.return_value = self.json_path + def test_the_hub_is_asked_anonymously_before_the_download(self): + self.download.return_value = self.json_path order = MagicMock() order.attach_mock(self.metadata, "metadata") - order.attach_mock(download, "download") + order.attach_mock(self.download, "download") resolve_system_prompt(_spec()) self.assertEqual([c[0] for c in order.mock_calls], ["metadata", "download"]) self.assertIs(self.metadata.call_args.kwargs["token"], False) - @patch("dw.tasks.text_generation.hf_hub_download") - def test_a_file_the_hub_will_not_serve_publicly_is_refused(self, download): + def test_a_file_the_hub_will_not_serve_publicly_is_refused(self): from huggingface_hub.errors import HfHubHTTPError class _Gated(HfHubHTTPError): @@ -682,39 +666,35 @@ def __init__(self, message): resolve_system_prompt(_spec()) self.assertIn("public file", str(caught.exception)) self.assertIs(caught.exception.__cause__, err) - download.assert_not_called() + self.download.assert_not_called() - @patch("dw.tasks.text_generation.hf_hub_download") - def test_an_oversize_file_is_refused_before_downloading(self, download): + def test_an_oversize_file_is_refused_before_downloading(self): self.metadata.return_value = MagicMock(size=1024 * 1024 + 1) with self.assertRaises(ValueError): resolve_system_prompt(_spec()) - download.assert_not_called() + self.download.assert_not_called() - @patch("dw.tasks.text_generation.hf_hub_download") - def test_a_file_of_unknown_size_is_refused_before_downloading(self, download): + def test_a_file_of_unknown_size_is_refused_before_downloading(self): self.metadata.return_value = MagicMock(size=None) with self.assertRaisesRegex(ValueError, "no size"): resolve_system_prompt(_spec()) - download.assert_not_called() + self.download.assert_not_called() def test_a_non_string_key_is_refused(self): for key in (["t2i"], {"a": 1}, 3): with self.assertRaises(ValueError): resolve_system_prompt(_spec(key=key)) - @patch("dw.tasks.text_generation.hf_hub_download") - def test_the_download_never_carries_the_server_token(self, download): - download.return_value = self.json_path + def test_the_download_never_carries_the_server_token(self): + self.download.return_value = self.json_path resolve_system_prompt(_spec()) - self.assertIs(download.call_args.kwargs["token"], False) + self.assertIs(self.download.call_args.kwargs["token"], False) - @patch("dw.tasks.text_generation.hf_hub_download") - def test_an_oversize_file_is_refused_unread(self, download): + def test_an_oversize_file_is_refused_unread(self): big = os.path.join(self.dir, "big.txt") with open(big, "wb") as f: f.truncate(1024 * 1024 + 1) - download.return_value = big + self.download.return_value = big with patch("builtins.open", side_effect=AssertionError("read")): with self.assertRaises(ValueError): resolve_system_prompt(_spec(file="big.txt", key=None)) @@ -734,13 +714,11 @@ def test_an_unknown_field_or_shape_is_refused(self): with self.assertRaises(ValueError): resolve_system_prompt(bad) - @patch("dw.tasks.text_generation.hf_pipeline") - @patch("dw.tasks.text_generation.hf_hub_download") - def test_generate_text_sends_the_resolved_text(self, download, mock_pipeline): - download.return_value = self.json_path + def test_generate_text_sends_the_resolved_text(self): + self.download.return_value = self.json_path pipe = MagicMock() pipe.return_value = [{"generated_text": "out"}] - mock_pipeline.return_value = pipe + self.mock_pipeline.return_value = pipe generate_text( "a dog", device="cpu", model_name="m/from-hub", system_prompt=_spec()