[model,test,docs] feat: add FLUX.3 adapter

[Jayce-Ping]

Summary

Add Flux3Adapter (model_type: "flux3") supporting Black Forest Labs' newly released FLUX.3 Image Klein family (black-forest-labs/flux-3-image-klein-base and black-forest-labs/flux-3-image-klein) from huggingface/diffusers@7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941 (huggingface/diffusers#15046).

Key Changes

• Submodule Bump (diffusers): Update the bundled diffusers submodule to 7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941 (Flux3Transformer2DModel, Flux3ModularPipeline, Flux3DistilledModularPipeline, Flux3ImageProcessor). The new pin requires Python 3.11+ and huggingface-hub>=1.32; the README install note and .agents/knowledge/dependencies.md say so, and QWEN_IMAGE_21_DIFFUSERS_COMMIT now names the same pin (the Qwen-Image 2.1 functions it documents as copied are unchanged between the two commits).
• Component Runtime (Flux3Adapter):
  • Wrap Flux3ModularPipeline / Flux3DistilledModularPipeline with ModularPipelineRuntime, rewrite local snapshot ComponentSpec paths when loading from disk, and declare component_load_dtype_defaults = {"transformers": torch.bfloat16, "text_encoders": torch.bfloat16, "vae": torch.float32}.
  • Opt into supports_fsdp2_cpu_efficient_loading = True and ddp_find_unused_parameters = True (required because mode_blocks.image_cond receives no gradient during text-only T2I steps).
  • Override _gradient_checkpointing_units() to return txt_mode_blocks -> mode_blocks.image -> mode_blocks.image_cond -> single_transformer_blocks in exact forward execution order.
• Prompt & Condition Image Encoding:
  • Implement Qwen3-VL text encoding with the official 80-token window right-padding recipe, explicit padding position_ids = 1, and 8-layer hidden-state stacking (4, 8, 12, 16, 20, 24, 28, 32) -> 20480 channels. encode_prompt defaults to the official max_sequence_length=8192, so preprocessing no longer truncates long captions at 512 tokens.
  • Preserve per-sample 80-token windows without cross-sample padding pollution when batch prompts differ in window length.
  • Encode multi-reference condition images with AutoencoderKLFlux2 (sample_mode="argmax"), image_rope_scale = 0.5, and negative l coordinates -(1 + idx) (-1, -2, ...).
• Unified Base (True CFG) & Step-Distilled (Flow-Map) Forward / Inference:
  • Automatically detect flow_map=True (is_distilled) to pass target_timestep = t_next / 1000 and narrow _resolve_pipeline_io_contract() to NegativePromptPolicy.UNSUPPORTED.
  • Normalize img_ids, txt_ids, and image_cond_ids to 2D (seq_len, 4) for Flux3PosEmbed.
• Offline SFT / DPO Support:
  • Extend src/flow_factory/models/flux/_output.py with encode_flux3_output_images (sample_mode="sample"), prepare_flux3_condition_latents, and standalone 2x2 patchify/unpatchify helpers.
  • Set offline_training_forward_overrides = MappingProxyType({"guidance_scale": 1.0}).
  • Register flux3 in the offline smoke catalog (text_to_image, image_to_image, multi_image_to_image).
• Examples & Docs:
  • Add examples/grpo/lora/flux3/default.yaml and update README.md, examples/README.md, guidance/datasets.md, and .agents/knowledge/architecture.md.

Fixes from GPU validation

Commit	Fix
[model] fix: default FLUX.3 encode_prompt to the official 8192-token limit	Preprocessing used the 512-token default; a 999-token caption got a 512-token window instead of the official 1040-token one.
[dataset,test] fix: register flux3 in the offline smoke catalog and README table check	tests/dataset/test_offline_smoke_profiles.py and tests/docs/test_sensenova_docs.py failed once flux3 was registered (both pass on main).
[model,docs,agents] fix: align Diffusers submodule references with the 7ecb1e1b pin	Stale Qwen-Image 2.1 pin and the undocumented Python 3.11 floor of the new submodule.

Verification

CPU (final head)

• Full suite: pytest tests -> 3000 passed, 50 skipped, 0 failed with the submodule installed editable (pip install -e ./diffusers).
• scripts/validate_agent_harness.py, scripts/audit_docstrings.py (0 errors), and Black/isort on every changed Python file pass.

GPU (8x H20 97,871 MiB, driver 535.247.01)

Environment: Python 3.12, torch 2.10.0+cu129, transformers 5.14.1, accelerate 1.14.0, peft 0.20.0, deepspeed 0.18.3, huggingface-hub 1.32.0, diffusers 0.41.0.dev0 = submodule 7ecb1e1b. Weights: flux-3-image-klein-base@3f7fa5ee and flux-3-image-klein@33733762.

Tested commit: 516f9653. The final head differs only in narrowing the 8192 default to encode_prompt (the four other entry points it also touched keep their previous defaults) and dropping the test that came with it. Training is unaffected: prompts are encoded during preprocessing through encode_prompt, rollouts reuse the cached embeddings, and the parity script passes max_sequence_length explicitly.

