[Jayce-Ping]

Summary

Training now feeds every dataset prompt to the model verbatim. Prompt enhancement moves out of the
adapters into a separate offline step, ff-prompt-enhance, which rewrites a dataset's prompts with an
enhancer model and writes a new dataset. Whether to enhance, with which model, and how many rewrites
per prompt are decided before training, so training needs no prompt-enhancement logic or checks.

The PR also ships GenEval prompts rewritten by the FLUX.3 and Qwen-Image-2.1 enhancers, with GRPO
examples that train on them.

Changes

• Adapters: remove FLUX.2 caption upsampling (caption_upsample_temperature) and LTX2 prompt
  enhancement (system_prompt, prompt_enhancement_seed, and the image argument of the LTX2 I2AV
  encode_prompt), plus their keys in the LTX2 GRPO examples.

• flow_factory.prompt_enhance: BasePromptEnhancer, a registry, and six enhancers:

  enhancer	Model
  flux3	FLUX.3 prompt upsampler black-forest-labs/flux-3-image-klein-pup (its t2i / i2i system prompts, reference images in order)
  flux2	FLUX.2-dev Mistral3 text encoder through Flux2Pipeline.upsample_prompt; only the text components are loaded
  ltx2	LTX-2 Gemma3 text encoder, or the LTX-2.5 Gemma4 prompt_enhancer, through LTX2Pipeline.enhance_prompt; needs a Gemma3 processor, e.g. diffusers/LTX-2.3-Diffusers
  qwen-image-2.1-t2i	Qwen/Qwen-Image-2.1-PE-T2I, prompt-only rows
  qwen-image-2.1-i2i	Qwen/Qwen-Image-2.1-PE-I2I, image-editing rows (one or more references)
  openai-compatible	Any vLLM / OpenAI-style chat-completions server

  The two Qwen-Image-2.1 enhancers follow the official
  prompt_rewrite recipe:

  • the checkpoint's system_prompt.txt;
  • images first, each downscaled to at most 1024×1024 pixels;
  • thinking enabled;
  • each task's sampling profile: T2I presence_penalty: 1.5 and max_new_tokens: 16256, edit no
    presence penalty and max_new_tokens: 24000. transformers has no presence penalty, so the
    enhancers apply a vLLM-style logits processor.

  The enhanced prompt is the answer's rewritten_prompt. The wh_ratio / ratio_follow canvas
  recommendation is dropped because training renders at the configured resolution. Two deviations
  from the official transformers script:

  • generation also stops at <|im_end|>, because the checkpoints' generation_config lists only
    <|endoftext|>;
  • a sampled rewrite that loops until max_new_tokens, or answers without a parseable
    rewritten_prompt, is resampled from the continuing RNG stream of its seed, up to three
    attempts. The official script keeps the raw text with parse_ok: false instead.

• Runner: each rank enhances a disjoint row shard. The output split has num_variants rows per
  source row, with variants adjacent; prompt is the enhanced text, source_prompt and pe_variant
  are added, and all other fields are kept. Media directories are symlinked, and
  prompt_enhance_manifest.json records the run.

  • Every (row, variant) gets the seed seed + variant * num_rows + row, independent of
    batch_size and the process count.
  • Per-rank part files make runs resumable. Each starts with a fingerprint of the generation settings
    and source rows; resuming under a different fingerprint fails, while raising num_variants
    extends a finished run.

• CLI: ff-prompt-enhance / flow-factory-prompt-enhance, sharing the ff-train launcher.

