diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index 6a855052..0d2bfb77 100644 --- a/.agents/knowledge/architecture.md +++ b/.agents/knowledge/architecture.md @@ -130,6 +130,7 @@ extensions are `GRPOGuardTrainer → GRPOTrainer`, `DPPOTrainer → GRPOTrainer` | `flux1-kontext` | `Flux1KontextAdapter` | Image-to-Image | | `flux2` | `Flux2Adapter` | Text-to-Image & Image(s)-to-Image | | `flux2-klein` | `Flux2KleinAdapter` | Text-to-Image & Image(s)-to-Image | +| `flux3` | `Flux3Adapter` | Text-to-Image & Image(s)-to-Image | | `qwen-image` | `QwenImageAdapter` | Text-to-Image | | `qwen-image-edit-plus` | `QwenImageEditPlusAdapter` | Image(s)-to-Image | | `qwen-image-2.1` | `QwenImage21Adapter` | Text-to-Image & Image(s)-to-Image | diff --git a/.agents/knowledge/dependencies.md b/.agents/knowledge/dependencies.md index c84df288..4b5d39ac 100644 --- a/.agents/knowledge/dependencies.md +++ b/.agents/knowledge/dependencies.md @@ -80,6 +80,8 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package ### diffusers - Use the released `diffusers>=0.40.0` package as the authoritative API. The repository submodule may be used for upstream development, but must not silently override the declared runtime dependency. +- `flux3` and `qwen-image-2.1` need the pinned submodule (`7ecb1e1b`), which requires Python 3.11+ + and `huggingface-hub>=1.32`; released diffusers keeps Python 3.10 support for the other adapters. - Model adapters depend on specific pipeline classes from diffusers. Major version bumps may rename or remove pipeline classes. - `load_pipeline()` in each adapter returns a `DiffusionPipeline`-compatible object; breaking changes in diffusers' pipeline API require adapter updates. diff --git a/.agents/knowledge/topics/train_inference_consistency.md b/.agents/knowledge/topics/train_inference_consistency.md index 3bbc5eaa..3ea8a776 100644 --- a/.agents/knowledge/topics/train_inference_consistency.md +++ b/.agents/knowledge/topics/train_inference_consistency.md @@ -36,7 +36,7 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch 4. **Model weight change**: EMA swap without restore between `inference()` and first `forward()` call. 5. **Scheduler state mismatch**: `step_index` not matching (e.g., dual-scheduler models). 6. **`num_inference_steps` changed**: invalidates sigma schedule, all trajectory timesteps are wrong. -7. **Batch/pack composition mismatch (pack-dependent adapters)**: For adapters whose batched `forward()` is *pack-composition-dependent* (e.g. Bagel NaViT sequence packing, where a sample's linear-projection matmuls run over the concatenated `sum_seqlen` of the whole pack), bf16 rounding depends on *which* samples share the pack. If a training micro-batch packs a different sample set than the corresponding rollout pack, the on-policy `forward()` is no longer bit-identical -> `ratio != 1` (~1e-4) even though every stored argument matches. Per-sample (B=1) adapters are immune. The trigger is the optimize-time sample shuffle reordering `samples` before chunking into micro-batches. +7. **Batch/pack composition mismatch (pack-dependent adapters)**: For adapters whose batched `forward()` is *pack-composition-dependent* (e.g. Bagel NaViT sequence packing, where a sample's linear-projection matmuls run over the concatenated `sum_seqlen` of the whole pack), bf16 rounding depends on *which* samples share the pack. If a training micro-batch packs a different sample set than the corresponding rollout pack, the on-policy `forward()` is no longer bit-identical -> `ratio != 1` (~1e-4) even though every stored argument matches. Per-sample (B=1) adapters are immune. The trigger is an optimize-time shuffle that reorders individual samples before chunking into micro-batches; adapters declare the dependence with `requires_preserved_replay_batch_composition`, and the optimize loop then shuffles whole rollout packs (see below). 8. **Stochastic conditioning encoder recomputed per forward**: When `forward()` *rebuilds* its conditioning from raw inputs each call (rather than replaying a stored embedding) and that encoder is non-deterministic, rollout and training diverge. Bagel I2I is the canonical case: the condition-image VAE (`DiagonalGaussian`, default `sample=True` -> `mean + std*randn`) is encoded **once** at rollout but **re-encoded every training `forward()`** (`_forward_rebuild` -> `_build_gen_context`), so each draws different noise -> different KV context -> on-policy `ratio != 1` (~2e-4). T2I is immune (text-only context, deterministic). Fix: make the condition encode deterministic (`vae.reg.sample = False`, posterior mean) or cache the rollout encoding and replay it. (Only affects `vae.encode` of conditions; generation uses init noise and `vae.decode` is unaffected.) **Bagel applies this fix**: `pipeline.vae.reg.sample = False` in `BagelAdapter.load_pipeline()`. 9. **Preference arms replayed with the wrong batch**: DPO shares forward-process noise across chosen/rejected states, but each arm still owns its conditioning batch. Both policy and reference forwards for the rejected state must receive `rejected_batch`; pairing it with `chosen_batch` evaluates the rejected trajectory under another conditioning context. 10. **Condition realization redrawn per candidate**: Geometry-bound VAE conditions and stochastic @@ -48,14 +48,18 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch boundary so mask/layout binding cannot drift between candidates. 11. **Model-specific velocity direction assumed by a trainer**: Standard flow adapters predict noise-ward velocity (`noise - clean`), while MiniMax H3 predicts data-ward velocity (`clean - noise`). Any `x0` target must use `adapter.project_velocity_to_clean_state()` rather than spelling `xt - sigma * velocity` inside a trainer. 12. **Different forward graph for the same inputs**: A prefix KV cache (prefill + cached decode) and a joint forward are mathematically equal but not bf16-equal. On Qwen-Image 2.1 they differ by 0.9–2.5% relative L2 in velocity (about 6e-6 in fp32), far outside `clip_range` and exact replay tolerances; grad mode alone does not change the result. Rollout and replay must build the same graph, so Qwen-Image 2.1 records `use_kv_cache` on the rollout sample and replay reads it back as a batch key. +13. **Execution path chosen from the batch's contents**: An adapter that runs a uniform batch batched and a mixed batch one sample at a time computes the same sample with different GEMM shapes whenever its rollout pack and its training micro-batch differ in that respect, and bf16 rounding then differs. FLUX.3 did this with its 80-token prompt windows: group-contiguous rollout packs repeat one prompt and ran batched, while shuffled micro-batches mix window counts and ran per sample, so about half the replays differed by up to 7e-6 in log-prob. Prompts long enough to span two windows trigger it, including some short captions once the chat template is added. FLUX.3 therefore declares `requires_preserved_replay_batch_composition`, so replay keeps its rollout packs (item 7). ## Pack-composition-dependent adapters: `shuffle_samples` -`samples` reaches `optimize()` in rollout-pack order (`generate_samples` appends each `inference` pack in order; `compute_advantages(store_to_samples=True)` attaches advantages in place without reordering), so `samples[i*B:(i+1)*B]` is exactly rollout pack `i`. Set `train.shuffle_samples: false` so the optimize loop keeps that order: each training micro-batch then packs the *same* samples as its rollout pack -> deterministic bf16 -> on-policy `ratio == 1`. +`samples` reaches `optimize()` in rollout-pack order (`generate_samples` appends each `inference` pack in order; `compute_advantages(store_to_samples=True)` attaches advantages in place without reordering), so `samples[i*B:(i+1)*B]` is exactly rollout pack `i`. `train.shuffle_samples` keeps that pack structure whenever the adapter needs it: -- Requires matched sampling and training `per_device_batch_size` (so contiguous chunks reproduce the rollout packs). These are the *same* field: rollout uses `training_args.per_device_batch_size` for the train DataLoader (`data_utils/loader.py`) and `optimize()` chunks by the same value, so they cannot drift unless a separate sampling batch size is introduced. -- Default is `True` (per-inner-epoch shuffle); only pack-dependent adapters (Bagel) need `false`. The off-policy decorrelation cost is minor because the rollout order is already sampler-randomized. `BagelAdapter` warns at init if `shuffle_samples` is left `True`. -- Wired via `BaseTrainer._order_samples_for_optimize(samples, inner_epoch)` (used by grpo/dppo/nft/awm/opd; dpo shuffles chosen/rejected pairs separately). +- `true` (default) with a batch-invariant adapter: every inner epoch permutes individual samples. +- `true` with an adapter that declares `requires_preserved_replay_batch_composition` (Bagel, FLUX.3): every inner epoch permutes whole rollout packs and keeps each pack's samples and order, so each training micro-batch replays exactly one rollout pack -> deterministic bf16 -> on-policy `ratio == 1`. For these adapters `generate_samples` always records the rollout packs in an `AcquisitionManifest`, and the ordering fails fast if a pack does not match `per_device_batch_size` or the samples changed after rollout (for example a custom `generate_samples()` that does not record boundaries). +- `false`: the rollout order is kept for every adapter. +- Rollout and optimize chunk by the same `training_args.per_device_batch_size` (the train DataLoader in `data_utils/loader.py` and `optimize()`), so a replayed pack has its rollout size unless a separate sampling batch size is introduced. +- Shuffling packs instead of samples costs little decorrelation: the rollout order is already sampler-randomized, and with one optimizer step per epoch the accumulated gradient covers the same samples in any order. +- Wired via `BaseTrainer._order_samples_for_optimize(samples, inner_epoch)` (used by grpo/dppo/nft/awm/opd; dpo shuffles chosen/rejected pairs separately because it compares policy and reference on the same micro-batch rather than against rollout log-probs). Reward/optimization overlap keeps requiring `shuffle_samples: false` and validates the same packs against its reward tiles. ## Where in Code @@ -84,6 +88,14 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch - **Lesson**: Never flip the grad mode above code that decides from it whether a graph may outlive the call. A graph built under forced grad without a later backward is a leak whenever activation checkpointing saves non-tensor inputs that point back into it. - **Related Constraint**: N/A +### Replay FLUX.3 rollout packs whole +- **Date**: 2026-10-18 +- **Symptom**: FLUX.3 GRPO with `per_device_batch_size: 8` logged on-policy ratios that rounded to 1.0000 but were not exactly 1 (`ratio_std` up to 1e-4), while Qwen-Image 2.1 under the same trainer stayed exactly 1. An exact-replay probe found 421 of 768 replayed log-probs off by up to 7e-6 with GenEval prompts rewritten by the FLUX.3 upsampler, and 140 of 768 with PickScore captions. +- **Root Cause**: `inference()` and `forward()` each chose batched execution for a batch whose prompts shared one 80-token window count and per-sample execution otherwise. Rollout packs of one repeated prompt ran at batch size 8; after the optimize-loop shuffle the same samples replayed at batch size 1, or the other way round. Keeping the rollout packs (`shuffle_samples: false`) made every replay exact, which isolated the pack composition. +- **Fix**: FLUX.3 declares `requires_preserved_replay_batch_composition`, and `shuffle_samples` now shuffles a pack-dependent adapter's recorded rollout packs as whole units instead of individual samples, so every micro-batch replays one rollout pack in rollout order. FLUX.3 keeps batched execution for uniform-window packs. Running every sample alone was also exact but cost about 13% rollout and 15% optimization time at `per_device_batch_size: 8`. +- **Lesson**: A GPU parity gate must include `per_device_batch_size > 1`, prompts of mixed lengths, and the default shuffle, and compare log-probs exactly rather than through the rounded ratio metric. Gates inherited from a `per_device_batch_size: 1` example never ran the batched path. +- **Related Constraint**: Constraint #7 + ## Cross-refs - UP: [`constraints.md` #7](../constraints.md#7-coupled-vs-decoupled-paradigm), [Architecture Execution Pipelines](../architecture.md#execution-pipelines) diff --git a/README.md b/README.md index 02fc59eb..f23bb6e5 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,17 @@ # 🔥 News +* **[2026-10-16]** **FLUX.3 Image Klein** support! Train both the base (`black-forest-labs/flux-3-image-klein-base`) + and step-distilled flow-map (`black-forest-labs/flux-3-image-klein`) text-to-image and multi-reference + image-to-image models with the [GRPO + LoRA recipe](examples/grpo/lora/flux3/default.yaml). + Until FLUX.3 reaches a released Diffusers package, install the pinned submodule + (it requires Python 3.11 or newer): +```bash +git submodule update --init +pip install -e ./diffusers +pip install -e . +``` + * **[2026-09-20]** **Qwen-Image 2.1** support! Train the unified text-to-image and ordered multi-image editing model with the [GRPO + LoRA recipe](examples/grpo/lora/qwen_image_2_1/default.yaml). Until Qwen-Image 2.1 reaches a released Diffusers package, install the pinned submodule: @@ -74,11 +85,13 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac Image(s)-to-ImageQwen-Image-Edit-250920Bqwen-image-edit-plus Qwen-Image-Edit-251120Bqwen-image-edit-plus - Text-to-Image & Image(s)-to-ImageFLUX.2-dev32Bflux2 + Text-to-Image & Image(s)-to-ImageFLUX.2-dev32Bflux2 FLUX.2-klein-4B4Bflux2-klein FLUX.2-klein-9B9Bflux2-klein FLUX.2-klein-base-4B4Bflux2-klein FLUX.2-klein-base-9B9Bflux2-klein + flux-3-image-klein-base5.5Bflux3 + flux-3-image-klein5.5Bflux3 BAGEL-7B-MoT14Bbagel SenseNova-U1 1.016Bsensenova SenseNova-U1 1.516Bsensenova @@ -107,7 +120,7 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac > To support new models, see [Guidance/New Model](guidance/new_model.md). > **Offline output support:** SFT and offline DPO currently support `sd3-5`, `flux1`, -> `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`, +> `flux1-kontext`, `flux2`, `flux2-klein`, `flux3`, `qwen-image`, `qwen-image-edit-plus`, > `qwen-image-2.1`, `z-image`, > `bagel`, `sensenova`, `wan2_t2v`, `wan2_i2v`, `ltx2_t2av`, `ltx2_i2av`, and all > MiniMax H3 workflows. Video/audio targets are encoded on demand and are never written to diff --git a/dataset/offline_smoke/profiles.py b/dataset/offline_smoke/profiles.py index 5bb0c9db..14d0181c 100644 --- a/dataset/offline_smoke/profiles.py +++ b/dataset/offline_smoke/profiles.py @@ -230,9 +230,11 @@ def _case( _MODEL_TYPES = { - "text_to_image": "sd3-5 flux1 flux2 flux2-klein qwen-image qwen-image-2.1 z-image bagel sensenova", - "image_to_image": "flux1-kontext flux2 flux2-klein qwen-image-edit-plus bagel sensenova", - "multi_image_to_image": "flux2 flux2-klein qwen-image-edit-plus bagel sensenova", + "text_to_image": ( + "sd3-5 flux1 flux2 flux2-klein flux3 qwen-image qwen-image-2.1 z-image bagel sensenova" + ), + "image_to_image": "flux1-kontext flux2 flux2-klein flux3 qwen-image-edit-plus bagel sensenova", + "multi_image_to_image": "flux2 flux2-klein flux3 qwen-image-edit-plus bagel sensenova", "text_to_video": "wan2_t2v", "first_frame_to_video": "wan2_i2v", "first_last_frame_to_video": "wan2_i2v", diff --git a/diffusers b/diffusers index 80c7ed26..7ecb1e1b 160000 --- a/diffusers +++ b/diffusers @@ -1 +1 @@ -Subproject commit 80c7ed262aeffbeb43ef13ae04baeb9b84515a69 +Subproject commit 7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941 diff --git a/examples/README.md b/examples/README.md index f0d98462..54815474 100644 --- a/examples/README.md +++ b/examples/README.md @@ -96,6 +96,14 @@ The same adapter also accepts prompt-only records. It currently requires the bun submodule, enforces `per_device_batch_size: 1`, and keeps prefix KV caching disabled for exact rollout/replay parity. +## FLUX.3 example + +- [`flux3` text-to-image & multi-reference editing GRPO](grpo/lora/flux3/default.yaml) + +Supports both `black-forest-labs/flux-3-image-klein-base` (true CFG) and +`black-forest-labs/flux-3-image-klein` (step-distilled flow map, no CFG) via the bundled +Diffusers submodule. + ## MiniMax H3 validation status The T2VA `debug.yaml` recipe is real-weight validated with the 61 GB checkpoint diff --git a/examples/grpo/lora/flux3/default.yaml b/examples/grpo/lora/flux3/default.yaml new file mode 100644 index 00000000..adac52ff --- /dev/null +++ b/examples/grpo/lora/flux3/default.yaml @@ -0,0 +1,113 @@ +# FLUX.3 currently requires the bundled Diffusers submodule: +# git submodule update --init +# pip install -e ./diffusers + +# Environment Configuration +launcher: "accelerate" +config_file: config/deepspeed/deepspeed_zero2.yaml # Switch to config/accelerate_configs/fsdp2.yaml if you have limited GPU memory. +num_processes: 8 +main_process_port: 29500 +mixed_precision: "bf16" + +# Data Configuration +data: + datasets: + - name: default + dataset_dir: "dataset/sharegpt4o_image_mini" + train: { weight: 1 } + eval: {} + preprocessing_batch_size: 8 + dataloader_num_workers: 16 + force_reprocess: true + cache_dir: "~/.cache/flow_factory/datasets" + sampler_type: "auto" + sampler_subgroup_size: null + +# Model Configuration +model: + finetune_type: "lora" + lora_rank: 64 + lora_alpha: 128 + target_modules: "default" + model_name_or_path: "black-forest-labs/flux-3-image-klein-base" # Options: black-forest-labs/flux-3-image-klein-base (CFG), black-forest-labs/flux-3-image-klein (step-distilled flow-map) + model_type: "flux3" + resume_path: null + resume_type: null + +log: + run_name: null + project: "Flow-Factory" + logging_backend: "wandb" + save_dir: "saves/" + save_freq: 20 + save_model_only: true + +# Training Configuration +train: + trainer_type: "grpo" + reward_optimization_overlap: false + reward_optimization_overlap_mode: "ready" + reward_optimization_overlap_poll_interval: 0.05 + advantage_aggregation: "gdpo" + clip_range: 1.0e-4 + adv_clip_range: 5.0 + kl_type: "v-based" + kl_beta: 0 + ref_param_device: "cuda" + + # Sampling settings + resolution: 384 + condition_image_size: 384 + num_inference_steps: 10 # Use 4-16 with guidance_scale: 1.0 for black-forest-labs/flux-3-image-klein + guidance_scale: 3.5 + + # Batch and sampling + per_device_batch_size: 1 + group_size: 16 + global_std: false + unique_sample_num_per_epoch: 48 + gradient_step_per_epoch: 2 + gradient_accumulation_steps: auto + + ema_decay: 0.9 + ema_update_interval: 4 + ema_device: "cuda" + enable_gradient_checkpointing: true + offload_samples_to_cpu: false + seed: 42 + +# Scheduler Configuration +scheduler: + dynamics_type: "Flow-SDE" + noise_level: 0.7 + num_sde_steps: 1 + sde_steps: [1, 2, 3] + seed: 42 + +# Evaluation Configuration +eval: + resolution: 1024 + condition_image_size: 1024 + per_device_batch_size: 1 + guidance_scale: 3.5 + num_inference_steps: 28 + eval_freq: 20 + seed: 42 + +# Reward Model Configuration +rewards: + - name: "pick_score" + reward_model: "PickScore" + weight: 1 + batch_size: 16 + device: "cuda" + dtype: bfloat16 + +# Optimizer Configuration +optimizers: + - name: default + learning_rate: 3.0e-4 + weight_decay: 1.0e-4 + betas: [0.9, 0.999] + eps: 1.0e-8 + max_grad_norm: 1.0 diff --git a/guidance/algorithms.md b/guidance/algorithms.md index 96542d4f..ed1155fb 100644 --- a/guidance/algorithms.md +++ b/guidance/algorithms.md @@ -172,7 +172,7 @@ The `kl_beta` parameter controls the coefficient of the KL divergence term. - **Full-Parameter Fine-Tuning**: The overhead is significant. You may want to set `ref_param_device` to `cpu` to save memory. - **No KL-Loss**: Setting `kl_beta` to `0` automatically disables this term and eliminates extra memory usage. -**Pack-composition-dependent adapters (e.g. Bagel)**: GRPO's on-policy step assumes `forward()` is bit-identical between rollout and training. Adapters that batch via sequence packing (NaViT) are not batch-invariant, so the default optimize-time sample shuffle breaks the on-policy `ratio == 1`. Set `train.shuffle_samples: false` (with matched sampling/training `per_device_batch_size`) for these adapters. See `workflow.md` Stage 6 and the train-inference consistency topic doc. +**Pack-composition-dependent adapters (e.g. Bagel, FLUX.3)**: GRPO's on-policy step assumes `forward()` is bit-identical between rollout and training. Adapters whose batched forward depends on which samples share a batch (NaViT sequence packing, or FLUX.3 batching only prompts that share an 80-token window count) are not batch-invariant. They declare `requires_preserved_replay_batch_composition`, and `train.shuffle_samples` then shuffles whole rollout packs instead of individual samples, so the on-policy `ratio == 1` holds under the default shuffle. See `workflow.md` Stage 6 and the train-inference consistency topic doc. #### GRPO-Guard diff --git a/guidance/datasets.md b/guidance/datasets.md index ab3aaf35..726fb708 100644 --- a/guidance/datasets.md +++ b/guidance/datasets.md @@ -365,7 +365,7 @@ output semantics. | Offline status | Model types | Notes | |---|---|---| -| Supported | `sd3-5`, `flux1`, `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`, `qwen-image-2.1`, `z-image`, `bagel`, `sensenova` | Image-output codecs with adapter-specific geometry and packing. Qwen-Image 2.1 keeps ordered condition blocks before the target block and runs offline batches one sample at a time. SenseNova uses the existing grouped `images` input with within-type order, not heterogeneous references. | +| Supported | `sd3-5`, `flux1`, `flux1-kontext`, `flux2`, `flux2-klein`, `flux3`, `qwen-image`, `qwen-image-edit-plus`, `qwen-image-2.1`, `z-image`, `bagel`, `sensenova` | Image-output codecs with adapter-specific geometry and packing. Qwen-Image 2.1 keeps ordered condition blocks before the target block and runs offline batches one sample at a time. SenseNova uses the existing grouped `images` input with within-type order, not heterogeneous references. | | Supported | `wan2_t2v` | Video targets require `fps`; the codec resamples to configured frames/rate and samples the Wan VAE posterior on the fly. | | Supported | `wan2_i2v` | Input media uses a checkpoint-specific contract. Expanded-timestep TI2V and standard Wan2.1 I2V accept exactly `first_frame`; dedicated Wan2.1 FLF2V requires both `first_frame` and `last_frame`; Wan2.2 I2V-A14B accepts `first_frame` plus an optional VAE-only `last_frame`. Condition pixels are cached at configured geometry and encoded with VAE posterior mode once per batch. Video targets require `fps`; offline execution is B=1. | | Supported | `ltx2_t2av`, `ltx2_i2av` | Every candidate is an exact ordered `(video, audio)` pair with required `fps` and `sample_rate`. Both streams are aligned to the official LTX2 clock and encoded/packed on the fly. I2AV requires the `first_frame` image slot, substitutes its posterior-mode first latent into each target, and excludes the pinned tokens with an active mask. | diff --git a/guidance/gpu_validation.md b/guidance/gpu_validation.md index 863ac848..14ebf98b 100644 --- a/guidance/gpu_validation.md +++ b/guidance/gpu_validation.md @@ -269,7 +269,7 @@ recorded in the resolved YAML; do not silently return to a large quality recipe. | Mode | Starting recipe | Train geometry | Steps | Algorithm-specific notes | |---|---|---|---|---| | `sd35-t2i` | `examples/grpo/lora/sd3_5/default.yaml` | `resolution: 256` | `num_inference_steps: 2` | Use the checked-in SD3.5 SFT/offline-DPO/TDM recipe as the algorithm overlay. | -| `bagel-mri2i` | `examples/grpo/lora/bagel/i2i.yaml` | `resolution: 256` | `num_inference_steps: 2` | Every row has exactly two ordered references; keep `shuffle_samples: false`. | +| `bagel-mri2i` | `examples/grpo/lora/bagel/i2i.yaml` | `resolution: 256` | `num_inference_steps: 2` | Every row has exactly two ordered references. | | `wan-t2v` | `examples/grpo/lora/wan21/t2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Target video has at least five frames and carries its source `fps`. | | `wan-i2v-first` | `examples/grpo/lora/wan22/i2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Use TI2V-5B and exactly one condition image. | | `wan-flf2v` | `examples/grpo/lora/wan21/i2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Use the dedicated FLF2V checkpoint with two ordered condition images. | diff --git a/guidance/workflow.md b/guidance/workflow.md index cd9a2df9..d6c641cb 100644 --- a/guidance/workflow.md +++ b/guidance/workflow.md @@ -625,7 +625,7 @@ For a pack-composition-dependent adapter such as Bagel at `per_device_batch_size whole work units, but validation rejects any unit that would split or repack one of those recorded batches. -> **`shuffle_samples` and on-policy ratio**: the optimize loop reorders `samples` each inner epoch (`train.shuffle_samples: true`, the default). For adapters whose batched `forward()` is *pack-composition-dependent* (e.g. Bagel NaViT packing), this makes a training micro-batch pack a different sample set than its rollout pack, so the on-policy `ratio != 1`. Set `train.shuffle_samples: false` for such adapters (with matched sampling/training `per_device_batch_size`) so each micro-batch reproduces its rollout pack. See the train-inference consistency topic doc. +> **`shuffle_samples` and on-policy ratio**: the optimize loop reorders `samples` each inner epoch (`train.shuffle_samples: true`, the default). Batch-invariant adapters shuffle individual samples. Adapters whose batched `forward()` is *pack-composition-dependent* (Bagel NaViT packing, FLUX.3 batching only prompts that share an 80-token window count) declare `requires_preserved_replay_batch_composition`; for them the loop shuffles whole rollout packs, keeping each pack's samples and order, so every micro-batch replays one rollout pack and the on-policy `ratio == 1`. `shuffle_samples: false` keeps the rollout order for every adapter. See the train-inference consistency topic doc. ### Algorithm-Specific Optimization diff --git a/src/flow_factory/hparams/model_args.py b/src/flow_factory/hparams/model_args.py index e022ed11..fdb4412a 100644 --- a/src/flow_factory/hparams/model_args.py +++ b/src/flow_factory/hparams/model_args.py @@ -165,6 +165,7 @@ class ModelArguments(ArgABC): "flux1-kontext", "flux2", "flux2-klein", + "flux3", "qwen-image", "qwen-image-edit-plus", "qwen-image-2.1", diff --git a/src/flow_factory/hparams/training_args/_base.py b/src/flow_factory/hparams/training_args/_base.py index 360ff6ef..a078fc59 100644 --- a/src/flow_factory/hparams/training_args/_base.py +++ b/src/flow_factory/hparams/training_args/_base.py @@ -320,12 +320,12 @@ def get_sampler_selection_contract(self) -> SamplerSelectionContract: default=True, metadata={ "help": ( - "Shuffle samples before each inner optimize epoch. Keep True normally. " - "Set False for adapters whose batched forward is pack-composition-dependent " - "(e.g. Bagel/NaViT sequence packing): then each training micro-batch packs the " - "same samples as its rollout pack, so the bf16 forward is bit-identical between " - "rollout and training and the on-policy ratio stays 1. Requires matched sampling " - "and training `per_device_batch_size`." + "Shuffle before each inner optimize epoch. Batch-invariant adapters shuffle " + "individual samples. Adapters whose batched forward depends on batch composition " + "(they declare `requires_preserved_replay_batch_composition`, e.g. Bagel and " + "FLUX.3) shuffle whole rollout packs instead, keeping each pack's samples and " + "order, so every training micro-batch replays one rollout pack and the on-policy " + "ratio stays exactly 1. False keeps the rollout order for every adapter." ) }, ) diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py index e2c69673..75d1fdcc 100644 --- a/src/flow_factory/models/abc.py +++ b/src/flow_factory/models/abc.py @@ -236,9 +236,11 @@ class BaseAdapter(ABC): supported_diffusers_cache_policies: ClassVar[Optional[frozenset[str]]] = None supports_fsdp2_cpu_efficient_loading: ClassVar[bool] = False # Opt in when replaying a sample under a different micro-batch composition - # changes its numerical result (for example Bagel's NaViT sequence packing). - # Reward/optimization overlap then records rollout packs and refuses to tile - # or replay them with different boundaries. + # changes its numerical result, i.e. the forward is not batch-invariant (for + # example Bagel's NaViT sequence packing). The generation loop then records the + # rollout packs: `train.shuffle_samples` shuffles whole packs instead of samples, + # and reward/optimization overlap refuses to tile or replay them with different + # boundaries. requires_preserved_replay_batch_composition: ClassVar[bool] = False # Opt in only when FSDP2 communication overlap exceeds the model's activation headroom. fsdp2_use_default_stream_unshard: ClassVar[bool] = False diff --git a/src/flow_factory/models/bagel/bagel.py b/src/flow_factory/models/bagel/bagel.py index c03db026..864cc471 100644 --- a/src/flow_factory/models/bagel/bagel.py +++ b/src/flow_factory/models/bagel/bagel.py @@ -233,16 +233,6 @@ def __init__(self, config: Arguments, accelerator: Accelerator): self.pipeline: BagelPseudoPipeline self.scheduler: FlowMatchEulerDiscreteSDEScheduler - # Bagel's batched forward is pack-composition-dependent (NaViT packing), so the - # optimize-time sample shuffle must be off to keep the on-policy ratio == 1. - if getattr(self.training_args, "shuffle_samples", True): - logger.warning( - "Bagel is not batch-invariant (NaViT packing) but training_args.shuffle_samples " - "is True: this reorders optimize micro-batches vs rollout packs and breaks " - "train-inference consistency (on-policy ratio != 1). Set `train.shuffle_samples: " - "false` (see the Bagel example configs / train_inference_consistency.md)." - ) - # ─────────────────── Tokenizer & Transforms ─────────────────── def _init_tokenizer_and_transforms(self): diff --git a/src/flow_factory/models/flux/_output.py b/src/flow_factory/models/flux/_output.py index 3ecf69e9..9eb4db6e 100644 --- a/src/flow_factory/models/flux/_output.py +++ b/src/flow_factory/models/flux/_output.py @@ -105,6 +105,91 @@ def prepare_flux1_output_latents( return packed, target_ids +def patchify_flux2_latents(latents: torch.Tensor) -> torch.Tensor: + """Convert 4D VAE latents ``(B, C, H, W)`` to 2x2 patchified ``(B, C*4, H//2, W//2)``.""" + if latents.ndim != 4: + raise ValueError( + f"expected 4D BCHW latents for 2x2 patchify, received shape {tuple(latents.shape)}" + ) + batch_size, channels, height, width = latents.shape + if height % 2 or width % 2: + raise ValueError( + f"expected even latent height/width for 2x2 patchify, received {(height, width)}" + ) + latents = latents.view(batch_size, channels, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 3, 5, 2, 4) + return latents.reshape(batch_size, channels * 4, height // 2, width // 2) + + +def unpatchify_flux2_latents(latents: torch.Tensor) -> torch.Tensor: + """Convert 2x2 patchified latents ``(B, C*4, H, W)`` back to ``(B, C, H*2, W*2)``.""" + if latents.ndim != 4: + raise ValueError( + f"expected 4D BCHW latents for 2x2 unpatchify, received shape {tuple(latents.shape)}" + ) + batch_size, channels, height, width = latents.shape + if channels % 4: + raise ValueError( + f"expected channel count divisible by 4 for 2x2 unpatchify, received {channels}" + ) + latents = latents.reshape(batch_size, channels // 4, 2, 2, height, width) + latents = latents.permute(0, 1, 4, 2, 5, 3) + return latents.reshape(batch_size, channels // 4, height * 2, width * 2) + + +def pack_flux3_latents(latents: torch.Tensor) -> torch.Tensor: + """Pack 4D patchified latents ``(B, C, H, W)`` into tokens ``(B, H*W, C)``.""" + if latents.ndim != 4: + raise ValueError( + f"expected 4D BCHW latents for FLUX.3 packing, received shape {tuple(latents.shape)}" + ) + return latents.flatten(2).transpose(1, 2) + + +def unpack_flux3_latents( + latents: torch.Tensor, + latent_height: int, + latent_width: int, +) -> torch.Tensor: + """Unpack FLUX.3 latent tokens ``(B, H*W, C)`` onto a 2D grid ``(B, C, H, W)``.""" + if latents.ndim != 3: + raise ValueError( + f"expected 3D (B, seq_len, C) latents for FLUX.3 unpacking, " + f"received shape {tuple(latents.shape)}" + ) + batch_size, seq_len, channels = latents.shape + if latent_height <= 0 or latent_width <= 0 or latent_height * latent_width != seq_len: + raise ValueError( + f"FLUX.3 latent token count {seq_len} does not match " + f"latent grid {latent_height} x {latent_width}" + ) + return latents.view(batch_size, latent_height, latent_width, channels).permute(0, 3, 1, 2) + + +def prepare_flux3_image_ids( + latent_height: int, + latent_width: int, + *, + image_rope_scale: float = 0.5, + l_coord: int = 0, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """Build 4D position IDs ``(t, h, w, l)`` of shape ``(H*W, 4)`` for FLUX.3 image tokens.""" + if latent_height <= 0 or latent_width <= 0: + raise ValueError( + f"expected positive FLUX.3 latent grid dimensions, received {(latent_height, latent_width)}" + ) + t_coords = torch.arange(1) + h_coords = torch.arange(latent_height) + w_coords = torch.arange(latent_width) + l_coords = torch.tensor([l_coord]) + ids = torch.cartesian_prod(t_coords, h_coords, w_coords, l_coords).to(dtype=torch.float32) + ids[:, 1:3] = ids[:, 1:3] * float(image_rope_scale) + if device is not None: + ids = ids.to(device) + return ids + + def encode_flux2_output_images( adapter: Any, pixel_values: torch.Tensor, @@ -135,6 +220,41 @@ def encode_flux2_output_images( ) +def encode_flux3_output_images( + adapter: Any, + pixel_values: torch.Tensor, + generator: Optional[torch.Generator], + *, + image_rope_scale: float = 0.5, +) -> EncodedImageTensor: + """Sample the FLUX.3 target posterior, normalize, and pack tokens with scaled RoPE IDs.""" + latents = encode_flux2_vae_image( + adapter, + pixel_values, + sample_mode="sample", + generator=generator, + ) + batch_size = latents.shape[0] + latent_height, latent_width = latents.shape[-2:] + latent_ids = ( + prepare_flux3_image_ids( + latent_height, + latent_width, + image_rope_scale=image_rope_scale, + l_coord=0, + device=adapter.device, + ) + .unsqueeze(0) + .repeat(batch_size, 1, 1) + ) + packed = pack_flux3_latents(latents).to(adapter.device) + return EncodedImageTensor( + latents=packed, + forward_context={"latent_ids": latent_ids}, + decode_context={"latent_ids": latent_ids}, + ) + + def encode_flux2_vae_image( adapter: Any, pixel_values: torch.Tensor, @@ -142,10 +262,10 @@ def encode_flux2_vae_image( sample_mode: VAELatentSampleMode, generator: Optional[torch.Generator] = None, ) -> torch.Tensor: - """Apply the exact FLUX.2 transform shared by condition and target roles. + """Apply the exact FLUX.2/FLUX.3 transform shared by condition and target roles. Args: - adapter: FLUX.2 adapter exposing the canonical VAE and patchify primitive. + adapter: FLUX.2 or FLUX.3 adapter exposing the canonical VAE and patchify primitive. pixel_values: Preprocessed BCHW image tensor. sample_mode: Explicit posterior selection for the calling pipeline role. generator: Generator forwarded unchanged for posterior sampling. @@ -165,7 +285,11 @@ def encode_flux2_vae_image( generator=generator, source=f"{type(adapter).__name__} FLUX.2 VAE encode", ) - latents = adapter.pipeline._patchify_latents(latents) + pipeline_patchify = getattr(getattr(adapter, "pipeline", None), "_patchify_latents", None) + if callable(pipeline_patchify): + latents = pipeline_patchify(latents) + else: + latents = patchify_flux2_latents(latents) if latents.ndim != 4: raise ValueError( f"{type(adapter).__name__} FLUX.2 patchify expected BCHW output, " @@ -259,10 +383,72 @@ def prepare_flux2_condition_latents( return packed, image_latent_ids +def prepare_flux3_condition_latents( + adapter: Any, + images: List[torch.Tensor], + *, + batch_size: int, + device: torch.device, + dtype: torch.dtype, + image_rope_scale: float = 0.5, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Compose FLUX.3 reference-image tokens and ``l = -(1 + idx)`` position IDs. + + Args: + adapter: FLUX.3 adapter exposing the canonical VAE. + images: One or more preprocessed BCHW condition tensors for one sample. + batch_size: Number of prompt rows that reuse the condition sequence. + device: Target device for condition tensors and identifiers. + dtype: Target compute dtype for packed condition tokens. + image_rope_scale: Spatial RoPE coordinate multiplier (0.5 in FLUX.3). + + Returns: + Packed condition latents ``(B, cond_seq_len, C)`` and position IDs ``(B, cond_seq_len, 4)``. + """ + if not images: + raise ValueError(f"{type(adapter).__name__} requires at least one condition image") + if type(batch_size) is not int or batch_size <= 0: + raise ValueError( + f"{type(adapter).__name__} condition batch_size must be a positive int, " + f"received {batch_size!r}" + ) + vae_dtype = getattr(adapter.vae, "dtype", dtype) + cond_tokens: List[torch.Tensor] = [] + cond_ids: List[torch.Tensor] = [] + for index, image in enumerate(images): + encoded = encode_flux2_vae_image( + adapter, + image.to(device=device, dtype=vae_dtype), + sample_mode="argmax", + ) + latent_height, latent_width = encoded.shape[-2:] + cond_tokens.append(pack_flux3_latents(encoded).squeeze(0)) + cond_ids.append( + prepare_flux3_image_ids( + latent_height, + latent_width, + image_rope_scale=image_rope_scale, + l_coord=-(1 + index), + device=device, + ) + ) + packed = torch.cat(cond_tokens, dim=0).unsqueeze(0).repeat(batch_size, 1, 1) + packed = packed.to(device=device, dtype=dtype) + image_latent_ids = torch.cat(cond_ids, dim=0).unsqueeze(0).repeat(batch_size, 1, 1).to(device) + return packed, image_latent_ids + + __all__ = [ "encode_flux1_vae_image", "encode_flux2_output_images", "encode_flux2_vae_image", + "encode_flux3_output_images", + "pack_flux3_latents", + "patchify_flux2_latents", "prepare_flux1_output_latents", "prepare_flux2_condition_latents", + "prepare_flux3_condition_latents", + "prepare_flux3_image_ids", + "unpack_flux3_latents", + "unpatchify_flux2_latents", ] diff --git a/src/flow_factory/models/flux/flux3.py b/src/flow_factory/models/flux/flux3.py new file mode 100644 index 00000000..9de0d293 --- /dev/null +++ b/src/flow_factory/models/flux/flux3.py @@ -0,0 +1,1573 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/models/flux/flux3.py +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Any, ClassVar, Dict, List, Literal, Mapping, Optional, Sequence, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn +from accelerate import Accelerator +from PIL import Image + +import diffusers +from diffusers.utils.torch_utils import randn_tensor + +try: + from diffusers import Flux3Transformer2DModel, ModularPipeline + from diffusers.modular_pipelines.flux3 import ( + Flux3AutoBlocks, + Flux3DistilledAutoBlocks, + Flux3DistilledModularPipeline, + Flux3ModularPipeline, + ) + from diffusers.modular_pipelines.flux3.image_processor import Flux3ImageProcessor + + _FLUX3_IMPORT_ERROR: Optional[Exception] = None +except (ImportError, RuntimeError) as exc: # pragma: no cover - exercised with older diffusers + Flux3Transformer2DModel = Any # type: ignore[misc,assignment] + ModularPipeline = Any # type: ignore[misc,assignment] + Flux3AutoBlocks = Any # type: ignore[misc,assignment] + Flux3DistilledAutoBlocks = Any # type: ignore[misc,assignment] + Flux3DistilledModularPipeline = Any # type: ignore[misc,assignment] + Flux3ModularPipeline = Any # type: ignore[misc,assignment] + Flux3ImageProcessor = Any # type: ignore[misc,assignment] + _FLUX3_IMPORT_ERROR = exc + +from ...contracts import InputMediaOrder, NegativePromptPolicy, PipelineIOContract +from ...hparams import Arguments +from ...samples import I2ISample +from ...scheduler import ( + FlowMatchEulerDiscreteSDEScheduler, + FlowMatchEulerDiscreteSDESchedulerOutput, + set_scheduler_timesteps, +) +from ...utils.image import ( + ImageBatch, + ImageSingle, + MultiImageBatch, + is_multi_image_batch, + standardize_image_batch, +) +from ...utils.logger_utils import setup_logger +from ...utils.trajectory_collector import ( + TrajectoryIndicesType, + create_callback_collector, + create_trajectory_collector, +) +from ..abc import BaseAdapter +from ..checkpointing import CheckpointUnit +from ..configured_image_output import ( + ConfiguredImageOutputAdapterMixin, + EncodedImageTensor, +) +from ..pipeline_contracts import image_output_contract +from ..precision import DTypePolicy +from ..runtime import ModularPipelineRuntime +from ._output import ( + encode_flux3_output_images, + pack_flux3_latents, + prepare_flux3_condition_latents, + prepare_flux3_image_ids, + unpack_flux3_latents, + unpatchify_flux2_latents, +) + +logger = setup_logger(__name__) + +FLUX3_DIFFUSERS_COMMIT = "7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941" +FLUX3_INSTALL = "git submodule update --init && pip install -e ./diffusers" +CONDITION_IMAGE_SIZE = (1024, 1024) +IMAGE_ROPE_SCALE = 0.5 +TEXT_PAD_MULTIPLE = 80 +# Matches the `Flux3TextEncoderStep` inference default, so long captions are not truncated. +DEFAULT_MAX_SEQUENCE_LENGTH = 8192 +DEFAULT_TEXT_ENCODER_OUT_LAYERS: Tuple[int, ...] = (4, 8, 12, 16, 20, 24, 28, 32) + + +@dataclass +class Flux3Sample(I2ISample): + """Output sample class for Flux3Adapter models.""" + + _shared_fields: ClassVar[frozenset[str]] = frozenset({}) + latent_ids: Optional[torch.Tensor] = None + text_ids: Optional[torch.Tensor] = None + negative_text_ids: Optional[torch.Tensor] = None + image_latents: Optional[torch.Tensor] = None + image_latent_ids: Optional[torch.Tensor] = None + + +class Flux3Adapter(ConfiguredImageOutputAdapterMixin, BaseAdapter): + """Adapter for FLUX.3 modular text-to-image and multi-reference image-to-image models. + + A batch whose prompts share one 80-token window count runs batched; any other batch, and + every image-conditioned batch, runs one sample at a time. A sample's bf16 result therefore + depends on the batch it is in, so the adapter declares + ``requires_preserved_replay_batch_composition`` and the optimize loop replays whole rollout + packs. + """ + + supports_diffusers_cache: ClassVar[bool] = False + supports_fsdp2_cpu_efficient_loading: ClassVar[bool] = True + ddp_find_unused_parameters: ClassVar[bool] = True + requires_preserved_replay_batch_composition: ClassVar[bool] = True + component_load_dtype_defaults: ClassVar[DTypePolicy] = MappingProxyType( + { + "transformers": torch.bfloat16, + "text_encoders": torch.bfloat16, + "vae": torch.float32, + } + ) + offline_training_forward_overrides = MappingProxyType({"guidance_scale": 1.0}) + pipeline_io_contract = image_output_contract( + negative_prompt=NegativePromptPolicy.OPTIONAL, + input_image_min_count=0, + input_image_max_count=None, + input_order=InputMediaOrder.WITHIN_TYPE, + ) + + def __init__(self, config: Arguments, accelerator: Accelerator): + if _FLUX3_IMPORT_ERROR is not None: + raise ImportError( + "Flux3Adapter requires the bundled Diffusers submodule at " + f"{FLUX3_DIFFUSERS_COMMIT}. Install it with `{FLUX3_INSTALL}`. " + f"The active diffusers version is {getattr(diffusers, '__version__', 'unknown')}." + ) from _FLUX3_IMPORT_ERROR + super().__init__(config, accelerator) + self.pipeline: Flux3ModularPipeline + self.scheduler: FlowMatchEulerDiscreteSDEScheduler + self._has_warned_inference_fallback = False + self._has_warned_forward_fallback = False + + def load_pipeline(self) -> Flux3ModularPipeline: + """Load the FLUX.3 modular pipeline specification from a local path or Hub repo.""" + if _FLUX3_IMPORT_ERROR is not None: + raise ImportError( + "Flux3Adapter requires the bundled Diffusers submodule at " + f"{FLUX3_DIFFUSERS_COMMIT}. Install it with `{FLUX3_INSTALL}`. " + f"The active diffusers version is {getattr(diffusers, '__version__', 'unknown')}." + ) from _FLUX3_IMPORT_ERROR + model_name_or_path = self.model_args.model_name_or_path + pipeline = ModularPipeline.from_pretrained(model_name_or_path) + if Path(model_name_or_path).is_dir(): + for name in pipeline.pretrained_component_names: + spec = pipeline._component_specs[name] + spec.pretrained_model_name_or_path = model_name_or_path + spec.revision = None + return pipeline + + def build_component_runtime(self) -> ModularPipelineRuntime: + """Build the modular runtime and materialize the target transformer.""" + runtime = ModularPipelineRuntime.from_adapter(self, self.load_pipeline()) + target_names = runtime.resolve_component_names(self.model_args.target_components) + runtime.materialize_components(target_names) + return runtime + + @property + def is_distilled(self) -> bool: + """Return whether the loaded checkpoint uses the step-distilled flow-map contract.""" + transformer = ( + self.get_component("transformer") if self.has_component("transformer") else None + ) + transformer_config = getattr(transformer, "config", None) + if getattr(transformer_config, "flow_map", None) is not None: + return bool(transformer_config.flow_map) + if Flux3DistilledModularPipeline is not Any and isinstance( + self.pipeline, Flux3DistilledModularPipeline + ): + return True + return "guider" not in self.component_runtime.declared_component_names + + def _resolve_pipeline_io_contract(self) -> Optional[PipelineIOContract]: + """Narrow negative-prompt support when a step-distilled flow-map checkpoint is loaded.""" + if self.is_distilled: + return image_output_contract( + negative_prompt=NegativePromptPolicy.UNSUPPORTED, + input_image_min_count=0, + input_image_max_count=None, + input_order=InputMediaOrder.WITHIN_TYPE, + ) + return type(self).pipeline_io_contract + + @property + def preprocessing_modules(self) -> List[str]: + return ["text_encoder", "tokenizer", "vae", "image_processor"] + + @property + def inference_modules(self) -> List[str]: + return ["transformer", "vae", "image_processor"] + + @property + def tokenizer(self) -> Any: + """Expose the tokenizer owned by the Qwen3-VL processor.""" + if self.get_component("tokenizer") is None and self.has_component("tokenizer"): + self.component_runtime.materialize_components(["tokenizer"]) + processor = self.get_component("tokenizer") + if processor is None: + raise ValueError("No tokenizer found in the FLUX.3 pipeline.") + return getattr(processor, "tokenizer", processor) + + @property + def processor(self) -> Any: + """Return the Qwen3-VL processor declared under the modular ``tokenizer`` slot.""" + if self.get_component("tokenizer") is None and self.has_component("tokenizer"): + self.component_runtime.materialize_components(["tokenizer"]) + processor = self.get_component("tokenizer") + if processor is None: + raise ValueError("No processor/tokenizer component found in the FLUX.3 pipeline.") + return processor + + @property + def default_target_modules(self) -> List[str]: + """Default LoRA target modules across all FLUX.3 mode and single-stream blocks.""" + return [ + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "ff.linear_in", + "ff.linear_out", + ] + + def _gradient_checkpointing_units( + self, + component_name: str, + component: torch.nn.Module, + ) -> List[CheckpointUnit]: + """Return FLUX.3 transformer blocks in forward execution order.""" + txt_blocks = getattr(component, "txt_mode_blocks", None) + mode_blocks = getattr(component, "mode_blocks", None) + single_blocks = getattr(component, "single_transformer_blocks", None) + image_blocks = ( + mode_blocks["image"] + if isinstance(mode_blocks, nn.ModuleDict) and "image" in mode_blocks + else None + ) + cond_blocks = ( + mode_blocks["image_cond"] + if isinstance(mode_blocks, nn.ModuleDict) and "image_cond" in mode_blocks + else None + ) + if ( + isinstance(txt_blocks, nn.ModuleList) + and isinstance(image_blocks, nn.ModuleList) + and isinstance(cond_blocks, nn.ModuleList) + and isinstance(single_blocks, nn.ModuleList) + ): + return [ + *[(f"txt_mode_blocks.{idx}", block) for idx, block in enumerate(txt_blocks)], + *[(f"mode_blocks.image.{idx}", block) for idx, block in enumerate(image_blocks)], + *[ + (f"mode_blocks.image_cond.{idx}", block) + for idx, block in enumerate(cond_blocks) + ], + *[ + (f"single_transformer_blocks.{idx}", block) + for idx, block in enumerate(single_blocks) + ], + ] + return super()._gradient_checkpointing_units(component_name, component) + + # ======================== Prompt Encoding ======================== + + @staticmethod + def _prepare_text_ids(seq_len: int, device: Optional[torch.device] = None) -> torch.Tensor: + """Prepare 4D position IDs ``(T, H, W, L)`` of shape ``(seq_len, 4)`` for text tokens.""" + if seq_len <= 0: + raise ValueError(f"expected positive text sequence length, received {seq_len}") + ids = torch.cartesian_prod( + torch.arange(1), + torch.arange(1), + torch.arange(1), + torch.arange(seq_len), + ).to(dtype=torch.float32) + if device is not None: + ids = ids.to(device) + return ids + + @staticmethod + def _encode_single_qwen3_vl_prompt( + text_encoder: Any, + processor: Any, + prompt: str, + *, + dtype: torch.dtype, + device: torch.device, + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH, + pad_multiple: int = TEXT_PAD_MULTIPLE, + text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Encode one prompt with the FLUX.3 80-token right-padded Qwen3-VL recipe.""" + if not isinstance(prompt, str): + raise TypeError( + f"expected prompt to be str, received {type(prompt).__name__}: {prompt!r}" + ) + if max_sequence_length <= 0: + raise ValueError( + f"expected positive max_sequence_length, received {max_sequence_length}" + ) + if pad_multiple <= 0: + raise ValueError(f"expected positive pad_multiple, received {pad_multiple}") + if not text_encoder_out_layers: + raise ValueError("expected non-empty text_encoder_out_layers") + + inner_tokenizer = getattr(processor, "tokenizer", processor) + chat_text = processor.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=False, + add_generation_prompt=True, + ) + unpadded_ids = inner_tokenizer( + chat_text, + return_tensors="pt", + truncation=True, + max_length=max_sequence_length, + )["input_ids"] + real_length = max(int(unpadded_ids.shape[1]), 1) + padded_length = min( + math.ceil(real_length / pad_multiple) * pad_multiple, + max_sequence_length, + ) + inputs = inner_tokenizer( + chat_text, + return_tensors="pt", + padding="max_length", + padding_side="right", + truncation=True, + max_length=padded_length, + ) + + input_ids = inputs["input_ids"].to(device) + attention_mask = inputs["attention_mask"].to(device) + # Qwen3-VL computes padding position IDs as 1 in transformers 4.57 (the FLUX.3 training + # environment) and 0 in 5.x. Pass explicit position IDs so padding tokens match upstream. + position_ids = torch.arange(padded_length, device=device)[None].masked_fill( + attention_mask == 0, 1 + ) + + output = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_hidden_states=True, + use_cache=False, + ) + + stacked = torch.stack([output.hidden_states[k] for k in text_encoder_out_layers], dim=1) + stacked = stacked.to(dtype=dtype, device=device) + batch_size, num_layers, seq_len, hidden_dim = stacked.shape + prompt_embeds = stacked.permute(0, 2, 1, 3).reshape( + batch_size, seq_len, num_layers * hidden_dim + ) + return input_ids.squeeze(0), prompt_embeds.squeeze(0) + + @classmethod + def _get_qwen3_vl_prompt_embeds( + cls, + text_encoder: Any, + processor: Any, + prompt: Union[str, Sequence[str]], + *, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH, + pad_multiple: int = TEXT_PAD_MULTIPLE, + text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, + ) -> Tuple[ + Union[torch.Tensor, List[torch.Tensor]], + Union[torch.Tensor, List[torch.Tensor]], + Union[torch.Tensor, List[torch.Tensor]], + ]: + """Encode one or more prompts while preserving each prompt's exact 80-token window.""" + resolved_dtype = text_encoder.dtype if dtype is None else dtype + resolved_device = text_encoder.device if device is None else device + prompts = [prompt] if isinstance(prompt, str) else list(prompt) + if not prompts: + raise ValueError("expected at least one prompt to encode") + + ids_list: List[torch.Tensor] = [] + embeds_list: List[torch.Tensor] = [] + text_ids_list: List[torch.Tensor] = [] + for single_prompt in prompts: + single_ids, single_embeds = cls._encode_single_qwen3_vl_prompt( + text_encoder=text_encoder, + processor=processor, + prompt=single_prompt, + dtype=resolved_dtype, + device=resolved_device, + max_sequence_length=max_sequence_length, + pad_multiple=pad_multiple, + text_encoder_out_layers=text_encoder_out_layers, + ) + ids_list.append(single_ids) + embeds_list.append(single_embeds) + text_ids_list.append( + cls._prepare_text_ids(single_embeds.shape[0], device=resolved_device) + ) + + lengths = {item.shape[0] for item in embeds_list} + if len(lengths) == 1: + return ( + torch.stack(ids_list, dim=0), + torch.stack(embeds_list, dim=0), + torch.stack(text_ids_list, dim=0), + ) + return ids_list, embeds_list, text_ids_list + + def encode_prompt( + self, + prompt: Union[str, List[str]], + negative_prompt: Optional[Union[str, List[str]]] = None, + guidance_scale: float = 3.5, + device: Optional[torch.device] = None, + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH, + text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, + pad_multiple: int = TEXT_PAD_MULTIPLE, + ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]: + """Encode prompt(s) and optional negative prompt(s) with Qwen3-VL.""" + if ( + self.get_component("text_encoder") is None or self.get_component("tokenizer") is None + ) and self.has_component("text_encoder"): + self.on_load_components(["text_encoder", "tokenizer"], device=device) + + text_encoder = self._require_component("text_encoder") + processor = self.processor + resolved_device = text_encoder.device if device is None else device + + if prompt is None: + prompt = "" + prompt_list = [prompt] if isinstance(prompt, str) else list(prompt) + + if self.is_distilled and negative_prompt is not None: + neg_values = ( + [negative_prompt] if isinstance(negative_prompt, str) else list(negative_prompt) + ) + if any(item is not None and item != "" for item in neg_values): + raise ValueError( + "Step-distilled FLUX.3 (flow_map=True) does not support negative_prompt; " + f"received {negative_prompt!r}." + ) + + do_classifier_free_guidance = (not self.is_distilled) and guidance_scale > 1.0 + + prompt_ids, prompt_embeds, text_ids = self._get_qwen3_vl_prompt_embeds( + text_encoder=text_encoder, + processor=processor, + prompt=prompt_list, + dtype=text_encoder.dtype, + device=resolved_device, + max_sequence_length=max_sequence_length, + pad_multiple=pad_multiple, + text_encoder_out_layers=text_encoder_out_layers, + ) + results: Dict[str, Union[torch.Tensor, List[torch.Tensor]]] = { + "prompt_ids": prompt_ids, + "prompt_embeds": prompt_embeds, + "text_ids": text_ids, + } + + if do_classifier_free_guidance: + negative_prompt = "" if negative_prompt is None else negative_prompt + neg_list = ( + [negative_prompt] if isinstance(negative_prompt, str) else list(negative_prompt) + ) + if len(neg_list) == 1 and len(prompt_list) > 1: + neg_list = neg_list * len(prompt_list) + if len(neg_list) != len(prompt_list): + raise ValueError( + f"The number of negative prompts ({len(neg_list)}) must match " + f"the number of prompts ({len(prompt_list)})." + ) + neg_ids, neg_embeds, neg_text_ids = self._get_qwen3_vl_prompt_embeds( + text_encoder=text_encoder, + processor=processor, + prompt=neg_list, + dtype=text_encoder.dtype, + device=resolved_device, + max_sequence_length=max_sequence_length, + pad_multiple=pad_multiple, + text_encoder_out_layers=text_encoder_out_layers, + ) + results.update( + { + "negative_prompt_ids": neg_ids, + "negative_prompt_embeds": neg_embeds, + "negative_text_ids": neg_text_ids, + } + ) + + return results + + # ======================== Image Encoding ======================== + + def encode_image( + self, + images: Union[ImageSingle, ImageBatch, MultiImageBatch], + condition_image_size: Union[int, Tuple[int, int]] = CONDITION_IMAGE_SIZE, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + generator: Optional[torch.Generator] = None, + ) -> Dict[str, Union[List[List[torch.Tensor]], List[Optional[torch.Tensor]]]]: + """Encode reference image(s) into normalized packed FLUX.3 conditioning tokens.""" + del generator + if self.get_component("vae") is None and self.has_component("vae"): + self.on_load_components(["vae", "image_processor"], device=device) + + vae = self._require_component("vae") + resolved_device = vae.device if device is None else device + transformer = ( + self.get_component("transformer") if self.has_component("transformer") else None + ) + resolved_dtype = ( + dtype if dtype is not None else getattr(transformer, "dtype", None) or vae.dtype + ) + + if not self._is_multi_images_batch(images): + images = [images] # type: ignore[list-item] + + standardized_batches = [ + ( + [] + if isinstance(imgs, list) and not imgs + else self._standardize_image_input(imgs, output_type="pil") + ) + for imgs in images # type: ignore[union-attr] + ] + + condition_image_tensors: List[List[torch.Tensor]] = [ + ( + self._resize_condition_images( + condition_images=imgs, + condition_image_size=condition_image_size, + ) + if imgs + else [] + ) + for imgs in standardized_batches + ] + + image_latents_list: List[Optional[torch.Tensor]] = [] + image_latent_ids_list: List[Optional[torch.Tensor]] = [] + for cond_img_tensors in condition_image_tensors: + if not cond_img_tensors: + image_latents_list.append(None) + image_latent_ids_list.append(None) + continue + image_latents, image_latent_ids = prepare_flux3_condition_latents( + self, + cond_img_tensors, + batch_size=1, + device=resolved_device, + dtype=resolved_dtype, + image_rope_scale=IMAGE_ROPE_SCALE, + ) + image_latents_list.append(image_latents.squeeze(0)) + image_latent_ids_list.append(image_latent_ids.squeeze(0)) + + postprocessed_condition_images: List[List[torch.Tensor]] = [ + [ + self.pipeline.image_processor.postprocess(img, output_type="pt")[0].to( + resolved_device + ) + for img in cond_img_tensors + ] + for cond_img_tensors in condition_image_tensors + ] + + return { + "condition_images": postprocessed_condition_images, + "image_latents": image_latents_list, + "image_latent_ids": image_latent_ids_list, + } + + def _resize_condition_images( + self, + condition_images: Union[Image.Image, List[Image.Image]], + condition_image_size: Union[int, Tuple[int, int]] = CONDITION_IMAGE_SIZE, + ) -> List[torch.Tensor]: + """Resize reference images when exceeding the area budget and center-crop to 16px grid.""" + if isinstance(condition_images, Image.Image): + condition_images = [condition_images] + + condition_images = self._standardize_image_input( + condition_images, + output_type="pil", + ) + + if isinstance(condition_image_size, int): + condition_image_size = (condition_image_size, condition_image_size) + elif not isinstance(condition_image_size, (tuple, list)) or len(condition_image_size) != 2: + raise ValueError( + f"expected condition_image_size as int or (height, width), " + f"received {condition_image_size!r}" + ) + max_area = int(condition_image_size[0]) * int(condition_image_size[1]) + if max_area <= 0: + raise ValueError( + f"expected positive condition_image_size area, received {condition_image_size!r}" + ) + + multiple_of = self.pipeline.vae_scale_factor * 2 + condition_image_tensors: List[torch.Tensor] = [] + for img in condition_images: + image_width, image_height = img.size + if image_width * image_height > max_area: + img = self.pipeline.image_processor._resize_to_target_area(img, max_area) + image_width, image_height = img.size + + cropped_width = max(multiple_of, (image_width // multiple_of) * multiple_of) + cropped_height = max(multiple_of, (image_height // multiple_of) * multiple_of) + tensor_img = self.pipeline.image_processor.preprocess( + img, + height=cropped_height, + width=cropped_width, + resize_mode="crop", + ) + condition_image_tensors.append(tensor_img) + + return condition_image_tensors + + @staticmethod + def _is_multi_images_batch(images: Optional[Union[ImageBatch, MultiImageBatch]]) -> bool: + return images is not None and is_multi_image_batch(images) + + @staticmethod + def _is_multi_image_latents( + image_latents: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]], + ) -> bool: + return ( + isinstance(image_latents, list) + and len(image_latents) > 0 + and all( + latent is None or (isinstance(latent, torch.Tensor) and latent.ndim == 2) + for latent in image_latents + ) + ) or (isinstance(image_latents, torch.Tensor) and image_latents.ndim == 3) + + def _standardize_image_input( + self, + images: Union[ImageSingle, ImageBatch], + output_type: Literal["pil", "pt", "np"] = "pil", + ) -> ImageBatch: + """Standardize image input to the requested container format.""" + if isinstance(images, Image.Image): + images = [images] + return standardize_image_batch(images, output_type=output_type) + + # ======================== Output Image Encoding & Decoding ======================== + + def _output_geometry_multiple(self) -> int: + """Require the 16-pixel grid imposed by the 8x VAE and 2x2 latent patchification.""" + return self.pipeline.vae_scale_factor * 2 + + def _encode_output_images( + self, + pixel_values: torch.Tensor, + condition: Mapping[str, Any], + generator: Optional[torch.Generator], + ) -> EncodedImageTensor: + """Sample and pack target images with the FLUX.3 latent recipe.""" + del condition + return encode_flux3_output_images( + self, + pixel_values, + generator, + image_rope_scale=IMAGE_ROPE_SCALE, + ) + + def _resolve_latent_grid( + self, + latents: torch.Tensor, + *, + latent_ids: Optional[torch.Tensor] = None, + height: Optional[int] = None, + width: Optional[int] = None, + ) -> Tuple[int, int]: + """Resolve ``(latent_height, latent_width)`` from explicit dimensions or position IDs.""" + scale = self.pipeline.vae_scale_factor * 2 + seq_len = latents.shape[1] + if height is not None and width is not None: + latent_height = int(height) // scale + latent_width = int(width) // scale + if latent_height * latent_width != seq_len: + raise ValueError( + f"FLUX.3 latents has {seq_len} tokens, which does not match " + f"height x width = {height} x {width} ({latent_height * latent_width} tokens)" + ) + return latent_height, latent_width + + if latent_ids is not None: + ids_2d = self._as_2d_position_ids(latent_ids, "latent_ids") + max_h = float(ids_2d[:, 1].max().item()) + max_w = float(ids_2d[:, 2].max().item()) + latent_height = int(round(max_h / IMAGE_ROPE_SCALE)) + 1 + latent_width = int(round(max_w / IMAGE_ROPE_SCALE)) + 1 + if latent_height * latent_width == seq_len: + return latent_height, latent_width + + side = int(math.isqrt(seq_len)) + if side * side == seq_len: + return side, side + raise ValueError( + f"Cannot infer FLUX.3 2D latent grid for seq_len={seq_len}; " + "pass explicit height/width or valid latent_ids." + ) + + def decode_latents( + self, + latents: torch.Tensor, + latent_ids: Optional[torch.Tensor] = None, + height: Optional[int] = None, + width: Optional[int] = None, + output_type: Literal["pil", "pt", "np"] = "pil", + ) -> Union[List[Image.Image], torch.Tensor, np.ndarray]: + """Unpack, BatchNorm-denormalize, unpatchify, and decode FLUX.3 image latents.""" + if self.get_component("vae") is None and self.has_component("vae"): + self.on_load_components(["vae", "image_processor"], device=latents.device) + + vae = self._require_component("vae") + latent_height, latent_width = self._resolve_latent_grid( + latents, + latent_ids=latent_ids, + height=height, + width=width, + ) + unpacked = unpack_flux3_latents(latents, latent_height, latent_width) + + latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(unpacked.device, unpacked.dtype) + latents_bn_std = torch.sqrt( + vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps + ).to(unpacked.device, unpacked.dtype) + unpacked = unpacked * latents_bn_std + latents_bn_mean + unpatched = unpatchify_flux2_latents(unpacked) + + decoded = vae.decode(unpatched.to(device=vae.device, dtype=vae.dtype), return_dict=False)[0] + return self.pipeline.image_processor.postprocess(decoded, output_type=output_type) + + # ======================== Inference ======================== + + @staticmethod + def _prepare_latents( + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + *, + vae_scale_factor: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Draw initial float32 noise on the 2D grid, cast, pack to tokens, and build RoPE IDs.""" + image_scale_factor = vae_scale_factor * 2 + if height % image_scale_factor != 0 or width % image_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {image_scale_factor}, " + f"received {height} and {width}." + ) + latent_height = height // image_scale_factor + latent_width = width // image_scale_factor + if latents is None: + shape = (batch_size, num_channels_latents, latent_height, latent_width) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"Passed a list of {len(generator)} generators for batch_size={batch_size}." + ) + latents = randn_tensor( + shape, + generator=generator, + device=device, + dtype=torch.float32, + ).to(dtype) + else: + latents = latents.to(device=device, dtype=dtype) + if latents.ndim == 4: + latents = pack_flux3_latents(latents) + latent_ids = ( + prepare_flux3_image_ids( + latent_height, + latent_width, + image_rope_scale=IMAGE_ROPE_SCALE, + l_coord=0, + device=device, + ) + .unsqueeze(0) + .repeat(batch_size, 1, 1) + ) + return latents, latent_ids + + @torch.no_grad() + def _inference( + self, + images: Optional[Union[ImageBatch, MultiImageBatch]] = None, + prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: Optional[Union[str, List[str]]] = None, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 50, + guidance_scale: float = 3.5, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + prompt_ids: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + text_ids: Optional[torch.Tensor] = None, + negative_prompt_ids: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_text_ids: Optional[torch.Tensor] = None, + condition_images: Optional[MultiImageBatch] = None, + image_latents: Optional[torch.Tensor] = None, + image_latent_ids: Optional[torch.Tensor] = None, + condition_image_size: Union[int, Tuple[int, int]] = CONDITION_IMAGE_SIZE, + attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH, + text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, + pad_multiple: int = TEXT_PAD_MULTIPLE, + compute_log_prob: bool = False, + extra_call_back_kwargs: List[str] = [], + trajectory_indices: TrajectoryIndicesType = "all", + ) -> List[Flux3Sample]: + """Run FLUX.3 rollout on a uniform-window T2I batch or a single sample.""" + device = self.device + transformer = self._require_component("transformer") + dtype = transformer.dtype + do_classifier_free_guidance = (not self.is_distilled) and guidance_scale > 1.0 + resolved_attention_kwargs = ( + attention_kwargs if attention_kwargs is not None else joint_attention_kwargs + ) + + # 1. Encode prompt + if isinstance(prompt, str): + prompt = [prompt] + + if (prompt_embeds is None or prompt_ids is None or text_ids is None) or ( + do_classifier_free_guidance + and ( + negative_prompt_embeds is None + or negative_prompt_ids is None + or negative_text_ids is None + ) + ): + prompt_encoding = self.encode_prompt( + prompt=prompt or [""], + negative_prompt=negative_prompt, + guidance_scale=guidance_scale, + device=device, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + pad_multiple=pad_multiple, + ) + prompt_ids = prompt_encoding["prompt_ids"] # type: ignore[assignment] + prompt_embeds = prompt_encoding["prompt_embeds"] # type: ignore[assignment] + text_ids = prompt_encoding["text_ids"] # type: ignore[assignment] + if do_classifier_free_guidance: + negative_prompt_ids = prompt_encoding["negative_prompt_ids"] # type: ignore[assignment] + negative_prompt_embeds = prompt_encoding["negative_prompt_embeds"] # type: ignore[assignment] + negative_text_ids = prompt_encoding["negative_text_ids"] # type: ignore[assignment] + else: + prompt_ids = prompt_ids.to(device) + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + text_ids = text_ids.to(device) + if negative_prompt_ids is not None: + negative_prompt_ids = negative_prompt_ids.to(device) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + if negative_text_ids is not None: + negative_text_ids = negative_text_ids.to(device) + + if ( + not isinstance(prompt_embeds, torch.Tensor) + or not isinstance(prompt_ids, torch.Tensor) + or not isinstance(text_ids, torch.Tensor) + ): + raise ValueError( + "Flux3Adapter._inference requires uniform stacked prompt tensors; " + "use Flux3Adapter.inference() for variable-window batches." + ) + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + if isinstance(negative_prompt_embeds, torch.Tensor): + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + + batch_size = prompt_embeds.shape[0] + + # 2. Encode condition images + images = ( + [images] if images is not None and not self._is_multi_images_batch(images) else images + ) + if images is not None and ( + condition_images is None or image_latents is None or image_latent_ids is None + ): + image_encoding = self.encode_image( + images=images, + condition_image_size=condition_image_size, + device=device, + dtype=dtype, + generator=generator[0] if isinstance(generator, list) else generator, + ) + condition_images = image_encoding["condition_images"] # type: ignore[assignment] + raw_latents = image_encoding["image_latents"] + raw_ids = image_encoding["image_latent_ids"] + if any(item is None for item in raw_latents): + image_latents = None + image_latent_ids = None + else: + image_latents = torch.stack( + [item for item in raw_latents if item is not None], dim=0 + ).to(device=device, dtype=dtype) + image_latent_ids = torch.stack( + [item for item in raw_ids if item is not None], dim=0 + ).to(device=device) + else: + image_latents = ( + image_latents.to(device=device, dtype=dtype) if image_latents is not None else None + ) + image_latent_ids = ( + image_latent_ids.to(device=device) if image_latent_ids is not None else None + ) + + # 3. Prepare initial latents + num_channels_latents = int(getattr(transformer.config, "in_channels", 128)) + latents, latent_ids = self._prepare_latents( + batch_size=batch_size, + num_channels_latents=num_channels_latents, + height=height, + width=width, + vae_scale_factor=self.pipeline.vae_scale_factor, + dtype=dtype, + device=device, + generator=generator, + ) + + # 4. Prepare timesteps (constant shift 6.93 on the scheduler) + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + timesteps = set_scheduler_timesteps( + scheduler=self.scheduler, + num_inference_steps=num_inference_steps, + sigmas=sigmas, + seq_len=latents.shape[1], + device=device, + ) + + # 5. Denoising loop + latent_collector = create_trajectory_collector(trajectory_indices, num_inference_steps) + latents = self.cast_latents(latents, default_dtype=dtype) + latent_collector.collect(latents, step_idx=0) + if compute_log_prob: + log_prob_collector = create_trajectory_collector( + trajectory_indices, num_inference_steps + ) + callback_collector = create_callback_collector(trajectory_indices, num_inference_steps) + + for i, t in enumerate(timesteps): + current_noise_level = self.scheduler.get_noise_level_for_timestep(t) + t_next = ( + timesteps[i + 1] + if i + 1 < len(timesteps) + else torch.tensor(0, device=device, dtype=timesteps.dtype) + ) + return_kwargs = list( + set(["next_latents", "log_prob", "velocity"] + extra_call_back_kwargs) + ) + current_compute_log_prob = compute_log_prob and current_noise_level > 0 + + output = self._forward( + t=t, + t_next=t_next, + latents=latents, + latent_ids=latent_ids, + prompt_embeds=prompt_embeds, + text_ids=text_ids, + image_latents=image_latents, + image_latent_ids=image_latent_ids, + negative_prompt_embeds=negative_prompt_embeds, + negative_text_ids=negative_text_ids, + guidance_scale=guidance_scale, + attention_kwargs=resolved_attention_kwargs, + compute_log_prob=current_compute_log_prob, + return_kwargs=return_kwargs, + noise_level=current_noise_level, + ) + + latents = self.cast_latents(output.next_latents, default_dtype=dtype) + latent_collector.collect(latents, i + 1) + if current_compute_log_prob: + log_prob_collector.collect(output.log_prob, i) + + callback_collector.collect_step( + step_idx=i, + output=output, + keys=extra_call_back_kwargs, + capturable={"noise_level": current_noise_level}, + ) + + # 6. Decode latents to images + decoded_images = self.decode_latents( + latents, + latent_ids=latent_ids, + height=height, + width=width, + output_type="pt", + ) + + # 7. Assemble Flux3Sample instances + extra_call_back_res = callback_collector.get_result() + callback_index_map = callback_collector.get_index_map() + all_latents = latent_collector.get_result() + latent_index_map = latent_collector.get_index_map() + all_log_probs = log_prob_collector.get_result() if compute_log_prob else None + log_prob_index_map = log_prob_collector.get_index_map() if compute_log_prob else None + + samples = [ + Flux3Sample( + timesteps=timesteps, + all_latents=( + torch.stack([lat[b] for lat in all_latents], dim=0) + if all_latents is not None + else None + ), + log_probs=( + torch.stack([lp[b] for lp in all_log_probs], dim=0) + if all_log_probs is not None + else None + ), + latent_index_map=latent_index_map, + log_prob_index_map=log_prob_index_map, + height=height, + width=width, + image=decoded_images[b], + latent_ids=latent_ids[b], + prompt=prompt[b] if isinstance(prompt, list) else prompt, + prompt_ids=prompt_ids[b], + prompt_embeds=prompt_embeds[b], + text_ids=text_ids[b], + negative_prompt=( + negative_prompt[b] if isinstance(negative_prompt, list) else negative_prompt + ), + negative_prompt_ids=( + negative_prompt_ids[b] if negative_prompt_ids is not None else None + ), + negative_prompt_embeds=( + negative_prompt_embeds[b] if negative_prompt_embeds is not None else None + ), + negative_text_ids=(negative_text_ids[b] if negative_text_ids is not None else None), + condition_images=condition_images[b] if condition_images is not None else None, + image_latents=image_latents[b] if image_latents is not None else None, + image_latent_ids=image_latent_ids[b] if image_latent_ids is not None else None, + extra_kwargs={ + **{k: v[b] for k, v in extra_call_back_res.items()}, + "callback_index_map": callback_index_map, + }, + ) + for b in range(batch_size) + ] + + if hasattr(self.pipeline, "maybe_free_model_hooks"): + self.pipeline.maybe_free_model_hooks() + + return samples + + @torch.no_grad() + def inference( + self, + images: Optional[MultiImageBatch] = None, + prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: Optional[Union[str, List[str]]] = None, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 50, + guidance_scale: float = 3.5, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + condition_image_size: Union[int, Tuple[int, int]] = CONDITION_IMAGE_SIZE, + attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH, + text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, + pad_multiple: int = TEXT_PAD_MULTIPLE, + prompt_ids: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + prompt_embeds: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + text_ids: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + negative_prompt_ids: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + negative_prompt_embeds: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + negative_text_ids: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + condition_images: Optional[MultiImageBatch] = None, + image_latents: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + image_latent_ids: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + compute_log_prob: bool = False, + extra_call_back_kwargs: List[str] = [], + trajectory_indices: TrajectoryIndicesType = "all", + ) -> List[Flux3Sample]: + """Run FLUX.3 inference; falls back to per-sample execution for I2I or ragged windows.""" + if isinstance(prompt, str): + prompt = [prompt] + if isinstance(negative_prompt, str): + negative_prompt = [negative_prompt] * (len(prompt) if prompt is not None else 1) + + is_nested_images = self._is_multi_images_batch(images) + is_nested_image_latents = self._is_multi_image_latents(image_latents) + is_ragged_prompts = isinstance(prompt_embeds, list) or isinstance( + negative_prompt_embeds, list + ) + + # If raw prompts are passed with batch_size > 1 and no pre-encoded embeds, check whether + # their 80-token windows differ by encoding once upfront. + if ( + prompt is not None + and len(prompt) > 1 + and prompt_embeds is None + and not (is_nested_images or is_nested_image_latents) + ): + encoded = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + guidance_scale=guidance_scale, + device=self.device, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + pad_multiple=pad_multiple, + ) + prompt_ids = encoded["prompt_ids"] + prompt_embeds = encoded["prompt_embeds"] + text_ids = encoded["text_ids"] + negative_prompt_ids = encoded.get("negative_prompt_ids") + negative_prompt_embeds = encoded.get("negative_prompt_embeds") + negative_text_ids = encoded.get("negative_text_ids") + is_ragged_prompts = isinstance(prompt_embeds, list) or isinstance( + negative_prompt_embeds, list + ) + + fall_back = is_nested_images or is_nested_image_latents or is_ragged_prompts + if not fall_back: + return self._inference( + images=images, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator, + condition_image_size=condition_image_size, + attention_kwargs=attention_kwargs, + joint_attention_kwargs=joint_attention_kwargs, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + pad_multiple=pad_multiple, + prompt_ids=prompt_ids, # type: ignore[arg-type] + prompt_embeds=prompt_embeds, # type: ignore[arg-type] + text_ids=text_ids, # type: ignore[arg-type] + negative_prompt_ids=negative_prompt_ids, # type: ignore[arg-type] + negative_prompt_embeds=negative_prompt_embeds, # type: ignore[arg-type] + negative_text_ids=negative_text_ids, # type: ignore[arg-type] + condition_images=condition_images, + image_latents=image_latents, # type: ignore[arg-type] + image_latent_ids=image_latent_ids, # type: ignore[arg-type] + compute_log_prob=compute_log_prob, + extra_call_back_kwargs=extra_call_back_kwargs, + trajectory_indices=trajectory_indices, + ) + + if not self._has_warned_inference_fallback: + logger.warning( + "FLUX.3 per-sample inference fallback activated (image conditioning or " + "variable 80-token prompt windows). This warning will only appear once." + ) + self._has_warned_inference_fallback = True + + if images is not None: + batch_size = len(images) + elif image_latents is not None: + batch_size = len(image_latents) + elif prompt_embeds is not None: + batch_size = len(prompt_embeds) + elif prompt is not None: + batch_size = len(prompt) + else: + raise ValueError("Flux3Adapter.inference requires prompt, prompt_embeds, or images.") + + samples: List[Flux3Sample] = [] + for idx in range(batch_size): + this_prompt = [prompt[idx]] if prompt is not None else None + this_neg_prompt = ( + [negative_prompt[idx]] if isinstance(negative_prompt, list) else negative_prompt + ) + this_prompt_ids = ( + prompt_ids[idx].unsqueeze(0) + if prompt_ids is not None and prompt_ids[idx] is not None + else None + ) + this_prompt_embeds = ( + prompt_embeds[idx].unsqueeze(0) + if prompt_embeds is not None and prompt_embeds[idx] is not None + else None + ) + this_text_ids = ( + text_ids[idx].unsqueeze(0) + if text_ids is not None and text_ids[idx] is not None + else None + ) + this_neg_prompt_ids = ( + negative_prompt_ids[idx].unsqueeze(0) + if negative_prompt_ids is not None and negative_prompt_ids[idx] is not None + else None + ) + this_neg_prompt_embeds = ( + negative_prompt_embeds[idx].unsqueeze(0) + if negative_prompt_embeds is not None and negative_prompt_embeds[idx] is not None + else None + ) + this_neg_text_ids = ( + negative_text_ids[idx].unsqueeze(0) + if negative_text_ids is not None and negative_text_ids[idx] is not None + else None + ) + this_images = images[idx] if images is not None else None + this_condition_images = ( + condition_images[idx : idx + 1] if condition_images is not None else None + ) + this_image_latents = ( + image_latents[idx].unsqueeze(0) + if image_latents is not None and image_latents[idx] is not None + else None + ) + this_image_latent_ids = ( + image_latent_ids[idx].unsqueeze(0) + if image_latent_ids is not None and image_latent_ids[idx] is not None + else None + ) + sample = self._inference( + images=this_images, + prompt=this_prompt, + negative_prompt=this_neg_prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator[idx] if isinstance(generator, list) else generator, + condition_image_size=condition_image_size, + attention_kwargs=attention_kwargs, + joint_attention_kwargs=joint_attention_kwargs, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + pad_multiple=pad_multiple, + prompt_ids=this_prompt_ids, + prompt_embeds=this_prompt_embeds, + text_ids=this_text_ids, + negative_prompt_ids=this_neg_prompt_ids, + negative_prompt_embeds=this_neg_prompt_embeds, + negative_text_ids=this_neg_text_ids, + condition_images=this_condition_images, + image_latents=this_image_latents, + image_latent_ids=this_image_latent_ids, + compute_log_prob=compute_log_prob, + extra_call_back_kwargs=extra_call_back_kwargs, + trajectory_indices=trajectory_indices, + ) + samples.extend(sample) + + return samples + + # ======================== Forward ======================== + + @staticmethod + def _as_2d_position_ids(ids: torch.Tensor, name: str) -> torch.Tensor: + """Normalize ``(seq_len, 4)`` or batched ``(B, seq_len, 4)`` position IDs to 2D.""" + if not isinstance(ids, torch.Tensor): + raise TypeError(f"expected {name} to be torch.Tensor, received {type(ids).__name__}") + if ids.ndim == 2 and ids.shape[-1] == 4: + return ids + if ids.ndim == 3 and ids.shape[-1] == 4: + return ids[0] + raise ValueError( + f"expected {name} to have shape (seq_len, 4) or (B, seq_len, 4), " + f"received {tuple(ids.shape)}" + ) + + def _forward( + self, + t: torch.Tensor, + latents: torch.Tensor, + latent_ids: torch.Tensor, + prompt_embeds: torch.Tensor, + text_ids: torch.Tensor, + image_latents: Optional[torch.Tensor] = None, + image_latent_ids: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_text_ids: Optional[torch.Tensor] = None, + guidance_scale: float = 3.5, + t_next: Optional[torch.Tensor] = None, + next_latents: Optional[torch.Tensor] = None, + noise_level: Optional[float] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + compute_log_prob: bool = True, + return_kwargs: Sequence[str] = ( + "velocity", + "next_latents", + "next_latents_mean", + "std_dev_t", + "dt", + "log_prob", + ), + ) -> FlowMatchEulerDiscreteSDESchedulerOutput: + """Core single-step FLUX.3 forward handling T2I, I2I, CFG, and step-distilled flow maps.""" + batch_size = latents.shape[0] + transformer = self.transformer + transformer_dtype = getattr(transformer, "dtype", latents.dtype) + resolved_attention_kwargs = ( + attention_kwargs if attention_kwargs is not None else joint_attention_kwargs + ) + + is_distilled = self.is_distilled + if not is_distilled and guidance_scale > 1.0 and negative_prompt_embeds is None: + logger.warning( + "Passed `guidance_scale=%s` > 1.0 to Flux3Adapter, but `negative_prompt_embeds` " + "is None. Falling back to no classifier-free guidance.", + guidance_scale, + ) + do_classifier_free_guidance = ( + not is_distilled and guidance_scale > 1.0 and negative_prompt_embeds is not None + ) + + # 1. Prepare inputs in transformer compute dtype and 2D position IDs + latent_model_input = latents.to(dtype=transformer_dtype) + cond_prompt_embeds = prompt_embeds.to(dtype=transformer_dtype) + img_ids_2d = self._as_2d_position_ids(latent_ids, "latent_ids").to(latents.device) + txt_ids_2d = self._as_2d_position_ids(text_ids, "text_ids").to(latents.device) + + image_cond_hidden_states: Optional[torch.Tensor] = None + image_cond_ids_2d: Optional[torch.Tensor] = None + if image_latents is not None: + if image_latent_ids is None: + raise ValueError("`image_latents` requires `image_latent_ids` in Flux3Adapter.") + image_cond_hidden_states = image_latents.to( + device=latents.device, dtype=transformer_dtype + ) + image_cond_ids_2d = self._as_2d_position_ids(image_latent_ids, "image_latent_ids").to( + latents.device + ) + + timestep = (t.expand(batch_size) / 1000.0).to( + device=latents.device, dtype=transformer_dtype + ) + target_timestep: Optional[torch.Tensor] = None + if is_distilled: + target_t = ( + torch.zeros(batch_size, device=latents.device, dtype=t.dtype) + if t_next is None + else t_next.expand(batch_size) + ) + target_timestep = (target_t / 1000.0).to(device=latents.device, dtype=transformer_dtype) + + # 2. Conditional forward pass + velocity = transformer( + hidden_states=latent_model_input, + encoder_hidden_states=cond_prompt_embeds, + timestep=timestep, + img_ids=img_ids_2d, + txt_ids=txt_ids_2d, + image_cond_hidden_states=image_cond_hidden_states, + image_cond_ids=image_cond_ids_2d, + target_timestep=target_timestep, + attention_kwargs=resolved_attention_kwargs, + return_dict=False, + )[0] + + # 3. Unconditional forward pass for true CFG + if do_classifier_free_guidance: + uncond_prompt_embeds = negative_prompt_embeds.to(dtype=transformer_dtype) + uncond_txt_ids_2d = ( + self._as_2d_position_ids(negative_text_ids, "negative_text_ids").to(latents.device) + if negative_text_ids is not None + else self._prepare_text_ids(uncond_prompt_embeds.shape[1], device=latents.device) + ) + neg_velocity = transformer( + hidden_states=latent_model_input, + encoder_hidden_states=uncond_prompt_embeds, + timestep=timestep, + img_ids=img_ids_2d, + txt_ids=uncond_txt_ids_2d, + image_cond_hidden_states=image_cond_hidden_states, + image_cond_ids=image_cond_ids_2d, + target_timestep=target_timestep, + attention_kwargs=resolved_attention_kwargs, + return_dict=False, + )[0] + velocity = neg_velocity + guidance_scale * (velocity - neg_velocity) + + # 4. Scheduler step + output = self.scheduler.step( + velocity=velocity, + timestep=t, + latents=latents, + timestep_next=t_next, + next_latents=next_latents, + compute_log_prob=compute_log_prob, + return_dict=True, + return_kwargs=list(return_kwargs), + noise_level=noise_level, + ) + return output + + def forward( + self, + t: torch.Tensor, + latents: torch.Tensor, + latent_ids: Union[torch.Tensor, List[torch.Tensor]], + prompt_embeds: Union[torch.Tensor, List[torch.Tensor]], + text_ids: Union[torch.Tensor, List[torch.Tensor]], + image_latents: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + image_latent_ids: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + negative_prompt_embeds: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + negative_text_ids: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + guidance_scale: float = 3.5, + t_next: Optional[torch.Tensor] = None, + next_latents: Optional[torch.Tensor] = None, + noise_level: Optional[float] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + compute_log_prob: bool = True, + return_kwargs: Sequence[str] = ( + "velocity", + "next_latents", + "next_latents_mean", + "std_dev_t", + "dt", + "log_prob", + ), + ) -> FlowMatchEulerDiscreteSDESchedulerOutput: + """General forward pass supporting T2I, multi-reference I2I, and ragged prompt windows.""" + has_image_cond = image_latents is not None and ( + isinstance(image_latents, torch.Tensor) + or any(item is not None for item in image_latents) + ) + has_ragged_prompts = ( + isinstance(prompt_embeds, list) + or isinstance(text_ids, list) + or isinstance(negative_prompt_embeds, list) + or isinstance(negative_text_ids, list) + ) + fall_back = has_image_cond or has_ragged_prompts + + if not fall_back: + return self._forward( + t=t, + latents=latents, + latent_ids=latent_ids, # type: ignore[arg-type] + prompt_embeds=prompt_embeds, # type: ignore[arg-type] + text_ids=text_ids, # type: ignore[arg-type] + image_latents=None, + image_latent_ids=None, + negative_prompt_embeds=negative_prompt_embeds, # type: ignore[arg-type] + negative_text_ids=negative_text_ids, # type: ignore[arg-type] + guidance_scale=guidance_scale, + t_next=t_next, + next_latents=next_latents, + noise_level=noise_level, + attention_kwargs=attention_kwargs, + joint_attention_kwargs=joint_attention_kwargs, + compute_log_prob=compute_log_prob, + return_kwargs=return_kwargs, + ) + + if not self._has_warned_forward_fallback: + logger.warning( + "FLUX.3: Image conditioning or variable prompt window detected in forward(). " + "Processing samples individually (warning shown once)." + ) + self._has_warned_forward_fallback = True + + batch_size = latents.shape[0] + outputs: List[FlowMatchEulerDiscreteSDESchedulerOutput] = [] + + for idx in range(batch_size): + single_t = t[idx].unsqueeze(0) if t.ndim > 0 and t.shape[0] == batch_size else t + single_t_next = ( + t_next[idx].unsqueeze(0) + if t_next is not None and t_next.ndim > 0 and t_next.shape[0] == batch_size + else t_next + ) + single_latents = latents[idx].unsqueeze(0) + single_latent_ids = ( + latent_ids[idx].unsqueeze(0) + if isinstance(latent_ids, list) or latent_ids.ndim == 3 + else latent_ids + ) + single_next_latents = ( + next_latents[idx].unsqueeze(0) if next_latents is not None else None + ) + single_prompt_embeds = prompt_embeds[idx].unsqueeze(0) + single_text_ids = ( + text_ids[idx].unsqueeze(0) + if isinstance(text_ids, list) or text_ids.ndim == 3 + else text_ids + ) + single_image_latents = ( + image_latents[idx].unsqueeze(0) + if image_latents is not None and image_latents[idx] is not None + else None + ) + single_image_latent_ids = ( + image_latent_ids[idx].unsqueeze(0) + if image_latent_ids is not None and image_latent_ids[idx] is not None + else None + ) + single_neg_prompt_embeds = ( + negative_prompt_embeds[idx].unsqueeze(0) + if negative_prompt_embeds is not None and negative_prompt_embeds[idx] is not None + else None + ) + single_neg_text_ids = ( + negative_text_ids[idx].unsqueeze(0) + if negative_text_ids is not None and negative_text_ids[idx] is not None + else None + ) + + out = self._forward( + t=single_t, + latents=single_latents, + latent_ids=single_latent_ids, + prompt_embeds=single_prompt_embeds, + text_ids=single_text_ids, + image_latents=single_image_latents, + image_latent_ids=single_image_latent_ids, + negative_prompt_embeds=single_neg_prompt_embeds, + negative_text_ids=single_neg_text_ids, + guidance_scale=guidance_scale, + t_next=single_t_next, + next_latents=single_next_latents, + noise_level=noise_level, + attention_kwargs=attention_kwargs, + joint_attention_kwargs=joint_attention_kwargs, + compute_log_prob=compute_log_prob, + return_kwargs=return_kwargs, + ) + outputs.append(out) + + outputs_dict = [o.to_dict() for o in outputs] + return FlowMatchEulerDiscreteSDESchedulerOutput.from_dict( + { + key: ( + torch.cat([item[key] for item in outputs_dict], dim=0) + if outputs_dict[0][key] is not None + else None + ) + for key in outputs_dict[0].keys() + } + ) diff --git a/src/flow_factory/models/qwen_image_21/output_codec.py b/src/flow_factory/models/qwen_image_21/output_codec.py index 8113ae22..3a5fd2b1 100644 --- a/src/flow_factory/models/qwen_image_21/output_codec.py +++ b/src/flow_factory/models/qwen_image_21/output_codec.py @@ -23,7 +23,7 @@ from ..configured_image_output import EncodedImageTensor -QWEN_IMAGE_21_DIFFUSERS_COMMIT = "80c7ed262aeffbeb43ef13ae04baeb9b84515a69" +QWEN_IMAGE_21_DIFFUSERS_COMMIT = "7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941" def _retrieve_vae_latents( diff --git a/src/flow_factory/models/registry.py b/src/flow_factory/models/registry.py index 6d42da5a..83e6a02d 100644 --- a/src/flow_factory/models/registry.py +++ b/src/flow_factory/models/registry.py @@ -30,6 +30,7 @@ "flux1-kontext": "flow_factory.models.flux.flux1_kontext.Flux1KontextAdapter", "flux2": "flow_factory.models.flux.flux2.Flux2Adapter", "flux2-klein": "flow_factory.models.flux.flux2_klein.Flux2KleinAdapter", + "flux3": "flow_factory.models.flux.flux3.Flux3Adapter", "qwen-image": "flow_factory.models.qwen_image.qwen_image.QwenImageAdapter", "qwen-image-edit-plus": "flow_factory.models.qwen_image.qwen_image_edit_plus.QwenImageEditPlusAdapter", "qwen-image-2.1": "flow_factory.models.qwen_image_21.qwen_image_21.QwenImage21Adapter", diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py index a9386dda..fd7281ed 100644 --- a/src/flow_factory/trainers/abc.py +++ b/src/flow_factory/trainers/abc.py @@ -181,6 +181,48 @@ def validate_samples(self, samples: Sequence[BaseSample], *, context: str) -> No ) +def _require_rollout_packs( + manifest: Optional[AcquisitionManifest], + samples: List[BaseSample], + per_device_batch_size: int, + purpose: str, +) -> Tuple[Tuple[int, ...], ...]: + """Return the recorded rollout packs of ``samples`` for a pack-composition-dependent adapter. + + Args: + manifest: Acquisition manifest recorded by the generation loop, if any. + samples: Rank-local samples, which must still be in rollout order. + per_device_batch_size: Optimize micro-batch size every rollout pack must match. + purpose: What needs the packs; it prefixes the missing-metadata error. + + Returns: + One tuple of sample indices per rollout micro-batch, in rollout order. + + Raises: + RuntimeError: If the generation loop recorded no rollout micro-batches. + ValueError: If a rollout micro-batch does not match ``per_device_batch_size`` or the + samples changed since rollout. + """ + if manifest is None or not manifest.rollout_batches: + raise RuntimeError( + f"{purpose} for a pack-composition-dependent adapter requires rollout " + "micro-batch metadata; custom generate_samples() implementations must " + "record the base generation-loop batch boundaries" + ) + + invalid_sizes = tuple( + len(batch) for batch in manifest.rollout_batches if len(batch) != per_device_batch_size + ) + if invalid_sizes: + raise ValueError( + "pack-composition-dependent replay requires every rollout micro-batch to match " + f"train.per_device_batch_size={per_device_batch_size}; invalid_sizes={invalid_sizes!r}" + ) + + manifest.validate_samples(samples) + return manifest.rollout_batches + + class BaseTrainer(MultiRoleCheckpointingMixin, MultiRoleBackendValidationMixin, ABC): """Abstract base class for Flow-Factory trainers.""" @@ -231,6 +273,7 @@ def __init__( self._acquisition_cycle_active = False self._acquisition_cycle_incomplete = False self._reward_overlap_acquisition_manifest: Optional[AcquisitionManifest] = None + self._rollout_acquisition_manifest: Optional[AcquisitionManifest] = None self.acquisition_driver: AcquisitionDriver = build_acquisition_driver( type(self).execution_contract ) @@ -2241,24 +2284,8 @@ def _validate_reward_overlap_replay_batch_composition( ) manifest = plan.manifest - if manifest is None or not manifest.rollout_batches: - raise RuntimeError( - "reward overlap for a pack-composition-dependent adapter requires rollout " - "micro-batch metadata; custom generate_samples() implementations must " - "record the base generation-loop batch boundaries" - ) - batch_size = self.training_args.per_device_batch_size - invalid_sizes = tuple( - len(batch) for batch in manifest.rollout_batches if len(batch) != batch_size - ) - if invalid_sizes: - raise ValueError( - "pack-composition-dependent replay requires every rollout micro-batch to " - f"match train.per_device_batch_size={batch_size}; invalid_sizes={invalid_sizes!r}" - ) - - manifest.validate_samples(samples) + _require_rollout_packs(manifest, samples, batch_size, "reward overlap") boundaries = manifest.rollout_boundaries for tile in plan.tiles: if tile.start not in boundaries or tile.stop not in boundaries: @@ -2865,17 +2892,29 @@ def _order_samples_for_optimize( ) -> List[BaseSample]: """Return the per-inner-epoch sample ordering for the optimize loop. - When ``training_args.shuffle_samples`` is False, the rollout-pack order is - preserved so each training micro-batch packs exactly the samples of its - corresponding rollout ``inference`` pack. For adapters whose batched forward - is pack-composition-dependent (e.g. Bagel/NaViT packing), this keeps the - bf16 forward bit-identical between rollout and training (on-policy ratio==1). + When ``training_args.shuffle_samples`` is False, the rollout order is kept, so each + training micro-batch packs exactly the samples of its rollout ``inference`` pack. + Otherwise a batch-invariant adapter's samples are shuffled individually. An adapter + that declares ``requires_preserved_replay_batch_composition`` (its batched forward + depends on which samples share a batch, e.g. Bagel's NaViT packing or FLUX.3's + window-dependent batching) has its recorded rollout packs shuffled as whole units: + every micro-batch then replays one rollout pack with its samples in rollout order, + so the bf16 forward stays bit-identical to the rollout (on-policy ratio == 1). """ if not self.training_args.shuffle_samples: return samples perm_gen = create_generator(self.training_args.seed, self.epoch, inner_epoch) - perm = torch.randperm(len(samples), generator=perm_gen) - return [samples[i] for i in perm] + if not self.adapter.requires_preserved_replay_batch_composition: + perm = torch.randperm(len(samples), generator=perm_gen) + return [samples[i] for i in perm] + packs = _require_rollout_packs( + self._rollout_acquisition_manifest, + samples, + self.training_args.per_device_batch_size, + "optimize-loop pack shuffling", + ) + pack_perm = torch.randperm(len(packs), generator=perm_gen) + return [samples[i] for pack_index in pack_perm.tolist() for i in packs[pack_index]] def _maybe_offload_samples_to_cpu(self, samples: List[BaseSample]) -> None: """Offload each sample's tensors to pinned CPU when offload is enabled. @@ -3161,8 +3200,12 @@ def generate_samples( reward_buffer is not None and getattr(self.training_args, "reward_optimization_overlap", False) ) + record_rollout_batches = ( + capture_overlap_manifest or self.adapter.requires_preserved_replay_batch_composition + ) recorded_rollout_batches: List[Tuple[int, ...]] = [] self._reward_overlap_acquisition_manifest = None + self._rollout_acquisition_manifest = None if reward_buffer is not None: reward_buffer.clear() @@ -3196,15 +3239,17 @@ def generate_samples( trajectory_indices=trajectory_indices, **extra_inference_kwargs, ) - if capture_overlap_manifest: + if record_rollout_batches: recorded_rollout_batches.append(tuple(id(sample) for sample in sample_batch)) samples.extend(sample_batch) - if capture_overlap_manifest: - self._reward_overlap_acquisition_manifest = AcquisitionManifest.from_samples( + if record_rollout_batches: + self._rollout_acquisition_manifest = AcquisitionManifest.from_samples( samples, rollout_batch_object_ids=recorded_rollout_batches, ) + if capture_overlap_manifest: + self._reward_overlap_acquisition_manifest = self._rollout_acquisition_manifest # Multi-source invariant: when more than one training source is # active, batches flow through `MultiSourceTrainDataLoader`, which diff --git a/tests/docs/test_flux3_docs.py b/tests/docs/test_flux3_docs.py new file mode 100644 index 00000000..e71196b5 --- /dev/null +++ b/tests/docs/test_flux3_docs.py @@ -0,0 +1,29 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] + + +def test_readme_lists_flux3_klein_models() -> None: + readme = (ROOT / "README.md").read_text(encoding="utf-8") + for repo in ( + "https://huggingface.co/black-forest-labs/flux-3-image-klein-base", + "https://huggingface.co/black-forest-labs/flux-3-image-klein", + ): + row_start = readme.index(repo) + model_row = readme[row_start : readme.index("", row_start)] + assert "5.5B" in model_row + assert "flux3" in model_row diff --git a/tests/docs/test_sensenova_docs.py b/tests/docs/test_sensenova_docs.py index abedc4a1..1c3a0e38 100644 --- a/tests/docs/test_sensenova_docs.py +++ b/tests/docs/test_sensenova_docs.py @@ -24,7 +24,7 @@ def _text(path: str) -> str: def test_main_readme_lists_sensenova_as_t2i_and_multi_reference_i2i() -> None: readme = _text("README.md") t2i_start = readme.index('Text-to-Image') - combined_start = readme.index('Text-to-Image & Image(s)-to-Image') + combined_start = readme.index('Text-to-Image & Image(s)-to-Image') video_start = readme.index('Text-to-Video') t2i_only_rows = readme[t2i_start:combined_start] diff --git a/tests/models/test_all_model_loaders.py b/tests/models/test_all_model_loaders.py index 31630831..9af909df 100644 --- a/tests/models/test_all_model_loaders.py +++ b/tests/models/test_all_model_loaders.py @@ -192,6 +192,36 @@ def test_bagel_adapter_imports_with_its_optional_kernel_contract( assert partial_inv_freq.shape == (2,) +def test_flux3_loader_builds_modular_pipeline_and_rewrites_local_specs( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + module = importlib.import_module("flow_factory.models.flux.flux3") + spec = SimpleNamespace( + pretrained_model_name_or_path="black-forest-labs/flux-3-image-klein-base", revision="main" + ) + pipeline = SimpleNamespace( + pretrained_component_names=["transformer"], + _component_specs={"transformer": spec}, + ) + calls = [] + + monkeypatch.setattr( + module.ModularPipeline, + "from_pretrained", + lambda path: calls.append(path) or pipeline, + ) + adapter = object.__new__(module.Flux3Adapter) + adapter.model_args = SimpleNamespace(model_name_or_path=str(tmp_path)) + + assert adapter.load_pipeline() is pipeline + assert calls == [str(tmp_path)] + assert spec.pretrained_model_name_or_path == str(tmp_path) + assert spec.revision is None + assert module.Flux3Adapter.supports_fsdp2_cpu_efficient_loading + assert module.Flux3Adapter.ddp_find_unused_parameters + + def test_all_registered_adapters_have_an_audited_load_pipeline() -> None: from flow_factory.models.registry import list_registered_models @@ -199,6 +229,7 @@ def test_all_registered_adapters_have_an_audited_load_pipeline() -> None: audited_paths = { f"{module_name}.{adapter_name}" for module_name, adapter_name in _CLASSIC_ADAPTERS } | { + "flow_factory.models.flux.flux3.Flux3Adapter", "flow_factory.models.bagel.bagel.BagelAdapter", "flow_factory.models.sensenova.sensenova.SenseNovaAdapter", "flow_factory.models.minimax_h3.adapters.MiniMaxH3T2VAAdapter", @@ -219,6 +250,7 @@ def test_model_type_literal_matches_registered_model_keys() -> None: def test_model_specific_load_dtype_defaults_are_explicit_and_narrow() -> None: + from flow_factory.models.flux.flux3 import Flux3Adapter from flow_factory.models.minimax_h3.adapters import ( MiniMaxH3FL2VAAdapter, MiniMaxH3Ref2VAAdapter, @@ -239,6 +271,7 @@ def test_model_specific_load_dtype_defaults_are_explicit_and_narrow() -> None: assert all( adapter.supports_fsdp2_cpu_efficient_loading for adapter in ( + Flux3Adapter, MiniMaxH3T2VAAdapter, MiniMaxH3FL2VAAdapter, MiniMaxH3Ref2VAAdapter, @@ -249,6 +282,7 @@ def test_model_specific_load_dtype_defaults_are_explicit_and_narrow() -> None: "text_encoders": torch.bfloat16, "vae": torch.float32, } + assert dict(Flux3Adapter.component_load_dtype_defaults) == expected_wan_defaults assert Wan2_T2V_Adapter.component_load_dtype_defaults == expected_wan_defaults assert Wan2_I2V_Adapter.component_load_dtype_defaults == { **expected_wan_defaults, diff --git a/tests/models/test_flux3.py b/tests/models/test_flux3.py new file mode 100644 index 00000000..7828b676 --- /dev/null +++ b/tests/models/test_flux3.py @@ -0,0 +1,795 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU contract, parity, and replay tests for the FLUX.3 adapter.""" + +from __future__ import annotations + +import math +from pathlib import Path +from types import SimpleNamespace +from typing import Any, List + +import numpy as np +import pytest +import torch +from PIL import Image + +from diffusers import ( + AutoencoderKLFlux2, + ClassifierFreeGuidance, + FlowMatchEulerDiscreteScheduler, + Flux3Transformer2DModel, +) +from diffusers.modular_pipelines import PipelineState +from diffusers.modular_pipelines.flux3 import ( + Flux3AutoBlocks, + Flux3DistilledAutoBlocks, + Flux3DistilledModularPipeline, + Flux3ModularPipeline, +) +from diffusers.modular_pipelines.flux3.before_denoise import ( + Flux3DistilledTextRoPEInputsStep, + Flux3PrepareImageCondLatentsStep, + Flux3PrepareLatentsStep, + Flux3SetTimestepsStep, + Flux3TextRoPEInputsStep, +) +from diffusers.modular_pipelines.flux3.decoders import Flux3DecodeStep, Flux3UnpackLatentsStep +from diffusers.modular_pipelines.flux3.denoise import Flux3DenoiseStep, Flux3DistilledDenoiseStep +from diffusers.modular_pipelines.flux3.encoders import ( + Flux3ProcessImagesInputStep, + Flux3TextEncoderStep, + Flux3VaeEncoderStep, +) +from diffusers.modular_pipelines.flux3.image_processor import Flux3ImageProcessor +from diffusers.modular_pipelines.flux3.inputs import ( + Flux3DistilledTextInputStep, + Flux3ImageCondInputStep, + Flux3TextInputStep, +) +from flow_factory.contracts import ( + BatchCapability, + GeometrySource, + InputMediaOrder, + NegativePromptPolicy, +) +from flow_factory.hparams import Arguments +from flow_factory.models.abc import BaseAdapter +from flow_factory.models.flux._output import encode_flux3_output_images +from flow_factory.models.flux.flux3 import ( + DEFAULT_MAX_SEQUENCE_LENGTH, + IMAGE_ROPE_SCALE, + TEXT_PAD_MULTIPLE, + Flux3Adapter, + Flux3Sample, +) +from flow_factory.models.registry import get_model_adapter_class +from flow_factory.models.runtime import ModularPipelineRuntime +from flow_factory.samples import BaseSample +from flow_factory.scheduler import FlowMatchEulerDiscreteSDEScheduler + + +class _FakeInnerTokenizer: + def __call__( + self, + text: str, + *, + return_tensors: str = "pt", + truncation: bool = True, + max_length: int | None = None, + padding: str | bool = False, + padding_side: str = "right", + ) -> dict[str, torch.Tensor]: + assert return_tensors == "pt" + assert truncation is True + raw_tokens = [min(ord(ch), 250) for ch in text][:max_length] + if not raw_tokens: + raw_tokens = [1] + if padding == "max_length": + assert padding_side == "right" + pad_count = max(0, max_length - len(raw_tokens)) + input_ids = raw_tokens + [0] * pad_count + attention_mask = [1] * len(raw_tokens) + [0] * pad_count + else: + input_ids = raw_tokens + attention_mask = [1] * len(raw_tokens) + return { + "input_ids": torch.tensor([input_ids], dtype=torch.long), + "attention_mask": torch.tensor([attention_mask], dtype=torch.long), + } + + +class _FakeProcessor: + def __init__(self) -> None: + self.tokenizer = _FakeInnerTokenizer() + self.chat_calls: list[str] = [] + + def apply_chat_template( + self, + messages: list[dict[str, str]], + *, + tokenize: bool = False, + add_generation_prompt: bool = True, + ) -> str: + assert tokenize is False + assert add_generation_prompt is True + content = messages[0]["content"] + self.chat_calls.append(content) + return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n" + + +class _FakeQwen3VL(torch.nn.Module): + def __init__(self, hidden_dim: int = 4, num_layers: int = 33) -> None: + super().__init__() + self.hidden_dim = hidden_dim + self.num_layers = num_layers + self.dummy_param = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32)) + self.calls: list[dict[str, Any]] = [] + + @property + def dtype(self) -> torch.dtype: + return self.dummy_param.dtype + + @property + def device(self) -> torch.device: + return self.dummy_param.device + + def forward( + self, + *, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + output_hidden_states: bool = True, + use_cache: bool = False, + ) -> Any: + assert output_hidden_states is True + assert use_cache is False + self.calls.append( + { + "input_ids": input_ids.clone(), + "attention_mask": attention_mask.clone(), + "position_ids": position_ids.clone(), + } + ) + batch_size, seq_len = input_ids.shape + base = ( + input_ids.to(torch.float32) * attention_mask.to(torch.float32) + + position_ids.to(torch.float32) * 0.01 + ) + base = base.unsqueeze(-1).repeat(1, 1, self.hidden_dim) + hidden_states = tuple(base + float(layer_idx) * 0.1 for layer_idx in range(self.num_layers)) + return SimpleNamespace(hidden_states=hidden_states) + + +def _build_tiny_flux3_components(*, flow_map: bool = False) -> dict[str, Any]: + torch.manual_seed(0) + transformer = Flux3Transformer2DModel( + in_channels=16, + cond_in_channels=16, + num_layers=1, + num_single_layers=1, + attention_head_dim=8, + num_attention_heads=2, + joint_attention_dim=16, + timestep_channels=32, + axes_dims_rope=(2, 2, 2, 2), + flow_map=flow_map, + ).eval() + vae = AutoencoderKLFlux2( + in_channels=3, + out_channels=3, + down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D"), + up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D"), + block_out_channels=(8, 16), + layers_per_block=1, + latent_channels=4, + norm_num_groups=4, + ).eval() + image_processor = Flux3ImageProcessor(vae_scale_factor=4, vae_latent_channels=4) + scheduler = FlowMatchEulerDiscreteScheduler(shift=6.93, use_dynamic_shifting=False) + text_encoder = _FakeQwen3VL(hidden_dim=2, num_layers=33).eval() + tokenizer = _FakeProcessor() + + blocks_cls = Flux3DistilledAutoBlocks if flow_map else Flux3AutoBlocks + pipeline: Flux3ModularPipeline = blocks_cls().init_pipeline() + update_kwargs: dict[str, Any] = { + "transformer": transformer, + "vae": vae, + "image_processor": image_processor, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + if not flow_map: + update_kwargs["guider"] = ClassifierFreeGuidance(guidance_scale=3.5) + pipeline.update_components(**update_kwargs) + return { + "pipeline": pipeline, + "transformer": transformer, + "vae": vae, + "image_processor": image_processor, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + + +def _build_flux3_adapter(*, flow_map: bool = False, noise_level: float = 0.0) -> Flux3Adapter: + bundle = _build_tiny_flux3_components(flow_map=flow_map) + adapter = object.__new__(Flux3Adapter) + adapter.config = SimpleNamespace( + mixed_precision="no", + scheduler_args=SimpleNamespace( + noise_level=noise_level, + sde_steps=None, + num_sde_steps=None, + seed=42, + dynamics_type="Flow-SDE" if noise_level > 0 else "ODE", + ), + ) + adapter.model_args = SimpleNamespace( + model_name_or_path="tiny-flux3", + target_components=["transformer"], + target_modules="default", + finetune_type="lora", + trainable_parameters_dtype=torch.float32, + component_load_dtypes=None, + frozen_parameters_dtype=None, + ) + adapter.training_args = SimpleNamespace( + height=32, + width=32, + latent_storage_dtype=None, + enable_gradient_checkpointing=False, + ) + adapter.eval_args = SimpleNamespace() + adapter.accelerator = SimpleNamespace( + device=torch.device("cpu"), + is_main_process=True, + is_local_main_process=True, + process_index=0, + distributed_type=None, + mixed_precision="no", + unwrap_model=lambda model: model, + wait_for_everyone=lambda: None, + ) + adapter._mode = "eval" if noise_level == 0.0 else "rollout" + adapter._component_load_dtype_manifest = Flux3Adapter.component_load_dtype_defaults + adapter._component_load_dtype_overrides = None + adapter.component_runtime = ModularPipelineRuntime.from_adapter(adapter, bundle["pipeline"]) + adapter.pipeline = bundle["pipeline"] + adapter._components = adapter.component_runtime.override_components + adapter.pipeline.scheduler = FlowMatchEulerDiscreteSDEScheduler( + shift=6.93, + use_dynamic_shifting=False, + noise_level=noise_level, + dynamics_type="Flow-SDE" if noise_level > 0 else "ODE", + seed=42, + ) + adapter.scheduler_group = adapter.build_scheduler_group() + adapter._effective_pipeline_io_contract = adapter._resolve_pipeline_io_contract() + adapter._condition_state_preparer = adapter._build_condition_state_preparer_declaration() + adapter._condition_state_encoding_modules = ( + adapter._validate_condition_state_preparer_lifecycle() + ) + adapter._output_state_codec = adapter._build_output_state_codec_declaration() + adapter._output_state_encoding_modules = adapter._validate_output_state_codec_lifecycle() + adapter._has_warned_inference_fallback = False + adapter._has_warned_forward_fallback = False + return adapter + + +def test_registry_inheritance_and_io_contracts() -> None: + assert get_model_adapter_class("flux3") is Flux3Adapter + Flux3Adapter.validate_offline_output_capability() + assert BaseAdapter in Flux3Adapter.__bases__ + assert not any( + base.__module__.startswith("flow_factory.models.flux.") and base is not Flux3Adapter + for base in Flux3Adapter.__mro__ + ) + assert not Flux3Adapter.supports_diffusers_cache + assert Flux3Adapter.supports_fsdp2_cpu_efficient_loading + assert Flux3Adapter.ddp_find_unused_parameters + + base_adapter = _build_flux3_adapter(flow_map=False) + assert not base_adapter.is_distilled + base_contract = base_adapter.effective_pipeline_io_contract + assert base_contract is not None + assert base_contract.negative_prompt is NegativePromptPolicy.OPTIONAL + assert base_contract.geometry_source is GeometrySource.CONFIGURED + assert base_contract.batch_capability is BatchCapability.UNIFORM + assert base_contract.input_media.order is InputMediaOrder.WITHIN_TYPE + + distilled_adapter = _build_flux3_adapter(flow_map=True) + assert distilled_adapter.is_distilled + distilled_contract = distilled_adapter.effective_pipeline_io_contract + assert distilled_contract is not None + assert distilled_contract.negative_prompt is NegativePromptPolicy.UNSUPPORTED + with pytest.raises(ValueError, match="does not support negative_prompt"): + distilled_adapter.encode_prompt("a fox", negative_prompt="blurry", guidance_scale=1.0) + + +def test_grpo_example_parses_through_production_config() -> None: + root = Path(__file__).resolve().parents[2] + config = Arguments.load_from_yaml(str(root / "examples/grpo/lora/flux3/default.yaml")) + + assert config.model_args.model_type == "flux3" + assert config.model_args.model_name_or_path == "black-forest-labs/flux-3-image-klein-base" + assert config.training_args.trainer_type == "grpo" + assert config.training_args.guidance_scale == 3.5 + assert config.eval_args.guidance_scale == 3.5 + + +def test_prompt_encoding_enforces_80_token_windows_and_pad_position_ids() -> None: + adapter = _build_flux3_adapter(flow_map=False) + encoded = adapter.encode_prompt( + prompt="a red fox", + guidance_scale=3.5, + text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), + pad_multiple=TEXT_PAD_MULTIPLE, + ) + + assert isinstance(encoded["prompt_ids"], torch.Tensor) + assert isinstance(encoded["prompt_embeds"], torch.Tensor) + assert isinstance(encoded["text_ids"], torch.Tensor) + assert encoded["prompt_ids"].shape == (1, 80) + assert encoded["prompt_embeds"].shape == (1, 80, 16) + assert encoded["text_ids"].shape == (1, 80, 4) + assert encoded["negative_prompt_embeds"].shape == (1, 80, 16) + assert encoded["negative_text_ids"].shape == (1, 80, 4) + + text_encoder: _FakeQwen3VL = adapter.text_encoder # type: ignore[assignment] + first_call = text_encoder.calls[0] + pad_positions = first_call["position_ids"][first_call["attention_mask"] == 0] + assert pad_positions.numel() > 0 + assert torch.all(pad_positions == 1) + + # Variable-length prompts with different 80-token windows remain un-merged so shorter + # prompts are never polluted with extra unmasked padding tokens. + short_prompt = "fox" + long_prompt = "word " * 15 # 75 chars + chat template -> (80, 160] window + ragged = adapter.encode_prompt( + prompt=[short_prompt, long_prompt], + guidance_scale=1.0, + text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), + ) + assert isinstance(ragged["prompt_embeds"], list) + assert ragged["prompt_embeds"][0].shape == (80, 16) + assert ragged["prompt_embeds"][1].shape == (160, 16) + + +def test_default_max_sequence_length_matches_upstream_and_keeps_long_prompts() -> None: + upstream_default = next( + param.default + for param in Flux3TextEncoderStep().inputs + if param.name == "max_sequence_length" + ) + assert DEFAULT_MAX_SEQUENCE_LENGTH == upstream_default == 8192 + + adapter = _build_flux3_adapter(flow_map=False) + # The fake tokenizer emits one token per character. + encoded = adapter.encode_prompt(prompt="x" * 600, guidance_scale=1.0) + + text_encoder: _FakeQwen3VL = adapter.text_encoder # type: ignore[assignment] + real_length = int(text_encoder.calls[-1]["attention_mask"].sum()) + assert real_length > 600 + expected_window = math.ceil(real_length / TEXT_PAD_MULTIPLE) * TEXT_PAD_MULTIPLE + assert encoded["prompt_embeds"].shape == (1, expected_window, 16) + assert encoded["text_ids"].shape == (1, expected_window, 4) + + # Raw-prompt rollouts (e.g. enable_preprocess: false) encode through inference(). + adapter.inference( + prompt=["x" * 600], + height=16, + width=16, + num_inference_steps=1, + guidance_scale=1.0, + generator=torch.Generator().manual_seed(0), + compute_log_prob=False, + ) + assert int(text_encoder.calls[-1]["attention_mask"].sum()) == real_length + + +def test_condition_image_encoding_assigns_negative_l_coords_and_preserves_empty_slots() -> None: + adapter = _build_flux3_adapter(flow_map=False) + ref0 = Image.new("RGB", (16, 8), (120, 40, 200)) + ref1 = Image.new("RGB", (8, 8), (10, 180, 90)) + + encoded = adapter.encode_image([[], [ref0, ref1]], condition_image_size=(64, 64)) + + assert encoded["condition_images"][0] == [] + assert encoded["image_latents"][0] is None + assert encoded["image_latent_ids"][0] is None + + # Tiny VAE has scale_factor=2, so 2x2 patchify grid is 4px: + # ref0 (16x8) -> 4x2 = 8 tokens at l=-1; ref1 (8x8) -> 2x2 = 4 tokens at l=-2. + latents_1 = encoded["image_latents"][1] + ids_1 = encoded["image_latent_ids"][1] + assert isinstance(latents_1, torch.Tensor) + assert isinstance(ids_1, torch.Tensor) + assert latents_1.shape == (12, 16) + assert ids_1.shape == (12, 4) + assert torch.all(ids_1[:8, 3] == -1.0) + assert torch.all(ids_1[8:, 3] == -2.0) + assert torch.allclose(ids_1[:8, 1].max(), torch.tensor(1.0 * IMAGE_ROPE_SCALE)) + assert torch.allclose(ids_1[:8, 2].max(), torch.tensor(3.0 * IMAGE_ROPE_SCALE)) + + +def test_output_codec_and_decode_latents_round_trip() -> None: + adapter = _build_flux3_adapter(flow_map=False) + pixels = torch.zeros(2, 3, 16, 16) + encoded = encode_flux3_output_images( + adapter, + pixels, + torch.Generator().manual_seed(7), + image_rope_scale=IMAGE_ROPE_SCALE, + ) + assert encoded.latents.shape == (2, 16, 16) + assert encoded.forward_context["latent_ids"].shape == (2, 16, 4) + + decoded_from_hw = adapter.decode_latents( + encoded.latents, + height=16, + width=16, + output_type="pt", + ) + decoded_from_ids = adapter.decode_latents( + encoded.latents, + latent_ids=encoded.forward_context["latent_ids"], + output_type="pt", + ) + assert decoded_from_hw.shape == (2, 3, 16, 16) + torch.testing.assert_close(decoded_from_hw, decoded_from_ids, rtol=0, atol=0) + + +@pytest.mark.parametrize("flow_map", [False, True]) +def test_end_to_end_parity_against_diffusers_modular_blocks(flow_map: bool) -> None: + """Verify exact numerical parity between Flux3Adapter.inference() and Diffusers blocks.""" + adapter = _build_flux3_adapter(flow_map=flow_map, noise_level=0.0) + pipe = adapter.pipeline + height, width = 16, 16 + num_inference_steps = 3 + guidance_scale = 1.0 if flow_map else 3.5 + layers = (4, 8, 12, 16, 20, 24, 28, 32) + ref_images = [ + Image.fromarray(np.full((16, 8, 3), 80, dtype=np.uint8)), + Image.fromarray(np.full((8, 8, 3), 160, dtype=np.uint8)), + ] + + # 1. Encode prompts with adapter + prompt_enc = adapter.encode_prompt( + prompt="a snowy forest", + negative_prompt="" if not flow_map else None, + guidance_scale=guidance_scale, + max_sequence_length=64, + text_encoder_out_layers=layers, + ) + + # 2. Run official Diffusers modular blocks on the same components & seed + ref_pipe = (Flux3DistilledAutoBlocks if flow_map else Flux3AutoBlocks)().init_pipeline() + ref_update = { + "transformer": adapter.transformer, + "vae": adapter.vae, + "image_processor": pipe.image_processor, + "scheduler": FlowMatchEulerDiscreteScheduler(shift=6.93, use_dynamic_shifting=False), + "text_encoder": adapter.text_encoder, + "tokenizer": adapter.processor, + } + if not flow_map: + ref_update["guider"] = ClassifierFreeGuidance(guidance_scale=guidance_scale) + ref_pipe.update_components(**ref_update) + + state = PipelineState( + values={ + "prompt_embeds": prompt_enc["prompt_embeds"].clone(), + "negative_prompt_embeds": ( + None if flow_map else prompt_enc["negative_prompt_embeds"].clone() + ), + "image": ref_images, + "height": height, + "width": width, + "num_images_per_prompt": 1, + "generator": torch.Generator().manual_seed(123), + "num_inference_steps": num_inference_steps, + "timesteps": None, + "sigmas": None, + "latents": None, + "attention_kwargs": None, + "output_type": "pt", + } + ) + steps: List[Any] = [ + Flux3ProcessImagesInputStep(), + Flux3VaeEncoderStep(), + Flux3DistilledTextInputStep() if flow_map else Flux3TextInputStep(), + Flux3ImageCondInputStep(default_height=height, default_width=width), + Flux3PrepareLatentsStep( + default_height=height, + default_width=width, + image_rope_scale=IMAGE_ROPE_SCALE, + ), + Flux3PrepareImageCondLatentsStep(image_rope_scale=IMAGE_ROPE_SCALE), + Flux3SetTimestepsStep(shift=6.93, num_inference_steps=num_inference_steps), + Flux3DistilledTextRoPEInputsStep() if flow_map else Flux3TextRoPEInputsStep(), + Flux3DistilledDenoiseStep() if flow_map else Flux3DenoiseStep(), + Flux3UnpackLatentsStep(), + Flux3DecodeStep(), + ] + for step in steps: + ref_pipe, state = step(ref_pipe, state) + expected_images = state.values["images"] + + # 3. Run Flux3Adapter.inference() with the same seed and condition images + samples = adapter.inference( + images=[ref_images], + prompt=["a snowy forest"], + negative_prompt=[""] if not flow_map else None, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + condition_image_size=(64, 64), + max_sequence_length=64, + text_encoder_out_layers=layers, + generator=torch.Generator().manual_seed(123), + compute_log_prob=False, + ) + + assert len(samples) == 1 + torch.testing.assert_close(samples[0].image, expected_images[0], rtol=0, atol=1e-6) + + +def test_sde_rollout_replay_parity_and_gradient_flow() -> None: + adapter = _build_flux3_adapter(flow_map=False, noise_level=0.7) + adapter.transformer.train() + ref_image = Image.fromarray(np.full((8, 8, 3), 100, dtype=np.uint8)) + + samples = adapter.inference( + images=[[ref_image]], + prompt=["relight "], + height=16, + width=16, + num_inference_steps=3, + guidance_scale=2.0, + condition_image_size=(64, 64), + max_sequence_length=64, + text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), + generator=torch.Generator().manual_seed(99), + compute_log_prob=True, + ) + sample = samples[0] + assert isinstance(sample, Flux3Sample) + assert sample.all_latents is not None + assert sample.log_probs is not None + assert sample.all_latents.shape == (4, 16, 16) + assert sample.log_probs.shape == (2,) + + batch = BaseSample.stack(samples) + step_idx = 0 + t = batch["timesteps"][:, step_idx] + t_next = batch["timesteps"][:, step_idx + 1] + latents = batch["all_latents"][:, step_idx] + next_latents = batch["all_latents"][:, step_idx + 1] + noise_level = adapter.scheduler.get_noise_level_for_timestep(t[0]) + + replay = adapter.forward( + t=t, + t_next=t_next, + latents=latents, + next_latents=next_latents, + latent_ids=batch["latent_ids"], + prompt_embeds=batch["prompt_embeds"], + text_ids=batch["text_ids"], + image_latents=batch["image_latents"], + image_latent_ids=batch["image_latent_ids"], + negative_prompt_embeds=batch["negative_prompt_embeds"], + negative_text_ids=batch["negative_text_ids"], + guidance_scale=2.0, + noise_level=noise_level, + compute_log_prob=True, + ) + + assert replay.log_prob is not None + torch.testing.assert_close( + replay.log_prob.detach(), + batch["log_probs"][:, step_idx], + rtol=0, + atol=1e-6, + ) + ratio = torch.exp(replay.log_prob - batch["log_probs"][:, step_idx]) + torch.testing.assert_close(ratio.detach(), torch.ones_like(ratio), rtol=0, atol=1e-6) + + loss = -replay.log_prob.mean() + loss.backward() + + assert adapter.transformer.txt_mode_blocks[0].attn.to_q.weight.grad is not None + assert adapter.transformer.mode_blocks["image"][0].attn.to_q.weight.grad is not None + assert adapter.transformer.mode_blocks["image_cond"][0].attn.to_q.weight.grad is not None + assert adapter.transformer.single_transformer_blocks[0].attn.to_q.weight.grad is not None + + +def test_rollout_packs_replay_exactly_through_their_rollout_path() -> None: + # Window-dependent batching makes the forward pack-dependent; the optimize loop then keeps + # every rollout pack intact, so replaying a pack must reproduce its rollout bit for bit. + assert Flux3Adapter.requires_preserved_replay_batch_composition is True + adapter = _build_flux3_adapter(flow_map=False, noise_level=0.7) + adapter.transformer.train() + transformer_batch_sizes: list[int] = [] + adapter.transformer.register_forward_pre_hook( + lambda module, args, kwargs: transformer_batch_sizes.append( + kwargs["hidden_states"].shape[0] + ), + with_kwargs=True, + ) + short_prompt, long_prompt = "fox", "word " * 15 # one and two 80-token windows + for pack, batch_size in (([short_prompt, short_prompt], 2), ([short_prompt, long_prompt], 1)): + transformer_batch_sizes.clear() + samples = adapter.inference( + prompt=pack, + height=16, + width=16, + num_inference_steps=3, + guidance_scale=2.0, + text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), + generator=torch.Generator().manual_seed(7), + compute_log_prob=True, + ) + assert set(transformer_batch_sizes) == {batch_size} + + transformer_batch_sizes.clear() + batch = BaseSample.stack(samples) + for step_idx in range(2): + t = batch["timesteps"][:, step_idx] + replay = adapter.forward( + t=t, + t_next=batch["timesteps"][:, step_idx + 1], + latents=batch["all_latents"][:, step_idx], + next_latents=batch["all_latents"][:, step_idx + 1], + latent_ids=batch["latent_ids"], + prompt_embeds=batch["prompt_embeds"], + text_ids=batch["text_ids"], + negative_prompt_embeds=batch["negative_prompt_embeds"], + negative_text_ids=batch["negative_text_ids"], + guidance_scale=2.0, + noise_level=adapter.scheduler.get_noise_level_for_timestep(t[0]), + compute_log_prob=True, + ) + assert torch.equal(replay.log_prob.detach(), batch["log_probs"][:, step_idx]) + assert set(transformer_batch_sizes) == {batch_size} + + +def test_gradient_checkpointing_order_and_lora_round_trip(tmp_path: Path) -> None: + adapter = _build_flux3_adapter(flow_map=False, noise_level=0.0) + units = adapter._gradient_checkpointing_units("transformer", adapter.transformer) + unit_names = [name for name, _ in units] + assert unit_names == [ + "txt_mode_blocks.0", + "mode_blocks.image.0", + "mode_blocks.image_cond.0", + "single_transformer_blocks.0", + ] + + adapter.model_args.lora_rank = 4 + adapter.model_args.lora_alpha = 8 + adapter.target_module_map = {"transformer": adapter.default_target_modules} + adapter.apply_lora( + target_modules=adapter.default_target_modules, + components=["transformer"], + ) + + trainable_params = adapter.get_trainable_parameters() + assert len(trainable_params) > 0 + with torch.no_grad(): + for param in trainable_params: + param.add_(0.05) + expected_snapshot = [param.detach().clone() for param in trainable_params] + + ckpt_dir = tmp_path / "flux3_lora_ckpt" + adapter.save_checkpoint(str(ckpt_dir), save_ema=False, model_only=True) + + with torch.no_grad(): + for param in trainable_params: + param.zero_() + adapter.load_checkpoint(str(ckpt_dir), resume_type="lora") + + restored_params = adapter.get_trainable_parameters() + for expected, restored in zip(expected_snapshot, restored_params, strict=True): + torch.testing.assert_close(restored, expected, rtol=0, atol=0) + + +def test_offline_sft_forward_state_and_load_plan_contracts() -> None: + from flow_factory.loading import ComponentRole, build_adapter_load_plan + + adapter = _build_flux3_adapter(flow_map=False, noise_level=0.0) + adapter.training_args.height = 16 + adapter.training_args.width = 16 + + # 1. Verify ModelLoadCoordinator plan partitions TARGET, HOST, and AUXILIARY + plan = build_adapter_load_plan(adapter) + assert plan.request_for_component("transformer").role is ComponentRole.TARGET + assert plan.request_for_component("tokenizer").role is ComponentRole.HOST + assert plan.request_for_component("scheduler").role is ComponentRole.HOST + assert plan.request_for_component("text_encoder").role is ComponentRole.AUXILIARY + assert plan.request_for_component("vae").role is ComponentRole.AUXILIARY + assert plan.request_for_component("image_processor").role is ComponentRole.HOST + + # 2. Verify offline output-state codec + forward_state velocity matching + target_image = Image.new("RGB", (16, 16), (60, 120, 180)) + media_batch = ( + (SimpleNamespace(type="image", payload=target_image, fps=None, sample_rate=None),), + ) + prompt_enc = adapter.encode_prompt( + prompt="a castle", + guidance_scale=1.0, + max_sequence_length=64, + text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), + ) + condition = { + "prompt_embeds": prompt_enc["prompt_embeds"], + "text_ids": prompt_enc["text_ids"], + } + prepared_condition = adapter.prepare_condition_state(condition) + encoded_output = adapter.encode_output_state( + media_batch, + prepared_condition, + torch.Generator().manual_seed(17), + ) + + times = adapter.build_training_component_times(torch.tensor([500.0])) + noised = adapter.add_forward_process_noise( + encoded_output.clean_state, + times, + generator=torch.Generator().manual_seed(23), + ) + forward_kwargs = { + **prepared_condition.condition, + **prepared_condition.forward_context, + **encoded_output.forward_context, + **dict(adapter.offline_training_forward_overrides), + } + step_out = adapter.forward_state( + batch=forward_kwargs, + state=noised.state, + times=times, + compute_log_prob=False, + return_fields=("velocity",), + ) + assert step_out.velocity is not None + per_elem_sq_error = { + name: (step_out.velocity.components[name] - noised.target_velocity.components[name]).pow(2) + for name in adapter.trajectory_component_order + } + reduced = adapter.reduce_flow_matching_objective_values(per_elem_sq_error, state=noised.state) + assert reduced.shape == (1,) + assert torch.isfinite(reduced).all() + + # 3. Verify Flux3Sample reconstruction_required_fields + sample = Flux3Sample( + prompt="a castle", + image=torch.zeros(3, 16, 16), + height=16, + width=16, + latent_ids=torch.zeros(16, 4), + text_ids=torch.zeros(80, 4), + ) + reconstructed = Flux3Sample( + **{ + field_name: getattr(sample, field_name) + for field_name in Flux3Sample.reconstruction_required_fields + if hasattr(sample, field_name) + } + ) + assert isinstance(reconstructed, Flux3Sample) diff --git a/tests/models/test_modern_image_output_codecs.py b/tests/models/test_modern_image_output_codecs.py index 16b8e6ff..d2711e2d 100644 --- a/tests/models/test_modern_image_output_codecs.py +++ b/tests/models/test_modern_image_output_codecs.py @@ -32,9 +32,11 @@ from flow_factory.models.flux._output import ( encode_flux2_vae_image, prepare_flux2_condition_latents, + prepare_flux3_condition_latents, ) from flow_factory.models.flux.flux2 import Flux2Adapter from flow_factory.models.flux.flux2_klein import Flux2KleinAdapter +from flow_factory.models.flux.flux3 import Flux3Adapter from flow_factory.models.qwen_image._output import encode_qwen_vae_image from flow_factory.models.qwen_image.qwen_image import QwenImageAdapter from flow_factory.models.qwen_image.qwen_image_edit_plus import QwenImageEditPlusAdapter @@ -180,6 +182,7 @@ def _pack_qwen_latents( [ Flux2Adapter, Flux2KleinAdapter, + Flux3Adapter, QwenImageAdapter, QwenImageEditPlusAdapter, QwenImage21Adapter, @@ -452,6 +455,12 @@ def test_qwen_edit_keeps_condition_argmax_and_target_first_ordered_shapes() -> N BatchCapability.UNIFORM, InputMediaOrder.WITHIN_TYPE, ), + ( + Flux3Adapter, + GeometrySource.CONFIGURED, + BatchCapability.UNIFORM, + InputMediaOrder.WITHIN_TYPE, + ), ( QwenImageAdapter, GeometrySource.CONFIGURED, diff --git a/tests/models/test_offline_training_guidance.py b/tests/models/test_offline_training_guidance.py index c06d7637..f17ade16 100644 --- a/tests/models/test_offline_training_guidance.py +++ b/tests/models/test_offline_training_guidance.py @@ -19,6 +19,7 @@ from flow_factory.models.flux.flux1_kontext import Flux1KontextAdapter from flow_factory.models.flux.flux2 import Flux2Adapter from flow_factory.models.flux.flux2_klein import Flux2KleinAdapter +from flow_factory.models.flux.flux3 import Flux3Adapter from flow_factory.models.qwen_image.qwen_image import QwenImageAdapter from flow_factory.models.qwen_image.qwen_image_edit_plus import QwenImageEditPlusAdapter from flow_factory.models.qwen_image_21.qwen_image_21 import QwenImage21Adapter @@ -39,6 +40,7 @@ def test_offline_training_forward_overrides_are_immutable() -> None: BaseAdapter, SD3_5Adapter, Flux2KleinAdapter, + Flux3Adapter, QwenImageAdapter, QwenImageEditPlusAdapter, QwenImage21Adapter, diff --git a/tests/trainers/test_optimize_sample_order.py b/tests/trainers/test_optimize_sample_order.py new file mode 100644 index 00000000..57a01c14 --- /dev/null +++ b/tests/trainers/test_optimize_sample_order.py @@ -0,0 +1,181 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Optimize-loop sample ordering for batch-invariant and pack-dependent adapters.""" + +from __future__ import annotations + +import contextlib +from types import SimpleNamespace +from typing import Any, Iterator, List, Optional + +import pytest +import torch + +from flow_factory.samples import AcquisitionManifest +from flow_factory.trainers.abc import BaseTrainer +from flow_factory.utils.base import create_generator + +SEED, EPOCH = 42, 3 + + +def _samples(num_packs: int, pack_size: int) -> List[SimpleNamespace]: + return [ + SimpleNamespace(unique_id=pack, name=f"pack{pack}-item{item}") + for pack in range(num_packs) + for item in range(pack_size) + ] + + +def _packs_manifest(samples: List[Any], pack_size: int) -> AcquisitionManifest: + return AcquisitionManifest.from_samples( + samples, + rollout_batch_object_ids=[ + [id(sample) for sample in samples[start : start + pack_size]] + for start in range(0, len(samples), pack_size) + ], + ) + + +def _trainer( + *, + pack_dependent: bool, + shuffle: bool = True, + per_device_batch_size: int = 2, + manifest: Optional[AcquisitionManifest] = None, +) -> SimpleNamespace: + return SimpleNamespace( + training_args=SimpleNamespace( + shuffle_samples=shuffle, seed=SEED, per_device_batch_size=per_device_batch_size + ), + epoch=EPOCH, + adapter=SimpleNamespace(requires_preserved_replay_batch_composition=pack_dependent), + _rollout_acquisition_manifest=manifest, + ) + + +def _names(samples: List[Any]) -> List[str]: + return [sample.name for sample in samples] + + +@pytest.mark.parametrize("pack_dependent", [False, True]) +def test_shuffle_samples_false_keeps_the_rollout_order(pack_dependent: bool) -> None: + samples = _samples(num_packs=4, pack_size=2) + trainer = _trainer(pack_dependent=pack_dependent, shuffle=False) + + assert BaseTrainer._order_samples_for_optimize(trainer, samples, 0) == samples + + +def test_batch_invariant_adapters_shuffle_individual_samples() -> None: + samples = _samples(num_packs=6, pack_size=2) + trainer = _trainer(pack_dependent=False) + + ordered = BaseTrainer._order_samples_for_optimize(trainer, samples, inner_epoch=1) + + permutation = torch.randperm(len(samples), generator=create_generator(SEED, EPOCH, 1)) + assert ordered == [samples[index] for index in permutation] + assert ordered != samples + + +def test_pack_dependent_adapters_shuffle_whole_rollout_packs() -> None: + pack_size = 2 + samples = _samples(num_packs=6, pack_size=pack_size) + trainer = _trainer(pack_dependent=True, manifest=_packs_manifest(samples, pack_size)) + + ordered = BaseTrainer._order_samples_for_optimize(trainer, samples, inner_epoch=1) + + micro_batches = [ + _names(ordered[start : start + pack_size]) for start in range(0, len(ordered), pack_size) + ] + rollout_packs = [ + _names(samples[start : start + pack_size]) for start in range(0, len(samples), pack_size) + ] + pack_permutation = torch.randperm( + len(rollout_packs), generator=create_generator(SEED, EPOCH, 1) + ) + assert micro_batches == [rollout_packs[index] for index in pack_permutation] + assert micro_batches != rollout_packs + assert BaseTrainer._order_samples_for_optimize(trainer, samples, inner_epoch=1) == ordered + + +def test_single_sample_packs_shuffle_like_individual_samples() -> None: + samples = _samples(num_packs=8, pack_size=1) + pack_dependent = _trainer( + pack_dependent=True, per_device_batch_size=1, manifest=_packs_manifest(samples, 1) + ) + batch_invariant = _trainer(pack_dependent=False, per_device_batch_size=1) + + assert BaseTrainer._order_samples_for_optimize( + pack_dependent, samples, 2 + ) == BaseTrainer._order_samples_for_optimize(batch_invariant, samples, 2) + + +def test_pack_shuffling_requires_matching_rollout_packs() -> None: + samples = _samples(num_packs=3, pack_size=2) + + with pytest.raises(RuntimeError, match="optimize-loop pack shuffling .* rollout micro-batch"): + BaseTrainer._order_samples_for_optimize(_trainer(pack_dependent=True), samples, 0) + + uneven = _trainer( + pack_dependent=True, + per_device_batch_size=2, + manifest=_packs_manifest(samples, pack_size=3), + ) + with pytest.raises(ValueError, match=r"train.per_device_batch_size=2; invalid_sizes=\(3, 3\)"): + BaseTrainer._order_samples_for_optimize(uneven, samples, 0) + + reordered = _trainer(pack_dependent=True, manifest=_packs_manifest(samples, pack_size=2)) + with pytest.raises(ValueError, match="samples changed after its manifest was created"): + BaseTrainer._order_samples_for_optimize(reordered, samples[::-1], 0) + + +class _GenerationHarness: + def __init__(self, pack_dependent: bool) -> None: + self.adapter = SimpleNamespace( + rollout=lambda: None, requires_preserved_replay_batch_composition=pack_dependent + ) + self.training_args = SimpleNamespace( + reward_optimization_overlap=False, num_batches_per_epoch=2 + ) + self.dataloader = [{"pack": 0}, {"pack": 1}] + self.epoch = 0 + self.show_progress_bar = False + self.train_dataloaders_by_source = {"default": self.dataloader} + + @contextlib.contextmanager + def _rollout_acceleration(self) -> Iterator[None]: + yield + + @contextlib.contextmanager + def autocast(self) -> Iterator[None]: + yield + + def sample_batch(self, batch: dict, **_: Any) -> List[SimpleNamespace]: + return [SimpleNamespace(unique_id=batch["pack"], source=None) for _ in range(2)] + + +@pytest.mark.parametrize("pack_dependent", [False, True]) +def test_generation_records_rollout_packs_for_pack_dependent_adapters( + pack_dependent: bool, +) -> None: + harness = _GenerationHarness(pack_dependent) + + samples = BaseTrainer.generate_samples(harness) + + assert harness._reward_overlap_acquisition_manifest is None + if pack_dependent: + assert harness._rollout_acquisition_manifest.rollout_batches == ((0, 1), (2, 3)) + harness._rollout_acquisition_manifest.validate_samples(samples) + else: + assert harness._rollout_acquisition_manifest is None