1. Real-weight 4-layer parity (1 GPU per variant; T2I and 2-reference I2I at 1024², base 28 steps CFG 3.5, distilled 16 steps)

• Config + components (L1/L2), both variants: scheduler (shift=6.93, no dynamic shifting), VAE scale, image-processor, and transformer configs identical; every transformer / Qwen3-VL / VAE weight bit-equal, including the 86 fp32 norm_q/norm_k tensors; prompt embeddings (80- and 160-token windows), a 999-token prompt at the new default (1040-token window), and the 8016 reference-image tokens are all bit-exact.
• End-to-end (L3/L4), strict mode, both variants: timesteps, initial noise, final latents, and decoded images are bit-exact (max abs diff 0) for T2I and I2I, without autocast and under bf16 autocast. Strict mode keeps trajectory latents in the transformer dtype (latent_storage_dtype: null) and makes the reference scheduler step with Flow-Factory's fp32 Euler arithmetic; everything else runs the official modular blocks.
  • In the autocast I2I case the reference images are encoded outside autocast, as dataset preprocessing does. Encoding inside autocast differs only because the shared FLUX.2/3 VAE helper normalizes in the VAE output dtype (as the official FLUX.2 pipelines do), while the official FLUX.3 step uses vae.dtype. The trainer never encodes conditions under autocast.
• Production defaults (informational): with fp16 trajectory storage and the shared scheduler's fp32 Euler (Diffusers computes dt * v in bf16), the first step already differs by one bf16 ulp and the trajectories drift: PSNR 37.1 / 43.6 dB (base T2I / +autocast), 53.4 / 43.8 dB (distilled T2I), 20.0 / 27.3 dB (base I2I), 30.1 / 29.9 dB (distilled I2I). Images keep the same subject and composition; the 28-step CFG trajectory amplifies the ulp differences. These settings are framework-wide, not FLUX.3-specific.

2. On-policy rollout/replay (real GRPO trainer, bf16 autocast + gradient checkpointing)

A per-rank probe on GRPOTrainer recorded every replay: max |exp(new_log_prob − old_log_prob) − 1| = 0 for every on-policy replay across all 21 GRPO runs (T2I/I2I, base/distilled, every backend, including the 384 first-step replays of the shipped example). The example's second gradient step per epoch reuses pre-update rollouts by design and stays within clip_range (max |ratio − 1| = 1.3e-5, clip fraction 0). norm_q/norm_k stay fp32 under DDP and FSDP2 (see the ZeRO-2 note below). Under DDP, LoRA gradients reach all four block streams in I2I; in T2I mode_blocks.image_cond correctly receives none.

3. Backend x algorithm smoke matrix (2 GPUs each unless noted; all exit 0, losses finite)

Run	ZeRO-2	DDP	FSDP2
GRPO LoRA, T2I (PickScore)	pass	pass (+ eval)	pass, also with fsdp_cpu_ram_efficient_loading: true
GRPO LoRA, multi-reference I2I (ShareGPT-4o-Image-Mini)	pass	pass	pass (+ eval)
GRPO LoRA, distilled flux-3-image-klein, T2I and I2I		pass
SFT LoRA, 1 epoch	pass	pass	pass
Offline-DPO LoRA, 1 epoch	pass	pass	pass
examples/grpo/lora/flux3/default.yaml as shipped (8 GPUs, ZeRO-2, eval at 1024² on 128 prompts, 1 epoch)	pass (eval PickScore 0.7920, 2 optimizer steps, 1,367 s)

4. Checkpoint round-trip

• save_model_only: true: the checkpoint reloaded into a fresh real-weight adapter matches all 516/516 saved LoRA tensors bit-exactly, and a resume_type: lora GRPO run trains on-policy from it.
• save_model_only: false, resuming from checkpoint-1 with resume_type: state (DDP, ZeRO-2, FSDP2 for T2I; DDP for I2I): the first resumed epoch reproduces the source run exactly (all logged metrics and per-rank rollout/replay log-probs bit-equal), so weights, RNG, and sampler state are restored. Under DDP and FSDP2, the optimizer state before the first resumed step is bit-equal to the source run's; under ZeRO-2 (where the step runs inside accelerator.backward), the Adam step counter continues from the checkpoint. Later epochs differ by design: two fresh runs from the same seed already diverge after the first optimizer step (non-deterministic GPU backward), as a control run showed.

Known framework-level behavior (not introduced here)

• Under DeepSpeed bf16, DeepSpeedEngine casts the whole module with module.bfloat16(), so Diffusers fp32 islands (norm_q/norm_k here; Wan and MiniMax H3 have the same) become bf16 and torch.rms_norm falls back from the fused kernel with a dtype-mismatch warning. Rollout and replay stay consistent (ratio exactly 1). Restoring fp32 islands after DeepSpeed preparation belongs in the trainer and is left for a separate PR.

[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.