• GenEval datasets and examples: both datasets cover the GenEval test split (553 rows) and the
  first 1024 train rows, which is the slice max_dataset_size: 1024 trains on. They keep the
  include / exclude / tag fields the reward scores, and each has a README with the enhancer
  revision.

  Dataset	Produced by	Training example
  dataset/geneval_flux3_pe	examples/prompt_enhance/flux3_geneval.yaml (upstream greedy decoding; one degenerate row resampled, see its README)	examples/grpo/lora/flux3/geneval.yaml (the FLUX.3 adapter from [model,test,docs] feat: add FLUX.3 adapter #247, now on main)
  dataset/geneval_qwen_image_21_pe	examples/prompt_enhance/qwen_image_2_1_geneval.yaml (official T2I sampling profile)	examples/grpo/lora/qwen_image_2_1/geneval.yaml

  Both training examples are examples/grpo/lora/sd3_5/geneval.yaml with only the model and the
  dataset swapped. The Qwen-Image-2.1 example also sets per_device_batch_size: 1, which its
  adapter requires.

• Docs and harness: guidance/prompt_enhance.md, examples/prompt_enhance/*.yaml, README,
  AGENTS.md, and architecture updates; a risk route and eval case for
  src/flow_factory/prompt_enhance/.

Fixes from GPU validation

Commit	Fix
[prompt_enhance,test] fix: left-pad FLUX.2 upsampling prompts	Flux2Pipeline.upsample_prompt pads every prompt to 2048 tokens before generate, and the FLUX.2-dev Mistral tokenizer pads on the right, so the first new token was predicted from a pad position. With transformers 5.14 the outputs ignored the prompt and echoed the system message's numbered rules. The enhancer now sets left padding.
[prompt_enhance,test,docs] fix: load the undeclared LTX-2.3 Diffusers processor	diffusers/LTX-2.3-Diffusers ships its Gemma3 processor under processor/ without declaring it in model_index.json, so the enhancer failed. It now loads processor/processor_config.json when the checkpoint ships it and still fails fast when it has neither (as Lightricks/LTX-2 does). The example points at diffusers/LTX-2.3-Diffusers.

Compatibility

Configs that still set caption_upsample_temperature, system_prompt, or prompt_enhancement_seed
keep loading. The keys land in extra_kwargs with the existing unknown-key warning and no longer have
an effect. To train on enhanced prompts, run ff-prompt-enhance and point dataset_dir at its output.
The removed FLUX.2 path appended the list returned by upsample_prompt into the prompt list, and it
inherited the right-padding problem above, so it did not produce usable prompts before this change
either.

Re-run pip install -e . to register the new console scripts.

Validation

CPU

• tests/prompt_enhance, 33 tests:
  • runner end to end, and GeneralDataset consuming the output unchanged;
  • two-rank sharding, resume, fingerprint mismatch, and per-row seeds;
  • enhancer wiring against fakes, with the upstream diffusers / transformers signatures pinned,
    including FLUX.2 left padding and loading an undeclared LTX2 processor/ subfolder;
  • for Qwen-Image-2.1: the official messages, sampling profile, and stop tokens; vLLM presence-penalty
    semantics; image downscaling; answer parsing; resampling within a row's seed; and rejecting rows
    of the other task;
  • CLI launch and parsing of the example configs.
• Full CPU suite on this branch after merging main (which now carries [model,test,docs] feat: add FLUX.3 adapter #247 and the Diffusers pin
  7ecb1e1b): 3042 passed, 50 skipped.
• Real-process runs:
  • a single process through prompt_enhance_cli;
  • three gloo ranks with a real barrier, whose merged output equals the single-process output;
  • the OpenAI-compatible enhancer against a local HTTP server (per-row seeds, auth header, PNG
    data-URL images).
• scripts/agent_scope.py: R2. Harness validator and docstring audit are clean.

GPU (8x H20, Python 3.12, torch 2.10.0+cu129, transformers 5.14.1)

Flux2Pipeline.upsample_prompt and LTX2Pipeline.enhance_prompt are identical at the repository's
Diffusers pin and at the FLUX.3 pin 7ecb1e1b; the runs used the latter. Every run went through
prompt_enhance_cli.

Run	Result
flux3, flux-3-image-klein-pup, 8 ranks, num_variants: 2, sampling, 16 image-editing + 16 prompt-only rows	Pass, 109 s. One part file per rank and split, 32 rows per split with adjacent ordered variants; 16/16 rows have distinct variants; no <think> or chat-template remnants and no truncated captions; image rows describe their references; the Chinese prompt yields a faithful English caption. 233-429 tokens.
Resume and fingerprint	Pass. A run killed after the train split and restarted reused the train parts, finished the test split, and produced output byte-identical to the uninterrupted run. Rerunning into the same output_dir with seed: 43 raised ValueError: ... was written with different ['seed'].
flux2, local FLUX.2-dev text encoder, 2 ranks	Pass after the left-padding fix (before it, captions were unrelated to their prompts). 24-472 tokens.
ltx2, diffusers/LTX-2.3-Diffusers	Pass after the processor fix, both from a local snapshot and from the Hub repo id. Text rows produce T2V captions and image rows I2V captions that describe the reference; the manifest records dedicated_prompt_enhancer: false. 70-158 tokens. Before the fix the same checkpoint failed fast at load.
qwen-image-2.1-t2i / -i2i against the official run_transformers.py reference, on the official example inputs with Flow-Factory's per-row seeds	6 of 7 rewrites are identical: all 3 T2I prompts (one Chinese), and 3 of 4 edits, including a 1088×1440 source that is downscaled and a Chinese two-image <image1> / <image2> composite. On the 7th, "Translate into Hindi." with a 3840×2160 source and seed 43, both implementations generate the same loop until the 24000-token limit. The reference keeps the raw thinking text with parse_ok: false; the enhancer logged the failed attempt, resampled, and wrote a grounded rewrite that replaces the image's Japanese title.
qwen-image-2.1-t2i, 8 ranks, 8 pickscore prompts × 2 variants	Pass, 176 s. Variants differ; English rewrites of 363-1131 words; quoted text such as "philo is a weird" survives verbatim; the manifest records the effective profile and the checkpoint system-prompt hash.
qwen-image-2.1-i2i, 8 ranks, 8 sharegpt4o_image_mini editing rows × 2 variants	Pass, 293 s. Variants differ; edit instructions of 114-443 words describe the actual reference contents; images is symlinked and the image field kept.
dataset/geneval_qwen_image_21_pe, 48 ranks on 6 nodes	1577 rows in 24 min, every row on its first attempt. Rewrites run 207-1039 words; some class names are paraphrased (hair dryer, sofa, doughnut), which the detector-based reward does not depend on.
FLUX.3 GenEval GRPO, examples/grpo/lora/flux3/geneval.yaml with #247, #248, and #249 merged; 8 GPUs, ZeRO-2, 10 epochs	Pass, about 5.6 min per epoch. Evaluation on the 553 test prompts is deterministic (0.7027 in two launches) and stays flat over this short run (0.7027 at epoch 0, 0.7021 at epoch 10; train reward 0.60-0.76 per epoch). This run used the FLUX.3 adapter from before #247's on-policy fix, so the first-step ratios only rounded to 1; with the fix they are bit-exact (see #247). No sample was clipped, and the LoRA moved: 456 of 516 tensors changed and lora_B reached 0.0021 after 20 steps.
Qwen-Image-2.1 GenEval GRPO, examples/grpo/lora/qwen_image_2_1/geneval.yaml; 4 GPUs, 1 epoch, a reduced per-epoch sample count	Pass. 34 GB per GPU without gradient checkpointing; ratio 1 on the first optimizer step and 0.9998-1 on the second; 131 s per epoch at 36 samples per GPU.

Every observed length fits the matching training adapter's default max_sequence_length (FLUX.3
8192, FLUX.2 512, LTX2 1024); the Qwen-Image-2.1 adapter does not truncate prompts. FLUX.2's longest
caption is close to its limit, as the guide warns. The Qwen-Image-2.1 rewriters generate about 34
tokens/s per H20 here, because the environment lacks the flash-linear-attention kernels and
transformers falls back to torch.

Not run:

• openai-compatible against a real vLLM server (covered by the local HTTP server above), and the
  official run_vllm.py path for the Qwen-Image-2.1 rewriters.
• The LTX-2.5 dedicated prompt_enhancer path: no Diffusers-format LTX-2.5 checkpoint is published.
• GRPO regression runs for FLUX.2 and LTX2. The removed adapter branches were disabled by default,
  and the LTX2 examples set them to null.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.